This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new b5b1810c89 feat: complete range repartition physical planning (#23617)
b5b1810c89 is described below
commit b5b1810c89308de3f9d4d37b769a70645b6b1825
Author: Saad Tajwar <[email protected]>
AuthorDate: Tue Jul 21 05:26:52 2026 -0400
feat: complete range repartition physical planning (#23617)
## Which issue does this PR close?
- Closes #23230
## Rationale for this change
After #23231 was merged in for supporting physical execution of the
range repartitioning scheme, we still had a few methods on physical
planning unimplemented, specifically `try_swapping_with_projection`,
`try_pushdown_sort`, `repartitioned` - this PR finishes the
implementation of those methods
## What changes are included in this PR?
- `try_swapping_with_projection`: similar to the `Hash` scheme, for
`Range` we call `update_expr` for each of the range key expressions to
attempt rewriting based on the projection expressions
- `try_pushdown_sort`: same as other variants, we delegate to the child
and wrap with a new `RepartitionExec`
- `repartitioned`: unable to support for Range, left comment in codebase
with explanation
## Are these changes tested?
Yes
## Are there any user-facing changes?
No
---
datafusion/physical-plan/src/repartition/mod.rs | 346 ++++++++++++++++++++++--
1 file changed, 329 insertions(+), 17 deletions(-)
diff --git a/datafusion/physical-plan/src/repartition/mod.rs
b/datafusion/physical-plan/src/repartition/mod.rs
index 8da20d3d23..12229c26b7 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) => {
+ // 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() {
+ 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,281 @@ mod tests {
Ok(())
}
+ #[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(
+ Arc::new(EmptyExec::new(Arc::clone(&schema))),
+ range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
+ )?);
+
+ let projection =
+ projection_on_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::<RepartitionExec>()
+ .expect("top node should be RepartitionExec");
+
+ assert!(swapped_repartition.input().is::<ProjectionExec>());
+ let range =
expect_range_partitioning(swapped_repartition.partitioning());
+ assert_eq!(range.ordering()[0].to_string(), "id@1 ASC");
+ assert_eq!(
+ range.split_points(),
+ &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])]
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn range_repartition_does_not_swap_when_projection_drops_key() ->
Result<()> {
+ // 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_on_columns(&(Arc::clone(&repartition) as _),
&["payload"])?;
+ assert!(
+ repartition
+ .try_swapping_with_projection(&projection)?
+ .is_none()
+ );
+
+ // 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_on_columns(&(Arc::clone(&repartition) as _), &["a",
"c"])?;
+ assert!(
+ repartition
+ .try_swapping_with_projection(&projection)?
+ .is_none()
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ 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 source with preserve_order: Range maintains input
order.
+ let source = Arc::new(ExactSortPushdownExec::new(
+ Arc::clone(&schema),
+ 2,
+ ordering.clone(),
+ ));
+ let repartition = Arc::new(
+ 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::<RepartitionExec>()
+ .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");
+ assert_eq!(
+ inner.properties().output_ordering().map(|o|
o.to_string()),
+ Some(ordering.to_string()),
+ "pushed repartition output ordering should match the
requested sort"
+ );
+ }
+ other => panic!("expected Exact sort pushdown, got {other:?}"),
+ }
+
+ Ok(())
+ }
+
+ #[test]
+ fn
range_repartition_try_pushdown_sort_unsupported_without_order_maintenance()
+ -> 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 source without preserve_order: Range does not
maintain order.
+ let source = Arc::new(ExactSortPushdownExec::new(
+ Arc::clone(&schema),
+ 2,
+ ordering.clone(),
+ ));
+ let repartition = Arc::new(RepartitionExec::try_new(
+ source,
+ range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?,
+ )?);
+ assert!(!repartition.maintains_input_order()[0]);
+
+ assert!(matches!(
+ repartition.try_pushdown_sort(ordering.as_ref())?,
+ SortOrderPushdownResult::Unsupported
+ ));
+
+ Ok(())
+ }
+
+ fn range_partitioning_on_columns(
+ schema: &SchemaRef,
+ key_columns: &[&str],
+ split_points: Vec<Vec<u32>>,
+ ) -> Result<Partitioning> {
+ let Some(ordering) = LexOrdering::new(
+ key_columns
+ .iter()
+ .map(|name| {
+ Ok(PhysicalSortExpr::new(
+ col(name, schema)?,
+ SortOptions::default(),
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?,
+ ) else {
+ return exec_err!("range ordering must not be empty");
+ };
+ Ok(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(),
+ )?))
+ }
+
+ fn projection_on_columns(
+ input: &Arc<dyn ExecutionPlan>,
+ names: &[&str],
+ ) -> Result<ProjectionExec> {
+ let exprs = names
+ .iter()
+ .map(|name| {
+ Ok(ProjectionExpr {
+ expr: col(name, &input.schema())?,
+ alias: (*name).to_string(),
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ ProjectionExec::try_new(exprs, Arc::clone(input))
+ }
+
+ fn expect_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 {
+ cache: Arc<PlanProperties>,
+ }
+
+ impl ExactSortPushdownExec {
+ fn new(schema: SchemaRef, num_partitions: usize, ordering:
LexOrdering) -> Self {
+ use crate::execution_plan::{Boundedness, EmissionType};
+ Self {
+ cache: Arc::new(PlanProperties::new(
+ EquivalenceProperties::new_with_orderings(schema,
[ordering]),
+ Partitioning::UnknownPartitioning(num_partitions),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ )),
+ }
+ }
+ }
+
+ 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<PlanProperties> {
+ &self.cache
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ Ok(Box::pin(EmptyRecordBatchStream::new(self.schema())))
+ }
+
+ fn try_pushdown_sort(
+ &self,
+ _order: &[PhysicalSortExpr],
+ ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> {
+ Ok(SortOrderPushdownResult::Exact {
+ inner: Arc::new(self.clone()),
+ })
+ }
+ }
+
#[tokio::test]
async fn test_repartition_with_coalescing() -> Result<()> {
let schema = test_schema(false);
@@ -3796,6 +4081,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<Schema> {
Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)]))
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
