wiedld commented on code in PR #15010:
URL: https://github.com/apache/datafusion/pull/15010#discussion_r1980470202
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -371,46 +371,95 @@ macro_rules! plans_matches_expected {
}
}
+fn test_suite_default_config_options() -> ConfigOptions {
+ let mut config = ConfigOptions::new();
+
+ // 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;
+
+ // Use a small batch size, to trigger RoundRobin in tests
+ config.execution.batch_size = 1;
+
+ config
+}
+
+/// How the optimizers are run.
+#[derive(PartialEq, Clone)]
+enum DoFirst {
+ /// Runs: (EnforceDistribution, EnforceDistribution, EnforceSorting)
+ Distribution,
+ /// Runs: (EnforceSorting, EnforceDistribution, EnforceDistribution)
+ Sorting,
+}
+
+#[derive(Clone)]
+struct TestConfig {
+ config: ConfigOptions,
+ optimizers_to_run: DoFirst,
+}
+
+impl TestConfig {
+ fn new(optimizers_to_run: DoFirst) -> Self {
+ Self {
+ config: test_suite_default_config_options(),
+ optimizers_to_run,
+ }
+ }
+
+ /// If preferred, will not repartition / resort data if it is already
sorted.
+ fn with_prefer_existing_sort(mut self) -> Self {
+ self.config.optimizer.prefer_existing_sort = true;
+ self
+ }
+
+ /// If preferred, will not attempt to convert Union to Interleave.
+ fn with_prefer_existing_union(mut self) -> Self {
+ self.config.optimizer.prefer_existing_union = true;
+ self
+ }
+
+ /// If preferred, will repartition file scans.
+ /// Accepts a minimum file size to repartition.
+ fn with_prefer_repartition_file_scans(mut self, file_min_size: usize) ->
Self {
+ self.config.optimizer.repartition_file_scans = true;
+ self.config.optimizer.repartition_file_min_size = file_min_size;
+ self
+ }
+
+ /// Set the preferred target partitions for query execution concurrency.
+ fn with_query_execution_partitions(mut self, target_partitions: usize) ->
Self {
+ self.config.execution.target_partitions = target_partitions;
+ self
+ }
+}
+
/// Runs the repartition optimizer and asserts the plan against the expected
/// Arguments
/// * `EXPECTED_LINES` - Expected output plan
/// * `PLAN` - Input plan
-/// * `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
+/// * `CONFIG` - [`TestConfig`]
macro_rules! assert_optimized {
- ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr) => {
- assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, false,
10, false, 1024, false);
- };
-
- ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr,
$PREFER_EXISTING_SORT: expr) => {
- assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST,
$PREFER_EXISTING_SORT, 10, false, 1024, false);
- };
-
- ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr,
$PREFER_EXISTING_SORT: expr, $PREFER_EXISTING_UNION: expr) => {
- assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST,
$PREFER_EXISTING_SORT, 10, false, 1024, $PREFER_EXISTING_UNION);
- };
-
- ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr,
$PREFER_EXISTING_SORT: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS:
expr, $REPARTITION_FILE_MIN_SIZE: expr) => {
- assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST,
$PREFER_EXISTING_SORT, $TARGET_PARTITIONS, $REPARTITION_FILE_SCANS,
$REPARTITION_FILE_MIN_SIZE, false);
+ ($EXPECTED_LINES: expr, $PLAN: expr) => {
+ assert_optimized!($EXPECTED_LINES, $PLAN,
&TestConfig::new(DoFirst::Distribution));
Review Comment:
Note: this option is technically not used, so I could remove it. It's
basically a 1-to-1 replacement of the "default options" provided for this macro
on main.
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
