Fixin some benchmarks

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

.github/workflows/sql-benchmarks.yml

@@ -169,7 +169,7 @@ jobs:
             target/release_debug/df-bench ${{ matrix.subcommand }} \
               --formats "$df_formats" \
               -i1 \
-              -d gh-json \
+              -d table \
               $opts
           fi
 
@@ -178,15 +178,15 @@ jobs:
             target/release_debug/ddb-bench ${{ matrix.subcommand }} \
               --formats "$ddb_formats" \
               -i1 \
-              -d gh-json \
+              -d table \
               $opts
           fi
 
           # Generate data with lance-bench (runs each query once)
           if [ "$has_lance" = "true" ] && [ -f "target/release_debug/lance-bench" ]; then
             target/release_debug/lance-bench ${{ matrix.subcommand }} \
               -i1 \
-              -d gh-json \
+              -d table \
               $opts
           fi
 
@@ -278,10 +278,9 @@ jobs:
           OTEL_EXPORTER_OTLP_HEADERS: "${{ (inputs.mode != 'pr' || github.event.pull_request.head.repo.fork == false) && secrets.OTEL_EXPORTER_OTLP_HEADERS || '' }}"
           OTEL_RESOURCE_ATTRIBUTES: "bench-name=${{ matrix.id }}"
         run: |
-          # Extract formats for each engine (filter out lance from df_formats)
-          df_formats=$(echo "${{ matrix.targets }}" | tr ',' '\n' | grep '^datafusion:' | grep -v ':lance$' | sed 's/datafusion://' | tr '\n' ',' | sed 's/,$//')
+          # Extract formats for each engine
+          df_formats=$(echo "${{ matrix.targets }}" | tr ',' '\n' | grep '^datafusion:' | sed 's/datafusion://' | tr '\n' ',' | sed 's/,$//')
           ddb_formats=$(echo "${{ matrix.targets }}" | tr ',' '\n' | grep '^duckdb:' | sed 's/duckdb://' | tr '\n' ',' | sed 's/,$//')
-          has_lance=$(echo "${{ matrix.targets }}" | grep -q 'datafusion:lance' && echo "true" || echo "false")
 
           # Build options string if scale_factor is set
           opts=""
@@ -310,17 +309,8 @@ jobs:
               -o ddb-results.json
           fi
 
-          # Run lance-bench with remote storage
-          if [ "$has_lance" = "true" ] && [ -f "target/release_debug/lance-bench" ]; then
-            target/release_debug/lance-bench ${{ matrix.subcommand }} \
-              --use-remote-data-dir ${{ matrix.remote_storage }} \
-              -d gh-json \
-              $opts \
-              -o lance-results.json
-          fi
-
           # Merge results
-          cat df-results.json ddb-results.json lance-results.json 2>/dev/null > results.json || true
+          cat df-results.json ddb-results.json 2>/dev/null > results.json || true
 
       - name: Install uv
         if: inputs.mode == 'pr'

Cargo.lock

@@ -9,14 +9,14 @@ use clap::Parser;
 use clap::value_parser;
 use ddb_bench::DuckDBCtx;
 use vortex_bench::BenchmarkArg;
-use vortex_bench::BenchmarkOutput;
 use vortex_bench::CompactionStrategy;
 use vortex_bench::Engine;
 use vortex_bench::Format;
 use vortex_bench::Opt;
 use vortex_bench::Opts;
 use vortex_bench::conversions::convert_parquet_to_vortex;
 use vortex_bench::create_benchmark;
+use vortex_bench::create_output_writer;
 use vortex_bench::display::DisplayFormat;
 use vortex_bench::runner::SqlBenchmarkRunner;
 use vortex_bench::runner::filter_queries;
@@ -132,8 +132,9 @@ async fn main() -> anyhow::Result<()> {
         |ctx, _query_idx, query| ctx.execute_query(query),
     )?;
 
-    let output = BenchmarkOutput::with_path(benchmark.dataset_name(), args.output_path);
-    runner.export_to(&args.display_format, output.create_writer()?)?;
+    let benchmark_id = format!("duckdb-{}", benchmark.dataset_name());
+    let writer = create_output_writer(&args.display_format, args.output_path, &benchmark_id)?;
+    runner.export_to(&args.display_format, writer)?;
 
     Ok(())
 }

benchmarks/ddb-bench/src/main.rs

@@ -15,6 +15,7 @@ version = { workspace = true }
 
 [dependencies]
 anyhow = { workspace = true }
+arrow-ipc.workspace = true
 clap = { workspace = true, features = ["derive"] }
 datafusion = { workspace = true, features = [
     "parquet",
@@ -24,6 +25,7 @@ datafusion = { workspace = true, features = [
 ] }
 datafusion-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+futures.workspace = true
 itertools.workspace = true
 object_store = { workspace = true, features = ["aws", "gcp"] }
 opentelemetry.workspace = true

benchmarks/df-bench/Cargo.toml

@@ -11,19 +11,22 @@ use datafusion::datasource::listing::ListingOptions;
 use datafusion::datasource::listing::ListingTable;
 use datafusion::datasource::listing::ListingTableConfig;
 use datafusion::datasource::listing::ListingTableUrl;
+use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion_physical_plan::ExecutionPlan;
 use df_bench::format_to_df_format;
+use futures::StreamExt;
+use tokio::fs::File;
 use vortex_bench::Benchmark;
 use vortex_bench::BenchmarkArg;
-use vortex_bench::BenchmarkOutput;
 use vortex_bench::CompactionStrategy;
 use vortex_bench::Engine;
 use vortex_bench::Format;
 use vortex_bench::Opt;
 use vortex_bench::Opts;
 use vortex_bench::conversions::convert_parquet_to_vortex;
 use vortex_bench::create_benchmark;
+use vortex_bench::create_output_writer;
 use vortex_bench::display::DisplayFormat;
 use vortex_bench::runner::SqlBenchmarkRunner;
 use vortex_bench::runner::filter_queries;
@@ -160,8 +163,9 @@ async fn main() -> anyhow::Result<()> {
         )
         .await?;
 
-    let output = BenchmarkOutput::with_path(benchmark.dataset_name(), args.output_path);
-    runner.export_to(&args.display_format, output.create_writer()?)?;
+    let benchmark_id = format!("datafusion-{}", benchmark.dataset_name());
+    let writer = create_output_writer(&args.display_format, args.output_path, &benchmark_id)?;
+    runner.export_to(&args.display_format, writer)?;
 
     Ok(())
 }
@@ -171,26 +175,88 @@ async fn register_benchmark_tables<B: Benchmark + ?Sized>(
     benchmark: &B,
     format: Format,
 ) -> anyhow::Result<()> {
-    let benchmark_base = benchmark.data_url().join(&format!("{}/", format.name()))?;
-    let file_format = format_to_df_format(format);
+    if matches!(format, Format::Arrow) {
+        // For Arrow format, load Arrow IPC files into in-memory tables
+        register_arrow_tables(session, benchmark).await
+    } else {
+        let benchmark_base = benchmark.data_url().join(&format!("{}/", format.name()))?;
+        let file_format = format_to_df_format(format);
+
+        for table in benchmark.table_specs().iter() {
+            let pattern = benchmark.pattern(table.name, format);
+            let table_url = ListingTableUrl::try_new(benchmark_base.clone(), pattern)?;
+
+            let mut config = ListingTableConfig::new(table_url).with_listing_options(
+                ListingOptions::new(file_format.clone())
+                    .with_session_config_options(session.state().config()),
+            );
+
+            config = match table.schema.as_ref() {
+                Some(schema) => config.with_schema(Arc::new(schema.clone())),
+                None => config.infer_schema(&session.state()).await?,
+            };
+
+            let listing_table = Arc::new(ListingTable::try_new(config)?);
+
+            session.register_table(table.name, listing_table)?;
+        }
 
-    for table in benchmark.table_specs().iter() {
-        let pattern = benchmark.pattern(table.name, format);
-        let table_url = ListingTableUrl::try_new(benchmark_base.clone(), pattern)?;
+        Ok(())
+    }
+}
 
-        let mut config = ListingTableConfig::new(table_url).with_listing_options(
-            ListingOptions::new(file_format.clone())
-                .with_session_config_options(session.state().config()),
-        );
+/// Load Arrow IPC files into in-memory DataFusion tables.
+async fn register_arrow_tables<B: Benchmark + ?Sized>(
+    session: &SessionContext,
+    benchmark: &B,
+) -> anyhow::Result<()> {
+    use datafusion::datasource::MemTable;
 
-        config = match table.schema.as_ref() {
-            Some(schema) => config.with_schema(Arc::new(schema.clone())),
-            None => config.infer_schema(&session.state()).await?,
-        };
+    let parquet_dir = benchmark
+        .data_url()
+        .to_file_path()
+        .map_err(|_| anyhow::anyhow!("Arrow format requires local file path"))?
+        .join(Format::Parquet.name());
 
-        let listing_table = Arc::new(ListingTable::try_new(config)?);
+    // Read all arrow files from the directory
+    let data_files = std::fs::read_dir(&parquet_dir)?.collect::<Result<Vec<_>, _>>()?;
 
-        session.register_table(table.name, listing_table)?;
+    for table in benchmark.table_specs().iter() {
+        let pattern = benchmark.pattern(table.name, Format::Parquet);
+
+        // Find files matching this table's pattern
+        let matching_files: Vec<_> = data_files
+            .iter()
+            .filter(|entry| {
+                let filename = entry.file_name();
+                let filename_str = filename.to_str().unwrap_or("");
+                match &pattern {
+                    Some(p) => p.matches(filename_str),
+                    None => filename_str == format!("{}.{}", table.name, Format::Parquet.ext()),
+                }
+            })
+            .collect();
+
+        // Load all matching files into memory
+        let mut all_batches = Vec::new();
+        let mut schema = None;
+
+        for dir_entry in matching_files {
+            let file = File::open(dir_entry.path()).await?;
+            let mut reader = ParquetRecordBatchStreamBuilder::new(file).await?.build()?;
+            if schema.is_none() {
+                schema = Some(reader.schema()).cloned();
+            }
+
+            while let Some(batch) = reader.next().await {
+                all_batches.push(batch?);
+            }
+        }
+
+        if let Some(schema) = schema {
+            let mem_table = MemTable::try_new(schema, vec![all_batches])?;
+            session.register_table(table.name, Arc::new(mem_table))?;
+        }
     }
 
     Ok(())

benchmarks/df-bench/src/main.rs

@@ -15,12 +15,12 @@ use lance_bench::convert::convert_parquet_to_lance;
 use tracing::info;
 use vortex_bench::Benchmark;
 use vortex_bench::BenchmarkArg;
-use vortex_bench::BenchmarkOutput;
 use vortex_bench::Engine;
 use vortex_bench::Format;
 use vortex_bench::Opt;
 use vortex_bench::Opts;
 use vortex_bench::create_benchmark;
+use vortex_bench::create_output_writer;
 use vortex_bench::display::DisplayFormat;
 use vortex_bench::runner::SqlBenchmarkRunner;
 use vortex_bench::runner::filter_queries;
@@ -117,8 +117,9 @@ async fn main() -> anyhow::Result<()> {
         )
         .await?;
 
-    let output = BenchmarkOutput::with_path(benchmark.dataset_name(), args.output_path);
-    runner.export_to(&args.display_format, output.create_writer()?)?;
+    let benchmark_id = format!("lance-{}", benchmark.dataset_name());
+    let writer = create_output_writer(&args.display_format, args.output_path, &benchmark_id)?;
+    runner.export_to(&args.display_format, writer)?;
 
     Ok(())
 }

benchmarks/lance-bench/src/main.rs

@@ -19,6 +19,7 @@ workspace = true
 [dependencies]
 anyhow = { workspace = true }
 arrow-array = { workspace = true }
+arrow-ipc = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-select = { workspace = true }
 async-trait = { workspace = true }

vortex-bench/Cargo.toml

@@ -7,6 +7,7 @@ use std::path::Path;
 use std::path::PathBuf;
 
 use arrow_array::RecordBatchReader;
+use arrow_ipc::writer::FileWriter;
 use futures::StreamExt;
 use futures::TryStreamExt;
 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
@@ -119,3 +120,73 @@ pub async fn convert_parquet_to_vortex(
         .await?;
     Ok(())
 }
+
+/// Convert all Parquet files in a directory to Arrow IPC format.
+///
+/// This function reads Parquet files from `{input_path}/parquet/` and writes
+/// Arrow IPC files to `{input_path}/arrow/`.
+///
+/// The conversion is idempotent - existing Arrow files will not be regenerated.
+pub async fn convert_parquet_to_arrow(input_path: &Path) -> anyhow::Result<()> {
+    let arrow_dir = input_path.join(Format::Arrow.name());
+    let parquet_path = input_path.join(Format::Parquet.name());
+    create_dir_all(&arrow_dir).await?;
+
+    let parquet_inputs = fs::read_dir(&parquet_path)?.collect::<std::io::Result<Vec<_>>>()?;
+
+    trace!(
+        "Found {} parquet files in {}",
+        parquet_inputs.len(),
+        parquet_path.to_str().unwrap()
+    );
+
+    let iter = parquet_inputs
+        .iter()
+        .filter(|entry| entry.path().extension().is_some_and(|e| e == "parquet"));
+
+    futures::stream::iter(iter)
+        .map(|dir_entry| {
+            let filename = {
+                let mut temp = dir_entry.path();
+                temp.set_extension("");
+                temp.file_name().unwrap().to_str().unwrap().to_string()
+            };
+            let parquet_file_path = parquet_path.join(format!("{filename}.parquet"));
+            let output_path = arrow_dir.join(format!("{filename}.{}", Format::Arrow.ext()));
+
+            tokio::spawn(
+                async move {
+                    idempotent_async(output_path.as_path(), move |arrow_file| async move {
+                        info!("Converting '{filename}' from Parquet to Arrow IPC");
+
+                        // Read Parquet file
+                        let parquet_file = File::open(&parquet_file_path)?;
+                        let reader =
+                            ParquetRecordBatchReaderBuilder::try_new(parquet_file)?.build()?;
+                        let schema = reader.schema();
+
+                        // Write Arrow IPC file
+                        let arrow_output = File::create(&arrow_file)?;
+                        let mut writer = FileWriter::try_new(arrow_output, &schema)?;
+
+                        for batch in reader {
+                            let batch = batch?;
+                            writer.write(&batch)?;
+                        }
+
+                        writer.finish()?;
+
+                        anyhow::Ok(())
+                    })
+                    .await
+                    .expect("Failed to write Arrow file")
+                }
+                .in_current_span(),
+            )
+        })
+        .buffer_unordered(16)
+        .try_collect::<Vec<_>>()
+        .await?;
+
+    Ok(())
+}

vortex-bench/src/conversions.rs

@@ -49,6 +49,7 @@ pub use benchmark::TableSpec;
 pub use datasets::BenchmarkDataset;
 pub use engines::df;
 pub use output::BenchmarkOutput;
+pub use output::create_output_writer;
 use vortex::VortexSessionDefault;
 pub use vortex::error::vortex_panic;
 use vortex::io::session::RuntimeSessionExt;