From b740eabd94ef60182b84f28515de03763d5c4d4b Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Tue, 8 Sep 2026 17:54:55 +0530 Subject: [PATCH] perf: reduce invariant-check code generation --- datafusion/physical-plan/Cargo.toml | 4 + .../physical-plan/benches/invariant_check.rs | 61 +++++ .../src/distribution_requirements.rs | 41 ++- .../physical-plan/src/execution_plan.rs | 234 ++++++++++++++++-- 4 files changed, 306 insertions(+), 34 deletions(-) create mode 100644 datafusion/physical-plan/benches/invariant_check.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 0aa22653b44ee..534cea8ea9cbb 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -109,6 +109,10 @@ tokio = { workspace = true, features = [ harness = false name = "partial_ordering" +[[bench]] +harness = false +name = "invariant_check" + [[bench]] harness = false name = "union_schema" diff --git a/datafusion/physical-plan/benches/invariant_check.rs b/datafusion/physical-plan/benches/invariant_check.rs new file mode 100644 index 0000000000000..65d9c3f23bedf --- /dev/null +++ b/datafusion/physical-plan/benches/invariant_check.rs @@ -0,0 +1,61 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::datatypes::{Schema, SchemaRef}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::execution_plan::{ + ExecutionPlan, InvariantLevel, check_default_invariants, +}; +use datafusion_physical_plan::union::UnionExec; + +fn empty_exec(schema: &SchemaRef) -> Arc { + Arc::new(EmptyExec::new(Arc::clone(schema))) +} + +fn bench_invariant_checks(c: &mut Criterion) { + let schema = Arc::new(Schema::empty()); + let leaf = EmptyExec::new(Arc::clone(&schema)); + let union = + UnionExec::try_new((0..4).map(|_| empty_exec(&schema)).collect::>()) + .unwrap(); + + let mut group = c.benchmark_group("check_default_invariants"); + group.bench_function("leaf", |b| { + b.iter(|| { + black_box(check_default_invariants( + black_box(&leaf), + InvariantLevel::Always, + )) + }); + }); + group.bench_function("four_children", |b| { + b.iter(|| { + black_box(check_default_invariants( + black_box(union.as_ref()), + InvariantLevel::Always, + )) + }); + }); + group.finish(); +} + +criterion_group!(benches, bench_invariant_checks); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs index 119b2357c02b4..f7a17ddc7b8d6 100644 --- a/datafusion/physical-plan/src/distribution_requirements.rs +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -17,6 +17,8 @@ //! Input distribution requirements for physical execution plans. +use std::sync::Arc; + use datafusion_common::{Result, internal_err}; use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; @@ -205,25 +207,46 @@ impl InputDistributionRequirements { } /// Validate the requirements against a plan's children. - pub(crate) fn check_invariants( + #[inline] + pub(crate) fn check_invariants( &self, - plan: &P, + plan_name: &str, + children: Vec<&Arc>, check: InvariantLevel, ) -> Result<()> { - let children = plan.children(); - self.validate_shape(plan.name(), children.len())?; + let children_len = children.len(); + if self.children.len() != children_len { + return self.validate_shape(plan_name, children_len); + } + + if let Some(co_partitioned) = &self.co_partitioned { + self.validate_shape(plan_name, children_len)?; + if matches!(check, InvariantLevel::Executable) { + return self.check_co_partitioning_invariant( + plan_name, + co_partitioned, + children, + ); + } + } + Ok(()) + } + + fn check_co_partitioning_invariant( + &self, + plan_name: &str, + co_partitioned: &[usize], + children: Vec<&Arc>, + ) -> Result<()> { let children = children .into_iter() .map(|child| child.as_ref()) .collect::>(); - if matches!(check, InvariantLevel::Executable) - && let Some(co_partitioned) = &self.co_partitioned - && !self.co_partitioning_satisfied(co_partitioned, &children) - { + if !self.co_partitioning_satisfied(co_partitioned, &children) { return internal_err!( "{} requires children {:?} to be co-partitioned", - plan.name(), + plan_name, co_partitioned ); } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 6cb3eeef8cefb..978059b9faf6b 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -1601,6 +1601,39 @@ impl PlanProperties { } } +#[derive(Debug)] +struct DefaultInvariantChecker<'a> { + plan_name: &'a str, +} + +impl<'a> DefaultInvariantChecker<'a> { + fn new(plan_name: &'a str) -> Self { + Self { plan_name } + } + + /// All dynamic expressions must have an expression id. + fn check_dynamic_expression_invariants( + &self, + dynamic_expressions: Vec>, + ) -> Result<()> { + let mut produced_ids = HashSet::new(); + for expr in dynamic_expressions { + let Some(expression_id) = expr.expression_id() else { + return internal_err!( + "{}::dynamic_expressions_produced returned an expression without an expression ID", + self.plan_name + ); + }; + assert_or_internal_err!( + produced_ids.insert(expression_id), + "{}::dynamic_expressions_produced returned duplicate expression ID {expression_id}", + self.plan_name + ); + } + Ok(()) + } +} + macro_rules! check_len { ($target:expr, $func_name:ident, $expected_len:expr) => { let actual_len = $target.$func_name().len(); @@ -1616,27 +1649,6 @@ macro_rules! check_len { }; } -/// All dynamic expressions must have an expression id. -fn check_dynamic_expression_invariants( - plan: &P, -) -> Result<()> { - let mut produced_ids = HashSet::new(); - for expr in plan.dynamic_expressions_produced() { - let Some(expression_id) = expr.expression_id() else { - return internal_err!( - "{}::dynamic_expressions_produced returned an expression without an expression ID", - plan.name() - ); - }; - assert_or_internal_err!( - produced_ids.insert(expression_id), - "{}::dynamic_expressions_produced returned duplicate expression ID {expression_id}", - plan.name() - ); - } - Ok(()) -} - /// Checks a set of invariants that apply to all ExecutionPlan implementations. /// Returns an error if the given node does not conform. pub fn check_default_invariants( @@ -1648,9 +1660,19 @@ pub fn check_default_invariants( check_len!(plan, maintains_input_order, children_len); check_len!(plan, required_input_ordering, children_len); check_len!(plan, benefits_from_input_partitioning, children_len); - plan.input_distribution_requirements() - .check_invariants(plan, check)?; - check_dynamic_expression_invariants(plan)?; + + let input_distribution_requirements = plan.input_distribution_requirements(); + let children = plan.children(); + let checker = DefaultInvariantChecker::new(plan.name()); + input_distribution_requirements.check_invariants( + checker.plan_name, + children, + check, + )?; + let dynamic_expressions = plan.dynamic_expressions_produced(); + if !dynamic_expressions.is_empty() { + checker.check_dynamic_expression_invariants(dynamic_expressions)?; + } Ok(()) } @@ -2094,16 +2116,31 @@ mod tests { #[derive(Debug)] pub struct EmptyExec { + children: Vec>, dynamic_expressions: Vec>, + maintains_input_order_len: Option, + required_input_ordering_len: Option, + benefits_from_input_partitioning_len: Option, + input_distribution_requirements: Option, } impl EmptyExec { pub fn new(_schema: SchemaRef) -> Self { Self { + children: vec![], dynamic_expressions: vec![], + maintains_input_order_len: None, + required_input_ordering_len: None, + benefits_from_input_partitioning_len: None, + input_distribution_requirements: None, } } + fn with_child(mut self, child: Arc) -> Self { + self.children.push(child); + self + } + fn with_dynamic_expressions( mut self, dynamic_expressions: Vec>, @@ -2111,6 +2148,37 @@ mod tests { self.dynamic_expressions = dynamic_expressions; self } + + fn with_maintains_input_order_len(mut self, len: usize) -> Self { + self.maintains_input_order_len = Some(len); + self + } + + fn with_required_input_ordering_len(mut self, len: usize) -> Self { + self.required_input_ordering_len = Some(len); + self + } + + fn with_benefits_from_input_partitioning_len(mut self, len: usize) -> Self { + self.benefits_from_input_partitioning_len = Some(len); + self + } + + fn with_input_distribution_requirements_len(mut self, len: usize) -> Self { + self.input_distribution_requirements = + Some(InputDistributionRequirements::new( + vec![Distribution::UnspecifiedDistribution; len], + )); + self + } + + fn with_input_distribution_requirements( + mut self, + requirements: InputDistributionRequirements, + ) -> Self { + self.input_distribution_requirements = Some(requirements); + self + } } impl DisplayAs for EmptyExec { @@ -2133,7 +2201,42 @@ mod tests { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() + } + + fn maintains_input_order(&self) -> Vec { + vec![ + false; + self.maintains_input_order_len + .unwrap_or(self.children.len()) + ] + } + + fn required_input_ordering(&self) -> Vec> { + vec![ + None; + self.required_input_ordering_len + .unwrap_or(self.children.len()) + ] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![ + true; + self.benefits_from_input_partitioning_len + .unwrap_or(self.children.len()) + ] + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + self.input_distribution_requirements + .clone() + .unwrap_or_else(|| { + InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution; + self.children.len() + ]) + }) } fn replace_children( @@ -2208,6 +2311,87 @@ mod tests { Ok(()) } + #[test] + fn test_default_invariant_lengths() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let child: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let make_plan = + || EmptyExec::new(Arc::clone(&schema)).with_child(Arc::clone(&child)); + + check_default_invariants(&make_plan(), InvariantLevel::Always)?; + + let cases = [ + ( + make_plan().with_maintains_input_order_len(0), + "Internal error: Assertion failed: actual_len == children_len (left: 0, right: 1): EmptyExec::maintains_input_order returned Vec with incorrect size: 0 != 1", + ), + ( + make_plan().with_required_input_ordering_len(0), + "Internal error: Assertion failed: actual_len == children_len (left: 0, right: 1): EmptyExec::required_input_ordering returned Vec with incorrect size: 0 != 1", + ), + ( + make_plan().with_benefits_from_input_partitioning_len(0), + "Internal error: Assertion failed: actual_len == children_len (left: 0, right: 1): EmptyExec::benefits_from_input_partitioning returned Vec with incorrect size: 0 != 1", + ), + ( + make_plan().with_input_distribution_requirements_len(0), + "Internal error: EmptyExec::input_distribution_requirements returned incorrect child count: 0 != 1", + ), + ]; + + for (plan, expected) in cases { + let error = check_default_invariants(&plan, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.starts_with(expected), "{error}"); + } + + let invalid_co_partitioning = make_plan() + .with_child(child) + .with_input_distribution_requirements( + InputDistributionRequirements::co_partitioned(vec![ + Distribution::UnspecifiedDistribution, + Distribution::UnspecifiedDistribution, + ]), + ); + let error = + check_default_invariants(&invalid_co_partitioning, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!( + error.starts_with( + "Internal error: EmptyExec has invalid co-partitioning requirement: child 0 has unspecified distribution" + ), + "{error}" + ); + + let left: Arc = + Arc::new(crate::empty::EmptyExec::new(Arc::clone(&schema))); + let right: Arc = + Arc::new(crate::empty::EmptyExec::new(Arc::clone(&schema))); + let valid_co_partitioning = EmptyExec::new(schema) + .with_child(left) + .with_child(right) + .with_input_distribution_requirements( + InputDistributionRequirements::co_partitioned(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]), + ); + check_default_invariants(&valid_co_partitioning, InvariantLevel::Executable)?; + + Ok(()) + } + + #[test] + fn test_default_invariants_accept_trait_object() -> Result<()> { + let plan = EmptyExec::new(Arc::new(Schema::empty())); + let plan: &dyn ExecutionPlan = &plan; + + check_default_invariants(plan, InvariantLevel::Always)?; + plan.check_invariants(InvariantLevel::Always) + } + #[derive(Debug)] pub struct RenamedEmptyExec;