Add partition attributes to sort node

unix280 · Dec 2, 2024 · 6cbc7c0 · 6cbc7c0
1 parent 4f159c0
commit 6cbc7c0
Show file tree

Hide file tree

Showing 18 changed files with 310 additions and 21 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/CanonicalPlanGenerator.java
@@ -751,7 +751,8 @@ public Optional<PlanNode> visitSort(SortNode node, Context context)
                 planNodeidAllocator.getNextId(),
                 source.get(),
                 getCanonicalOrderingScheme(node.getOrderingScheme(), context.getExpressions()),
-                node.isPartial());
+                node.isPartial(),
+                node.getPartitionBy());
         context.addPlan(node, new CanonicalPlan(canonicalPlan, strategy));
         return Optional.of(canonicalPlan);
     }

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java
@@ -1260,7 +1260,7 @@ private PlanBuilder sort(PlanBuilder subPlan, Optional<OrderBy> orderBy, List<Ex
         OrderingScheme orderingScheme = toOrderingScheme(
                 orderByExpressions.stream().map(subPlan::translate).collect(toImmutableList()),
                 orderBy.get().getSortItems().stream().map(PlannerUtils::toSortOrder).collect(toImmutableList()));
-        planNode = new SortNode(getSourceLocation(orderBy.get()), idAllocator.getNextId(), subPlan.getRoot(), orderingScheme, false);
+        planNode = new SortNode(getSourceLocation(orderBy.get()), idAllocator.getNextId(), subPlan.getRoot(), orderingScheme, false, ImmutableList.of());
 
         return subPlan.withNewRoot(planNode);
     }

diff --git a/.../main/java/com/facebook/presto/sql/planner/iterative/rule/RemoveRedundantSortColumns.java b/.../main/java/com/facebook/presto/sql/planner/iterative/rule/RemoveRedundantSortColumns.java
@@ -50,6 +50,6 @@ public Result apply(SortNode node, Captures captures, Context context)
             return Result.empty();
         }
 
-        return Result.ofPlanNode(new SortNode(node.getSourceLocation(), node.getId(), node.getStatsEquivalentPlanNode(), node.getSource(), newOrderingScheme, node.isPartial()));
+        return Result.ofPlanNode(new SortNode(node.getSourceLocation(), node.getId(), node.getStatsEquivalentPlanNode(), node.getSource(), newOrderingScheme, node.isPartial(), node.getPartitionBy()));
     }
 }
diff --git a/...in/java/com/facebook/presto/sql/planner/iterative/rule/SimplifySortWithConstantInput.java b/...in/java/com/facebook/presto/sql/planner/iterative/rule/SimplifySortWithConstantInput.java
@@ -68,7 +68,7 @@ public Result apply(SortNode node, Captures captures, Context context)
                 return Result.ofPlanNode(projectNode);
             }
             OrderingScheme orderExcludeConstantVariable = new OrderingScheme(newOrderBy);
-            return Result.ofPlanNode(new SortNode(node.getSourceLocation(), context.getIdAllocator().getNextId(), projectNode, orderExcludeConstantVariable, node.isPartial()));
+            return Result.ofPlanNode(new SortNode(node.getSourceLocation(), context.getIdAllocator().getNextId(), projectNode, orderExcludeConstantVariable, node.isPartial(), node.getPartitionBy()));
         }
         return Result.empty();
     }

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
@@ -530,6 +530,44 @@ public PlanWithProperties visitTopN(TopNNode node, PreferredProperties preferred
 
         @Override
         public PlanWithProperties visitSort(SortNode node, PreferredProperties preferredProperties)
+        {
+            if (!node.getPartitionBy().isEmpty()) {
+                return planSortWithPartition(node, preferredProperties);
+            }
+            return planSortWithoutPartition(node, preferredProperties);
+        }
+
+        private PlanWithProperties planSortWithPartition(SortNode node, PreferredProperties preferredProperties)
+        {
+            List<LocalProperty<VariableReferenceExpression>> desiredProperties = new ArrayList<>();
+            checkArgument(!node.getPartitionBy().isEmpty());
+            desiredProperties.add(new GroupingProperty<>(node.getPartitionBy()));
+
+            node.getOrderingScheme().getOrderByVariables().stream()
+                    .map(variable -> new SortingProperty<>(variable, node.getOrderingScheme().getOrdering(variable)))
+                    .forEach(desiredProperties::add);
+
+            PlanWithProperties child = planChild(
+                    node,
+                    PreferredProperties.partitionedWithLocal(ImmutableSet.copyOf(node.getPartitionBy()), desiredProperties)
+                            .mergeWithParent(preferredProperties, !isExactPartitioningPreferred(session)));
+
+            if (!isStreamPartitionedOn(child.getProperties(), node.getPartitionBy()) &&
+                    !isNodePartitionedOn(child.getProperties(), node.getPartitionBy())) {
+                child = withDerivedProperties(
+                        partitionedExchange(
+                                idAllocator.getNextId(),
+                                selectExchangeScopeForPartitionedRemoteExchange(child.getNode(), false),
+                                child.getNode(),
+                                createPartitioning(node.getPartitionBy()),
+                                Optional.empty()),
+                        child.getProperties());
+            }
+
+            return rebaseAndDeriveProperties(node, child);
+        }
+
+        private PlanWithProperties planSortWithoutPartition(SortNode node, PreferredProperties preferredProperties)
         {
             PlanWithProperties child = planChild(node, PreferredProperties.undistributed());
 
@@ -561,7 +599,8 @@ public PlanWithProperties visitSort(SortNode node, PreferredProperties preferred
                                         idAllocator.getNextId(),
                                         source,
                                         node.getOrderingScheme(),
-                                        true),
+                                        true,
+                                        node.getPartitionBy()),
                                 node.getOrderingScheme()),
                         child.getProperties());
             }

diff --git a/...o-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java b/...o-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddLocalExchanges.java
@@ -191,6 +191,30 @@ public PlanWithProperties visitExplainAnalyze(ExplainAnalyzeNode node, StreamPre
         @Override
         public PlanWithProperties visitSort(SortNode node, StreamPreferredProperties parentPreferences)
         {
+            if (!node.getPartitionBy().isEmpty()) {
+                return planSortWithPartition(node, parentPreferences);
+            }
+            return planSortWithoutPartition(node, parentPreferences);
+        }
+
+        private PlanWithProperties planSortWithPartition(SortNode node, StreamPreferredProperties parentPreferences)
+        {
+            checkArgument(!node.getPartitionBy().isEmpty());
+            StreamPreferredProperties childRequirements = parentPreferences
+                    .constrainTo(node.getSource().getOutputVariables())
+                    .withDefaultParallelism(session)
+                    .withPartitioning(node.getPartitionBy());
+
+            PlanWithProperties child = planAndEnforce(node.getSource(), childRequirements, childRequirements);
+
+            SortNode result = new SortNode(node.getSourceLocation(), idAllocator.getNextId(), child.getNode(), node.getOrderingScheme(), node.isPartial(), node.getPartitionBy());
+
+            return deriveProperties(result, child.getProperties());
+        }
+
+        private PlanWithProperties planSortWithoutPartition(SortNode node, StreamPreferredProperties parentPreferences)
+        {
+            checkArgument(node.getPartitionBy().isEmpty());
             // Remove sort if the child is already sorted and in a single stream
             // TODO: extract to its own optimization after AddLocalExchanges once the
             // constraint optimization framework is in a better state to be extended

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/LimitPushDown.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/LimitPushDown.java
@@ -207,7 +207,7 @@ public PlanNode visitSort(SortNode node, RewriteContext<LimitContext> context)
             }
             else if (rewrittenSource != node.getSource()) {
                 planChanged = true;
-                return new SortNode(node.getSourceLocation(), node.getId(), rewrittenSource, node.getOrderingScheme(), node.isPartial());
+                return new SortNode(node.getSourceLocation(), node.getId(), rewrittenSource, node.getOrderingScheme(), node.isPartial(), node.getPartitionBy());
             }
             return node;
         }

diff --git a/...src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java b/...src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java
@@ -714,7 +714,7 @@ public PlanNode visitSort(SortNode node, RewriteContext<Set<VariableReferenceExp
 
             PlanNode source = context.rewrite(node.getSource(), expectedInputs);
 
-            return new SortNode(node.getSourceLocation(), node.getId(), node.getStatsEquivalentPlanNode(), source, node.getOrderingScheme(), node.isPartial());
+            return new SortNode(node.getSourceLocation(), node.getId(), node.getStatsEquivalentPlanNode(), source, node.getOrderingScheme(), node.isPartial(), node.getPartitionBy());
         }
 
         @Override

diff --git a/.../src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/.../src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
@@ -563,7 +563,7 @@ public PlanNode visitSort(SortNode node, RewriteContext<Void> context)
         {
             PlanNode source = context.rewrite(node.getSource());
 
-            return new SortNode(node.getSourceLocation(), node.getId(), source, canonicalizeAndDistinct(node.getOrderingScheme()), node.isPartial());
+            return new SortNode(node.getSourceLocation(), node.getId(), source, canonicalizeAndDistinct(node.getOrderingScheme()), node.isPartial(), node.getPartitionBy());
         }
 
         @Override

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java
@@ -1104,9 +1104,11 @@ public Void visitSort(SortNode node, Void context)
         {
             Iterable<String> keys = Iterables.transform(node.getOrderingScheme().getOrderByVariables(), input -> input + " " + node.getOrderingScheme().getOrdering(input));
 
-            addNode(node,
-                    format("%sSort", node.isPartial() ? "Partial" : ""),
-                    format("[%s]", Joiner.on(", ").join(keys)));
+            String detail = format("[%s]", Joiner.on(", ").join(keys));
+            if (!node.getPartitionBy().isEmpty()) {
+                detail = format("%s[Partition by %s]", detail, Joiner.on(", ").join(node.getPartitionBy()));
+            }
+            addNode(node, format("%sSort", node.isPartial() ? "Partial" : ""), detail);
 
             return processChildren(node, context);
         }

diff --git a/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java b/presto-main/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java
@@ -354,6 +354,8 @@ public class LocalQueryRunner
 
     private final ReadWriteLock lock = new ReentrantReadWriteLock();
 
+    private List<PlanOptimizer> additionalOptimizer = ImmutableList.of();
+
     public LocalQueryRunner(Session defaultSession)
     {
         this(defaultSession, new FeaturesConfig(), new FunctionsConfig(), new NodeSpillConfig(), false, false);
@@ -1097,12 +1099,21 @@ public Plan createPlan(Session session, @Language("SQL") String sql, Optimizer.P
         return createPlan(session, sql, getPlanOptimizers(forceSingleNode), stage, warningCollector);
     }
 
+    public void setAdditionalOptimizer(List<PlanOptimizer> additionalOptimizer)
+    {
+        this.additionalOptimizer = additionalOptimizer;
+    }
+
     public List<PlanOptimizer> getPlanOptimizers(boolean forceSingleNode)
     {
         FeaturesConfig featuresConfig = new FeaturesConfig()
                 .setDistributedIndexJoinsEnabled(false)
                 .setOptimizeHashGeneration(true);
-        return new PlanOptimizers(
+        ImmutableList.Builder<PlanOptimizer> planOptimizers = ImmutableList.builder();
+        if (!additionalOptimizer.isEmpty()) {
+            planOptimizers.addAll(additionalOptimizer);
+        }
+        planOptimizers.addAll(new PlanOptimizers(
                 metadata,
                 sqlParser,
                 forceSingleNode,
@@ -1116,7 +1127,8 @@ public List<PlanOptimizer> getPlanOptimizers(boolean forceSingleNode)
                 new CostComparator(featuresConfig),
                 taskCountEstimator,
                 partitioningProviderManager,
-                featuresConfig).getPlanningTimeOptimizers();
+                featuresConfig).getPlanningTimeOptimizers());
+        return planOptimizers.build();
     }
 
     public Plan createPlan(Session session, @Language("SQL") String sql, List<PlanOptimizer> optimizers, WarningCollector warningCollector)

diff --git a/...o-main/src/test/java/com/facebook/presto/sql/planner/TestEffectivePredicateExtractor.java b/...o-main/src/test/java/com/facebook/presto/sql/planner/TestEffectivePredicateExtractor.java
@@ -318,7 +318,8 @@ public void testSort()
                                 equals(BV, CV),
                                 lessThan(CV, bigintLiteral(10)))),
                 new OrderingScheme(ImmutableList.of(new Ordering(AV, SortOrder.ASC_NULLS_LAST))),
-                false);
+                false,
+                ImmutableList.of());
 
         RowExpression effectivePredicate = effectivePredicateExtractor.extract(node);
 

diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java
@@ -91,7 +91,7 @@ public ObjectMapper getObjectMapper()
         return objectMapper;
     }
 
-    private static ObjectMapper createObjectMapper()
+    protected static ObjectMapper createObjectMapper()
     {
         TestingTypeManager typeManager = new TestingTypeManager();
         TestingBlockEncodingSerde blockEncodingSerde = new TestingBlockEncodingSerde();

diff --git a/...o-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java b/...o-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java
@@ -246,7 +246,8 @@ public SortNode sort(List<VariableReferenceExpression> orderBy, PlanNode source)
                 idAllocator.getNextId(),
                 source,
                 new OrderingScheme(orderBy.stream().map(variable -> new Ordering(variable, SortOrder.ASC_NULLS_FIRST)).collect(toImmutableList())),
-                false);
+                false,
+                ImmutableList.of());
     }
 
     public OffsetNode offset(long rowCount, PlanNode source)

diff --git a/...c/test/java/com/facebook/presto/sql/planner/optimizations/TestAddPartitionToSortRule.java b/...c/test/java/com/facebook/presto/sql/planner/optimizations/TestAddPartitionToSortRule.java
@@ -0,0 +1,48 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sql.planner.optimizations;
+
+import com.facebook.presto.matching.Captures;
+import com.facebook.presto.matching.Pattern;
+import com.facebook.presto.spi.plan.SortNode;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
+import com.facebook.presto.sql.planner.iterative.Rule;
+import com.google.common.collect.ImmutableList;
+
+import java.util.Optional;
+
+import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.matching.Pattern.typeOf;
+import static com.google.common.base.Preconditions.checkState;
+
+public class TestAddPartitionToSortRule
+        implements Rule<SortNode>
+{
+    @Override
+    public Pattern<SortNode> getPattern()
+    {
+        return typeOf(SortNode.class);
+    }
+
+    @Override
+    public Result apply(SortNode node, Captures captures, Context context)
+    {
+        if (!node.getPartitionBy().isEmpty()) {
+            return Result.empty();
+        }
+        Optional<VariableReferenceExpression> partition = node.getSource().getOutputVariables().stream().filter(x -> x.getType().equals(BIGINT)).findFirst();
+        checkState(partition.isPresent());
+        return Result.ofPlanNode(new SortNode(node.getSourceLocation(), context.getIdAllocator().getNextId(), node.getSource(), node.getOrderingScheme(), node.isPartial(), ImmutableList.of(partition.get())));
+    }
+}
diff --git a/...test/java/com/facebook/presto/sql/planner/optimizations/TestSortWithinPartitionPlans.java b/...test/java/com/facebook/presto/sql/planner/optimizations/TestSortWithinPartitionPlans.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sql.planner.optimizations;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.sql.planner.RuleStatsRecorder;
+import com.facebook.presto.sql.planner.assertions.BasePlanTest;
+import com.facebook.presto.sql.planner.iterative.IterativeOptimizer;
+import com.facebook.presto.testing.LocalQueryRunner;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import org.testng.annotations.Test;
+
+import static com.facebook.presto.SystemSessionProperties.TASK_CONCURRENCY;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.anyTree;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.exchange;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.sort;
+import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.tableScan;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Scope.LOCAL;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Scope.REMOTE_STREAMING;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Type.GATHER;
+import static com.facebook.presto.sql.planner.plan.ExchangeNode.Type.REPARTITION;
+
+public class TestSortWithinPartitionPlans
+        extends BasePlanTest
+{
+    @Test
+    public void testSortWithPartition()
+    {
+        LocalQueryRunner localQueryRunner = getQueryRunner();
+        localQueryRunner.setAdditionalOptimizer(ImmutableList.of(new IterativeOptimizer(
+                localQueryRunner.getMetadata(),
+                new RuleStatsRecorder(),
+                localQueryRunner.getStatsCalculator(),
+                localQueryRunner.getEstimatedExchangesCostCalculator(),
+                ImmutableSet.of(new TestAddPartitionToSortRule()))));
+        Session session = Session.builder(this.getQueryRunner().getDefaultSession())
+                .setSystemProperty(TASK_CONCURRENCY, "2")
+                .build();
+
+        assertDistributedPlan("SELECT partkey, discount from lineitem order by discount",
+                session,
+                anyTree(
+                        exchange(REMOTE_STREAMING, GATHER, ImmutableList.of(),
+                                sort(
+                                        anyTree(
+                                                exchange(LOCAL, REPARTITION,
+                                                        exchange(REMOTE_STREAMING, REPARTITION,
+                                                                anyTree(
+                                                                        tableScan("lineitem")))))))));
+    }
+}