diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 01ebb21ad163..3511609e2e9b 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -477,7 +477,10 @@ impl ExecutionPlan for UnionExec { /// Combines multiple input streams by interleaving them. /// -/// This only works if all inputs have the same hash-partitioning. +/// All inputs must share an identical [`Partitioning::Hash`] or [`Partitioning::Range`] so that +/// partition `k` covers the same data across every input. Each output partition is the +/// interleaving of the same-indexed partition from all inputs: +/// `output[k] = input[0][k] + input[1][k] + ... + input[n-1][k]` /// /// # Data Flow /// ```text @@ -522,7 +525,7 @@ impl InterleaveExec { pub fn try_new(inputs: Vec>) -> Result { assert_or_internal_err!( can_interleave(inputs.iter()), - "Not all InterleaveExec children have a consistent hash partitioning" + "Not all InterleaveExec children have a consistent hash or range partitioning" ); let cache = Self::compute_properties(&inputs)?; Ok(InterleaveExec { @@ -682,8 +685,12 @@ impl ExecutionPlan for InterleaveExec { } } -/// If all the input partitions have the same Hash partition spec with the first_input_partition -/// The InterleaveExec is partition aware. +/// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] +/// spec, making them safe to interleave. Two inputs are interleave-compatible when partition +/// `k` covers the identical key range or hash bucket across every input. +/// +/// Note: compatibility is checked sequentially against the first input, so +/// `InputDistributionRequirements::co_partitioned` is not needed here. /// /// It might be too strict here in the case that the input partition specs are compatible but not exactly the same. /// For example one input partition has the partition spec Hash('a','b','c') and @@ -696,7 +703,7 @@ pub fn can_interleave>>( }; let reference = first.borrow().output_partitioning(); - matches!(reference, Partitioning::Hash(_, _)) + matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) && inputs .map(|plan| plan.borrow().output_partitioning().clone()) .all(|partition| partition == *reference) @@ -844,10 +851,13 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::DataType; + use datafusion_common::SplitPoint; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_physical_expr::RangePartitioning; use datafusion_physical_expr::equivalence::convert_to_orderings; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; // Generate a schema which consists of 7 columns (a, b, c, d, e, f, g) fn create_test_schema() -> Result { @@ -1288,6 +1298,124 @@ mod tests { ); } + fn make_hash_exec( + schema: &SchemaRef, + hash_cols: Vec<&str>, + buckets: usize, + ) -> Result> { + let exprs = hash_cols + .iter() + .map(|c| col(c, schema)) + .collect::>>()?; + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Hash(exprs, buckets), + )?)) + } + + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec, + sort_options: SortOptions, + ) -> Result> { + let sort_expr = + PhysicalSortExpr::new(col(schema.field(0).name(), schema)?, sort_options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let split_points = split_values + .into_iter() + .map(|v| SplitPoint::new(vec![ScalarValue::Int32(Some(v))])) + .collect(); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) + } + + #[test] + fn test_can_interleave_matrix() -> Result<()> { + let name_column = "name"; + let age_column = "age"; + let schema = Arc::new(Schema::new(vec![ + Field::new(name_column, DataType::Int32, true), + Field::new(age_column, DataType::Int32, true), + ])); + + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + struct Case { + inputs: Vec>, + expected: bool, + label: &'static str, + } + + let cases = vec![ + // compatible + Case { + label: "matching hash on single column", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + Case { + label: "matching hash on multiple columns", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + ], + }, + Case { + label: "matching range same splits and order", + expected: true, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending)?, + ], + }, + // incompatible + Case { + label: "subset range partition", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 15], ascending)?, + ], + }, + Case { + label: "range different split points", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 30], ascending)?, + ], + }, + Case { + label: "mixed range and hash", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + ]; + + for case in cases { + assert_eq!( + can_interleave(case.inputs.iter()), + case.expected, + "{}", + case.label + ); + } + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index ac92e10a8ea2..f91b296a2a4f 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -782,8 +782,8 @@ reset datafusion.optimizer.preserve_file_partitions; ########## # TEST 22: Union of Range Partitioned Inputs -# Each input exposes Range partitioning on range_key. These changes do not add a -# cross-child Range relationship for UNION ALL. +# Each input exposes the same Range partitioning on range_key, so the optimizer +# converts UnionExec to InterleaveExec to avoid redundant repartitioning. ########## query TT @@ -792,7 +792,7 @@ UNION ALL SELECT range_key, value FROM range_partitioned; ---- physical_plan -01)UnionExec +01)InterleaveExec 02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false 03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false @@ -1203,3 +1203,140 @@ reset datafusion.explain.physical_plan_only; statement ok reset datafusion.optimizer.enable_window_topn; + +########## +# TEST 34: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# In a three-way union, two inputs share the same Range split points [10,20,30] +# while the third has a partially-overlapping but different set [15,20,30]. +# can_interleave requires ALL inputs to match, so UnionExec is kept. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +1 10 +5 50 +5 50 +5 50 +10 100 +10 100 +10 100 +15 150 +15 150 +15 150 +20 200 +20 200 +20 200 +25 250 +25 250 +25 250 +30 300 +30 300 +30 300 +35 350 +35 350 +35 350 + +########## +# TEST 35: Incompatible Range Split Points Falls Back to UnionExec +# Two range-partitioned inputs with different split points cannot be interleaved, +# so the optimizer keeps UnionExec instead of converting to InterleaveExec. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +########## +# TEST 36: InterleaveExec Propagates Range Partitioning to Aggregate +# InterleaveExec outputs the same Range partitioning as its compatible inputs, +# allowing a downstream aggregate on range_key to run SinglePartitioned without +# a Hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(value)] +02)--InterleaveExec +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key ORDER BY range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +statement ok +reset datafusion.explain.physical_plan_only;