Refactor run function to aggregate batches and update command-line arguments

Author: kosiew

benchmarks/collect_gil_bench.py

@@ -21,14 +21,16 @@
 import time
 
 import pyarrow as pa
-from datafusion import SessionContext
+from datafusion import SessionContext, col
+from datafusion import functions as f
 
 
 def run(
     n_batches: int = 8,
     batch_size: int = 1_000_000,
     n_partitions: int | None = None,
 ) -> None:
+    """Aggregate column 'a' across partitions and report runtime."""
     ctx = SessionContext()
     batches = []
     for i in range(n_batches):
@@ -46,20 +48,36 @@ def run(
     df = ctx.create_dataframe(partitions)
 
     start = time.perf_counter()
-    df.collect()
+    df.aggregate([], [f.sum(col("a"))]).collect()
     duration = time.perf_counter() - start
-    print(f"{n_batches} batches collected in {duration:.3f}s")
+    print(f"{n_batches} batches aggregated in {duration:.3f}s")
 
 
 if __name__ == "__main__":
     import argparse
 
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--batches",
+        type=int,
+        default=8,
+        help="number of input batches to generate",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1_000_000,
+        help="number of rows per batch",
+    )
     parser.add_argument(
         "--partitions",
         type=int,
         default=None,
         help="number of partitions to create (defaults to one per batch)",
     )
     args = parser.parse_args()
-    run(n_partitions=args.partitions)
+    run(
+        n_batches=args.batches,
+        batch_size=args.batch_size,
+        n_partitions=args.partitions,
+    )