gene-bordegaray commented on code in PR #23617:
URL: https://github.com/apache/datafusion/pull/23617#discussion_r3597935486
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1561,10 +1561,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::new(
Review Comment:
can we use try_new
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -1603,12 +1622,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(_)
Review Comment:
can we remove this no-op match now since all are exhausted?
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2558,6 +2574,214 @@ mod tests {
Ok(())
}
+ #[test]
+ fn range_repartition_swaps_with_projection_keeping_key() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
Review Comment:
there is a lot of repeated boiler plate with schema set up and constructin.
Can we extract / try to reuse some helpers to reduce this?
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2558,6 +2574,214 @@ mod tests {
Ok(())
}
+ #[test]
+ fn range_repartition_swaps_with_projection_keeping_key() -> Result<()> {
Review Comment:
I noticed no test for the sort pushdown, can we add one 👍
should have a sort maintaining range repartition then and assert and make
sure result maintains the ordering
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
