Refactor test suite in EnforceDistribution, to use standard test config. #15010

wiedld · 2025-03-05T00:23:10Z

Which issue does this PR close?

Rationale for this change

Make it easier to determine the differences in test configuration, for each test case.
Also clearly defines what are the defaults used for the test suite.

What changes are included in this PR?

Define TestConfig, which requires one-of DoFirst optimizer run configs. Most test cases uses this test config with the default settings.

Provide a set of interfaces to define which configurations get tweaked.

Are these changes tested?

Yes

Are there any user-facing changes?

No

wiedld · 2025-03-05T00:24:56Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

+    // By default, will not repartition / resort data if it is already sorted.
+    config.optimizer.prefer_existing_sort = false;
+
+    // By default, will attempt to convert Union to Interleave.
+    config.optimizer.prefer_existing_union = false;
+
+    // By default, will not repartition file scans.
+    config.optimizer.repartition_file_scans = false;
+    config.optimizer.repartition_file_min_size = 1024;
+
+    // By default, set query execution concurrency to 10.
+    config.execution.target_partitions = 10;


These match the defaults on main:

datafusion/datafusion/core/tests/physical_optimizer/enforce_distribution.rs

Line 388 in c61e7e5

assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, false, 10, false, 1024, false);

wiedld · 2025-03-05T00:25:40Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

+    // Use a small batch size, to trigger RoundRobin in tests
+    config.execution.batch_size = 1;


This matches main:

datafusion/datafusion/core/tests/physical_optimizer/enforce_distribution.rs

Lines 412 to 413 in c61e7e5

// Use a small batch size, to trigger RoundRobin in tests

config.execution.batch_size = 1;

…est config builder

wiedld · 2025-03-05T00:31:35Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

-/// * `FIRST_ENFORCE_DIST` -
-///     true: (EnforceDistribution, EnforceDistribution,  EnforceSorting)
-///     false: else runs (EnforceSorting, EnforceDistribution, EnforceDistribution)
-/// * `PREFER_EXISTING_SORT` (optional) - if true, will not repartition / resort data if it is already sorted
-/// * `TARGET_PARTITIONS` (optional) - number of partitions to repartition to
-/// * `REPARTITION_FILE_SCANS` (optional) - if true, will repartition file scans
-/// * `REPARTITION_FILE_MIN_SIZE` (optional) - minimum file size to repartition
-/// * `PREFER_EXISTING_UNION` (optional) - if true, will not attempt to convert Union to Interleave


These are now replaced by the required DoFirst (for one-of 2 possible optimizer run patterns),
then all the optional settings.

I very much like how it is much more explicit

alamb

Thank you very much @wiedld -- I think this is a major step forward in test readability as all the relevant options are now explicit in the tests, rather than being true / false parameters where interpreting them requires looking at a macro definition.

cc @berkaysynnada and @mustafasrepo / @akurmustafa

alamb · 2025-03-05T11:26:56Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

-/// * `FIRST_ENFORCE_DIST` -
-///     true: (EnforceDistribution, EnforceDistribution,  EnforceSorting)
-///     false: else runs (EnforceSorting, EnforceDistribution, EnforceDistribution)
-/// * `PREFER_EXISTING_SORT` (optional) - if true, will not repartition / resort data if it is already sorted
-/// * `TARGET_PARTITIONS` (optional) - number of partitions to repartition to
-/// * `REPARTITION_FILE_SCANS` (optional) - if true, will repartition file scans
-/// * `REPARTITION_FILE_MIN_SIZE` (optional) - minimum file size to repartition
-/// * `PREFER_EXISTING_UNION` (optional) - if true, will not attempt to convert Union to Interleave


I very much like how it is much more explicit

alamb · 2025-03-05T11:30:35Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

+                assert_optimized!(
+                    expected,
+                    top_join.clone(),
+                    &TestConfig::new(DoFirst::Distribution)
+                );


Taking this a step more, what do you think of using a function on TestConfig so this looks like:

Suggested change

assert_optimized!(

expected,

top_join.clone(),

&TestConfig::new(DoFirst::Distribution)

);

TestConfig::new(DoFirst::Distribution)

.assert_optimized)

&expected,

&top_join)

);

I like this idea. I was already planning for the next followup PR to change the macro to a functional call.

Here is the draft of the followup PR: influxdata#61

Just to get an idea of where this is heading.

berkaysynnada · 2025-03-05T21:38:03Z

I'll review this tomorrow as well

alamb · 2025-03-06T00:46:06Z

BTW @blaginin is proposing to add insta in this PR:

Add insta / snapshot testing to CLI & set up AWS mock #13672

If we merge that some of these tests may become much eaiser to maintain / change 🤔

alamb · 2025-03-06T17:50:04Z

Let's plan to merge this PR tomorrow / after @berkaysynnada has had a chance to review

berkaysynnada

I really like the idea. Interpreting these tests becomes much clearer. Thank you @wiedld. I just have two questions, but perhaps I'm missing something as I'm a bit tired today 😅

berkaysynnada · 2025-03-06T22:32:16Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

-        true,
-        repartition_size,
-        false
+        &test_config.clone().with_query_execution_partitions(3)


Are we missing prefer_existing_sort and repartition_file_scans here?

I put the common config a few lines up here: https://github.com/apache/datafusion/pull/15010/files/5de55b3a8bf1858206dcecc0d3b1907228c970a4#r1984443646. Then only modified the unique config (with_query_execution_partitions) per test case.

berkaysynnada · 2025-03-06T22:32:31Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

-        true,
-        repartition_size,
-        false
+        &test_config.with_query_execution_partitions(8)


wiedld · 2025-03-07T04:49:36Z

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

+    let test_config = TestConfig::new(DoFirst::Distribution)
+        .with_prefer_existing_sort()
+        .with_prefer_repartition_file_scans(1);


Here's the common config.

wiedld · 2025-03-07T04:52:09Z

I just have two questions, but perhaps I'm missing something as I'm a bit tired today 😅

That's my everyday. 😆
Thanks for taking the time to review.

alamb · 2025-03-07T13:57:44Z

Thank you @berkaysynnada and @wiedld

…est config builder (apache#15010)

github-actions bot added the core Core DataFusion crate label Mar 5, 2025

wiedld commented Mar 5, 2025

View reviewed changes

refactor(15003): refactor test suite in EnforceDistribution, to use t…

82b0342

…est config builder

wiedld force-pushed the 15003/switch-to-builder branch from 9af23e0 to 82b0342 Compare March 5, 2025 00:28

wiedld commented Mar 5, 2025

View reviewed changes

wiedld marked this pull request as ready for review March 5, 2025 00:55

wiedld mentioned this pull request Mar 5, 2025

Improve EnforceDistribution testings. #15003

Open

9 tasks

alamb mentioned this pull request Mar 5, 2025

Weekly Plan (Andrew Lamb) March 3, 2025 #14978

Open

12 tasks

alamb approved these changes Mar 5, 2025

View reviewed changes

Merge branch 'main' into 15003/switch-to-builder

5de55b3

wiedld mentioned this pull request Mar 5, 2025

Refactor EnforceDistribution test cases to demonstrate dependencies across optimizer runs. influxdata/arrow-datafusion#61

Closed

berkaysynnada approved these changes Mar 6, 2025

View reviewed changes

wiedld commented Mar 7, 2025

View reviewed changes

alamb merged commit 297af95 into apache:main Mar 7, 2025
24 checks passed

alamb mentioned this pull request Mar 7, 2025

Refactor EnforceDistribution test cases to demonstrate dependencies across optimizer runs. #15074

Open

danila-b pushed a commit to danila-b/datafusion that referenced this pull request Mar 8, 2025

refactor(15003): refactor test suite in EnforceDistribution, to use t…

ece7b4a

…est config builder (apache#15010)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor test suite in EnforceDistribution, to use standard test config. #15010

Refactor test suite in EnforceDistribution, to use standard test config. #15010

wiedld commented Mar 5, 2025

wiedld Mar 5, 2025

wiedld Mar 5, 2025

wiedld Mar 5, 2025

alamb Mar 5, 2025

alamb left a comment

alamb Mar 5, 2025

alamb Mar 5, 2025

wiedld Mar 5, 2025

wiedld Mar 5, 2025

berkaysynnada commented Mar 5, 2025

alamb commented Mar 6, 2025

alamb commented Mar 6, 2025 •

edited

Loading

berkaysynnada left a comment

berkaysynnada Mar 6, 2025

wiedld Mar 7, 2025

berkaysynnada Mar 6, 2025

wiedld Mar 7, 2025

wiedld commented Mar 7, 2025

alamb commented Mar 7, 2025

		// Use a small batch size, to trigger RoundRobin in tests
		config.execution.batch_size = 1;

Refactor test suite in EnforceDistribution, to use standard test config. #15010

Refactor test suite in EnforceDistribution, to use standard test config. #15010

Conversation

wiedld commented Mar 5, 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?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

berkaysynnada commented Mar 5, 2025

alamb commented Mar 6, 2025

alamb commented Mar 6, 2025 • edited Loading

berkaysynnada left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wiedld commented Mar 7, 2025

alamb commented Mar 7, 2025

alamb commented Mar 6, 2025 •

edited

Loading