From c0f9ed3b1926d3a96b1c7693cb3e105d25132be0 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 14 Jul 2026 19:38:04 -0400 Subject: [PATCH 1/9] feat: finishing range repartition physical plan --- datafusion/physical-plan/src/repartition/mod.rs | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 6587157421946..2e8c3912c0ded 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1603,12 +1603,8 @@ impl ExecutionPlan for RepartitionExec { return Ok(SortOrderPushdownResult::Unsupported); } match self.partitioning() { - Partitioning::Range(_) => { - // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/23230 - return Ok(SortOrderPushdownResult::Unsupported); - } - Partitioning::RoundRobinBatch(_) + Partitioning::Range(_) + | Partitioning::RoundRobinBatch(_) | Partitioning::Hash(_, _) | Partitioning::UnknownPartitioning(_) => {} } @@ -1634,12 +1630,11 @@ impl ExecutionPlan for RepartitionExec { new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), - UnknownPartitioning(_) => UnknownPartitioning(target_partitions), Range(_) => { - // Range repartition optimizations are tracked in - // https://github.com/apache/datafusion/issues/23230 + // Number of partitions is constrained by the split points and cannot be changed return Ok(None); } + UnknownPartitioning(_) => UnknownPartitioning(target_partitions), }; Ok(Some(Arc::new(Self { input: Arc::clone(&self.input), From 935183531b1925837b0fd429e76c245f150395e5 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 14 Jul 2026 20:29:07 -0400 Subject: [PATCH 2/9] working on swapping with projection pass --- .../physical-plan/src/repartition/mod.rs | 33 ++++++++++++++++--- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 2e8c3912c0ded..18ef9916ebd84 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1561,10 +1561,35 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Hash(new_partitions, *size) } - Partitioning::Range(_) => { - // Range partitioning optimizer propagation is tracked in - // https://github.com/apache/datafusion/issues/23230 - return Ok(None); + Partitioning::Range(range_partitioning) => { + let updated_exprs = range_partitioning + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .filter_map(|expr| { + update_expr(&expr, projection.expr(), false).ok()? + }) + // .filter_map(|expr| expr.transpose()?) + .collect::>(); + + if updated_exprs.len() != range_partitioning.ordering().len() { + return Ok(None); + } + + let sort_exprs = range_partitioning + .ordering() + .iter() + .zip(updated_exprs) + .map(|(sort_expr, expr)| { + PhysicalSortExpr::new(expr, sort_expr.options) + }) + .collect::>(); + let ordering = LexOrdering::new(sort_exprs)?; + + Partitioning::Range(RangePartitioning::new( + ordering, + range_partitioning.split_points().clone(), + )) } others => others.clone(), }; From e39d31ff87c58ac0f8d97d42929c2dcbfc7c94aa Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Wed, 15 Jul 2026 12:06:18 -0400 Subject: [PATCH 3/9] try_swapping_with_projection complete for range repartitioning plan --- .../physical-plan/src/repartition/mod.rs | 236 ++++++++++++++++-- 1 file changed, 214 insertions(+), 22 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 18ef9916ebd84..bbd21defdadb5 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1562,33 +1562,27 @@ impl ExecutionPlan for RepartitionExec { Partitioning::Hash(new_partitions, *size) } Partitioning::Range(range_partitioning) => { - let updated_exprs = range_partitioning - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .filter_map(|expr| { - update_expr(&expr, projection.expr(), false).ok()? - }) - // .filter_map(|expr| expr.transpose()?) - .collect::>(); - - if updated_exprs.len() != range_partitioning.ordering().len() { - return Ok(None); + // Rewriting the range key expressions in ordering based on the projection expressions + let mut sort_exprs = + Vec::with_capacity(range_partitioning.ordering().len()); + for sort_expr in range_partitioning.ordering() { + let Some(new_expr) = + update_expr(&sort_expr.expr, projection.expr(), false)? + else { + return Ok(None); + }; + sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options)); } - let sort_exprs = range_partitioning - .ordering() - .iter() - .zip(updated_exprs) - .map(|(sort_expr, expr)| { - PhysicalSortExpr::new(expr, sort_expr.options) - }) - .collect::>(); - let ordering = LexOrdering::new(sort_exprs)?; + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; Partitioning::Range(RangePartitioning::new( ordering, - range_partitioning.split_points().clone(), + range_partitioning.split_points().to_vec(), )) } others => others.clone(), @@ -2118,6 +2112,8 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -2578,6 +2574,202 @@ mod tests { Ok(()) } + #[test] + fn range_repartition_swaps_with_projection_keeping_key() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + + // Narrow to [id, payload], keeping the range key. + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: col("id", &repartition.schema())?, + alias: "id".to_string(), + }, + ProjectionExpr { + expr: col("payload", &repartition.schema())?, + alias: "payload".to_string(), + }, + ], + Arc::clone(&repartition) as _, + )?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::() + .expect("top node should be RepartitionExec"); + + assert!( + swapped_repartition.input().is::(), + "projection should be pushed below the repartition" + ); + assert_eq!(swapped.schema().fields().len(), 2); + assert_eq!(swapped.schema().field(0).name(), "id"); + assert_eq!(swapped.schema().field(1).name(), "payload"); + + match swapped_repartition.partitioning() { + Partitioning::Range(range) => { + assert_eq!(range.partition_count(), 3); + assert_eq!(range.ordering().len(), 1); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + assert_eq!( + range.split_points(), + &[ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ] + ); + } + other => panic!("expected Range partitioning, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_swaps_with_projection_reorders_key() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], + )?); + let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + + // Reorder so the range key moves from @0 to @1. + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: col("payload", &repartition.schema())?, + alias: "payload".to_string(), + }, + ProjectionExpr { + expr: col("id", &repartition.schema())?, + alias: "id".to_string(), + }, + ], + Arc::clone(&repartition) as _, + )?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::() + .expect("top node should be RepartitionExec"); + + match swapped_repartition.partitioning() { + Partitioning::Range(range) => { + assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); + } + other => panic!("expected Range partitioning, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], + )?); + let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + + // Drop the range key column. + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: col("region", &repartition.schema())?, + alias: "region".to_string(), + }, + ProjectionExpr { + expr: col("payload", &repartition.schema())?, + alias: "payload".to_string(), + }, + ], + Arc::clone(&repartition) as _, + )?; + + let result = repartition.try_swapping_with_projection(&projection)?; + assert!( + result.is_none(), + "swap should fail when projection drops the range key" + ); + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_compound_key_part( + ) -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ])], + )?); + let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + + // Keep `a` and `c`, but drop `b` which is part of the range key. + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: col("a", &repartition.schema())?, + alias: "a".to_string(), + }, + ProjectionExpr { + expr: col("c", &repartition.schema())?, + alias: "c".to_string(), + }, + ], + Arc::clone(&repartition) as _, + )?; + + let result = repartition.try_swapping_with_projection(&projection)?; + assert!( + result.is_none(), + "swap should fail when projection drops any part of a compound range key" + ); + + Ok(()) + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { let schema = test_schema(false); From 3649ad206d76a251155429d034ee94d53a96b7dd Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Wed, 15 Jul 2026 12:13:19 -0400 Subject: [PATCH 4/9] test added for repartitioned --- .../physical-plan/src/repartition/mod.rs | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index bbd21defdadb5..5df94fdfacffe 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -3873,6 +3873,33 @@ mod test { Ok(()) } + #[test] + fn test_range_repartitioned_returns_none() -> Result<()> { + let schema = test_schema(); + let source = memory_exec(&schema); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("c0", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(source, partitioning)?; + + // Range partition count is fixed by split points, so repartitioned() + // cannot change it to an arbitrary target. + let result = exec.repartitioned(10, &Default::default())?; + assert!( + result.is_none(), + "range repartitioning should not support changing partition count" + ); + Ok(()) + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } From a10af3d27c0b5206e3d3ff93a5fd6a1a744bdc4f Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Wed, 15 Jul 2026 12:57:27 -0400 Subject: [PATCH 5/9] formatting --- .../physical-plan/src/repartition/mod.rs | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 5df94fdfacffe..db91e8d902fdc 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2583,7 +2583,11 @@ mod tests { ])); let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + [PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )] + .into(), vec![ SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), @@ -2649,7 +2653,11 @@ mod tests { ])); let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + [PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )] + .into(), vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], )?); let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); @@ -2695,7 +2703,11 @@ mod tests { ])); let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new(col("id", &schema)?, SortOptions::default())].into(), + [PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )] + .into(), vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], )?); let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); @@ -2725,8 +2737,8 @@ mod tests { } #[test] - fn range_repartition_does_not_swap_when_projection_drops_compound_key_part( - ) -> Result<()> { + fn range_repartition_does_not_swap_when_projection_drops_compound_key_part() + -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::UInt32, false), From 8551849ab876045d605fb0bbb275c5dcf3b8b762 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Thu, 16 Jul 2026 16:20:08 -0400 Subject: [PATCH 6/9] addressing comments --- .../physical-plan/src/repartition/mod.rs | 414 ++++++++++++------ 1 file changed, 288 insertions(+), 126 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index db91e8d902fdc..f711a0eae9bca 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1580,10 +1580,10 @@ impl ExecutionPlan for RepartitionExec { ); }; - Partitioning::Range(RangePartitioning::new( + Partitioning::Range(RangePartitioning::try_new( ordering, range_partitioning.split_points().to_vec(), - )) + )?) } others => others.clone(), }; @@ -1621,12 +1621,6 @@ impl ExecutionPlan for RepartitionExec { if !self.maintains_input_order()[0] { return Ok(SortOrderPushdownResult::Unsupported); } - match self.partitioning() { - Partitioning::Range(_) - | Partitioning::RoundRobinBatch(_) - | Partitioning::Hash(_, _) - | Partitioning::UnknownPartitioning(_) => {} - } // Delegate to the child and wrap with a new RepartitionExec self.input.try_pushdown_sort(order)?.try_map(|new_input| { @@ -2576,39 +2570,13 @@ mod tests { #[test] fn range_repartition_swaps_with_projection_keeping_key() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt32, false), - Field::new("region", DataType::Utf8, false), - Field::new("payload", DataType::UInt32, false), - ])); - let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new( - col("id", &schema)?, - SortOptions::default(), - )] - .into(), - vec![ - SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), - SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), - ], - )?); - let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + let schema = id_region_payload_schema(); + let repartition = + range_repartition_on_columns(&schema, &["id"], vec![vec![10], vec![20]])?; // Narrow to [id, payload], keeping the range key. - let projection = ProjectionExec::try_new( - vec![ - ProjectionExpr { - expr: col("id", &repartition.schema())?, - alias: "id".to_string(), - }, - ProjectionExpr { - expr: col("payload", &repartition.schema())?, - alias: "payload".to_string(), - }, - ], - Arc::clone(&repartition) as _, - )?; + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["id", "payload"])?; let swapped = repartition .try_swapping_with_projection(&projection)? @@ -2646,36 +2614,12 @@ mod tests { #[test] fn range_repartition_swaps_with_projection_reorders_key() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt32, false), - Field::new("region", DataType::Utf8, false), - Field::new("payload", DataType::UInt32, false), - ])); - let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new( - col("id", &schema)?, - SortOptions::default(), - )] - .into(), - vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], - )?); - let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + let schema = id_region_payload_schema(); + let repartition = range_repartition_on_columns(&schema, &["id"], vec![vec![10]])?; // Reorder so the range key moves from @0 to @1. - let projection = ProjectionExec::try_new( - vec![ - ProjectionExpr { - expr: col("payload", &repartition.schema())?, - alias: "payload".to_string(), - }, - ProjectionExpr { - expr: col("id", &repartition.schema())?, - alias: "id".to_string(), - }, - ], - Arc::clone(&repartition) as _, - )?; + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; let swapped = repartition .try_swapping_with_projection(&projection)? @@ -2696,36 +2640,12 @@ mod tests { #[test] fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt32, false), - Field::new("region", DataType::Utf8, false), - Field::new("payload", DataType::UInt32, false), - ])); - let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new( - col("id", &schema)?, - SortOptions::default(), - )] - .into(), - vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], - )?); - let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + let schema = id_region_payload_schema(); + let repartition = range_repartition_on_columns(&schema, &["id"], vec![vec![10]])?; // Drop the range key column. - let projection = ProjectionExec::try_new( - vec![ - ProjectionExpr { - expr: col("region", &repartition.schema())?, - alias: "region".to_string(), - }, - ProjectionExpr { - expr: col("payload", &repartition.schema())?, - alias: "payload".to_string(), - }, - ], - Arc::clone(&repartition) as _, - )?; + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["region", "payload"])?; let result = repartition.try_swapping_with_projection(&projection)?; assert!( @@ -2739,39 +2659,13 @@ mod tests { #[test] fn range_repartition_does_not_swap_when_projection_drops_compound_key_part() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::UInt32, false), - Field::new("b", DataType::UInt32, false), - Field::new("c", DataType::UInt32, false), - ])); - let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - [ - PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), - PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), - ] - .into(), - vec![SplitPoint::new(vec![ - ScalarValue::UInt32(Some(10)), - ScalarValue::UInt32(Some(1)), - ])], - )?); - let repartition = Arc::new(RepartitionExec::try_new(input, partitioning)?); + let schema = u32_fields_schema(&["a", "b", "c"]); + let repartition = + range_repartition_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?; // Keep `a` and `c`, but drop `b` which is part of the range key. - let projection = ProjectionExec::try_new( - vec![ - ProjectionExpr { - expr: col("a", &repartition.schema())?, - alias: "a".to_string(), - }, - ProjectionExpr { - expr: col("c", &repartition.schema())?, - alias: "c".to_string(), - }, - ], - Arc::clone(&repartition) as _, - )?; + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; let result = repartition.try_swapping_with_projection(&projection)?; assert!( @@ -2782,6 +2676,141 @@ mod tests { Ok(()) } + #[test] + fn range_repartition_try_pushdown_sort_maintains_ordering() -> Result<()> { + let schema = u32_fields_schema(&["id"]); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition sorted source so Range repartition must use + // preserve_order to maintain input ordering. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + ordering.clone(), + vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], + )?); + let repartition = Arc::new( + RepartitionExec::try_new(source, partitioning)?.with_preserve_order(), + ); + assert!( + repartition.maintains_input_order()[0], + "test setup requires an order-maintaining range repartition" + ); + assert_eq!( + repartition + .properties() + .output_ordering() + .map(|o| o.to_string()), + Some(ordering.to_string()) + ); + + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::() + .expect("pushdown should keep a RepartitionExec wrapper"); + assert!( + pushed.maintains_input_order()[0], + "pushed plan should still maintain input order" + ); + assert!(pushed.preserve_order()); + assert_eq!( + pushed.properties().output_ordering().map(|o| o.to_string()), + Some(ordering.to_string()) + ); + match pushed.partitioning() { + Partitioning::Range(range) => { + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + } + other => panic!("expected Range partitioning, got {other:?}"), + } + } + other => panic!( + "expected Exact sort pushdown through range repartition, got {other:?}" + ), + } + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() + -> Result<()> { + let schema = u32_fields_schema(&["id"]); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source without preserve_order: Range does not maintain order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + ordering.clone(), + vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], + )?); + let repartition = Arc::new(RepartitionExec::try_new(source, partitioning)?); + assert!(!repartition.maintains_input_order()[0]); + + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Unsupported + )); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_preserves_sortedness_within_partitions() -> Result<()> { + let schema = u32_fields_schema(&["id"]); + // Sorted single-partition input spans multiple range buckets. + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![1, 5, 10, 15, 20, 25]))], + )?; + let input = Arc::new(MockExec::new(vec![Ok(batch)], Arc::clone(&schema))); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(input, partitioning)?; + assert!(exec.maintains_input_order()[0]); + + let task_ctx = Arc::new(TaskContext::default()); + for partition in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(partition, Arc::clone(&task_ctx))?; + let mut values = Vec::new(); + while let Some(result) = stream.next().await { + let batch = result?; + values.extend(as_uint32_array(batch.column(0))?.values().iter().copied()); + } + assert!( + values.windows(2).all(|w| w[0] <= w[1]), + "partition {partition} lost ascending order: {values:?}" + ); + } + + Ok(()) + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { let schema = test_schema(false); @@ -2808,6 +2837,139 @@ mod tests { Ok(()) } + fn id_region_payload_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])) + } + + fn u32_fields_schema(names: &[&str]) -> SchemaRef { + Arc::new(Schema::new( + names + .iter() + .map(|name| Field::new(*name, DataType::UInt32, false)) + .collect::>(), + )) + } + + fn range_repartition_on_columns( + schema: &SchemaRef, + key_columns: &[&str], + split_points: Vec>, + ) -> Result> { + let ordering = LexOrdering::new( + key_columns + .iter() + .map(|name| { + Ok(PhysicalSortExpr::new( + col(name, schema)?, + SortOptions::default(), + )) + }) + .collect::>>()?, + ) + .expect("range ordering must not be empty"); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points + .into_iter() + .map(|values| { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::UInt32(Some(value))) + .collect(), + ) + }) + .collect(), + )?); + let input = Arc::new(EmptyExec::new(Arc::clone(schema))); + Ok(Arc::new(RepartitionExec::try_new(input, partitioning)?)) + } + + fn projection_columns( + input: &Arc, + names: &[&str], + ) -> Result { + let exprs = names + .iter() + .map(|name| { + Ok(ProjectionExpr { + expr: col(name, &input.schema())?, + alias: (*name).to_string(), + }) + }) + .collect::>>()?; + ProjectionExec::try_new(exprs, Arc::clone(input)) + } + + /// Test source that claims Exact support for any sort pushdown request. + #[derive(Debug, Clone)] + struct ExactSortPushdownExec { + cache: Arc, + } + + impl ExactSortPushdownExec { + fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { + use crate::execution_plan::{Boundedness, EmissionType}; + let eq_properties = + EquivalenceProperties::new_with_orderings(schema, [ordering]); + let cache = Arc::new(PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { cache } + } + } + + impl DisplayAs for ExactSortPushdownExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExactSortPushdownExec") + } + } + + impl ExecutionPlan for ExactSortPushdownExec { + fn name(&self) -> &str { + "ExactSortPushdownExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(self.schema()))) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(self.clone()), + }) + } + } + fn test_schema(nullable: bool) -> Arc { Arc::new(Schema::new(vec![Field::new( "c0", From af756be4b7a4e6c24aa5ba4dcc5776f42b8fb6c4 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Fri, 17 Jul 2026 09:57:11 -0400 Subject: [PATCH 7/9] test changes --- .../physical-plan/src/repartition/mod.rs | 336 ++++++------------ 1 file changed, 108 insertions(+), 228 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index bfdfab3209c84..fa2396422ce0a 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2615,57 +2615,19 @@ mod tests { } #[test] - fn range_repartition_swaps_with_projection_keeping_key() -> Result<()> { - let schema = id_region_payload_schema(); - let repartition = - range_repartition_on_columns(&schema, &["id"], vec![vec![10], vec![20]])?; - - // Narrow to [id, payload], keeping the range key. - let projection = - projection_columns(&(Arc::clone(&repartition) as _), &["id", "payload"])?; - - let swapped = repartition - .try_swapping_with_projection(&projection)? - .expect("swap should succeed when projection keeps the range key"); - let swapped_repartition = swapped - .downcast_ref::() - .expect("top node should be RepartitionExec"); - - assert!( - swapped_repartition.input().is::(), - "projection should be pushed below the repartition" - ); - assert_eq!(swapped.schema().fields().len(), 2); - assert_eq!(swapped.schema().field(0).name(), "id"); - assert_eq!(swapped.schema().field(1).name(), "payload"); - - match swapped_repartition.partitioning() { - Partitioning::Range(range) => { - assert_eq!(range.partition_count(), 3); - assert_eq!(range.ordering().len(), 1); - assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); - assert_eq!( - range.split_points(), - &[ - SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), - SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), - ] - ); - } - other => panic!("expected Range partitioning, got {other:?}"), - } - - Ok(()) - } - - #[test] - fn range_repartition_swaps_with_projection_reorders_key() -> Result<()> { - let schema = id_region_payload_schema(); - let repartition = range_repartition_on_columns(&schema, &["id"], vec![vec![10]])?; + fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); // Reorder so the range key moves from @0 to @1. let projection = - projection_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; let swapped = repartition .try_swapping_with_projection(&projection)? @@ -2674,114 +2636,85 @@ mod tests { .downcast_ref::() .expect("top node should be RepartitionExec"); - match swapped_repartition.partitioning() { - Partitioning::Range(range) => { - assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); - } - other => panic!("expected Range partitioning, got {other:?}"), - } + assert!(swapped_repartition.input().is::()); + assert_eq!( + as_range_partitioning(swapped_repartition.partitioning()).ordering()[0] + .to_string(), + "id@1 ASC" + ); Ok(()) } #[test] fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { - let schema = id_region_payload_schema(); - let repartition = range_repartition_on_columns(&schema, &["id"], vec![vec![10]])?; - - // Drop the range key column. + // Drop a simple range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); let projection = - projection_columns(&(Arc::clone(&repartition) as _), &["region", "payload"])?; - - let result = repartition.try_swapping_with_projection(&projection)?; + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?; assert!( - result.is_none(), - "swap should fail when projection drops the range key" + repartition + .try_swapping_with_projection(&projection)? + .is_none() ); - Ok(()) - } - - #[test] - fn range_repartition_does_not_swap_when_projection_drops_compound_key_part() - -> Result<()> { - let schema = u32_fields_schema(&["a", "b", "c"]); - let repartition = - range_repartition_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?; - - // Keep `a` and `c`, but drop `b` which is part of the range key. + // Drop part of a compound range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?, + )?); let projection = - projection_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; - - let result = repartition.try_swapping_with_projection(&projection)?; + projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; assert!( - result.is_none(), - "swap should fail when projection drops any part of a compound range key" + repartition + .try_swapping_with_projection(&projection)? + .is_none() ); Ok(()) } #[test] - fn range_repartition_try_pushdown_sort_maintains_ordering() -> Result<()> { - let schema = u32_fields_schema(&["id"]); + fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); let ordering = LexOrdering::new([PhysicalSortExpr::new( col("id", &schema)?, SortOptions::default(), )]) .expect("ordering must not be empty"); - // Multi-partition sorted source so Range repartition must use - // preserve_order to maintain input ordering. + // Multi-partition source with preserve_order: Range maintains input order. let source = Arc::new(ExactSortPushdownExec::new( Arc::clone(&schema), 2, ordering.clone(), )); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - ordering.clone(), - vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], - )?); let repartition = Arc::new( - RepartitionExec::try_new(source, partitioning)?.with_preserve_order(), - ); - assert!( - repartition.maintains_input_order()[0], - "test setup requires an order-maintaining range repartition" - ); - assert_eq!( - repartition - .properties() - .output_ordering() - .map(|o| o.to_string()), - Some(ordering.to_string()) + RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )? + .with_preserve_order(), ); + assert!(repartition.maintains_input_order()[0]); - match repartition.try_pushdown_sort(ordering.as_ref())? { - SortOrderPushdownResult::Exact { inner } => { - let pushed = inner - .downcast_ref::() - .expect("pushdown should keep a RepartitionExec wrapper"); - assert!( - pushed.maintains_input_order()[0], - "pushed plan should still maintain input order" - ); - assert!(pushed.preserve_order()); - assert_eq!( - pushed.properties().output_ordering().map(|o| o.to_string()), - Some(ordering.to_string()) - ); - match pushed.partitioning() { - Partitioning::Range(range) => { - assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); - } - other => panic!("expected Range partitioning, got {other:?}"), - } - } - other => panic!( - "expected Exact sort pushdown through range repartition, got {other:?}" - ), - } + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Exact { .. } + )); Ok(()) } @@ -2789,7 +2722,8 @@ mod tests { #[test] fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() -> Result<()> { - let schema = u32_fields_schema(&["id"]); + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); let ordering = LexOrdering::new([PhysicalSortExpr::new( col("id", &schema)?, SortOptions::default(), @@ -2802,11 +2736,10 @@ mod tests { 2, ordering.clone(), )); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - ordering.clone(), - vec![SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])], + let repartition = Arc::new(RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, )?); - let repartition = Arc::new(RepartitionExec::try_new(source, partitioning)?); assert!(!repartition.maintains_input_order()[0]); assert!(matches!( @@ -2817,94 +2750,11 @@ mod tests { Ok(()) } - #[tokio::test] - async fn range_repartition_preserves_sortedness_within_partitions() -> Result<()> { - let schema = u32_fields_schema(&["id"]); - // Sorted single-partition input spans multiple range buckets. - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(UInt32Array::from(vec![1, 5, 10, 15, 20, 25]))], - )?; - let input = Arc::new(MockExec::new(vec![Ok(batch)], Arc::clone(&schema))); - let partitioning = Partitioning::Range(RangePartitioning::try_new( - [PhysicalSortExpr::new( - col("id", &schema)?, - SortOptions::default(), - )] - .into(), - vec![ - SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), - SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), - ], - )?); - let exec = RepartitionExec::try_new(input, partitioning)?; - assert!(exec.maintains_input_order()[0]); - - let task_ctx = Arc::new(TaskContext::default()); - for partition in 0..exec.partitioning().partition_count() { - let mut stream = exec.execute(partition, Arc::clone(&task_ctx))?; - let mut values = Vec::new(); - while let Some(result) = stream.next().await { - let batch = result?; - values.extend(as_uint32_array(batch.column(0))?.values().iter().copied()); - } - assert!( - values.windows(2).all(|w| w[0] <= w[1]), - "partition {partition} lost ascending order: {values:?}" - ); - } - - Ok(()) - } - - #[tokio::test] - async fn test_repartition_with_coalescing() -> Result<()> { - let schema = test_schema(false); - // create 50 batches, each having 8 rows - let partition = create_vec_batches(50); - let partitions = vec![partition.clone(), partition.clone()]; - let partitioning = Partitioning::RoundRobinBatch(1); - - let session_config = SessionConfig::new().with_batch_size(200); - let task_ctx = TaskContext::default().with_session_config(session_config); - let task_ctx = Arc::new(task_ctx); - - // create physical plan - let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; - let exec = RepartitionExec::try_new(exec, partitioning)?; - - for i in 0..exec.partitioning().partition_count() { - let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; - while let Some(result) = stream.next().await { - let batch = result?; - assert_eq!(200, batch.num_rows()); - } - } - Ok(()) - } - - fn id_region_payload_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("id", DataType::UInt32, false), - Field::new("region", DataType::Utf8, false), - Field::new("payload", DataType::UInt32, false), - ])) - } - - fn u32_fields_schema(names: &[&str]) -> SchemaRef { - Arc::new(Schema::new( - names - .iter() - .map(|name| Field::new(*name, DataType::UInt32, false)) - .collect::>(), - )) - } - - fn range_repartition_on_columns( + fn range_partitioning_on_columns( schema: &SchemaRef, key_columns: &[&str], split_points: Vec>, - ) -> Result> { + ) -> Result { let ordering = LexOrdering::new( key_columns .iter() @@ -2917,7 +2767,7 @@ mod tests { .collect::>>()?, ) .expect("range ordering must not be empty"); - let partitioning = Partitioning::Range(RangePartitioning::try_new( + Ok(Partitioning::Range(RangePartitioning::try_new( ordering, split_points .into_iter() @@ -2930,12 +2780,10 @@ mod tests { ) }) .collect(), - )?); - let input = Arc::new(EmptyExec::new(Arc::clone(schema))); - Ok(Arc::new(RepartitionExec::try_new(input, partitioning)?)) + )?)) } - fn projection_columns( + fn projection_on_columns( input: &Arc, names: &[&str], ) -> Result { @@ -2951,6 +2799,13 @@ mod tests { ProjectionExec::try_new(exprs, Arc::clone(input)) } + fn as_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + match partitioning { + Partitioning::Range(range) => range, + other => panic!("expected Range partitioning, got {other:?}"), + } + } + /// Test source that claims Exact support for any sort pushdown request. #[derive(Debug, Clone)] struct ExactSortPushdownExec { @@ -2960,15 +2815,14 @@ mod tests { impl ExactSortPushdownExec { fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { use crate::execution_plan::{Boundedness, EmissionType}; - let eq_properties = - EquivalenceProperties::new_with_orderings(schema, [ordering]); - let cache = Arc::new(PlanProperties::new( - eq_properties, - Partitioning::UnknownPartitioning(num_partitions), - EmissionType::Incremental, - Boundedness::Bounded, - )); - Self { cache } + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, [ordering]), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } } } @@ -3016,6 +2870,32 @@ mod tests { } } + #[tokio::test] + async fn test_repartition_with_coalescing() -> Result<()> { + let schema = test_schema(false); + // create 50 batches, each having 8 rows + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone()]; + let partitioning = Partitioning::RoundRobinBatch(1); + + let session_config = SessionConfig::new().with_batch_size(200); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + assert_eq!(200, batch.num_rows()); + } + } + Ok(()) + } + fn test_schema(nullable: bool) -> Arc { Arc::new(Schema::new(vec![Field::new( "c0", From 549c50bab98547d8d527cfdd1eeb88cda5635930 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Fri, 17 Jul 2026 09:58:06 -0400 Subject: [PATCH 8/9] test fix --- datafusion/physical-plan/src/repartition/mod.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index fa2396422ce0a..44a7724c6d900 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2616,8 +2616,11 @@ mod tests { #[test] fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + // Three columns so the projection both narrows the schema (required for + // swap) and moves the range key from @0 to @1. let schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), Field::new("payload", DataType::UInt32, false), ])); let repartition = Arc::new(RepartitionExec::try_new( @@ -2625,7 +2628,6 @@ mod tests { range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, )?); - // Reorder so the range key moves from @0 to @1. let projection = projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; From b654808f1d1a249ad817ca5d56a71ea57f42e5fd Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Fri, 17 Jul 2026 11:19:57 -0400 Subject: [PATCH 9/9] tests addressed --- .../physical-plan/src/repartition/mod.rs | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 44a7724c6d900..8cb74fdc3855f 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1608,7 +1608,7 @@ impl ExecutionPlan for RepartitionExec { Partitioning::Hash(new_partitions, *size) } Partitioning::Range(range_partitioning) => { - // Rewriting the range key expressions in ordering based on the projection expressions + // Rewrite range key expressions through the projection. let mut sort_exprs = Vec::with_capacity(range_partitioning.ordering().len()); for sort_expr in range_partitioning.ordering() { @@ -2639,10 +2639,11 @@ mod tests { .expect("top node should be RepartitionExec"); assert!(swapped_repartition.input().is::()); + let range = expect_range_partitioning(swapped_repartition.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); assert_eq!( - as_range_partitioning(swapped_repartition.partitioning()).ordering()[0] - .to_string(), - "id@1 ASC" + range.split_points(), + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] ); Ok(()) @@ -2713,10 +2714,20 @@ mod tests { ); assert!(repartition.maintains_input_order()[0]); - assert!(matches!( - repartition.try_pushdown_sort(ordering.as_ref())?, - SortOrderPushdownResult::Exact { .. } - )); + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::() + .expect("pushdown should keep RepartitionExec"); + + assert!(pushed.preserve_order()); + assert!(pushed.maintains_input_order()[0]); + + let range = expect_range_partitioning(pushed.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + } + other => panic!("expected Exact sort pushdown, got {other:?}"), + } Ok(()) } @@ -2757,7 +2768,7 @@ mod tests { key_columns: &[&str], split_points: Vec>, ) -> Result { - let ordering = LexOrdering::new( + let Some(ordering) = LexOrdering::new( key_columns .iter() .map(|name| { @@ -2767,8 +2778,9 @@ mod tests { )) }) .collect::>>()?, - ) - .expect("range ordering must not be empty"); + ) else { + return exec_err!("range ordering must not be empty"); + }; Ok(Partitioning::Range(RangePartitioning::try_new( ordering, split_points @@ -2801,7 +2813,7 @@ mod tests { ProjectionExec::try_new(exprs, Arc::clone(input)) } - fn as_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { match partitioning { Partitioning::Range(range) => range, other => panic!("expected Range partitioning, got {other:?}"),