From 18f57b7122d8e1621f90cb94c6c85717fd1b0b35 Mon Sep 17 00:00:00 2001 From: Edson Petry Date: Tue, 14 Jul 2026 15:43:03 -0400 Subject: [PATCH 1/2] fix: preserve range partitioning through joins --- datafusion/physical-plan/src/joins/utils.rs | 65 +++++++++++++++++-- .../test_files/range_partitioning.slt | 30 ++++----- 2 files changed, 72 insertions(+), 23 deletions(-) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 2a7759a8abeec..d068bead9f35f 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/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; @@ -146,9 +147,16 @@ 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) + .expect("offsetting a range ordering keeps it non-empty"); + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) } result => result.clone(), }; @@ -2468,7 +2476,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 +4139,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![ diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa5..2580aa80b1085 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/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=, 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=, 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=, 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=, 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 +05)--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 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=, 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=, 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=, 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 l.range_key, SUM(l.value + r.value) From ad0aca8119b9240691ae3850d94d1aa02dc66c32 Mon Sep 17 00:00:00 2001 From: Edson Petry Date: Wed, 15 Jul 2026 03:28:31 -0400 Subject: [PATCH 2/2] fix: return error for empty range ordering --- datafusion/physical-plan/src/joins/utils.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index d068bead9f35f..90b39f7ada122 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -69,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; @@ -151,8 +151,11 @@ pub fn adjust_right_output_partitioning( range.ordering().iter().cloned(), left_columns_len as _, )?; - let ordering = LexOrdering::new(ordering) - .expect("offsetting a range ordering keeps it non-empty"); + 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(),