bench: keep the expect_plan check out of the measured region

benchmarks/sql_benchmarks/README.md

@@ -326,7 +326,9 @@ DROP TABLE test;
 <td>
 
 The expect_plan directive will check the physical plan for the string provided on the same line. This
-can be used to validate that a particular join was used. <br/> <br/> Example:<br/>
+can be used to validate that a particular join was used. The plan is rendered as <code>EXPLAIN</code>
+displays it, and the check runs once per benchmark, not once per iteration, so it adds no cost to the
+measured region. <br/> <br/> Example:<br/>
 <blockquote>expect_plan NestedLoopJoinExec</blockquote>
 
 </td>

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q02.benchmark

@@ -14,7 +14,7 @@ WHERE o.id NOT IN (SELECT i.id_n1 FROM large_inner i);
 true
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q2: uncorrelated NOT IN, 1% NULL on the subquery side.

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q03.benchmark

@@ -14,7 +14,7 @@ WHERE o.id_n50 NOT IN (SELECT i.id FROM large_inner i);
 true
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q3: uncorrelated NOT IN, 50% NULL on the outer side.

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q04.benchmark

@@ -19,7 +19,7 @@ WHERE o.id_n0 NOT IN (SELECT i.id_n0 FROM small_inner i WHERE i.z < o.z);
 true
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q4: non-equality-correlated NOT IN, nullable keys that hold no NULL.

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark

@@ -4,7 +4,7 @@ group null_aware_join
 load sql_benchmarks/null_aware_join/init/load.sql
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q5: non-equality-correlated NOT IN, 1% NULL on the outer side.

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark

@@ -4,7 +4,7 @@ group null_aware_join
 load sql_benchmarks/null_aware_join/init/load.sql
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q6: non-equality-correlated NOT IN, 50% NULL on the outer side.

benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark

@@ -4,7 +4,7 @@ group null_aware_join
 load sql_benchmarks/null_aware_join/init/load.sql
 
 expect_plan HashJoinExec
-expect_plan null_aware: true
+expect_plan null_aware
 
 run
 -- Q7: non-equality-correlated NOT IN, 50% NULL on the subquery side.

benchmarks/src/sql_benchmark.rs

@@ -21,7 +21,8 @@ use arrow::error::ArrowError;
 use arrow::util::display::{ArrayFormatter, FormatOptions};
 use datafusion::dataframe::DataFrameWriteOptions;
 use datafusion::datasource::MemTable;
-use datafusion::physical_plan::execute_stream;
+use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use datafusion::physical_plan::{ExecutionPlan, execute_stream};
 use datafusion::prelude::{CsvReadOptions, DataFrame, SessionContext};
 use datafusion_common::config::CsvOptions;
 use datafusion_common::{DataFusionError, Result, exec_datafusion_err};
@@ -63,6 +64,11 @@ pub struct SqlBenchmark {
     assert_queries: Vec<BenchmarkQuery>,
     /// Flag indicating whether the benchmark has been fully loaded
     is_loaded: bool,
+    /// Flag indicating whether the `expect` strings were checked against the
+    /// physical plan. The plan does not change between iterations, so the
+    /// check runs on the first iteration only and is not repeated in the
+    /// measured region.
+    plans_validated: bool,
     /// Stores the last run results if needed so they can be compared or persisted.
     last_results: Option<Vec<RecordBatch>>,
     /// echo statements
@@ -99,6 +105,7 @@ impl SqlBenchmark {
             benchmark_path: full_path.to_path_buf(),
             replacement_mapping,
             expect: vec![],
+            plans_validated: false,
             queries: HashMap::new(),
             result_queries: vec![],
             assert_queries: vec![],
@@ -239,7 +246,7 @@ impl SqlBenchmark {
                         );
 
                         let df = ctx.sql(query).await?;
-                        if !self.expect.is_empty() {
+                        if !self.expect.is_empty() && !self.plans_validated {
                             let physical_plan = df.create_physical_plan().await?;
                             self.validate_expected_plan(&physical_plan)?;
                         }
@@ -270,7 +277,11 @@ impl SqlBenchmark {
                         );
 
                         let row_count = self
-                            .execute_sql_without_result_buffering(query, ctx)
+                            .execute_sql_without_result_buffering(
+                                query,
+                                ctx,
+                                !self.plans_validated,
+                            )
                             .await?;
 
                         if is_result_statement(query) {
@@ -285,6 +296,10 @@ impl SqlBenchmark {
 
         debug!("Results have {result_count} rows");
 
+        // 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;
+
         // Store results for verification
         self.last_results = Some(result);
 
@@ -572,12 +587,22 @@ impl SqlBenchmark {
         Ok(())
     }
 
-    fn validate_expected_plan(&self, physical_plan: &impl Debug) -> Result<()> {
+    /// Checks the `expect` strings against the plan as `EXPLAIN` displays it.
+    ///
+    /// `{:#?}` would dump every `RecordBatch` an in-memory source holds, which
+    /// is megabytes for a benchmark that builds its tables with `CREATE TABLE
+    /// ... AS SELECT`.
+    fn validate_expected_plan(
+        &self,
+        physical_plan: &Arc<dyn ExecutionPlan>,
+    ) -> Result<()> {
         if self.expect.is_empty() {
             return Ok(());
         }
 
-        let plan_string = format!("{physical_plan:#?}");
+        let plan_string = DisplayableExecutionPlan::new(physical_plan.as_ref())
+            .indent(true)
+            .to_string();
 
         for exp_str in &self.expect {
             if !plan_string.contains(exp_str) {
@@ -594,13 +619,16 @@ impl SqlBenchmark {
         &self,
         sql: &str,
         ctx: &SessionContext,
+        validate_plan: bool,
     ) -> Result<usize> {
         let mut row_count = 0;
 
         let df = ctx.sql(sql).await?;
         let physical_plan = df.create_physical_plan().await?;
 
-        self.validate_expected_plan(&physical_plan)?;
+        if validate_plan {
+            self.validate_expected_plan(&physical_plan)?;
+        }
         let mut stream = execute_stream(physical_plan, ctx.task_ctx())?;
 
         while let Some(batch) = stream.next().await {
@@ -3333,6 +3361,31 @@ SELECT 1;
         );
     }
 
+    #[tokio::test]
+    async fn run_checks_expect_plan_once_per_benchmark() {
+        let ctx = SessionContext::new();
+        let benchmark_text = "expect_plan PlaceholderRowExec\nrun\nSELECT 1\n";
+
+        let mut benchmark = parse_benchmark(benchmark_text)
+            .await
+            .expect("benchmark should parse");
+        assert!(!benchmark.plans_validated);
+
+        benchmark
+            .run(&ctx, false)
+            .await
+            .expect("first run should accept the matching plan");
+        assert!(
+            benchmark.plans_validated,
+            "the first run should mark the plans as checked"
+        );
+
+        benchmark
+            .run(&ctx, false)
+            .await
+            .expect("later runs should not repeat the check");
+    }
+
     #[tokio::test]
     async fn run_accepts_matching_expect_plan_for_buffered_and_streaming_modes() {
         let ctx = SessionContext::new();