[feature](Nereids): push down topN through join (#24720)

Push TopN through Join. JoinType just can be left/right outer join or cross join, because data of their one child can't be filtered. new TopN is (original limit + original offset, 0) as limit and offset. (cherry picked from commit 3c9ff7a)
apache · xiaokang · Nov 9, 2023 · Nov 9, 2023 · Nov 9, 2023 · 7ac2c58ebfa5213eaa8e85964366a6a8b3173ee3
commit 7ac2c58ebfa5213eaa8e85964366a6a8b3173ee3
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java
@@ -71,6 +71,7 @@
 import org.apache.doris.nereids.rules.rewrite.InferFilterNotNull;
 import org.apache.doris.nereids.rules.rewrite.InferJoinNotNull;
 import org.apache.doris.nereids.rules.rewrite.InferPredicates;
+import org.apache.doris.nereids.rules.rewrite.LimitSortToTopN;
 import org.apache.doris.nereids.rules.rewrite.MergeFilters;
 import org.apache.doris.nereids.rules.rewrite.MergeOneRowRelationIntoUnion;
 import org.apache.doris.nereids.rules.rewrite.MergeProjects;
@@ -89,9 +90,9 @@
 import org.apache.doris.nereids.rules.rewrite.PushProjectThroughUnion;
 import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughProject;
 import org.apache.doris.nereids.rules.rewrite.PushdownLimit;
+import org.apache.doris.nereids.rules.rewrite.PushdownTopNThroughJoin;
 import org.apache.doris.nereids.rules.rewrite.PushdownTopNThroughWindow;
 import org.apache.doris.nereids.rules.rewrite.ReorderJoin;
-import org.apache.doris.nereids.rules.rewrite.ReplaceLimitNode;
 import org.apache.doris.nereids.rules.rewrite.RewriteCteChildren;
 import org.apache.doris.nereids.rules.rewrite.SemiJoinCommute;
 import org.apache.doris.nereids.rules.rewrite.SimplifyAggGroupBy;
@@ -263,14 +264,14 @@ public class Rewriter extends AbstractBatchJobExecutor {
             // ),
 
             topic("Limit optimization",
-                    topDown(
-                            // TODO: the logical plan should not contains any phase information,
-                            //       we should refactor like AggregateStrategies, e.g. LimitStrategies,
-                            //       generate one PhysicalLimit if current distribution is gather or two
-                            //       PhysicalLimits with gather exchange
-                            new ReplaceLimitNode(),
-                            new SplitLimit(),
-                            new PushdownLimit(),
+                    // TODO: the logical plan should not contains any phase information,
+                    //       we should refactor like AggregateStrategies, e.g. LimitStrategies,
+                    //       generate one PhysicalLimit if current distribution is gather or two
+                    //       PhysicalLimits with gather exchange
+                    topDown(new LimitSortToTopN()),
+                    topDown(new SplitLimit()),
+                    topDown(new PushdownLimit(),
+                            new PushdownTopNThroughJoin(),
                             new PushdownTopNThroughWindow(),
                             new CreatePartitionTopNFromWindow()
                     )
@@ -301,15 +302,15 @@ public class Rewriter extends AbstractBatchJobExecutor {
             topic("topn optimize",
                     topDown(new DeferMaterializeTopNResult())
             ),
+            topic("eliminate",
+                    // SORT_PRUNING should be applied after mergeLimit
+                    custom(RuleType.ELIMINATE_SORT, EliminateSort::new),
+                    bottomUp(new EliminateEmptyRelation())
+            ),
             // this rule batch must keep at the end of rewrite to do some plan check
             topic("Final rewrite and check",
                     custom(RuleType.ENSURE_PROJECT_ON_TOP_JOIN, EnsureProjectOnTopJoin::new),
-                    topDown(
-                            new PushdownFilterThroughProject(),
-                            new MergeProjects()
-                    ),
-                    // SORT_PRUNING should be applied after mergeLimit
-                    custom(RuleType.ELIMINATE_SORT, EliminateSort::new),
+                    topDown(new PushdownFilterThroughProject(), new MergeProjects()),
                     custom(RuleType.ADJUST_CONJUNCTS_RETURN_TYPE, AdjustConjunctsReturnType::new),
                     bottomUp(
                             new ExpressionRewrite(CheckLegalityAfterRewrite.INSTANCE),
@@ -323,10 +324,6 @@ public class Rewriter extends AbstractBatchJobExecutor {
                             new CollectFilterAboveConsumer(),
                             new CollectProjectAboveConsumer()
                     )
-            ),
-
-            topic("eliminate empty relation",
-                bottomUp(new EliminateEmptyRelation())
             )
     );
 

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java
@@ -162,8 +162,9 @@ public enum RuleType {
     COLUMN_PRUNING(RuleTypeClass.REWRITE),
     ELIMINATE_SORT(RuleTypeClass.REWRITE),
 
-    PUSHDOWN_TOP_N_THROUGH_PROJECTION_WINDOW(RuleTypeClass.REWRITE),
-    PUSHDOWN_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),
+    PUSHDOWN_MIN_MAX_THROUGH_JOIN(RuleTypeClass.REWRITE),
+    PUSHDOWN_SUM_THROUGH_JOIN(RuleTypeClass.REWRITE),
+    PUSHDOWN_COUNT_THROUGH_JOIN(RuleTypeClass.REWRITE),
 
     TRANSPOSE_LOGICAL_SEMI_JOIN_LOGICAL_JOIN(RuleTypeClass.REWRITE),
     TRANSPOSE_LOGICAL_SEMI_JOIN_LOGICAL_JOIN_PROJECT(RuleTypeClass.REWRITE),
@@ -243,7 +244,12 @@ public enum RuleType {
     PUSH_LIMIT_THROUGH_PROJECT_WINDOW(RuleTypeClass.REWRITE),
     PUSH_LIMIT_THROUGH_UNION(RuleTypeClass.REWRITE),
     PUSH_LIMIT_THROUGH_WINDOW(RuleTypeClass.REWRITE),
-    PUSH_LIMIT_INTO_SORT(RuleTypeClass.REWRITE),
+    LIMIT_SORT_TO_TOP_N(RuleTypeClass.REWRITE),
+    // topN push down
+    PUSH_TOP_N_THROUGH_JOIN(RuleTypeClass.REWRITE),
+    PUSH_TOP_N_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),
+    PUSH_TOP_N_THROUGH_PROJECT_WINDOW(RuleTypeClass.REWRITE),
+    PUSH_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),
     // adjust nullable
     ADJUST_NULLABLE(RuleTypeClass.REWRITE),
     ADJUST_CONJUNCTS_RETURN_TYPE(RuleTypeClass.REWRITE),

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateLimit.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateLimit.java
@@ -19,18 +19,35 @@
 
 import org.apache.doris.nereids.rules.Rule;
 import org.apache.doris.nereids.rules.RuleType;
+import org.apache.doris.nereids.trees.UnaryNode;
+import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
 import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
 
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
 /**
  * Eliminate limit = 0.
  */
-public class EliminateLimit extends OneRewriteRuleFactory {
+public class EliminateLimit implements RewriteRuleFactory {
+
     @Override
-    public Rule build() {
-        return logicalLimit()
-                .when(limit -> limit.getLimit() == 0)
-                .thenApply(ctx -> new LogicalEmptyRelation(ctx.statementContext.getNextRelationId(),
-                        ctx.root.getOutput()))
-                .toRule(RuleType.ELIMINATE_LIMIT);
+    public List<Rule> buildRules() {
+        return ImmutableList.of(
+                logicalLimit()
+                        .when(limit -> limit.getLimit() == 0)
+                        .thenApply(ctx -> new LogicalEmptyRelation(ctx.statementContext.getNextRelationId(),
+                                ctx.root.getOutput()))
+                        .toRule(RuleType.ELIMINATE_LIMIT),
+                logicalLimit(logicalOneRowRelation())
+                        .then(limit -> limit.getLimit() > 0 && limit.getOffset() == 0
+                                ? limit.child() : new LogicalEmptyRelation(StatementScopeIdGenerator.newRelationId(),
+                                limit.child().getOutput()))
+                        .toRule(RuleType.ELIMINATE_LIMIT_ON_ONE_ROW_RELATION),
+                logicalLimit(logicalEmptyRelation())
+                        .then(UnaryNode::child)
+                        .toRule(RuleType.ELIMINATE_LIMIT_ON_EMPTY_RELATION)
+        );
     }
 }
diff --git a/...reids/rules/rewrite/ReplaceLimitNode.java → ...ereids/rules/rewrite/LimitSortToTopN.java b/...reids/rules/rewrite/ReplaceLimitNode.java → ...ereids/rules/rewrite/LimitSortToTopN.java
@@ -19,10 +19,7 @@
 
 import org.apache.doris.nereids.rules.Rule;
 import org.apache.doris.nereids.rules.RuleType;
-import org.apache.doris.nereids.trees.UnaryNode;
-import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
 import org.apache.doris.nereids.trees.plans.Plan;
-import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
 import org.apache.doris.nereids.trees.plans.logical.LogicalSort;
 import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
@@ -35,7 +32,7 @@
 /**
  * rule to eliminate limit node by replace to other nodes.
  */
-public class ReplaceLimitNode implements RewriteRuleFactory {
+public class LimitSortToTopN implements RewriteRuleFactory {
     @Override
     public List<Rule> buildRules() {
         return ImmutableList.of(
@@ -47,8 +44,8 @@ public List<Rule> buildRules() {
                                     limit.getLimit(),
                                     limit.getOffset(),
                                     sort.child(0));
-                        }).toRule(RuleType.PUSH_LIMIT_INTO_SORT),
-                //limit->proj->sort ==> proj->topN
+                        }).toRule(RuleType.LIMIT_SORT_TO_TOP_N),
+                // limit -> proj -> sort ==> proj -> topN
                 logicalLimit(logicalProject(logicalSort()))
                         .then(limit -> {
                             LogicalProject project = limit.child();
@@ -58,15 +55,7 @@ public List<Rule> buildRules() {
                                     limit.getOffset(),
                                     sort.child(0));
                             return project.withChildren(Lists.newArrayList(topN));
-                        }).toRule(RuleType.PUSH_LIMIT_INTO_SORT),
-                logicalLimit(logicalOneRowRelation())
-                        .then(limit -> limit.getLimit() > 0 && limit.getOffset() == 0
-                                ? limit.child() : new LogicalEmptyRelation(StatementScopeIdGenerator.newRelationId(),
-                                limit.child().getOutput()))
-                        .toRule(RuleType.ELIMINATE_LIMIT_ON_ONE_ROW_RELATION),
-                logicalLimit(logicalEmptyRelation())
-                        .then(UnaryNode::child)
-                        .toRule(RuleType.ELIMINATE_LIMIT_ON_EMPTY_RELATION)
+                        }).toRule(RuleType.LIMIT_SORT_TO_TOP_N)
         );
     }
 }
diff --git a/...re/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownFilterThroughProject.java b/...re/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownFilterThroughProject.java
@@ -41,27 +41,26 @@ public class PushdownFilterThroughProject implements RewriteRuleFactory {
     @Override
     public List<Rule> buildRules() {
         return ImmutableList.of(
-            RuleType.PUSHDOWN_FILTER_THROUGH_PROJECT.build(logicalFilter(logicalProject())
-                    .whenNot(filter -> filter.child().getProjects().stream().anyMatch(
-                            expr -> expr.anyMatch(WindowExpression.class::isInstance)))
-                    .then(PushdownFilterThroughProject::pushdownFilterThroughProject)),
-            // filter(project(limit)) will change to filter(limit(project)) by PushdownProjectThroughLimit,
-            // then we should change filter(limit(project)) to project(filter(limit))
-            RuleType.PUSHDOWN_FILTER_THROUGH_PROJECT_UNDER_LIMIT
-                    .build(logicalFilter(logicalLimit(logicalProject()))
-                            .whenNot(filter -> filter.child().child().getProjects().stream()
-                                    .anyMatch(expr -> expr
-                                            .anyMatch(WindowExpression.class::isInstance)))
-                            .then(filter -> {
-                                LogicalLimit<LogicalProject<Plan>> limit = filter.child();
-                                LogicalProject<Plan> project = limit.child();
+                logicalFilter(logicalProject())
+                        .whenNot(filter -> filter.child().getProjects().stream().anyMatch(
+                                expr -> expr.anyMatch(WindowExpression.class::isInstance)))
+                        .then(PushdownFilterThroughProject::pushdownFilterThroughProject)
+                        .toRule(RuleType.PUSHDOWN_FILTER_THROUGH_PROJECT),
+                // filter(project(limit)) will change to filter(limit(project)) by PushdownProjectThroughLimit,
+                // then we should change filter(limit(project)) to project(filter(limit))
+                logicalFilter(logicalLimit(logicalProject()))
+                        .whenNot(filter -> filter.child().child().getProjects().stream()
+                                .anyMatch(expr -> expr.anyMatch(WindowExpression.class::isInstance)))
+                        .then(filter -> {
+                            LogicalLimit<LogicalProject<Plan>> limit = filter.child();
+                            LogicalProject<Plan> project = limit.child();
 
-                                return project.withProjectsAndChild(project.getProjects(),
-                                        new LogicalFilter<>(
-                                                ExpressionUtils.replace(filter.getConjuncts(),
-                                                        project.getAliasToProducer()),
-                                                limit.withChildren(project.child())));
-                            }))
+                            return project.withProjectsAndChild(project.getProjects(),
+                                    new LogicalFilter<>(
+                                            ExpressionUtils.replace(filter.getConjuncts(),
+                                                    project.getAliasToProducer()),
+                                            limit.withChildren(project.child())));
+                        }).toRule(RuleType.PUSHDOWN_FILTER_THROUGH_PROJECT_UNDER_LIMIT)
         );
     }
 

diff --git a/...ore/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownProjectThroughLimit.java b/...ore/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownProjectThroughLimit.java
@@ -24,6 +24,7 @@
 import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
 
 /**
+ * <pre>
  * Before:
  *          project
  *             │
@@ -42,6 +43,7 @@
  *             │
  *             ▼
  *          plan node
+ * </pre>
  */
 public class PushdownProjectThroughLimit extends OneRewriteRuleFactory {
 
@@ -50,9 +52,7 @@ public Rule build() {
         return logicalProject(logicalLimit()).thenApply(ctx -> {
             LogicalProject<LogicalLimit<Plan>> logicalProject = ctx.root;
             LogicalLimit<Plan> logicalLimit = logicalProject.child();
-            return new LogicalLimit<>(logicalLimit.getLimit(), logicalLimit.getOffset(),
-                    logicalLimit.getPhase(), logicalProject.withProjectsAndChild(logicalProject.getProjects(),
-                    logicalLimit.child()));
+            return logicalLimit.withChildren(logicalProject.withChildren(logicalLimit.child()));
         }).toRule(RuleType.PUSHDOWN_PROJECT_THROUGH_LIMIT);
     }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownTopNThroughJoin.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownTopNThroughJoin.java
@@ -0,0 +1,118 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.properties.OrderKey;
+import org.apache.doris.nereids.rules.Rule;
+import org.apache.doris.nereids.rules.RuleType;
+import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
+import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
+import org.apache.doris.nereids.util.Utils;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Push down TopN through Outer Join into left child .....
+ */
+public class PushdownTopNThroughJoin implements RewriteRuleFactory {
+
+    @Override
+    public List<Rule> buildRules() {
+        return ImmutableList.of(
+                // topN -> join
+                logicalTopN(logicalJoin())
+                        // TODO: complex orderby
+                        .when(topN -> topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                                .allMatch(Slot.class::isInstance))
+                        .then(topN -> {
+                            LogicalJoin<Plan, Plan> join = topN.child();
+                            Plan newJoin = pushLimitThroughJoin(topN, join);
+                            if (newJoin == null || topN.child().children().equals(newJoin.children())) {
+                                return null;
+                            }
+                            return topN.withChildren(newJoin);
+                        })
+                        .toRule(RuleType.PUSH_TOP_N_THROUGH_JOIN),
+
+                // topN -> project -> join
+                logicalTopN(logicalProject(logicalJoin()))
+                        .when(topN -> topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                                .allMatch(Slot.class::isInstance))
+                        .then(topN -> {
+                            LogicalProject<LogicalJoin<Plan, Plan>> project = topN.child();
+                            LogicalJoin<Plan, Plan> join = project.child();
+
+                            // If orderby exprs aren't all in the output of the project, we can't push down.
+                            // topN(order by: slot(a+1))
+                            // - project(a+1, b)
+                            // TODO: in the future, we also can push down it.
+                            Set<Slot> outputSet = project.child().getOutputSet();
+                            if (!topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                                    .flatMap(e -> e.getInputSlots().stream())
+                                    .allMatch(outputSet::contains)) {
+                                return null;
+                            }
+
+                            Plan newJoin = pushLimitThroughJoin(topN, join);
+                            if (newJoin == null || join.children().equals(newJoin.children())) {
+                                return null;
+                            }
+                            return topN.withChildren(project.withChildren(newJoin));
+                        }).toRule(RuleType.PUSH_TOP_N_THROUGH_PROJECT_JOIN)
+        );
+    }
+
+    private Plan pushLimitThroughJoin(LogicalTopN<? extends Plan> topN, LogicalJoin<Plan, Plan> join) {
+        switch (join.getJoinType()) {
+            case LEFT_OUTER_JOIN:
+                Set<Slot> rightOutputSet = join.right().getOutputSet();
+                if (topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                        .anyMatch(e -> Utils.isIntersecting(rightOutputSet, e.getInputSlots()))) {
+                    return null;
+                }
+                return join.withChildren(topN.withChildren(join.left()), join.right());
+            case RIGHT_OUTER_JOIN:
+                Set<Slot> leftOutputSet = join.left().getOutputSet();
+                if (topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                        .anyMatch(e -> Utils.isIntersecting(leftOutputSet, e.getInputSlots()))) {
+                    return null;
+                }
+                return join.withChildren(join.left(), topN.withChildren(join.right()));
+            case CROSS_JOIN:
+                List<Slot> orderbySlots = topN.getOrderKeys().stream().map(OrderKey::getExpr)
+                        .flatMap(e -> e.getInputSlots().stream()).collect(Collectors.toList());
+                if (join.left().getOutputSet().containsAll(orderbySlots)) {
+                    return join.withChildren(topN.withChildren(join.left()), join.right());
+                } else if (join.right().getOutputSet().containsAll(orderbySlots)) {
+                    return join.withChildren(join.left(), topN.withChildren(join.right()));
+                } else {
+                    return null;
+                }
+            default:
+                // don't push limit.
+                return null;
+        }
+    }
+}