diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 5858deb83c83c..956e7f8d4b530 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/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" diff --git a/datafusion/physical-plan/benches/union_schema.rs b/datafusion/physical-plan/benches/union_schema.rs new file mode 100644 index 0000000000000..d25fe2c0bdee8 --- /dev/null +++ b/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` +//! - `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 { + (0..METADATA_PER_FIELD) + .map(|m| (format!("key_{m}"), format!("value_{tag}_{i}_{m}"))) + .collect() +} + +fn flat_schema() -> Schema { + let fields: Vec = (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 = (0..NUM_FIELDS) + .map(|i| { + let children: Vec = (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 = + 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 { + Arc::new(EmptyExec::new(schema)) +} + +fn children_shared_arc(schema: &Schema, n: usize) -> Vec> { + let schema: SchemaRef = Arc::new(schema.clone()); + (0..n).map(|_| child(Arc::clone(&schema))).collect() +} + +fn children_content_equal(schema: &Schema, n: usize) -> Vec> { + (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> { + let mut fields: Vec = + 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> { + 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 = + 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); diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index c95678dac9cdd..a4eeb8b498560 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -257,6 +257,26 @@ impl ExecutionPlan for UnionExec { self: Arc, children: Vec>, ) -> Result> { + // 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]) -> Result { 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,