Refactor: Update enforce_sorting tests to use insta snapshots for easier updates #17900

alamb · 2025-10-03T14:22:04Z

Which issue does this PR close?

part of Migrate core tests to insta #15791
replaces some parts of this PR: Migrate core test to insta part 3 #16978

Rationale for this change

While working on the arrow 57 upgrade,

[WIP] Upgrade to arrow/parquet 57.0.0 #17888

I had to change quite few tests in enforce_sorting and the fact they are
not yet ported to insta means it was very painful to do.

What changes are included in this PR?

rewrite enforce_sorting so it uses insta for plan snapshots rather than hard coded vec<&str>

Also avoid the use of a macro for easier maintenance (and hopefully faster compile time)

Are these changes tested?

Yes, by existing tests

Are there any user-facing changes?

alamb · 2025-10-03T14:27:08Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

    let join = hash_join_exec(left_input, right_input, on, None, &JoinType::Inner)?;
    let physical_plan = sort_exec([sort_expr("a", &join.schema())].into(), join);

+    let test = EnforceSortingTest {


You can see that the plans did not change by turning on whitespace blind diff: https://github.com/apache/datafusion/pull/17900/files?w=1

alamb · 2025-10-03T14:27:41Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

+        let expected_optimized: Vec<&str> =optimized_plan_string.trim().lines().collect();
+
+        // return a string with both input and optimized plan
+        format!(


This somewhat crazy formatting is designed to make it easier to review the diff on github

…ier updates

alamb · 2025-10-03T16:18:28Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

+
+    /// Runs the enforce sorting test and returns a string with the input and
+    /// optimized plan as strings for snapshot comparison using insta
+    fn run(&self) -> String {


this mimics / replicates the logic of assert_optimized_plan!

I was not able to remove that macro yet because there are several parameterized plans that use it that I have to migrate somehow separately

blaginin · 2025-10-04T17:56:21Z

Since I was also working on that, I'll push a few commits to this branch to speed up the merge. Feel free to drop them if they don't fit

blaginin · 2025-10-04T17:56:56Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

    // 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:


if those two are the same, i think there's no need to duplicate the plans - so i assert just one

Also i removed expect_no_change, because it was duplicated and actually didn't catch some cases (see 0fd01c7)

this is very clever -- thank you

blaginin · 2025-10-04T17:58:28Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

+    Input Plan:
+    SortExec: expr=[a@2 ASC], preserve_partitioning=[false]
+      HashJoinExec: mode=Partitioned, join_type=Inner, on=[(col_a@0, c@2)]
+        DataSourceExec: partitions=1, partition_sizes=[0]
+        DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet


this is definitely just a personal preference but I find merged strings easier to follow than array representation with [] and "".

Before it would make sense, because it'd allow you to copy and paste the plan - but now this is managed by insta so no need use debug representation in my opinion

I agree -- I was using the debug representation to try and make review easier (to see that the plans hadn't changed)

🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Split the rstest parameterized test into two separate functions: - test_with_lost_ordering_unbounded - test_with_lost_ordering_bounded Each tests both repartition_sorts modes (false and true) using insta snapshots. Replaced the incomplete existing test_with_lost_ordering_bounded which only tested one mode. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

blaginin · 2025-10-04T19:46:02Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -2525,1213 +2696,1255 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> {
        "avg".to_string(),
        function_arg_unordered,
    );
-    struct TestCase<'a> {


for this test I asked claude to write a script that confirms before and after are the same:
https://gist.github.com/blaginin/606a2b6d9ba48627d0bdf5b9aa6d1c57

i manually changed something in "after" and it have caught it. also i randomly verified some cases

blaginin · 2025-10-04T20:19:55Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

+        .with_expected_description(
+            "// Should adjust the requirement in the third input of the union so\n\
+        // that it is not unnecessarily fine.",


Feel free to ignore, but I think it's not the best idea to mix comments and snapshots - it'll now be repeated twice in the code, making it messier, and it won't be clear for newcomers which one to update

I personally don’t see an issue with just adding the comment after the assert. I did that separately in: blaginin@ecdf6c3 in case you’ll like it🙂

I like it! Applied

blaginin

thank you!! wrote some comments, and you may need to check the changes i pushed here 😅

…tion)

…_insta

alamb

Thank you @blaginin -- I think this test is now much easier to deal with ❤️

alamb · 2025-10-06T11:11:26Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

    // 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:


this is very clever -- thank you

alamb · 2025-10-06T11:11:46Z

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

+        .with_expected_description(
+            "// Should adjust the requirement in the third input of the union so\n\
+        // that it is not unnecessarily fine.",


I like it! Applied

alamb · 2025-10-06T11:20:03Z

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


alamb · 2025-10-06T12:23:55Z

Thank you very much for helping to push this along @blaginin 🙏

github-actions bot added the core Core DataFusion crate label Oct 3, 2025

alamb commented Oct 3, 2025

View reviewed changes

alamb force-pushed the alamb/enforce_sorting_insta branch from 9311077 to ecfc0f8 Compare October 3, 2025 14:59

Refactor: Update enforce_sorting tests to use insta snapshots for eas…

1cac670

…ier updates

alamb force-pushed the alamb/enforce_sorting_insta branch from ecfc0f8 to 1cac670 Compare October 3, 2025 15:17

alamb changed the title ~~Refactor: Update enforce_sorting tests to use insta snapshots for easier updates~~ Refactor: Update enforce_sorting tests to use insta snapshots for easier updates (part 1) Oct 3, 2025

alamb marked this pull request as ready for review October 3, 2025 15:26

alamb mentioned this pull request Oct 3, 2025

Migrate core test to insta part 3 #16978

Draft

blaginin self-requested a review October 3, 2025 15:36

alamb added 2 commits October 3, 2025 11:44

port some more

d3a30c6

fmt

c0a6fb1

alamb commented Oct 3, 2025

View reviewed changes

Cleanup representation

8e33dd4

blaginin reviewed Oct 4, 2025

View reviewed changes

blaginin and others added 5 commits October 4, 2025 19:18

migrate test_sort_merge_join_order_by_left

9739758

migrate test_sort_merge_join_order_by_right

93a62a9

🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Start migrating test_window_partial_constant_and_set_monotonicity

214ce0d

Migrate the rest test_window_partial_constant_and_set_monotonicity

209afc0

blaginin reviewed Oct 4, 2025

View reviewed changes

Clippy

6d6e8b7

blaginin reviewed Oct 4, 2025

View reviewed changes

blaginin approved these changes Oct 4, 2025

View reviewed changes

blaginin changed the title ~~Refactor: Update enforce_sorting tests to use insta snapshots for easier updates (part 1)~~ Refactor: Update enforce_sorting tests to use insta snapshots for easier updates Oct 4, 2025

blaginin added 2 commits October 4, 2025 21:26

Kill expect_no_change (we now have Input / Optimized Plan unifica…

0fd01c7

…tion)

Fmt

d1e337d

blaginin mentioned this pull request Oct 4, 2025

Refactor: Split test_window_partial_constant_and_set_monotonicity from enforce_sorting #17903

Closed

Merge more plans with the same before and output

4ac1e56

alamb and others added 2 commits October 6, 2025 07:12

Merge remote-tracking branch 'apache/main' into alamb/enforce_sorting…

dc22509

…_insta

Kill expected_description

d4fcfbd

alamb commented Oct 6, 2025

View reviewed changes

alamb added this pull request to the merge queue Oct 6, 2025

Merged via the queue into apache:main with commit e03ba61 Oct 6, 2025
28 checks passed

Refactor: Update enforce_sorting tests to use insta snapshots for easier updates #17900

Refactor: Update enforce_sorting tests to use insta snapshots for easier updates #17900

Uh oh!

Conversation

alamb commented Oct 3, 2025

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blaginin commented Oct 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blaginin Oct 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blaginin Oct 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blaginin Oct 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

blaginin left a comment

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Oct 6, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

blaginin Oct 4, 2025 •

edited

Loading

blaginin Oct 4, 2025 •

edited

Loading

blaginin Oct 4, 2025 •

edited

Loading