alamb commented on code in PR #17900:
URL: https://github.com/apache/datafusion/pull/17900#discussion_r2405744240
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -291,15 +434,19 @@ async fn test_union_inputs_different_sorted() ->
Result<()> {
let physical_plan = sort_preserving_merge_exec(ordering, union);
// one input to the union is already sorted, one is not.
- let expected_input = [
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " UnionExec",
- " DataSourceExec: file_groups={1 group: [[x]]},
projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0
ASC, non_nullable_col@1 ASC], file_type=parquet",
- " SortExec: expr=[nullable_col@0 ASC],
preserve_partitioning=[false]",
- " DataSourceExec: file_groups={1 group: [[x]]},
projection=[nullable_col, non_nullable_col], file_type=parquet",
- ];
+ let test = EnforceSortingTest::new(physical_plan)
+ .with_repartition_sorts(true)
+ .with_expected_description("// should not add a sort at the output of
the union, input plan should not be changed")
+ .with_expect_no_change(true);
+ assert_snapshot!(test.run(), @r"
// should not add a sort at the output of the union, input plan should not
be changed
- assert_optimized!(expected_input, expected_input, physical_plan, true);
+ Input / Optimized Plan:
Review Comment:
this is very clever -- thank you
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -2480,6 +2547,73 @@ async fn
test_window_partial_constant_and_set_monotonicity() -> Result<()> {
.into();
let source = parquet_exec_with_sort(input_schema.clone(), vec![ordering])
as _;
+ // Macro for testing window function optimization with snapshots
Review Comment:
👍
##########
datafusion/core/tests/physical_optimizer/enforce_sorting.rs:
##########
@@ -505,29 +571,35 @@ async fn test_union_inputs_different_sorted6() ->
Result<()> {
// At the same time, this ordering requirement is unnecessarily fine.
// The final plan should be valid AND the ordering of the third child
// shouldn't be finer than necessary.
- let expected_input = [
- "SortPreservingMergeExec: [nullable_col@0 ASC]",
- " UnionExec",
- " SortExec: expr=[nullable_col@0 ASC],
preserve_partitioning=[false]",
- " DataSourceExec: file_groups={1 group: [[x]]},
projection=[nullable_col, non_nullable_col], file_type=parquet",
- " DataSourceExec: file_groups={1 group: [[x]]},
projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0
ASC], file_type=parquet",
- " SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1
ASC]",
- " RepartitionExec: partitioning=RoundRobinBatch(10),
input_partitions=1",
- " DataSourceExec: file_groups={1 group: [[x]]},
projection=[nullable_col, non_nullable_col], file_type=parquet",
- ];
+ let test = EnforceSortingTest::new(physical_plan)
+ .with_repartition_sorts(true)
+ .with_expected_description(
+ "// Should adjust the requirement in the third input of the union
so\n\
+ // that it is not unnecessarily fine.",
Review Comment:
I like it! Applied
--
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]
