From ae18f6b31327e486c4dd621e9a2809539287b6ab Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:48:21 -0500 Subject: [PATCH] bench: keep the expect_plan check out of the measured region `expect_plan` was checked on every iteration, inside the timed region, and it rendered the plan with `{:#?}`. A derived `Debug` prints every `RecordBatch` an in-memory source holds, so for a benchmark whose tables come from `CREATE TABLE ... AS SELECT` the dump is megabytes per iteration. In the `null_aware_join` suite this was most of the measured time. The check now renders the plan as `EXPLAIN` displays it, and runs on the first iteration only, because a plan does not change between iterations. A failure also prints an 8-line plan instead of a 1.8 MB dump. `null_aware_join`, release, apache/main, median of 30 iterations: | Query | before | after | |-------|--------|-------| | Q01 | 17.6 ms | 4.6 ms | | Q02 | 15.3 ms | 2.6 ms | | Q03 | 14.7 ms | 2.1 ms | | Q04 | 0.9 ms | 0.4 ms | | Q05 | 0.9 ms | 0.4 ms | | Q06 | 0.8 ms | 0.3 ms | | Q07 | 0.9 ms | 0.3 ms | | Q08 | 1.2 ms | 0.5 ms | The display form spells the flag as `null_aware`, not `null_aware: true`, so the six `null_aware_join` checks change with it. The other 89 `expect_plan` strings are operator names and are not affected. Co-Authored-By: Claude Opus 5 --- benchmarks/sql_benchmarks/README.md | 4 +- .../null_aware_join/benchmarks/q02.benchmark | 2 +- .../null_aware_join/benchmarks/q03.benchmark | 2 +- .../null_aware_join/benchmarks/q04.benchmark | 2 +- .../null_aware_join/benchmarks/q05.benchmark | 2 +- .../null_aware_join/benchmarks/q06.benchmark | 2 +- .../null_aware_join/benchmarks/q07.benchmark | 2 +- benchmarks/src/sql_benchmark.rs | 65 +++++++++++++++++-- 8 files changed, 68 insertions(+), 13 deletions(-) diff --git a/benchmarks/sql_benchmarks/README.md b/benchmarks/sql_benchmarks/README.md index fb09c95daa12f..0faa630f93106 100644 --- a/benchmarks/sql_benchmarks/README.md +++ b/benchmarks/sql_benchmarks/README.md @@ -326,7 +326,9 @@ DROP TABLE test; 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.

Example:
+can be used to validate that a particular join was used. The plan is rendered as EXPLAIN +displays it, and the check runs once per benchmark, not once per iteration, so it adds no cost to the +measured region.

Example:
expect_plan NestedLoopJoinExec
diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q02.benchmark index ba8554b1cb823..af86fa6471216 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q02.benchmark +++ b/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. diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q03.benchmark index 28f1ed3410a8a..a3676df23776b 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q03.benchmark +++ b/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. diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q04.benchmark index c72b5bceed74f..03ee21bf66944 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q04.benchmark +++ b/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. diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark index 2d0a04077d5cf..b50bbe1e7398c 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q05.benchmark +++ b/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. diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark index 668476ded9571..004f2cf628272 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q06.benchmark +++ b/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. diff --git a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark index 38bf7e9048bc9..d1aec98aa39ad 100644 --- a/benchmarks/sql_benchmarks/null_aware_join/benchmarks/q07.benchmark +++ b/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. diff --git a/benchmarks/src/sql_benchmark.rs b/benchmarks/src/sql_benchmark.rs index 0c0c9375e4283..3736f0311b026 100644 --- a/benchmarks/src/sql_benchmark.rs +++ b/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, /// 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>, /// 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, + ) -> 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 { 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();