Update app.py

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 import matplotlib.pyplot as plt
 from PIL import Image
 import io
+import os
 
 duckdb_con = duckdb.connect(database=":memory:")
 
@@ -33,6 +34,7 @@ def generate_data(n_rows: int, n_groups: int = 50) -> pd.DataFrame:
 # ----------------------------------------------------------
 
 def time_function(fn, repeats=3):
+    repeats = int(repeats)
     times = []
     for _ in range(repeats):
         start = time.perf_counter()
@@ -43,45 +45,37 @@ def time_function(fn, repeats=3):
 
 
 # ----------------------------------------------------------
-# Benchmark Operations
+# Benchmark Operations (Compute + I/O)
 # ----------------------------------------------------------
 
+# ---- Filter ----
 def bench_filter(df, repeats=3):
     def pandas_op():
         _ = df[(df["value1"] > 0.5) & (df["category"] == df["category"].iloc[0])]
 
     def duckdb_op():
         duckdb_con.register("df", df)
-        q = f"""
-        SELECT *
-        FROM df
-        WHERE value1 > 0.5
-        AND category = '{df['category'].iloc[0]}'
-        """
-        _ = duckdb_con.execute(q).fetchdf()
+        duckdb_con.execute(f"""
+            SELECT *
+            FROM df
+            WHERE value1 > 0.5
+              AND category='{df['category'].iloc[0]}'
+        """).fetchdf()
 
     p_mean, p_std, p_all = time_function(pandas_op, repeats)
     d_mean, d_std, d_all = time_function(duckdb_op, repeats)
 
-    return {
-        "operation": "Filter rows with comparison",
-        "pandas_mean_s": p_mean,
-        "pandas_std_s": p_std,
-        "duckdb_mean_s": d_mean,
-        "duckdb_std_s": d_std,
-        "speedup": p_mean / d_mean if d_mean > 0 else None,
-        "raw_pandas_runs": p_all,
-        "raw_duckdb_runs": d_all,
-    }
+    return build_result("Filter rows", p_mean, p_std, p_all, d_mean, d_std, d_all)
 
 
+# ---- Groupby ----
 def bench_groupby(df, repeats=3):
     def pandas_op():
         _ = df.groupby("category")[["value1", "value2"]].mean()
 
     def duckdb_op():
         duckdb_con.register("df", df)
-        _ = duckdb_con.execute("""
+        duckdb_con.execute("""
             SELECT category, AVG(value1), AVG(value2)
             FROM df GROUP BY category
         """).fetchdf()
@@ -89,18 +83,10 @@ def bench_groupby(df, repeats=3):
     p_mean, p_std, p_all = time_function(pandas_op, repeats)
     d_mean, d_std, d_all = time_function(duckdb_op, repeats)
 
-    return {
-        "operation": "Groupby mean",
-        "pandas_mean_s": p_mean,
-        "pandas_std_s": p_std,
-        "duckdb_mean_s": d_mean,
-        "duckdb_std_s": d_std,
-        "speedup": p_mean / d_mean if d_mean > 0 else None,
-        "raw_pandas_runs": p_all,
-        "raw_duckdb_runs": d_all,
-    }
+    return build_result("Groupby mean", p_mean, p_std, p_all, d_mean, d_std, d_all)
 
 
+# ---- Join ----
 def bench_join(df, repeats=3):
     categories = df["category"].unique()
     rng = np.random.default_rng(123)
@@ -114,7 +100,7 @@ def bench_join(df, repeats=3):
     def duckdb_op():
         duckdb_con.register("df", df)
         duckdb_con.register("dim_df", dim_df)
-        _ = duckdb_con.execute("""
+        duckdb_con.execute("""
             SELECT d.*, dim.weight
             FROM df d
             LEFT JOIN dim_df dim
@@ -124,13 +110,66 @@ def bench_join(df, repeats=3):
     p_mean, p_std, p_all = time_function(pandas_op, repeats)
     d_mean, d_std, d_all = time_function(duckdb_op, repeats)
 
+    return build_result("Join on category", p_mean, p_std, p_all, d_mean, d_std, d_all)
+
+
+# ---- Read CSV ----
+def bench_read_csv(temp_csv_path, repeats=3):
+    def pandas_op():
+        _ = pd.read_csv(temp_csv_path)
+
+    def duckdb_op():
+        _ = duckdb.read_csv_auto(temp_csv_path)
+
+    p_mean, p_std, p_all = time_function(pandas_op, repeats)
+    d_mean, d_std, d_all = time_function(duckdb_op, repeats)
+
+    return build_result("Read CSV", p_mean, p_std, p_all, d_mean, d_std, d_all)
+
+
+# ---- Read Parquet ----
+def bench_read_parquet(temp_parquet_path, repeats=3):
+    def pandas_op():
+        _ = pd.read_parquet(temp_parquet_path)
+
+    def duckdb_op():
+        _ = duckdb.read_parquet(temp_parquet_path)
+
+    p_mean, p_std, p_all = time_function(pandas_op, repeats)
+    d_mean, d_std, d_all = time_function(duckdb_op, repeats)
+
+    return build_result("Read Parquet", p_mean, p_std, p_all, d_mean, d_std, d_all)
+
+
+# ---- Write Parquet ----
+def bench_write_parquet(df, repeats=3):
+    def pandas_op():
+        df.to_parquet("temp_pd.parquet")
+
+    def duckdb_op():
+        duckdb_con.register("df", df)
+        duckdb_con.execute("COPY df TO 'temp_duck.parquet' (FORMAT PARQUET)")
+
+    p_mean, p_std, p_all = time_function(pandas_op, repeats)
+    d_mean, d_std, d_all = time_function(duckdb_op, repeats)
+
+    return build_result("Write Parquet", p_mean, p_std, p_all, d_mean, d_std, d_all)
+
+
+# ----------------------------------------------------------
+# Shared result formatting
+# ----------------------------------------------------------
+
+def build_result(op_name, p_mean, p_std, p_all, d_mean, d_std, d_all):
+    speedup = p_mean / d_mean if d_mean > 0 else None
+
     return {
-        "operation": "Left Join",
+        "operation": op_name,
         "pandas_mean_s": p_mean,
         "pandas_std_s": p_std,
         "duckdb_mean_s": d_mean,
         "duckdb_std_s": d_std,
-        "speedup": p_mean / d_mean if d_mean > 0 else None,
+        "speedup": speedup,
         "raw_pandas_runs": p_all,
         "raw_duckdb_runs": d_all,
     }
@@ -140,18 +179,19 @@ def bench_join(df, repeats=3):
 # Benchmark Dispatcher
 # ----------------------------------------------------------
 
-def run_benchmark(operation, df, repeats):
-    if operation == "Filter":
-        return bench_filter(df, repeats)
-    if operation == "Groupby":
-        return bench_groupby(df, repeats)
-    if operation == "Join":
-        return bench_join(df, repeats)
-    return None
+def run_benchmark(operation, df=None, repeats=3):
+    repeats = int(repeats)
+
+    if operation == "Filter": return bench_filter(df, repeats)
+    if operation == "Groupby": return bench_groupby(df, repeats)
+    if operation == "Join": return bench_join(df, repeats)
+    if operation == "Write Parquet": return bench_write_parquet(df, repeats)
+
+    raise ValueError(f"Unsupported operation: {operation}")
 
 
 # ----------------------------------------------------------
-# Chart generator (PIL Image) — FIXED
+# Chart generator (PIL Image)
 # ----------------------------------------------------------
 
 def generate_chart(result):
@@ -162,7 +202,7 @@ def generate_chart(result):
 
     ax.bar(engines, times)
     ax.set_ylabel("Time (seconds)")
-    ax.set_title("Pandas vs DuckDB Performance")
+    ax.set_title(result["operation"])
 
     buf = io.BytesIO()
     plt.tight_layout()
@@ -174,7 +214,7 @@ def generate_chart(result):
 
 
 # ----------------------------------------------------------
-# Markdown Formatting
+# Markdown result
 # ----------------------------------------------------------
 
 def format_result(result):
@@ -204,6 +244,20 @@ def format_result(result):
     return md
 
 
+# ----------------------------------------------------------
+# Helper to load custom dataset
+# ----------------------------------------------------------
+
+def load_custom_dataset(file):
+    if file.name.endswith(".csv"):
+        return pd.read_csv(file.name)
+    if file.name.endswith(".parquet"):
+        return pd.read_parquet(file.name)
+    if file.name.endswith(".arrow"):
+        return pd.read_feather(file.name)
+    raise ValueError("Unsupported file format")
+
+
 # ----------------------------------------------------------
 # Gradio App
 # ----------------------------------------------------------
@@ -216,62 +270,74 @@ with gr.Blocks(title="DuckDB vs Pandas Benchmark", theme=theme) as demo:
 
     with gr.Tabs():
 
-        # ---------------------- Synthetic Mode ----------------------
+        # ==================================================
+        # 🔥 Synthetic Mode
+        # ==================================================
         with gr.Tab("🔥 Synthetic Dataset Benchmarks"):
-            dataset_size = gr.Radio(["100k", "500k", "2M"], label="Dataset Size", value="100k")
-            operation_synth = gr.Radio(["Filter", "Groupby", "Join"], label="Operation")
+
+            dataset_size = gr.Radio(["100k", "500k", "2M"], value="100k", label="Dataset Size")
+
+            operation_synth = gr.Radio(
+                ["Filter", "Groupby", "Join", "Write Parquet"],
+                label="Operation",
+                value="Filter"
+            )
+
             repeats_synth = gr.Slider(1, 7, value=3, label="Repeats")
-            synth_btn = gr.Button("🚀 Run Benchmark")
 
-            synth_output = gr.Markdown()
-            synth_chart = gr.Image(label="Performance Chart")
+            btn_synth = gr.Button("🚀 Run Benchmark")
+
+            out_md_synth = gr.Markdown()
+            out_chart_synth = gr.Image()
 
             def synthetic_runner(size, operation, repeats):
+                repeats = int(repeats)
                 n = {"100k": 100_000, "500k": 500_000, "2M": 2_000_000}[size]
+
                 df = generate_data(n)
                 result = run_benchmark(operation, df, repeats)
                 chart = generate_chart(result)
+
                 return format_result(result), chart
 
-            synth_btn.click(
+            btn_synth.click(
                 synthetic_runner,
                 [dataset_size, operation_synth, repeats_synth],
-                [synth_output, synth_chart],
+                [out_md_synth, out_chart_synth],
             )
 
-        # ---------------------- Custom Dataset Mode ----------------------
+
+        # ==================================================
+        # 📁 Custom Dataset Mode
+        # ==================================================
         with gr.Tab("📁 Custom Dataset Upload"):
 
-            file_input = gr.File(label="Upload a CSV / Parquet / Arrow file")
-            operation_custom = gr.Radio(["Filter", "Groupby", "Join"], label="Operation")
+            file_in = gr.File(label="Upload CSV / Parquet / Arrow")
+            
+            operation_custom = gr.Radio(
+                ["Filter", "Groupby", "Join", "Write Parquet"],
+                label="Operation",
+                value="Filter"
+            )
+
             repeats_custom = gr.Slider(1, 7, value=3, label="Repeats")
 
-            custom_btn = gr.Button("Run on Uploaded Dataset")
-            custom_out = gr.Markdown()
-            custom_chart = gr.Image(label="Performance Chart")
+            btn_custom = gr.Button("Run on Uploaded Dataset")
 
-            def load_custom_dataset(file):
-                if file.name.endswith(".csv"):
-                    return pd.read_csv(file.name)
-                elif file.name.endswith(".parquet"):
-                    return pd.read_parquet(file.name)
-                elif file.name.endswith(".arrow"):
-                    return pd.read_feather(file.name)
-                else:
-                    raise ValueError("Unsupported format")
+            out_md_custom = gr.Markdown()
+            out_chart_custom = gr.Image()
 
             def custom_runner(file, operation, repeats):
+                repeats = int(repeats)
                 df = load_custom_dataset(file)
                 result = run_benchmark(operation, df, repeats)
-                chart = generate_chart(result)
-                return format_result(result), chart
+                return format_result(result), generate_chart(result)
 
-            custom_btn.click(
+            btn_custom.click(
                 custom_runner,
-                [file_input, operation_custom, repeats_custom],
-                [custom_out, custom_chart],
+                [file_in, operation_custom, repeats_custom],
+                [out_md_custom, out_chart_custom],
             )
 
-
 if __name__ == "__main__":
     demo.launch()