chore[bench]: move tpcds into new benchmark format (#3870)

Signed-off-by: Joe Isaacs <[email protected]>

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

bench-vortex/src/benchmark_driver.rs

@@ -8,11 +8,12 @@ use std::path::PathBuf;
 use anyhow::Result;
 use indicatif::ProgressBar;
 use log::warn;
+use vortex::error::VortexExpect;
 use vortex_datafusion::metrics::VortexMetricsFinder;
 
 use crate::benchmark_trait::Benchmark;
 use crate::display::DisplayFormat;
-use crate::engines::{EngineCtx, benchmark_datafusion_query, benchmark_duckdb_query};
+use crate::engines::{EngineCtx, benchmark_datafusion_query};
 use crate::measurements::{MemoryMeasurement, QueryMeasurement};
 use crate::memory::BenchmarkMemoryTracker;
 use crate::metrics::{MetricsSetExt, export_plan_spans};
@@ -143,13 +144,12 @@ fn execute_queries<B: Benchmark>(
             tracker.start_query();
         }
 
-        match engine_ctx {
+        let row_count = match engine_ctx {
             EngineCtx::DataFusion(ctx) => {
                 let (runs, (row_count, execution_plan)) = runtime.block_on(async {
                     benchmark_datafusion_query(iterations, || async {
-                        let (batches, plan) = df::execute_query(&ctx.session, query_string)
-                            .await
-                            .unwrap_or_else(|err| {
+                        let (batches, plan) =
+                            ctx.execute_query(query_string).await.unwrap_or_else(|err| {
                                 vortex_panic!("query: {query_idx} failed with: {err}")
                             });
                         let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum();
@@ -158,16 +158,6 @@ fn execute_queries<B: Benchmark>(
                     .await
                 });
 
-                // Validate row count if expected counts are provided
-                if let Some(expected_counts) = expected_row_counts {
-                    if query_idx < expected_counts.len() {
-                        assert_eq!(
-                            row_count, expected_counts[query_idx],
-                            "Row count mismatch for query {query_idx} - datafusion:{format}",
-                        );
-                    }
-                }
-
                 ctx.execution_plans
                     .push((query_idx, execution_plan.clone()));
 
@@ -193,19 +183,24 @@ fn execute_queries<B: Benchmark>(
                     storage: url_scheme_to_storage(benchmark.data_url())?,
                     runs,
                 });
+
+                row_count
             }
             EngineCtx::DuckDB(ctx) => {
-                let (runs, row_count) =
-                    benchmark_duckdb_query(query_idx, query_string, iterations, ctx);
-
-                // Validate row count if expected counts are provided
-                if let Some(expected_counts) = expected_row_counts {
-                    if query_idx < expected_counts.len() {
-                        assert_eq!(
-                            row_count, expected_counts[query_idx],
-                            "Row count mismatch for query {query_idx} - duckdb:{format}",
-                        );
-                    }
+                let mut runs = Vec::with_capacity(iterations);
+                let mut row_count = None;
+
+                for _ in 0..iterations {
+                    let (duration, current_row_count) =
+                        ctx.execute_query(query_string).unwrap_or_else(|err| {
+                            vortex_panic!("query: {query_idx} failed with: {err}")
+                        });
+
+                    runs.push(duration);
+                    row_count.inspect(|rc| {
+                        assert_eq!(*rc, current_row_count, "each row count must match")
+                    });
+                    row_count = Some(current_row_count);
                 }
 
                 query_measurements.push(QueryMeasurement {
@@ -215,6 +210,18 @@ fn execute_queries<B: Benchmark>(
                     storage: url_scheme_to_storage(benchmark.data_url())?,
                     runs,
                 });
+
+                row_count.vortex_expect("cannot have zero runs")
+            }
+        };
+
+        // Validate row count if expected counts are provided
+        if let Some(expected_counts) = expected_row_counts {
+            if query_idx < expected_counts.len() {
+                assert_eq!(
+                    row_count, expected_counts[query_idx],
+                    "Row count mismatch for query {query_idx} - duckdb:{format}",
+                );
             }
         }
 

bench-vortex/src/bin/query_bench.rs

@@ -5,10 +5,10 @@ use std::path::PathBuf;
 
 use bench_vortex::Target;
 use bench_vortex::benchmark_driver::{DriverConfig, run_benchmark};
-use bench_vortex::clickbench::Flavor;
-use bench_vortex::clickbench_benchmark::ClickBenchBenchmark;
+use bench_vortex::clickbench::{ClickBenchBenchmark, Flavor};
 use bench_vortex::display::DisplayFormat;
-use bench_vortex::tpch_benchmark::TpcHBenchmark;
+use bench_vortex::tpcds::TpcDsBenchmark;
+use bench_vortex::tpch::tpch_benchmark::TpcHBenchmark;
 use clap::{Parser, Subcommand, value_parser};
 
 #[derive(Parser, Debug)]
@@ -27,6 +27,10 @@ enum Commands {
     /// Run TPC-H queries
     #[command(name = "tpch")]
     TpcH(TpcHArgs),
+
+    /// Run TPC-DS queries
+    #[command(name = "tpcds")]
+    TpcDS(TpcDSArgs),
 }
 
 /// Common arguments shared across benchmarks
@@ -123,6 +127,26 @@ struct TpcHArgs {
     scale_factor: String,
 }
 
+#[derive(Parser, Debug)]
+struct TpcDSArgs {
+    #[command(flatten)]
+    common: CommonArgs,
+
+    #[arg(long, value_delimiter = ',', value_parser = value_parser!(Target),
+        default_values = vec![
+            "datafusion:parquet",
+            "datafusion:vortex",
+            "duckdb:parquet",
+            "duckdb:vortex",
+            "duckdb:duckdb"
+        ]
+    )]
+    targets: Vec<Target>,
+
+    #[arg(long, default_value = "1.0", value_parser=validate_scale_factor)]
+    scale_factor: String,
+}
+
 fn validate_scale_factor(val: &str) -> Result<String, String> {
     match val.parse::<f32>() {
         Ok(n) if [0.01, 0.1, 1., 10., 100., 1000.].contains(&n) => {
@@ -150,6 +174,7 @@ fn main() -> anyhow::Result<()> {
     match args.command {
         Commands::ClickBench(clickbench_args) => run_clickbench(clickbench_args),
         Commands::TpcH(tpch_args) => run_tpch(tpch_args),
+        Commands::TpcDS(tpcds_args) => run_tpcds(tpcds_args),
     }
 }
 
@@ -212,3 +237,31 @@ fn run_tpch(args: TpcHArgs) -> anyhow::Result<()> {
 
     Ok(())
 }
+
+fn run_tpcds(args: TpcDSArgs) -> anyhow::Result<()> {
+    // Create benchmark instance
+    let benchmark = TpcDsBenchmark::new(args.scale_factor, args.common.use_remote_data_dir)?;
+
+    // Configure driver
+    let config = DriverConfig {
+        targets: args.targets,
+        iterations: args.common.iterations,
+        threads: args.common.threads,
+        verbose: args.common.verbose,
+        display_format: args.common.display_format,
+        disable_datafusion_cache: args.common.disable_datafusion_cache,
+        queries: args.common.queries,
+        exclude_queries: args.common.exclude_queries,
+        output_path: args.common.output_path,
+        emit_plan: args.common.emit_plan,
+        export_spans: args.common.export_spans,
+        show_metrics: args.common.show_metrics,
+        hide_progress_bar: args.common.hide_progress_bar,
+        track_memory: args.common.track_memory,
+    };
+
+    // Run benchmark using the trait system
+    run_benchmark(benchmark, config)?;
+
+    Ok(())
+}