From 3a4a64e86fcbd877571500f49532bee06cd7f6a8 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 15 Jul 2026 16:31:34 -0400 Subject: [PATCH 1/5] fix merge conflict --- datafusion/physical-plan/src/union.rs | 77 +++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 01ebb21ad1635..c2528e61072e8 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,9 @@ 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. /// /// 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 +700,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 +848,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 +1295,66 @@ mod tests { ); } + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec, + descending: bool, + ) -> Result> { + let sort_expr = PhysicalSortExpr::new( + col("a", schema)?, + SortOptions { + descending, + nulls_first: false, + }, + ); + 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 partitioning = + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new(base, partitioning)?)) + } + + #[test] + fn test_can_interleave_range_matching() -> Result<()> { + let schema = create_test_schema()?; + let range_partitioning_a = make_range_exec(&schema, vec![10, 20], false)?; + let range_partitioning_b = make_range_exec(&schema, vec![10, 20], false)?; + assert!(can_interleave( + [&range_partitioning_a, &range_partitioning_b].into_iter() + )); + Ok(()) + } + + #[test] + fn test_can_interleave_range_mismatched_split_points() -> Result<()> { + let schema = create_test_schema()?; + let range_partitioning_a = make_range_exec(&schema, vec![10, 20], false)?; + let range_partitioning_b = make_range_exec(&schema, vec![10, 30], false)?; + assert!(!can_interleave( + [&range_partitioning_a, &range_partitioning_b].into_iter() + )); + Ok(()) + } + + #[test] + fn test_can_interleave_range_mixed_with_hash() -> Result<()> { + let schema = create_test_schema()?; + let range_partitioning = make_range_exec(&schema, vec![10, 20], false)?; + let hash_partitioning: Arc = + Arc::new(RepartitionExec::try_new( + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?), + Partitioning::Hash(vec![col("a", &schema)?], 3), + )?); + assert!(!can_interleave( + [&range_partitioning, &hash_partitioning].into_iter() + )); + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; From 8242bfa6f2ba86336ef35ea49f9f47265269e8e7 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 15 Jul 2026 16:37:51 -0400 Subject: [PATCH 2/5] fix merge conflict upstream --- datafusion/sqllogictest/test_files/range_partitioning.slt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index ac92e10a8ea22..d43ca2a58878d 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 From 0dc1ba6c4fd57dbbe22453f39b91cb8b6b30a6d0 Mon Sep 17 00:00:00 2001 From: Richard Date: Thu, 16 Jul 2026 10:52:09 -0400 Subject: [PATCH 3/5] add more test cases --- datafusion/physical-plan/src/union.rs | 175 +++++++++++++++++++------- 1 file changed, 133 insertions(+), 42 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index c2528e61072e8..64f717c74d210 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -1295,63 +1295,154 @@ 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, - descending: bool, + sort_options: SortOptions, ) -> Result> { - let sort_expr = PhysicalSortExpr::new( - col("a", schema)?, - SortOptions { - descending, - nulls_first: false, - }, - ); + 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 partitioning = - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?); let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); - Ok(Arc::new(RepartitionExec::try_new(base, partitioning)?)) + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) } #[test] - fn test_can_interleave_range_matching() -> Result<()> { - let schema = create_test_schema()?; - let range_partitioning_a = make_range_exec(&schema, vec![10, 20], false)?; - let range_partitioning_b = make_range_exec(&schema, vec![10, 20], false)?; - assert!(can_interleave( - [&range_partitioning_a, &range_partitioning_b].into_iter() - )); - Ok(()) - } + 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), + ])); - #[test] - fn test_can_interleave_range_mismatched_split_points() -> Result<()> { - let schema = create_test_schema()?; - let range_partitioning_a = make_range_exec(&schema, vec![10, 20], false)?; - let range_partitioning_b = make_range_exec(&schema, vec![10, 30], false)?; - assert!(!can_interleave( - [&range_partitioning_a, &range_partitioning_b].into_iter() - )); - Ok(()) - } + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + let descending = SortOptions { + descending: true, + nulls_first: false, + }; + let ascending_nulls_first = SortOptions { + descending: false, + nulls_first: true, + }; - #[test] - fn test_can_interleave_range_mixed_with_hash() -> Result<()> { - let schema = create_test_schema()?; - let range_partitioning = make_range_exec(&schema, vec![10, 20], false)?; - let hash_partitioning: Arc = - Arc::new(RepartitionExec::try_new( - Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?), - Partitioning::Hash(vec![col("a", &schema)?], 3), - )?); - assert!(!can_interleave( - [&range_partitioning, &hash_partitioning].into_iter() - )); + 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: "hash different columns", + expected: false, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![age_column], 3)?, + ], + }, + Case { + label: "hash same column different bucket count", + expected: false, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 4)?, + ], + }, + 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: "range ascending vs descending", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![20, 10], descending)?, + ], + }, + Case { + label: "range nulls_last vs nulls_first", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending_nulls_first)?, + ], + }, + 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(()) } From 9ee194d1b29eea1548a9279549004b31b84e8a08 Mon Sep 17 00:00:00 2001 From: Richard Date: Thu, 16 Jul 2026 11:06:32 -0400 Subject: [PATCH 4/5] add more test for .slt files --- datafusion/physical-plan/src/union.rs | 8 + .../test_files/range_partitioning.slt | 137 ++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 64f717c74d210..a3ce719a367ff 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -1385,6 +1385,14 @@ mod tests { ], }, // 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: "hash different columns", expected: false, diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index d43ca2a58878d..f91b296a2a4f8 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -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; From f6f7a6452e83d4dd6e30f85c411638af5ff89906 Mon Sep 17 00:00:00 2001 From: Richard Date: Thu, 16 Jul 2026 13:02:10 -0400 Subject: [PATCH 5/5] PR revisions --- datafusion/physical-plan/src/union.rs | 44 ++------------------------- 1 file changed, 3 insertions(+), 41 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index a3ce719a367ff..3511609e2e9b1 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -689,6 +689,9 @@ impl ExecutionPlan for InterleaveExec { /// 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 /// other has the partition spec Hash('a'), It is safe to derive the out partition with the spec Hash('a','b','c'). @@ -1343,15 +1346,6 @@ mod tests { descending: false, nulls_first: false, }; - let descending = SortOptions { - descending: true, - nulls_first: false, - }; - let ascending_nulls_first = SortOptions { - descending: false, - nulls_first: true, - }; - struct Case { inputs: Vec>, expected: bool, @@ -1393,22 +1387,6 @@ mod tests { make_range_exec(&schema, vec![10, 15], ascending)?, ], }, - Case { - label: "hash different columns", - expected: false, - inputs: vec![ - make_hash_exec(&schema, vec![name_column], 3)?, - make_hash_exec(&schema, vec![age_column], 3)?, - ], - }, - Case { - label: "hash same column different bucket count", - expected: false, - inputs: vec![ - make_hash_exec(&schema, vec![name_column], 3)?, - make_hash_exec(&schema, vec![name_column], 4)?, - ], - }, Case { label: "range different split points", expected: false, @@ -1417,22 +1395,6 @@ mod tests { make_range_exec(&schema, vec![10, 30], ascending)?, ], }, - Case { - label: "range ascending vs descending", - expected: false, - inputs: vec![ - make_range_exec(&schema, vec![10, 20], ascending)?, - make_range_exec(&schema, vec![20, 10], descending)?, - ], - }, - Case { - label: "range nulls_last vs nulls_first", - expected: false, - inputs: vec![ - make_range_exec(&schema, vec![10, 20], ascending)?, - make_range_exec(&schema, vec![10, 20], ascending_nulls_first)?, - ], - }, Case { label: "mixed range and hash", expected: false,