fix: fast paths for UnionExec schema recomputation (O(n²) planning on wide unions)

datafusion/physical-plan/Cargo.toml

@@ -88,6 +88,10 @@ tokio = { workspace = true, features = [
 harness = false
 name = "partial_ordering"
 
+[[bench]]
+harness = false
+name = "union_schema"
+
 [[bench]]
 harness = false
 name = "spill_io"

datafusion/physical-plan/benches/union_schema.rs

@@ -0,0 +1,180 @@
+// 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.
+
+//! Benchmark for `UnionExec` construction cost as a function of child count.
+//!
+//! Scenarios (run against a flat and a nested/struct schema):
+//! - `shared_arc`: every child returns the same `Arc<Schema>`
+//! - `content_equal`: pointer-distinct but identical schemas per child
+//! - `last_differs`: identical except the last child's deepest field (worst
+//!   case for any equality fast path: the scan is wasted, then the full
+//!   merge runs)
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion_physical_plan::empty::EmptyExec;
+use datafusion_physical_plan::union::UnionExec;
+use datafusion_physical_plan::ExecutionPlan;
+
+const NUM_FIELDS: usize = 10;
+const NESTED_CHILDREN: usize = 5;
+const METADATA_PER_FIELD: usize = 2;
+
+fn metadata(tag: &str, i: usize) -> HashMap<String, String> {
+    (0..METADATA_PER_FIELD)
+        .map(|m| (format!("key_{m}"), format!("value_{tag}_{i}_{m}")))
+        .collect()
+}
+
+fn flat_schema() -> Schema {
+    let fields: Vec<Field> = (0..NUM_FIELDS)
+        .map(|i| {
+            Field::new(format!("col_{i}"), DataType::Int64, true)
+                .with_metadata(metadata("f", i))
+        })
+        .collect();
+    Schema::new(fields)
+}
+
+fn nested_schema() -> Schema {
+    let fields: Vec<Field> = (0..NUM_FIELDS)
+        .map(|i| {
+            let children: Vec<Field> = (0..NESTED_CHILDREN)
+                .map(|c| {
+                    Field::new(format!("sub_{i}_{c}"), DataType::Int64, true)
+                        .with_metadata(metadata("n", i * NESTED_CHILDREN + c))
+                })
+                .collect();
+            Field::new(format!("col_{i}"), DataType::Struct(children.into()), true)
+                .with_metadata(metadata("s", i))
+        })
+        .collect();
+    Schema::new(fields)
+}
+
+/// Clone `schema` with extra metadata on its last field. Divergence must stay
+/// at the top level: differing nested fields change the field `DataType`
+/// itself, which `UnionExec::try_new` rejects ("Schemas have to be aligned").
+fn divergent(schema: &Schema) -> Schema {
+    let mut fields: Vec<Field> =
+        schema.fields().iter().map(|f| f.as_ref().clone()).collect();
+    let last = fields.pop().unwrap();
+    let mut md = last.metadata().clone();
+    md.insert("divergent".to_string(), "true".to_string());
+    fields.push(last.with_metadata(md));
+    Schema::new(fields)
+}
+
+fn child(schema: SchemaRef) -> Arc<dyn ExecutionPlan> {
+    Arc::new(EmptyExec::new(schema))
+}
+
+fn children_shared_arc(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
+    let schema: SchemaRef = Arc::new(schema.clone());
+    (0..n).map(|_| child(Arc::clone(&schema))).collect()
+}
+
+fn children_content_equal(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
+    (0..n).map(|_| child(Arc::new(schema.clone()))).collect()
+}
+
+/// First child keeps `schema`; the rest rename the first field — the common
+/// real-world unequal union (`SELECT a .. UNION ALL SELECT b ..`), where any
+/// equality check fails immediately.
+fn children_names_differ(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
+    let mut fields: Vec<Field> =
+        schema.fields().iter().map(|f| f.as_ref().clone()).collect();
+    let first = fields.remove(0);
+    let renamed = first.clone().with_name(format!("renamed_{}", first.name()));
+    fields.insert(0, renamed);
+    let alt = Schema::new(fields);
+    let mut children = vec![child(Arc::new(schema.clone()))];
+    children.extend((1..n).map(|_| child(Arc::new(alt.clone()))));
+    children
+}
+
+fn children_last_differs(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
+    let mut children = children_content_equal(schema, n - 1);
+    children.push(child(Arc::new(divergent(schema))));
+    children
+}
+
+fn bench_union_construction(c: &mut Criterion) {
+    for (suffix, schema, sizes) in [
+        ("", flat_schema(), &[100usize, 1000, 4000][..]),
+        ("_nested", nested_schema(), &[1000, 4000][..]),
+    ] {
+        let mut group = c.benchmark_group(format!("union_exec_try_new{suffix}"));
+        for &n in sizes {
+            let shared = children_shared_arc(&schema, n);
+            let content = children_content_equal(&schema, n);
+            let differs = children_last_differs(&schema, n);
+
+            group.bench_with_input(
+                BenchmarkId::new("shared_arc", n),
+                &shared,
+                |b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
+            );
+            group.bench_with_input(
+                BenchmarkId::new("content_equal", n),
+                &content,
+                |b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
+            );
+            group.bench_with_input(
+                BenchmarkId::new("last_differs", n),
+                &differs,
+                |b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
+            );
+            let names = children_names_differ(&schema, n);
+            group.bench_with_input(
+                BenchmarkId::new("names_differ", n),
+                &names,
+                |b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
+            );
+        }
+        group.finish();
+    }
+
+    // Reconstruct the union once per child, as optimizer rewrites do.
+    let mut group = c.benchmark_group("union_exec_rebuild_per_child");
+    let schema = flat_schema();
+    for n in [100, 1000] {
+        for (label, children) in [
+            ("content_equal", children_content_equal(&schema, n)),
+            ("last_differs", children_last_differs(&schema, n)),
+        ] {
+            let union: Arc<dyn ExecutionPlan> =
+                UnionExec::try_new(children.clone()).unwrap();
+            group.bench_with_input(BenchmarkId::new(label, n), &n, |b, _| {
+                b.iter(|| {
+                    let mut plan = Arc::clone(&union);
+                    for _ in 0..n {
+                        plan = plan.with_new_children(children.clone()).unwrap();
+                    }
+                    plan
+                })
+            });
+        }
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_union_construction);
+criterion_main!(benches);

datafusion/physical-plan/src/union.rs

@@ -257,6 +257,26 @@ impl ExecutionPlan for UnionExec {
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        // Fast path: if the children's schemas are unchanged the union schema
+        // cannot change, so skip the O(children x fields) `union_schema`
+        // recomputation. Plan properties can change even when schemas do not,
+        // so they are always recomputed.
+        if children.len() == self.inputs.len()
+            && children.len() >= 2
+            && children.iter().zip(self.inputs.iter()).all(|(new, old)| {
+                let new_schema = new.schema();
+                let old_schema = old.schema();
+                Arc::ptr_eq(&new_schema, &old_schema) || new_schema == old_schema
+            })
+        {
+            let schema = self.schema();
+            let cache = Self::compute_properties(&children, schema)?;
+            return Ok(Arc::new(UnionExec {
+                inputs: children,
+                metrics: ExecutionPlanMetricsSet::new(),
+                cache,
+            }));
+        }
         UnionExec::try_new(children)
     }
 
@@ -584,6 +604,17 @@ fn union_schema(inputs: &[Arc<dyn ExecutionPlan>]) -> Result<SchemaRef> {
 
     let first_schema = inputs[0].schema();
 
+    // Merging N identical schemas is the identity operation: if every input
+    // schema is pointer- or content-equal to the first (`Schema` equality
+    // covers every property the merge below reads), return it directly
+    // instead of paying the per-field metadata merge.
+    if inputs.iter().all(|input| {
+        let schema = input.schema();
+        Arc::ptr_eq(&schema, &first_schema) || schema == first_schema
+    }) {
+        return Ok(first_schema);
+    }
+
     let fields = (0..first_schema.fields().len())
         .map(|i| {
             // We take the name from the left side of the union to match how names are coerced during logical planning,