Avoid adding unnecessary exchanges under Window nodes

presto-main/src/main/java/io/prestosql/sql/planner/optimizations/AddExchanges.java

@@ -300,7 +300,8 @@ public PlanWithProperties visitWindow(WindowNode node, PreferredProperties prefe
                     PreferredProperties.partitionedWithLocal(ImmutableSet.copyOf(node.getPartitionBy()), desiredProperties)
                             .mergeWithParent(preferredProperties));
 
-            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy())) {
+            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy()) &&
+                    !child.getProperties().isNodePartitionedOn(node.getPartitionBy())) {
                 if (node.getPartitionBy().isEmpty()) {
                     child = withDerivedProperties(
                             gatheringExchange(idAllocator.getNextId(), REMOTE, child.getNode()),
@@ -337,7 +338,8 @@ public PlanWithProperties visitRowNumber(RowNumberNode node, PreferredProperties
                             .mergeWithParent(preferredProperties));
 
             // TODO: add config option/session property to force parallel plan if child is unpartitioned and window has a PARTITION BY clause
-            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy())) {
+            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy())
+                    && !child.getProperties().isNodePartitionedOn(node.getPartitionBy())) {
                 child = withDerivedProperties(
                         partitionedExchange(
                                 idAllocator.getNextId(),
@@ -370,7 +372,8 @@ public PlanWithProperties visitTopNRowNumber(TopNRowNumberNode node, PreferredPr
             }
 
             PlanWithProperties child = planChild(node, preferredChildProperties);
-            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy())) {
+            if (!child.getProperties().isStreamPartitionedOn(node.getPartitionBy())
+                    && !child.getProperties().isNodePartitionedOn(node.getPartitionBy())) {
                 // add exchange + push function to child
                 child = withDerivedProperties(
                         new TopNRowNumberNode(

presto-main/src/test/java/io/prestosql/sql/planner/assertions/PlanMatchPattern.java

@@ -264,11 +264,25 @@ public static ExpectedValueProvider<WindowNode.Frame> windowFrame(
                 endValue.map(SymbolAlias::new));
     }
 
-    public static PlanMatchPattern window(Consumer<WindowMatcher.Builder> windowMatcherBuilderConsumer, PlanMatchPattern source)
+    public static PlanMatchPattern window(Consumer<WindowMatcher.Builder> handler, PlanMatchPattern source)
     {
-        WindowMatcher.Builder windowMatcherBuilder = new WindowMatcher.Builder(source);
-        windowMatcherBuilderConsumer.accept(windowMatcherBuilder);
-        return windowMatcherBuilder.build();
+        WindowMatcher.Builder builder = new WindowMatcher.Builder(source);
+        handler.accept(builder);
+        return builder.build();
+    }
+
+    public static PlanMatchPattern rowNumber(Consumer<RowNumberMatcher.Builder> handler, PlanMatchPattern source)
+    {
+        RowNumberMatcher.Builder builder = new RowNumberMatcher.Builder(source);
+        handler.accept(builder);
+        return builder.build();
+    }
+
+    public static PlanMatchPattern topNRowNumber(Consumer<TopNRowNumberMatcher.Builder> handler, PlanMatchPattern source)
+    {
+        TopNRowNumberMatcher.Builder builder = new TopNRowNumberMatcher.Builder(source);
+        handler.accept(builder);
+        return builder.build();
     }
 
     public static PlanMatchPattern sort(PlanMatchPattern source)

presto-main/src/test/java/io/prestosql/sql/planner/assertions/RowNumberMatcher.java

@@ -0,0 +1,166 @@
+/*
+ * 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 io.prestosql.sql.planner.assertions;
+
+import io.prestosql.Session;
+import io.prestosql.cost.StatsProvider;
+import io.prestosql.metadata.Metadata;
+import io.prestosql.sql.planner.Symbol;
+import io.prestosql.sql.planner.plan.PlanNode;
+import io.prestosql.sql.planner.plan.RowNumberNode;
+
+import java.util.List;
+import java.util.Optional;
+
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static io.prestosql.sql.planner.assertions.MatchResult.NO_MATCH;
+import static io.prestosql.sql.planner.assertions.MatchResult.match;
+import static io.prestosql.sql.planner.assertions.PlanMatchPattern.node;
+import static java.util.Objects.requireNonNull;
+
+public class RowNumberMatcher
+        implements Matcher
+{
+    private final Optional<List<SymbolAlias>> partitionBy;
+    private final Optional<Optional<Integer>> maxRowCountPerPartition;
+    private final Optional<SymbolAlias> rowNumberSymbol;
+    private final Optional<Optional<SymbolAlias>> hashSymbol;
+
+    private RowNumberMatcher(
+            Optional<List<SymbolAlias>> partitionBy,
+            Optional<Optional<Integer>> maxRowCountPerPartition,
+            Optional<SymbolAlias> rowNumberSymbol,
+            Optional<Optional<SymbolAlias>> hashSymbol)
+    {
+        this.partitionBy = requireNonNull(partitionBy, "partitionBy is null");
+        this.maxRowCountPerPartition = requireNonNull(maxRowCountPerPartition, "maxRowCountPerPartition is null");
+        this.rowNumberSymbol = requireNonNull(rowNumberSymbol, "rowNumberSymbol is null");
+        this.hashSymbol = requireNonNull(hashSymbol, "hashSymbol is null");
+    }
+
+    @Override
+    public boolean shapeMatches(PlanNode node)
+    {
+        return node instanceof RowNumberNode;
+    }
+
+    @Override
+    public MatchResult detailMatches(PlanNode node, StatsProvider stats, Session session, Metadata metadata, SymbolAliases symbolAliases)
+    {
+        checkState(shapeMatches(node), "Plan testing framework error: shapeMatches returned false in detailMatches in %s", this.getClass().getName());
+
+        RowNumberNode rowNumberNode = (RowNumberNode) node;
+
+        if (partitionBy.isPresent()) {
+            List<Symbol> expected = partitionBy.get().stream()
+                    .map(alias -> alias.toSymbol(symbolAliases))
+                    .collect(toImmutableList());
+
+            if (!expected.equals(rowNumberNode.getPartitionBy())) {
+                return NO_MATCH;
+            }
+        }
+
+        if (rowNumberSymbol.isPresent()) {
+            Symbol expected = rowNumberSymbol.get().toSymbol(symbolAliases);
+            if (!expected.equals(rowNumberNode.getRowNumberSymbol())) {
+                return NO_MATCH;
+            }
+        }
+
+        if (maxRowCountPerPartition.isPresent()) {
+            if (!maxRowCountPerPartition.get().equals(rowNumberNode.getMaxRowCountPerPartition())) {
+                return NO_MATCH;
+            }
+        }
+
+        if (hashSymbol.isPresent()) {
+            Optional<Symbol> expected = hashSymbol.get().map(alias -> alias.toSymbol(symbolAliases));
+            if (!expected.equals(rowNumberNode.getHashSymbol())) {
+                return NO_MATCH;
+            }
+        }
+
+        return match();
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("partitionBy", partitionBy)
+                .add("maxRowCountPerPartition", maxRowCountPerPartition)
+                .add("rowNumberSymbol", rowNumberSymbol)
+                .add("hashSymbol", hashSymbol)
+                .toString();
+    }
+
+    /**
+     * By default, matches any RowNumberNode.  Users add additional constraints by
+     * calling the various member functions of the Builder, typically named according
+     * to the field names of RowNumberNode.
+     */
+    public static class Builder
+    {
+        private final PlanMatchPattern source;
+        private Optional<List<SymbolAlias>> partitionBy = Optional.empty();
+        private Optional<Optional<Integer>> maxRowCountPerPartition = Optional.empty();
+        private Optional<SymbolAlias> rowNumberSymbol = Optional.empty();
+        private Optional<Optional<SymbolAlias>> hashSymbol = Optional.empty();
+
+        Builder(PlanMatchPattern source)
+        {
+            this.source = requireNonNull(source, "source is null");
+        }
+
+        public Builder partitionBy(List<String> partitionBy)
+        {
+            requireNonNull(partitionBy, "partitionBy is null");
+            this.partitionBy = Optional.of(partitionBy.stream()
+                    .map(SymbolAlias::new)
+                    .collect(toImmutableList()));
+            return this;
+        }
+
+        public Builder maxRowCountPerPartition(Optional<Integer> maxRowCountPerPartition)
+        {
+            this.maxRowCountPerPartition = Optional.of(requireNonNull(maxRowCountPerPartition, "maxRowCountPerPartition is null"));
+            return this;
+        }
+
+        public Builder rowNumberSymbol(SymbolAlias rowNumberSymbol)
+        {
+            this.rowNumberSymbol = Optional.of(requireNonNull(rowNumberSymbol, "rowNumberSymbol is null"));
+            return this;
+        }
+
+        public Builder hashSymbol(Optional<SymbolAlias> hashSymbol)
+        {
+            this.hashSymbol = Optional.of(requireNonNull(hashSymbol, "hashSymbol is null"));
+            return this;
+        }
+
+        PlanMatchPattern build()
+        {
+            return node(RowNumberNode.class, source).with(
+                    new RowNumberMatcher(
+                            partitionBy,
+                            maxRowCountPerPartition,
+                            rowNumberSymbol,
+                            hashSymbol));
+        }
+    }
+}