diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index a5bf68a63c38..dafde268ba73 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1782,16 +1782,12 @@ mod test { /// `$PLAN`: the plan to optimized /// macro_rules! assert_plan { - ($EXPECTED_PLAN_LINES: expr, $PLAN: expr) => { + ($PLAN: expr, @ $EXPECTED: expr) => { let formatted = crate::displayable($PLAN).indent(true).to_string(); - let actual: Vec<&str> = formatted.trim().lines().collect(); - let expected_plan_lines: Vec<&str> = $EXPECTED_PLAN_LINES - .iter().map(|s| *s).collect(); - - assert_eq!( - expected_plan_lines, actual, - "\n**Original Plan Mismatch\n\nexpected:\n\n{expected_plan_lines:#?}\nactual:\n\n{actual:#?}\n\n" + insta::assert_snapshot!( + formatted, + @$EXPECTED ); }; } @@ -1808,13 +1804,12 @@ mod test { .with_preserve_order(); // Repartition should preserve order - let expected_plan = [ - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c0@0 ASC", - " UnionExec", - " DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC", - " DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC", - ]; - assert_plan!(expected_plan, &exec); + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c0@0 ASC + UnionExec + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); Ok(()) } @@ -1824,16 +1819,15 @@ mod test { let sort_exprs = sort_exprs(&schema); let source = sorted_memory_exec(&schema, sort_exprs); // output is sorted, but has only a single partition, so no need to sort - let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10)) - .unwrap() + let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10))? .with_preserve_order(); // Repartition should not preserve order - let expected_plan = [ - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - " DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC", - ]; - assert_plan!(expected_plan, &exec); + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); + Ok(()) } @@ -1848,13 +1842,12 @@ mod test { .with_preserve_order(); // Repartition should not preserve order, as there is no order to preserve - let expected_plan = [ - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - " UnionExec", - " DataSourceExec: partitions=1, partition_sizes=[0]", - " DataSourceExec: partitions=1, partition_sizes=[0]", - ]; - assert_plan!(expected_plan, &exec); + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 + UnionExec + DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] + "); Ok(()) } @@ -1869,11 +1862,10 @@ mod test { .unwrap(); // Repartition should not preserve order - let expected_plan = [ - "RepartitionExec: partitioning=RoundRobinBatch(20), input_partitions=1", - " DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC", - ]; - assert_plan!(expected_plan, exec.as_ref()); + assert_plan!(exec.as_ref(), @r" + RepartitionExec: partitioning=RoundRobinBatch(20), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); Ok(()) }