diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 8da20d3d23d9..bfdfab3209c8 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1607,10 +1607,29 @@ 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) => { + // 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 Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; + + Partitioning::Range(RangePartitioning::try_new( + ordering, + range_partitioning.split_points().to_vec(), + )?) } others => others.clone(), }; @@ -1648,16 +1667,6 @@ impl ExecutionPlan for RepartitionExec { if !self.maintains_input_order()[0] { 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::Hash(_, _) - | Partitioning::UnknownPartitioning(_) => {} - } // Delegate to the child and wrap with a new RepartitionExec self.input.try_pushdown_sort(order)?.try_map(|new_input| { @@ -1680,12 +1689,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), @@ -2144,6 +2152,8 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -2604,6 +2614,249 @@ mod tests { Ok(()) } + #[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]])?; + + // Reorder so the range key moves from @0 to @1. + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + + 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 = id_region_payload_schema(); + let repartition = range_repartition_on_columns(&schema, &["id"], vec![vec![10]])?; + + // Drop the range key column. + let projection = + projection_columns(&(Arc::clone(&repartition) as _), &["region", "payload"])?; + + 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 = 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 = + projection_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; + + 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(()) + } + + #[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); @@ -2630,6 +2883,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", @@ -3796,6 +4182,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)])) }