Reintroduce emit-plan for TPC-H benchmark (#3124)

just seems it got lost somewhere with all the recent changes?

Author: AdamGS

bench-vortex/src/bin/clickbench.rs

@@ -85,7 +85,7 @@ struct DataFusionCtx {
     )>,
 
     session: SessionContext,
-    emit_execution_plan: bool,
+    emit_plan: bool,
 }
 
 struct DuckDBCtx {
@@ -108,12 +108,12 @@ enum EngineCtx {
 }
 
 impl EngineCtx {
-    fn new_with_datafusion(session_ctx: SessionContext, emit_execution_plan: bool) -> Self {
+    fn new_with_datafusion(session_ctx: SessionContext, emit_plan: bool) -> Self {
         EngineCtx::DataFusion(DataFusionCtx {
             execution_plans: Vec::new(),
             metrics: Vec::new(),
             session: session_ctx,
-            emit_execution_plan,
+            emit_plan,
         })
     }
 
@@ -438,37 +438,37 @@ fn execute_queries(
 ) -> Vec<QueryMeasurement> {
     let mut query_measurements = Vec::default();
 
-    for (query_idx, query_string) in queries {
+    for &(query_idx, ref query_string) in queries.iter() {
         match engine_ctx {
             EngineCtx::DataFusion(ctx) => {
                 let (fastest_run, execution_plan) = benchmark_datafusion_query(
-                    *query_idx,
+                    query_idx,
                     query_string,
                     iterations,
                     &ctx.session,
                     tokio_runtime,
                 );
 
                 ctx.execution_plans
-                    .push((*query_idx, execution_plan.clone()));
+                    .push((query_idx, execution_plan.clone()));
 
-                if ctx.emit_execution_plan {
+                if ctx.emit_plan {
                     df::write_execution_plan(
-                        *query_idx,
+                        query_idx,
                         file_format,
                         CLICKBENCH_DATASET,
-                        &execution_plan,
+                        execution_plan.as_ref(),
                     );
                 }
 
                 ctx.metrics.push((
-                    *query_idx,
+                    query_idx,
                     file_format,
                     VortexMetricsFinder::find_all(execution_plan.as_ref()),
                 ));
 
                 query_measurements.push(QueryMeasurement {
-                    query_idx: *query_idx,
+                    query_idx,
                     target: Target::new(Engine::DataFusion, file_format),
                     storage: STORAGE_NVME.to_owned(),
                     fastest_run,
@@ -478,14 +478,14 @@ fn execute_queries(
 
             EngineCtx::DuckDB(args) => {
                 let fastest_run = benchmark_duckdb_query(
-                    *query_idx,
+                    query_idx,
                     query_string,
                     iterations,
                     &DuckDBExecutor::new(args.duckdb_path.clone(), args.duckdb_file(file_format)),
                 );
 
                 query_measurements.push(QueryMeasurement {
-                    query_idx: *query_idx,
+                    query_idx,
                     target: Target::new(Engine::DuckDB, file_format),
                     storage: STORAGE_NVME.to_owned(),
                     fastest_run,

bench-vortex/src/bin/tpch.rs

@@ -1,7 +1,9 @@
+use std::path::PathBuf;
 use std::process::ExitCode;
 use std::time::{Duration, Instant};
 
 use bench_vortex::ddb::{DuckDBExecutor, register_tables};
+use bench_vortex::df::write_execution_plan;
 use bench_vortex::display::{DisplayFormat, RatioMode, print_measurements_json, render_table};
 use bench_vortex::measurements::QueryMeasurement;
 use bench_vortex::metrics::{MetricsSetExt, export_plan_spans};
@@ -43,7 +45,7 @@ struct Args {
     )]
     targets: Vec<Target>,
     #[arg(long)]
-    duckdb_path: Option<std::path::PathBuf>,
+    duckdb_path: Option<PathBuf>,
     #[arg(short, long, value_delimiter = ',')]
     queries: Option<Vec<usize>>,
     #[arg(short, long, value_delimiter = ',')]
@@ -180,6 +182,7 @@ fn main() -> ExitCode {
         url,
         args.all_metrics,
         args.export_spans,
+        args.emit_plan,
         &args.duckdb_path,
     ))
 }
@@ -298,7 +301,8 @@ async fn bench_main(
     url: Url,
     display_all_metrics: bool,
     export_spans: bool,
-    duckdb_path: &Option<std::path::PathBuf>,
+    emit_plan: bool,
+    duckdb_path: &Option<PathBuf>,
 ) -> ExitCode {
     let expected_row_counts = if scale_factor == 1 {
         EXPECTED_ROW_COUNTS_SF1
@@ -379,6 +383,11 @@ async fn bench_main(
                             ],
                         );
                     }
+
+                    if emit_plan {
+                        write_execution_plan(query_idx, format, TPCH_DATASET, plan.as_ref());
+                    }
+
                     plans.push((query_idx, plan.clone()));
 
                     let storage = match bench_vortex::utils::url_scheme_to_storage(&url) {

bench-vortex/src/engines/df/mod.rs

@@ -138,7 +138,7 @@ pub fn write_execution_plan(
     query_idx: usize,
     format: crate::Format,
     dataset_name: &str,
-    execution_plan: &Arc<dyn ExecutionPlan>,
+    execution_plan: &dyn ExecutionPlan,
 ) {
     fs::write(
         format!("{dataset_name}_{format}_q{query_idx:02}.plan"),
@@ -150,7 +150,7 @@ pub fn write_execution_plan(
         format!("{dataset_name}_{format}_q{query_idx:02}.short.plan"),
         format!(
             "{}",
-            DisplayableExecutionPlan::with_full_metrics(execution_plan.as_ref())
+            DisplayableExecutionPlan::with_full_metrics(execution_plan)
                 .set_show_schema(true)
                 .set_show_statistics(true)
                 .indent(true)