gene-bordegaray commented on code in PR #23617:
URL: https://github.com/apache/datafusion/pull/23617#discussion_r3604636233
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2604,6 +2614,276 @@ 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())? {
Review Comment:
non-blocking: maybe one last assert comparing the sort of the pushed
repartiton and required 👍
--
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]
