feat(benchmarks): add plan diagnostics to SQL benchmark runners

benchmarks/src/bin/benchmark_runner.rs

@@ -1965,6 +1965,86 @@ description = "Run query one against CSV data."
         assert!(message.contains("alpha"), "{message}");
     }
 
+    #[test]
+    fn sql_runner_debug_plans_include_execution_metrics() {
+        const CHILD_ENV: &str = "DATAFUSION_SQL_DEBUG_TEST_MODE";
+        if let Ok(mode) = std::env::var(CHILD_ENV) {
+            let temp = suite_root();
+            write_benchmark(
+                temp.path(),
+                "alpha/benchmarks/q01.benchmark",
+                "name Q01\n\nrun\nSELECT value FROM (VALUES (3), (1), (2)) AS t(value) ORDER BY value\n",
+            );
+            let mut args = vec!["benchmark_runner", "alpha", "--query", "1"];
+            if mode == "criterion" {
+                args.push("--criterion");
+            } else {
+                args.extend(["--iterations", "1"]);
+            }
+            if mode == "buffered" {
+                args.extend(["--result-mode", "persist"]);
+            }
+            match parse_cli_from(args, temp.path()).unwrap() {
+                CliAction::Simple(config) => {
+                    tokio::runtime::Runtime::new()
+                        .unwrap()
+                        .block_on(run_simple_benchmarks(temp.path(), config))
+                        .unwrap();
+                }
+                CliAction::Criterion { config, .. } => {
+                    let output = tempfile::tempdir().unwrap();
+                    let mut criterion = Criterion::default()
+                        .output_directory(output.path())
+                        .warm_up_time(std::time::Duration::from_millis(1))
+                        .measurement_time(std::time::Duration::from_millis(10))
+                        .without_plots();
+                    run_criterion_benchmarks_impl(temp.path(), &config, &mut criterion)
+                        .unwrap();
+                }
+                _ => panic!("expected a benchmark run"),
+            }
+            return;
+        }
+
+        for mode in ["simple", "criterion", "buffered"] {
+            for debug in [false, true] {
+                let output = std::process::Command::new(std::env::current_exe().unwrap())
+                    .args([
+                        "--exact",
+                        "tests::sql_runner_debug_plans_include_execution_metrics",
+                        "--nocapture",
+                    ])
+                    // Other CLI tests mutate environment variables concurrently.
+                    .env_clear()
+                    .env(CHILD_ENV, mode)
+                    .env("DEBUG", debug.to_string())
+                    .output()
+                    .unwrap();
+                let stdout = String::from_utf8(output.stdout).unwrap();
+                assert!(
+                    output.status.success(),
+                    "{mode}, debug={debug}: {stdout}\n{}",
+                    String::from_utf8_lossy(&output.stderr)
+                );
+                for heading in [
+                    "=== Logical plan ===",
+                    "=== Optimized logical plan ===",
+                    "=== Physical plan ===",
+                    "=== Physical plan with metrics ===",
+                ] {
+                    assert_eq!(stdout.contains(heading), debug, "{mode}: {stdout}");
+                }
+                if debug {
+                    let metrics = stdout
+                        .split("=== Physical plan with metrics ===")
+                        .nth(1)
+                        .unwrap();
+                    assert!(metrics.contains("output_rows=3"), "{mode}: {stdout}");
+                }
+            }
+        }
+    }
+
     #[test]
     fn criterion_runner_saves_named_baseline() {
         let temp = tempfile::tempdir().unwrap();

benchmarks/src/sql_benchmark.rs

@@ -22,7 +22,7 @@ use arrow::util::display::{ArrayFormatter, FormatOptions};
 use datafusion::dataframe::DataFrameWriteOptions;
 use datafusion::datasource::MemTable;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
-use datafusion::physical_plan::{ExecutionPlan, execute_stream};
+use datafusion::physical_plan::{ExecutionPlan, collect, execute_stream};
 use datafusion::prelude::{CsvReadOptions, DataFrame, SessionContext};
 use datafusion_common::config::CsvOptions;
 use datafusion_common::{DataFusionError, Result, exec_datafusion_err};
@@ -71,6 +71,10 @@ pub struct SqlBenchmark {
     plans_validated: bool,
     /// Stores the last run results if needed so they can be compared or persisted.
     last_results: Option<Vec<RecordBatch>>,
+    /// Print plans and execution metrics for each run query.
+    debug: bool,
+    /// Whether diagnostics have been printed for a successful run.
+    diagnostics_printed: bool,
     /// echo statements
     echo: Vec<String>,
 }
@@ -111,6 +115,8 @@ impl SqlBenchmark {
             assert_queries: vec![],
             is_loaded: false,
             last_results: None,
+            debug: false,
+            diagnostics_printed: false,
             echo: vec![],
         };
         insert_replacement(
@@ -125,6 +131,13 @@ impl SqlBenchmark {
         Ok(bm)
     }
 
+    /// Enable plan diagnostics for run queries, including post-execution metrics.
+    /// Output occurs on the first successful run, including preparation or warm-up,
+    /// and affects that execution's timing.
+    pub fn set_debug(&mut self, debug: bool) {
+        self.debug = debug;
+    }
+
     /// Initializes the benchmark by executing `load` and `init` queries.
     ///
     /// Registers any required tables or sets up state in the provided
@@ -246,13 +259,18 @@ impl SqlBenchmark {
                         );
 
                         let df = ctx.sql(query).await?;
-                        if !self.expect.is_empty() && !self.plans_validated {
-                            let physical_plan = df.create_physical_plan().await?;
+                        let result_schema = Arc::new(df.schema().as_arrow().clone());
+                        let physical_plan = self.create_physical_plan(df).await?;
+
+                        if !self.plans_validated {
                             self.validate_expected_plan(&physical_plan)?;
                         }
 
-                        let result_schema = Arc::new(df.schema().as_arrow().clone());
-                        let mut batches = df.collect().await?;
+                        let mut batches =
+                            collect(Arc::clone(&physical_plan), ctx.task_ctx()).await?;
+
+                        self.print_plan_metrics(&physical_plan);
+
                         // save the output for select/with queries
                         if is_result_statement(query) {
                             if batches.is_empty() {
@@ -299,6 +317,7 @@ impl SqlBenchmark {
         // Every `run` query was planned and checked above, and a plan does not
         // change between iterations, so later iterations skip the check.
         self.plans_validated = true;
+        self.diagnostics_printed |= self.debug;
 
         // Store results for verification
         self.last_results = Some(result);
@@ -615,6 +634,36 @@ impl SqlBenchmark {
         Ok(())
     }
 
+    async fn create_physical_plan(
+        &self,
+        df: DataFrame,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if !self.debug || self.diagnostics_printed {
+            return df.create_physical_plan().await;
+        }
+
+        let (state, plan) = df.into_parts();
+        println!("=== Logical plan ===\n{plan}\n");
+        let plan = state.optimize(&plan)?;
+        println!("=== Optimized logical plan ===\n{plan}\n");
+        let physical_plan = state.create_physical_plan(&plan).await?;
+        println!(
+            "=== Physical plan ===\n{}\n",
+            DisplayableExecutionPlan::new(physical_plan.as_ref()).indent(true)
+        );
+        Ok(physical_plan)
+    }
+
+    fn print_plan_metrics(&self, physical_plan: &Arc<dyn ExecutionPlan>) {
+        if self.debug && !self.diagnostics_printed {
+            println!(
+                "=== Physical plan with metrics ===\n{}\n",
+                DisplayableExecutionPlan::with_metrics(physical_plan.as_ref())
+                    .indent(true)
+            );
+        }
+    }
+
     async fn execute_sql_without_result_buffering(
         &self,
         sql: &str,
@@ -624,12 +673,12 @@ impl SqlBenchmark {
         let mut row_count = 0;
 
         let df = ctx.sql(sql).await?;
-        let physical_plan = df.create_physical_plan().await?;
+        let physical_plan = self.create_physical_plan(df).await?;
 
         if validate_plan {
             self.validate_expected_plan(&physical_plan)?;
         }
-        let mut stream = execute_stream(physical_plan, ctx.task_ctx())?;
+        let mut stream = execute_stream(Arc::clone(&physical_plan), ctx.task_ctx())?;
 
         while let Some(batch) = stream.next().await {
             row_count += batch?.num_rows();
@@ -638,6 +687,7 @@ impl SqlBenchmark {
             // to reduce memory pressure
         }
 
+        self.print_plan_metrics(&physical_plan);
         Ok(row_count)
     }
 

benchmarks/src/sql_benchmark_runner.rs

@@ -678,6 +678,7 @@ pub async fn prepare_benchmark(
     benchmark: &mut SqlBenchmark,
     config: &SqlRunConfig,
 ) -> Result<()> {
+    benchmark.set_debug(config.common.debug);
     benchmark.initialize(ctx).await?;
     record_pool_peak_after_init(ctx)?;
     benchmark.assert(ctx).await?;