allow interleaveExec to support Range partioning

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<Arc<dyn ExecutionPlan>>) -> Result<Self> {
         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<T: Borrow<Arc<dyn ExecutionPlan>>>(
     };
 
     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<SchemaRef> {
@@ -1288,6 +1298,124 @@ mod tests {
         );
     }
 
+    fn make_hash_exec(
+        schema: &SchemaRef,
+        hash_cols: Vec<&str>,
+        buckets: usize,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let exprs = hash_cols
+            .iter()
+            .map(|c| col(c, schema))
+            .collect::<Result<Vec<_>>>()?;
+        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<i32>,
+        sort_options: SortOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        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<Arc<dyn ExecutionPlan>>,
+            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()?;

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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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=<slt:ignore>, 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;