Skip to content

Commit eb2cdf3

Browse files
committed
Add union_schema construction benchmark
Benchmark `UnionExec::try_new` construction cost as a function of child count, over both a flat and a nested/struct schema. Covers the shared-Arc and content-equal fast-path cases, the adversarial last-differs case (scan wasted, then full merge), and the names-differ case where equality fails immediately. Signed-off-by: Reid Kaufmann <[email protected]>
1 parent 8fd9e55 commit eb2cdf3

2 files changed

Lines changed: 163 additions & 0 deletions

File tree

‎datafusion/physical-plan/Cargo.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,10 @@ tokio = { workspace = true, features = [
108108
harness = false
109109
name = "partial_ordering"
110110

111+
[[bench]]
112+
harness = false
113+
name = "union_schema"
114+
111115
[[bench]]
112116
harness = false
113117
name = "spill_io"
Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Benchmark for `UnionExec` construction cost as a function of child count.
19+
//!
20+
//! Scenarios (run against a flat and a nested/struct schema):
21+
//! - `shared_arc`: every child returns the same `Arc<Schema>`
22+
//! - `content_equal`: pointer-distinct but identical schemas per child
23+
//! - `last_differs`: all children equal except extra metadata on the last
24+
//! child's last top-level field. Worst case for the equality fast path:
25+
//! the scan runs, fails on the final child, then the full merge runs anyway.
26+
//! - `names_differ`: the first field is renamed in every child but the first,
27+
//! the common real-world unequal union where equality fails immediately.
28+
29+
use std::collections::HashMap;
30+
use std::sync::Arc;
31+
32+
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
33+
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
34+
use datafusion_physical_plan::ExecutionPlan;
35+
use datafusion_physical_plan::empty::EmptyExec;
36+
use datafusion_physical_plan::union::UnionExec;
37+
38+
const NUM_FIELDS: usize = 10;
39+
const NESTED_CHILDREN: usize = 5;
40+
const METADATA_PER_FIELD: usize = 2;
41+
42+
fn metadata(tag: &str, i: usize) -> HashMap<String, String> {
43+
(0..METADATA_PER_FIELD)
44+
.map(|m| (format!("key_{m}"), format!("value_{tag}_{i}_{m}")))
45+
.collect()
46+
}
47+
48+
fn flat_schema() -> Schema {
49+
let fields: Vec<Field> = (0..NUM_FIELDS)
50+
.map(|i| {
51+
Field::new(format!("col_{i}"), DataType::Int64, true)
52+
.with_metadata(metadata("f", i))
53+
})
54+
.collect();
55+
Schema::new(fields)
56+
}
57+
58+
fn nested_schema() -> Schema {
59+
let fields: Vec<Field> = (0..NUM_FIELDS)
60+
.map(|i| {
61+
let children: Vec<Field> = (0..NESTED_CHILDREN)
62+
.map(|c| {
63+
Field::new(format!("sub_{i}_{c}"), DataType::Int64, true)
64+
.with_metadata(metadata("n", i * NESTED_CHILDREN + c))
65+
})
66+
.collect();
67+
Field::new(format!("col_{i}"), DataType::Struct(children.into()), true)
68+
.with_metadata(metadata("s", i))
69+
})
70+
.collect();
71+
Schema::new(fields)
72+
}
73+
74+
/// Clone `schema` with extra metadata on its last field. Divergence must stay
75+
/// at the top level: differing nested fields change the field `DataType`
76+
/// itself, which `UnionExec::try_new` rejects ("Schemas have to be aligned").
77+
fn divergent(schema: &Schema) -> Schema {
78+
let mut fields: Vec<Field> =
79+
schema.fields().iter().map(|f| f.as_ref().clone()).collect();
80+
let last = fields.pop().unwrap();
81+
let mut md = last.metadata().clone();
82+
md.insert("divergent".to_string(), "true".to_string());
83+
fields.push(last.with_metadata(md));
84+
Schema::new(fields)
85+
}
86+
87+
fn child(schema: SchemaRef) -> Arc<dyn ExecutionPlan> {
88+
Arc::new(EmptyExec::new(schema))
89+
}
90+
91+
fn children_shared_arc(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
92+
let schema: SchemaRef = Arc::new(schema.clone());
93+
(0..n).map(|_| child(Arc::clone(&schema))).collect()
94+
}
95+
96+
fn children_content_equal(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
97+
(0..n).map(|_| child(Arc::new(schema.clone()))).collect()
98+
}
99+
100+
/// First child keeps `schema`; the rest rename the first field — the common
101+
/// real-world unequal union (`SELECT a .. UNION ALL SELECT b ..`), where any
102+
/// equality check fails immediately.
103+
fn children_names_differ(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
104+
let mut fields: Vec<Field> =
105+
schema.fields().iter().map(|f| f.as_ref().clone()).collect();
106+
let first = fields.remove(0);
107+
let renamed = first.clone().with_name(format!("renamed_{}", first.name()));
108+
fields.insert(0, renamed);
109+
let alt = Schema::new(fields);
110+
let mut children = vec![child(Arc::new(schema.clone()))];
111+
children.extend((1..n).map(|_| child(Arc::new(alt.clone()))));
112+
children
113+
}
114+
115+
fn children_last_differs(schema: &Schema, n: usize) -> Vec<Arc<dyn ExecutionPlan>> {
116+
let mut children = children_content_equal(schema, n - 1);
117+
children.push(child(Arc::new(divergent(schema))));
118+
children
119+
}
120+
121+
fn bench_union_construction(c: &mut Criterion) {
122+
for (suffix, schema, sizes) in [
123+
("", flat_schema(), &[100usize, 1000, 4000][..]),
124+
("_nested", nested_schema(), &[1000, 4000][..]),
125+
] {
126+
let mut group = c.benchmark_group(format!("union_exec_try_new{suffix}"));
127+
for &n in sizes {
128+
let shared = children_shared_arc(&schema, n);
129+
let content = children_content_equal(&schema, n);
130+
let differs = children_last_differs(&schema, n);
131+
132+
group.bench_with_input(
133+
BenchmarkId::new("shared_arc", n),
134+
&shared,
135+
|b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
136+
);
137+
group.bench_with_input(
138+
BenchmarkId::new("content_equal", n),
139+
&content,
140+
|b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
141+
);
142+
group.bench_with_input(
143+
BenchmarkId::new("last_differs", n),
144+
&differs,
145+
|b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
146+
);
147+
let names = children_names_differ(&schema, n);
148+
group.bench_with_input(
149+
BenchmarkId::new("names_differ", n),
150+
&names,
151+
|b, ch| b.iter(|| UnionExec::try_new(ch.clone()).unwrap()),
152+
);
153+
}
154+
group.finish();
155+
}
156+
}
157+
158+
criterion_group!(benches, bench_union_construction);
159+
criterion_main!(benches);

0 commit comments

Comments
 (0)