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