fix: preserve range partitioning through joins (#23584)

## Which issue does this PR close?

- Closes #23450.

## Rationale for this change

#23184 allowed compatible range-partitioned inputs to satisfy
partitioned inner joins without hash repartitioning. However, join
output still downgraded `Partitioning::Range` to `UnknownPartitioning`,
so downstream joins and aggregates could not reuse the preserved range
layout and inserted avoidable hash repartitions.

## What changes are included in this PR?

- Preserve `Partitioning::Range` in `adjust_right_output_partitioning`.
- Shift right-side range-ordering column indexes into the joined schema
while retaining split points and sort options.
- Add unit coverage for compound range keys and non-default sort
options.
- Update SQL logic test plans for nested range joins and a range join
feeding an aggregate.

## Are these changes tested?

Yes.

- `./dev/rust_lint.sh`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `cargo test -p datafusion-physical-plan
test_adjust_right_output_partitioning_preserves_range`
- `cargo test --profile=ci --test sqllogictests --
range_partitioning.slt`
- `ulimit -n 10240 && RUST_BACKTRACE=1 cargo test --profile ci --exclude
datafusion-examples --exclude datafusion-benchmarks --exclude
datafusion-cli --workspace --lib --tests --bins --features
avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`

## Are there any user-facing changes?

Yes. Physical plans can preserve proven range partitioning through
joins, allowing compatible downstream joins and aggregates to avoid
unnecessary hash repartitioning. There are no public API changes.

Author: EdsonPetry

datafusion/physical-plan/src/joins/utils.rs

@@ -33,7 +33,8 @@ use crate::metrics::{
 };
 use crate::projection::{ProjectionExec, ProjectionExpr};
 use crate::{
-    ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics,
+    ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning,
+    RangePartitioning, Statistics,
 };
 // compatibility
 pub use super::join_filter::JoinFilter;
@@ -68,7 +69,7 @@ use datafusion_common::stats::Precision;
 use datafusion_common::utils::normalize_float_zero;
 use datafusion_common::{
     DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult,
-    not_impl_err, plan_err,
+    internal_datafusion_err, not_impl_err, plan_err,
 };
 use datafusion_expr::Operator;
 use datafusion_expr::interval_arithmetic::Interval;
@@ -146,9 +147,19 @@ pub fn adjust_right_output_partitioning(
             Partitioning::Hash(new_exprs, *size)
         }
         Partitioning::Range(range) => {
-            // Range partitioning optimizer propagation is tracked in
-            // https://github.com/apache/datafusion/issues/22395
-            Partitioning::UnknownPartitioning(range.partition_count())
+            let ordering = add_offset_to_physical_sort_exprs(
+                range.ordering().iter().cloned(),
+                left_columns_len as _,
+            )?;
+            let ordering = LexOrdering::new(ordering).ok_or_else(|| {
+                internal_datafusion_err!(
+                    "Offsetting range partitioning produced an empty ordering"
+                )
+            })?;
+            Partitioning::Range(RangePartitioning::new(
+                ordering,
+                range.split_points().to_vec(),
+            ))
         }
         result => result.clone(),
     };
@@ -2468,7 +2479,7 @@ mod tests {
     use arrow::datatypes::{DataType, Fields};
     use arrow::error::{ArrowError, Result as ArrowResult};
     use datafusion_common::stats::Precision::{Absent, Exact, Inexact};
-    use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err};
+    use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err};
     use datafusion_physical_expr::PhysicalSortExpr;
 
     use rstest::rstest;
@@ -4131,6 +4142,53 @@ mod tests {
         assert_eq!(result.column_statistics[1].byte_size, Inexact(256));
     }
 
+    #[test]
+    fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> {
+        let split_points = vec![
+            SplitPoint::new(vec![
+                ScalarValue::Int32(Some(10)),
+                ScalarValue::Int32(Some(100)),
+            ]),
+            SplitPoint::new(vec![
+                ScalarValue::Int32(Some(20)),
+                ScalarValue::Int32(Some(50)),
+            ]),
+        ];
+        let range = RangePartitioning::try_new(
+            LexOrdering::new([
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("a", 0)),
+                    SortOptions::new(false, true),
+                ),
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("b", 2)),
+                    SortOptions::new(true, false),
+                ),
+            ])
+            .unwrap(),
+            split_points.clone(),
+        )?;
+
+        let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?;
+        let expected = Partitioning::Range(RangePartitioning::new(
+            LexOrdering::new([
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("a", 3)),
+                    SortOptions::new(false, true),
+                ),
+                PhysicalSortExpr::new(
+                    Arc::new(Column::new("b", 5)),
+                    SortOptions::new(true, false),
+                ),
+            ])
+            .unwrap(),
+            split_points,
+        ));
+
+        assert_eq!(adjusted, expected);
+        Ok(())
+    }
+
     #[test]
     fn test_calculate_join_output_ordering() -> Result<()> {
         let left_ordering = LexOrdering::new(vec![

datafusion/sqllogictest/test_files/range_partitioning.slt

@@ -506,9 +506,8 @@ set datafusion.optimizer.preserve_file_partitions = 0;
 
 ##########
 # TEST 15: Nested Range Joins
-# Compatible Range partitioning satisfies the lower join inputs. The upper join
-# still repairs the intermediate join output with Hash repartitioning because
-# HashJoinExec does not currently expose Range output partitioning.
+# Compatible Range partitioning is preserved through the lower join, allowing
+# the upper join to consume it without Hash repartitioning either input.
 ##########
 
 query TT
@@ -519,12 +518,10 @@ JOIN range_partitioned s ON r.range_key = s.range_key;
 ----
 physical_plan
 01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@2, range_key@0)], projection=[range_key@0, value@1, value@3, value@5]
-02)--RepartitionExec: partitioning=Hash([range_key@2], 4), input_partitions=4
-03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)]
-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
-05)------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
-06)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
-07)----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
+02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)]
+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
+05)--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 IIII
 SELECT l.range_key, l.value, r.value, s.value
@@ -599,9 +596,8 @@ ORDER BY l.range_key;
 
 ##########
 # TEST 17: Range Join Feeds Aggregate
-# The join inputs avoid Hash repartitioning, but the aggregate above the join
-# still repartitions because HashJoinExec does not currently expose Range
-# output partitioning.
+# The join preserves compatible Range partitioning on range_key, allowing the
+# aggregate above it to avoid Hash repartitioning.
 ##########
 
 query TT
@@ -611,12 +607,10 @@ JOIN range_partitioned r ON l.range_key = r.range_key
 GROUP BY l.range_key;
 ----
 physical_plan
-01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)]
-02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
-03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)]
-04)------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3]
-05)--------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
-06)--------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
+01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)]
+02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3]
+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 l.range_key, SUM(l.value + r.value)