perf: reduce invariant-check code generation

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"

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);

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
             );
         }