Insta for `enforce_distrubution` #18185

blaginin · 2025-10-20T20:44:02Z

Part of #15791 .

Surpasses part of #16978

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

- smj_join_key_ordering - reorder_join_keys_to_left_input (inline snapshots with filter) - parallelization_ignores_limit - parallelization_prior_to_sort_preserving_merge - parallelization_sort_preserving_merge_with_union - parallelization_does_not_benefit Replace old-style expected string arrays with insta inline snapshot assertions. For reorder_join_keys_to_left_input, use regex filter to replace all join_type values with '...' and add separate assertion for top join type. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated the following tests to use inline insta snapshots: - repartition_sorted_limit - repartition_sorted_limit_with_filter - repartition_ignores_limit - repartition_ignores_union - repartition_through_sort_preserving_merge - repartition_ignores_sort_preserving_merge - repartition_ignores_sort_preserving_merge_with_union - repartition_does_not_destroy_sort - repartition_does_not_destroy_sort_more_complex - repartition_ignores_transitively_with_projection - repartition_transitively_past_sort_with_projection All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated the following tests to use inline insta snapshots: - remove_redundant_roundrobins - remove_unnecessary_spm_after_filter - preserve_ordering_through_repartition - no_need_for_sort_after_filter - do_not_preserve_ordering_through_repartition3 - do_not_add_unnecessary_hash - do_not_add_unnecessary_hash2 - optimize_away_unnecessary_repartition All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated the following tests to use inline insta snapshots: - optimize_away_unnecessary_repartition2 - parallelization_ignores_transitively_with_projection_csv All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated the following tests to use inline insta snapshots: - parallelization_single_partition - parallelization_multiple_files All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

…tputs Migrated the following tests to use inline insta snapshots: - repartition_transitively_with_projection - repartition_transitively_past_sort_with_filter - repartition_transitively_past_sort_with_projection_and_filter These tests have different outputs for DISTRIB_DISTRIB_SORT vs SORT_DISTRIB_DISTRIB, so each gets two separate inline snapshots. All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated test_distribute_sort_parquet to use inline insta snapshots. This test uses different optimizer configurations: - Initial plan before optimization - After Run::Distribution - After Run::Distribution + Run::Sorting All three states now use inline snapshots instead of string arrays. All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated test_distribute_sort_memtable to use inline insta snapshot. This test checks the final optimized plan after creating and executing a SQL query on a MemTable. All tests pass without --accept flag. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

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

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

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

This test uses a loop to test different compression types with conditional logic for expected outputs. Used insta::allow_duplicates! to handle inline snapshots within the loop. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

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

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

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

Add hide_first() helper function that replaces the first occurrence of a regex pattern in a plan, returning the captured value and modified string. Use this to test join key reordering across multiple join types while keeping snapshots readable. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Converted the complex multi_smj_joins test to use insta snapshot testing with filters. This test loops over 6 join types and tests multiple configurations with different optimizer runs. Key changes: - Wrapped loop with insta::allow_duplicates! - Added Settings filter to replace join_type={actual_value} with join_type=... - Converted all test_config.run() calls to run2() + assert_plan!() with inline snapshots - Preserved all existing comments explaining test behavior - Removed unused join_plan helper functions and dynamic string variables 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Converted the reorder_join_keys_to_right_input test to use insta snapshot testing with the hide_first helper function. This test loops over 8 join types and verifies that join key ordering is adjusted based on children inputs. Key changes: - Wrapped loop with insta::allow_duplicates! - Used hide_first() to replace only the first join_type occurrence with "..." - Converted assert_plan_txt! to insta::assert_snapshot! with inline snapshot - Preserved comment explaining the test behavior - Removed dynamic top_join_plan string variable 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrated the last two tests using assert_plan_txt! macro to insta snapshots and removed the macro entirely. Tests migrated: - do_not_put_sort_when_input_is_invalid - put_sort_when_input_is_valid Key changes: - Converted assert_plan_txt! calls to assert_plan! with inline snapshots - Preserved all comments explaining expected behavior - Removed assert_plan_txt! macro definition - Removed unused get_plan_string import The assert_plan_txt! macro is now completely eliminated from the codebase. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

blaginin · 2025-10-21T11:02:58Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

+        let mut settings = Settings::clone_current();
+        settings.add_filter(&format!("join_type={join_type}"), "join_type=...");
+
+        #[rustfmt::skip]


Screen.Recording.2025-10-21.at.11.25.18.mov

alamb · 2025-10-22T17:28:01Z

So close...

blaginin · 2025-10-22T18:58:39Z

This happened to be a bigger change than I expected, there are some really big and complicated tests, and if we just compare with a separate snapshot on every case, we'll add +1K more lines (like in the original PR). I'm doing something more complicated from the code point of view, but on the positive side the file size is the same.

I also think we may want to split some tests / and the file itself, but this may be the one to do on top of this PR

alamb · 2025-10-22T22:08:45Z

I think a sequence of PRs sounds like a good plan to me

blaginin · 2025-10-23T18:47:59Z

Sure, moved all easy cases to #18248

blaginin and others added 30 commits October 16, 2025 10:58

add assert_plan

096b453

Move multi_hash_joins

2592835

Migrate multi_joins_after_alias test to insta snapshots

b5f6160

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrate multi_joins_after_multi_alias test to insta snapshots

406f590

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrate join_after_agg_alias test to insta snapshots

99fc2b0

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrate hash_join_key_ordering test to insta snapshots

d114351

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Migrate multi_hash_join_key_ordering test to insta snapshots

5704e73

Replace old-style expected string arrays with insta snapshot assertions. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>

Merge branch 'main' into move-enforce-distrivution

a6c6f2a

Migrate more tests

80b964f

More

9fd6982

Migrate optimize_away_unnecessary_repartition2 to insta snapshots

a672a00

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

Migrate parallelization_ignores_transitively_with_projection_csv to i…

54a2b40

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

Migrate remove_redundant_roundrobins to insta snapshots

a9783e3

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

Migrate optimize_away_unnecessary_repartition to insta snapshots

9b73d35

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

Kill plans_matches_expected

6e9dab7

Migrate merge_does_not_need_sort to insta snapshots

6cec480

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

Migrate parallelization_two_partitions to insta snapshots

8439cbb

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

Migrate parallelization_two_partitions_into_four to insta snapshots

c80a32d

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

Migrate parallelization_sorted_limit to insta snapshots

d6da5ef

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

Migrate parallelization_limit_with_filter to insta snapshots

ceef90f

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

Migrate parallelization_ignores_transitively_with_projection_parquet …

2a2b86c

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

blaginin and others added 9 commits October 20, 2025 20:23

Migrate parallelization_union_inputs to insta snapshots

199a84d

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

Migrate smj_join_key_ordering to insta snapshots

83e1e52

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

kill old run

5f9ce0a

fmt

c02b9ed

clippy

18d9254

blaginin changed the title ~~Move enforce distrivution~~ Insta for enforce_distrubution Oct 20, 2025

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

blaginin added 2 commits October 21, 2025 09:56

typo and fmt

dbae19d

Fmt

2d67ef8

blaginin commented Oct 21, 2025

View reviewed changes

blaginin added a commit to blaginin/datafusion that referenced this pull request Oct 23, 2025

enforce_distribution.rs is done in apache#18185

37ed67d

blaginin mentioned this pull request Oct 23, 2025

Insta for enforce_distrubution (easy ones) #18248

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Insta for `enforce_distrubution` #18185

Insta for `enforce_distrubution` #18185

Uh oh!

blaginin commented Oct 20, 2025

Uh oh!

blaginin Oct 21, 2025

Uh oh!

alamb commented Oct 22, 2025

Uh oh!

blaginin commented Oct 22, 2025 •

edited

Loading

Uh oh!

alamb commented Oct 22, 2025

Uh oh!

blaginin commented Oct 23, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Insta for enforce_distrubution #18185

Are you sure you want to change the base?

Insta for enforce_distrubution #18185

Uh oh!

Conversation

blaginin commented Oct 20, 2025

Uh oh!

blaginin Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

alamb commented Oct 22, 2025

Uh oh!

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

Uh oh!

Uh oh!

alamb commented Oct 22, 2025

Uh oh!

blaginin commented Oct 23, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Insta for `enforce_distrubution` #18185

Insta for `enforce_distrubution` #18185

blaginin commented Oct 22, 2025 •

edited

Loading