This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 423df6f4fc perf: reduce invariant-check code generation (#25073)
423df6f4fc is described below
commit 423df6f4fc6a6eac6d0d1457166a166dafe8030d
Author: Kumar Ujjawal <[email protected]>
AuthorDate: Thu Sep 10 04:42:58 2026 +0000
perf: reduce invariant-check code generation (#25073)
## Which issue does this PR close?
- Part of #13815.
## Rationale for this change
`check_default_invariants` is generic over each `ExecutionPlan`
implementation. The compiler therefore generated many copies of its
dynamic-expression and input-distribution validation logic.
These checks depend only on erased values such as the plan name,
children, and dynamic expressions. Sharing the heavier validation paths
reduces generated code without changing invariant semantics or the
public API.
Local release measurements against `upstream/main` at `92746a993`:
| Measurement | `upstream/main` | This PR | Change |
|---|---:|---:|---:|
| Targeted invariant-check LLVM IR lines | 68,856 | 47,082 | -31.6% |
| Total physical-plan LLVM IR lines | 2,560,976 | 2,513,101 | -1.87% |
| Release `rlib` size | 56,223,208 bytes | 55,689,856 bytes | -0.95% |
| Linked `partial_ordering` benchmark binary | 5,882,960 bytes |
5,882,960 bytes | No change |
The linked benchmark binary was unchanged because its linker removed
unused code. This PR does not claim a measured end-to-end WASM binary
reduction.
The added Criterion benchmark showed no runtime regression:
| Case | `upstream/main` | This PR |
|---|---:|---:|
| Leaf plan | 13.016–13.030 ns | 12.939–12.951 ns |
| Four-child plan | 150.24–151.02 ns | 142.75–145.68 ns |
## What changes are included in this PR?
- Move dynamic-expression ID validation behind a non-generic helper.
- Make input-distribution invariant validation non-generic.
- Keep the inexpensive per-plan length checks inline.
- Skip dynamic-expression validation when a plan produces no dynamic
expressions.
- Preserve the public `check_default_invariants` signature and existing
error behavior.
- Add unit tests for malformed vectors, dynamic-expression IDs,
co-partitioning requirements, and trait-object callers.
- Add a Criterion benchmark for leaf and four-child plans.
## What is the testing strategy for this PR?
The new unit tests cover:
- incorrect invariant-vector lengths and their error messages;
- missing and duplicate dynamic-expression IDs;
- malformed and valid co-partitioning requirements;
- calls through concrete plans and `dyn ExecutionPlan`.
The following checks pass:
```shell
cargo fmt --all
cargo clippy --all-targets --all-features -- -D warnings
RUST_BACKTRACE=1 cargo test --profile ci \
--exclude datafusion-examples \
--exclude datafusion-benchmarks \
--exclude datafusion-cli \
--workspace --lib --tests --bins \
--features
avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption
```
The code-size comparison and runtime measurements use the new
`invariant_check` Criterion benchmark and release builds of
`datafusion-physical-plan`.
## Are there any user-facing changes?
No.
---
datafusion/physical-plan/Cargo.toml | 4 +
.../physical-plan/benches/invariant_check.rs | 61 ++++++
.../physical-plan/src/distribution_requirements.rs | 41 +++-
datafusion/physical-plan/src/execution_plan.rs | 234 ++++++++++++++++++---
4 files changed, 306 insertions(+), 34 deletions(-)
diff --git a/datafusion/physical-plan/Cargo.toml
b/datafusion/physical-plan/Cargo.toml
index 0aa22653b4..534cea8ea9 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 0000000000..65d9c3f23b
--- /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<dyn ExecutionPlan> {
+ 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::<Vec<_>>())
+ .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 119b2357c0..f7a17ddc7b 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<P: ExecutionPlan + ?Sized>(
+ #[inline]
+ pub(crate) fn check_invariants(
&self,
- plan: &P,
+ plan_name: &str,
+ children: Vec<&Arc<dyn ExecutionPlan>>,
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<dyn ExecutionPlan>>,
+ ) -> Result<()> {
let children = children
.into_iter()
.map(|child| child.as_ref())
.collect::<Vec<_>>();
- 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 6cb3eeef8c..978059b9fa 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<Arc<dyn PhysicalExpr>>,
+ ) -> 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<P: ExecutionPlan + ?Sized>(
- 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<P: ExecutionPlan + ?Sized>(
@@ -1648,9 +1660,19 @@ pub fn check_default_invariants<P: ExecutionPlan +
?Sized>(
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<Arc<dyn ExecutionPlan>>,
dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>,
+ maintains_input_order_len: Option<usize>,
+ required_input_ordering_len: Option<usize>,
+ benefits_from_input_partitioning_len: Option<usize>,
+ input_distribution_requirements: Option<InputDistributionRequirements>,
}
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<dyn ExecutionPlan>) -> Self {
+ self.children.push(child);
+ self
+ }
+
fn with_dynamic_expressions(
mut self,
dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>,
@@ -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<dyn ExecutionPlan>> {
- vec![]
+ self.children.iter().collect()
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![
+ false;
+ self.maintains_input_order_len
+ .unwrap_or(self.children.len())
+ ]
+ }
+
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>>
{
+ vec![
+ None;
+ self.required_input_ordering_len
+ .unwrap_or(self.children.len())
+ ]
+ }
+
+ fn benefits_from_input_partitioning(&self) -> Vec<bool> {
+ 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<dyn ExecutionPlan> =
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<dyn ExecutionPlan> =
+ Arc::new(crate::empty::EmptyExec::new(Arc::clone(&schema)));
+ let right: Arc<dyn ExecutionPlan> =
+ 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
