gene-bordegaray commented on code in PR #23617:
URL: https://github.com/apache/datafusion/pull/23617#discussion_r3603893451
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2604,6 +2614,264 @@ 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>());
+ 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<()> {
+ // 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]);
+
+ assert!(matches!(
Review Comment:
lets make this sronger:
```rust
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]);
et range = expect_range_partitioning(pushed.partitioning());
assert_eq!(range.ordering()[0].to_string(), "id@0 ASC");
}
other => panic!("expected Exact sort pushdown, got {other:?}"),
}
```
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2604,6 +2614,264 @@ 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>());
+ 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<()> {
+ // 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]);
+
+ assert!(matches!(
+ repartition.try_pushdown_sort(ordering.as_ref())?,
+ SortOrderPushdownResult::Exact { .. }
+ ));
+
+ 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 ordering = LexOrdering::new(
+ key_columns
+ .iter()
+ .map(|name| {
+ Ok(PhysicalSortExpr::new(
+ col(name, schema)?,
+ SortOptions::default(),
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?,
+ )
+ .expect("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 as_range_partitioning(partitioning: &Partitioning) ->
&RangePartitioning {
Review Comment:
I think this would be better as `expect_range_parittioning()` 👍
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2604,6 +2614,264 @@ 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>());
+ 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<()> {
+ // 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]);
+
+ assert!(matches!(
+ repartition.try_pushdown_sort(ordering.as_ref())?,
+ SortOrderPushdownResult::Exact { .. }
+ ));
+
+ 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 ordering = LexOrdering::new(
+ key_columns
+ .iter()
+ .map(|name| {
+ Ok(PhysicalSortExpr::new(
+ col(name, schema)?,
+ SortOptions::default(),
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?,
+ )
+ .expect("range ordering must not be empty");
Review Comment:
lets error here since we return a Result
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2604,6 +2614,264 @@ 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>());
+ assert_eq!(
Review Comment:
coudl we also asset split points 👍
##########
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
Review Comment:
nit:
```
// Rewrite range key expressions through the projection.
```
--
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]
