Remove prefer partitioning rules

Author: gaurav8297

core/trino-main/src/main/java/io/trino/SystemSessionProperties.java

@@ -80,7 +80,6 @@ public final class SystemSessionProperties
     public static final String QUERY_MAX_STAGE_COUNT = "query_max_stage_count";
     public static final String REDISTRIBUTE_WRITES = "redistribute_writes";
     public static final String USE_PREFERRED_WRITE_PARTITIONING = "use_preferred_write_partitioning";
-    public static final String PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS = "preferred_write_partitioning_min_number_of_partitions";
     public static final String SCALE_WRITERS = "scale_writers";
     public static final String TASK_SCALE_WRITERS_ENABLED = "task_scale_writers_enabled";
     public static final String MAX_WRITER_TASKS_COUNT = "max_writer_tasks_count";
@@ -306,16 +305,6 @@ public SystemSessionProperties(
                         "Use preferred write partitioning",
                         optimizerConfig.isUsePreferredWritePartitioning(),
                         false),
-                integerProperty(
-                        PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS,
-                        "Use preferred write partitioning when the number of written partitions exceeds the configured threshold",
-                        optimizerConfig.getPreferredWritePartitioningMinNumberOfPartitions(),
-                        value -> {
-                            if (value < 1) {
-                                throw new TrinoException(INVALID_SESSION_PROPERTY, format("%s must be greater than or equal to 1: %s", PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS, value));
-                            }
-                        },
-                        false),
                 booleanProperty(
                         SCALE_WRITERS,
                         "Scale out writers based on throughput (use minimum necessary)",
@@ -1093,11 +1082,6 @@ public static boolean isUsePreferredWritePartitioning(Session session)
         return session.getSystemProperty(USE_PREFERRED_WRITE_PARTITIONING, Boolean.class);
     }
 
-    public static int getPreferredWritePartitioningMinNumberOfPartitions(Session session)
-    {
-        return session.getSystemProperty(PREFERRED_WRITE_PARTITIONING_MIN_NUMBER_OF_PARTITIONS, Integer.class);
-    }
-
     public static boolean isScaleWriters(Session session)
     {
         return session.getSystemProperty(SCALE_WRITERS, Boolean.class);

core/trino-main/src/main/java/io/trino/sql/planner/LocalExecutionPlanner.java

@@ -3285,7 +3285,7 @@ public PhysicalOperation visitRefreshMaterializedView(RefreshMaterializedViewNod
         public PhysicalOperation visitTableWriter(TableWriterNode node, LocalExecutionPlanContext context)
         {
             // Set table writer count
-            int maxWriterCount = getWriterCount(session, node.getPartitioningScheme(), node.getPreferredPartitioningScheme(), node.getSource());
+            int maxWriterCount = getWriterCount(session, node.getPartitioningScheme(), node.getSource());
             context.setDriverInstanceCount(maxWriterCount);
             context.taskContext.setMaxWriterCount(maxWriterCount);
 
@@ -3443,7 +3443,7 @@ public PhysicalOperation visitSimpleTableExecuteNode(SimpleTableExecuteNode node
         public PhysicalOperation visitTableExecute(TableExecuteNode node, LocalExecutionPlanContext context)
         {
             // Set table writer count
-            int maxWriterCount = getWriterCount(session, node.getPartitioningScheme(), node.getPreferredPartitioningScheme(), node.getSource());
+            int maxWriterCount = getWriterCount(session, node.getPartitioningScheme(), node.getSource());
             context.setDriverInstanceCount(maxWriterCount);
             context.taskContext.setMaxWriterCount(maxWriterCount);
 
@@ -3470,7 +3470,7 @@ public PhysicalOperation visitTableExecute(TableExecuteNode node, LocalExecution
             return new PhysicalOperation(operatorFactory, outputMapping.buildOrThrow(), context, source);
         }
 
-        private int getWriterCount(Session session, Optional<PartitioningScheme> partitioningScheme, Optional<PartitioningScheme> preferredPartitioningScheme, PlanNode source)
+        private int getWriterCount(Session session, Optional<PartitioningScheme> partitioningScheme, PlanNode source)
         {
             // This check is required because we don't know which writer count to use when exchange is
             // single distribution. It could be possible that when scaling is enabled, a single distribution is
@@ -3480,18 +3480,12 @@ private int getWriterCount(Session session, Optional<PartitioningScheme> partiti
                 return 1;
             }
 
-            if (isLocalScaledWriterExchange(source)) {
-                return partitioningScheme.or(() -> preferredPartitioningScheme)
-                        // The default value of partitioned writer count is 32 which is high enough to use it
-                        // for both cases when scaling is enabled or not. Additionally, it doesn't lead to too many
-                        // small files since when scaling is disabled only single writer will handle a single partition.
-                        .map(scheme -> getTaskPartitionedWriterCount(session))
-                        .orElseGet(() -> getTaskScaleWritersMaxWriterCount(session));
-            }
-
+            // The default value of partitioned writer count is 32 which is high enough to use it
+            // for both cases when scaling is enabled or not. Additionally, it doesn't lead to too many
+            // small files since when scaling is disabled only single writer will handle a single partition.
             return partitioningScheme
                     .map(scheme -> getTaskPartitionedWriterCount(session))
-                    .orElseGet(() -> getTaskWriterCount(session));
+                    .orElseGet(() -> isLocalScaledWriterExchange(source) ? getTaskScaleWritersMaxWriterCount(session) : getTaskWriterCount(session));
         }
 
         private boolean isSingleGatheringExchange(PlanNode node)

core/trino-main/src/main/java/io/trino/sql/planner/LogicalPlanner.java

@@ -130,6 +130,7 @@
 import static io.trino.SystemSessionProperties.getMaxWriterTaskCount;
 import static io.trino.SystemSessionProperties.getRetryPolicy;
 import static io.trino.SystemSessionProperties.isCollectPlanStatisticsForAllQueries;
+import static io.trino.SystemSessionProperties.isUsePreferredWritePartitioning;
 import static io.trino.metadata.MetadataUtil.createQualifiedObjectName;
 import static io.trino.spi.StandardErrorCode.CATALOG_NOT_FOUND;
 import static io.trino.spi.StandardErrorCode.CONSTRAINT_VIOLATION;
@@ -660,7 +661,6 @@ private RelationPlan createTableWriterPlan(
             TableStatisticsMetadata statisticsMetadata)
     {
         Optional<PartitioningScheme> partitioningScheme = Optional.empty();
-        Optional<PartitioningScheme> preferredPartitioningScheme = Optional.empty();
 
         int maxWriterTasks = target.getMaxWriterTasks(plannerContext.getMetadata(), session).orElse(getMaxWriterTaskCount(session));
         Optional<Integer> maxWritersNodesCount = getRetryPolicy(session) != RetryPolicy.TASK
@@ -683,9 +683,9 @@ private RelationPlan createTableWriterPlan(
                         Partitioning.create(partitioningHandle.get(), partitionFunctionArguments),
                         outputLayout));
             }
-            else {
+            else if (isUsePreferredWritePartitioning(session)) {
                 // empty connector partitioning handle means evenly partitioning on partitioning columns
-                preferredPartitioningScheme = Optional.of(new PartitioningScheme(
+                partitioningScheme = Optional.of(new PartitioningScheme(
                         Partitioning.create(FIXED_HASH_DISTRIBUTION, partitionFunctionArguments),
                         outputLayout,
                         Optional.empty(),
@@ -721,7 +721,6 @@ private RelationPlan createTableWriterPlan(
                             symbols,
                             columnNames,
                             partitioningScheme,
-                            preferredPartitioningScheme,
                             Optional.of(partialAggregation),
                             Optional.of(result.getDescriptor().map(aggregations.getMappings()::get))),
                     target,
@@ -743,7 +742,6 @@ private RelationPlan createTableWriterPlan(
                         symbols,
                         columnNames,
                         partitioningScheme,
-                        preferredPartitioningScheme,
                         Optional.empty(),
                         Optional.empty()),
                 target,
@@ -956,7 +954,6 @@ private RelationPlan createTableExecutePlan(Analysis analysis, TableExecute stat
 
         // todo extract common method to be used here and in createTableWriterPlan()
         Optional<PartitioningScheme> partitioningScheme = Optional.empty();
-        Optional<PartitioningScheme> preferredPartitioningScheme = Optional.empty();
         if (layout.isPresent()) {
             List<Symbol> partitionFunctionArguments = new ArrayList<>();
             layout.get().getPartitionColumns().stream()
@@ -973,13 +970,13 @@ private RelationPlan createTableExecutePlan(Analysis analysis, TableExecute stat
                         Partitioning.create(partitioningHandle.get(), partitionFunctionArguments),
                         outputLayout));
             }
-            else {
+            else if (isUsePreferredWritePartitioning(session)) {
                 // empty connector partitioning handle means evenly partitioning on partitioning columns
                 int maxWriterTasks = tableExecuteTarget.getMaxWriterTasks(plannerContext.getMetadata(), session).orElse(getMaxWriterTaskCount(session));
                 Optional<Integer> maxWritersNodesCount = getRetryPolicy(session) != RetryPolicy.TASK
                         ? Optional.of(Math.min(maxWriterTasks, getMaxWriterTaskCount(session)))
                         : Optional.empty();
-                preferredPartitioningScheme = Optional.of(new PartitioningScheme(
+                partitioningScheme = Optional.of(new PartitioningScheme(
                         Partitioning.create(FIXED_HASH_DISTRIBUTION, partitionFunctionArguments),
                         outputLayout,
                         Optional.empty(),
@@ -1000,8 +997,7 @@ private RelationPlan createTableExecutePlan(Analysis analysis, TableExecute stat
                         symbolAllocator.newSymbol("fragment", VARBINARY),
                         symbols,
                         columnNames,
-                        partitioningScheme,
-                        preferredPartitioningScheme),
+                        partitioningScheme),
                 tableExecuteTarget,
                 symbolAllocator.newSymbol("rows", BIGINT),
                 Optional.empty(),

core/trino-main/src/main/java/io/trino/sql/planner/OptimizerConfig.java

@@ -30,7 +30,7 @@
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MINUTES;
 
-@DefunctConfig("adaptive-partial-aggregation.min-rows")
+@DefunctConfig({"adaptive-partial-aggregation.min-rows", "preferred-write-partitioning-min-number-of-partitions"})
 public class OptimizerConfig
 {
     private double cpuCostWeight = 75;
@@ -57,7 +57,6 @@ public class OptimizerConfig
     private boolean distributedSort = true;
 
     private boolean usePreferredWritePartitioning = true;
-    private int preferredWritePartitioningMinNumberOfPartitions = 50;
 
     private Duration iterativeOptimizerTimeout = new Duration(3, MINUTES); // by default let optimizer wait a long time in case it retrieves some data from ConnectorMetadata
 
@@ -373,20 +372,6 @@ public OptimizerConfig setUsePreferredWritePartitioning(boolean usePreferredWrit
         return this;
     }
 
-    @Min(1)
-    public int getPreferredWritePartitioningMinNumberOfPartitions()
-    {
-        return preferredWritePartitioningMinNumberOfPartitions;
-    }
-
-    @Config("preferred-write-partitioning-min-number-of-partitions")
-    @ConfigDescription("Use preferred write partitioning when the number of written partitions exceeds the configured threshold")
-    public OptimizerConfig setPreferredWritePartitioningMinNumberOfPartitions(int preferredWritePartitioningMinNumberOfPartitions)
-    {
-        this.preferredWritePartitioningMinNumberOfPartitions = preferredWritePartitioningMinNumberOfPartitions;
-        return this;
-    }
-
     public Duration getIterativeOptimizerTimeout()
     {
         return iterativeOptimizerTimeout;

core/trino-main/src/main/java/io/trino/sql/planner/PlanOptimizers.java

@@ -34,8 +34,6 @@
 import io.trino.sql.planner.iterative.rule.AddDynamicFilterSource;
 import io.trino.sql.planner.iterative.rule.AddExchangesBelowPartialAggregationOverGroupIdRuleSet;
 import io.trino.sql.planner.iterative.rule.AddIntermediateAggregations;
-import io.trino.sql.planner.iterative.rule.ApplyPreferredTableExecutePartitioning;
-import io.trino.sql.planner.iterative.rule.ApplyPreferredTableWriterPartitioning;
 import io.trino.sql.planner.iterative.rule.ApplyTableScanRedirection;
 import io.trino.sql.planner.iterative.rule.ArraySortAfterArrayDistinct;
 import io.trino.sql.planner.iterative.rule.CanonicalizeExpressions;
@@ -757,16 +755,6 @@ public PlanOptimizers(
                         statsCalculator,
                         costCalculator,
                         ImmutableSet.of(new RemoveRedundantIdentityProjections())),
-                // Prefer write partitioning rule requires accurate stats.
-                // Run it before reorder joins which also depends on accurate stats.
-                new IterativeOptimizer(
-                        plannerContext,
-                        ruleStats,
-                        statsCalculator,
-                        costCalculator,
-                        ImmutableSet.of(
-                                new ApplyPreferredTableWriterPartitioning(),
-                                new ApplyPreferredTableExecutePartitioning())),
                 // Because ReorderJoins runs only once,
                 // PredicatePushDown, columnPruningOptimizer and RemoveRedundantIdentityProjections
                 // need to run beforehand in order to produce an optimal join order