[opt](nereids)push down assert one row join

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -132,6 +132,7 @@
 import org.apache.doris.nereids.rules.rewrite.PushDownEncodeSlot;
 import org.apache.doris.nereids.rules.rewrite.PushDownFilterIntoSchemaScan;
 import org.apache.doris.nereids.rules.rewrite.PushDownFilterThroughProject;
+import org.apache.doris.nereids.rules.rewrite.PushDownJoinOnAssertNumRows;
 import org.apache.doris.nereids.rules.rewrite.PushDownLimit;
 import org.apache.doris.nereids.rules.rewrite.PushDownLimitDistinctThroughJoin;
 import org.apache.doris.nereids.rules.rewrite.PushDownLimitDistinctThroughUnion;
@@ -736,6 +737,7 @@ public class Rewriter extends AbstractBatchJobExecutor {
                 ),
                 topic("set initial join order",
                         bottomUp(ImmutableList.of(new InitJoinOrder())),
+                        bottomUp(ImmutableList.of(new PushDownJoinOnAssertNumRows(), new MergeProjectable())),
                         topDown(new SkewJoin())),
                 topic("agg rewrite",
                     // these rules should be put after mv optimization to avoid mv matching fail

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java

@@ -176,7 +176,7 @@ public class RuleSet {
             new PushDownFilterThroughPartitionTopN(),
             ExpressionNormalizationAndOptimization.NO_MIN_MAX_RANGE_INSTANCE,
             new EliminateFilter()
-    );
+            );
 
     public static final List<Rule> IMPLEMENTATION_RULES = planRuleFactories()
             .add(new LogicalCTEProducerToPhysicalCTEProducer())

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java

@@ -233,6 +233,7 @@ public enum RuleType {
     LOGICAL_SEMI_JOIN_COMMUTE(RuleTypeClass.REWRITE),
     TRANSPOSE_LOGICAL_SEMI_JOIN_AGG(RuleTypeClass.REWRITE),
     TRANSPOSE_LOGICAL_SEMI_JOIN_AGG_PROJECT(RuleTypeClass.REWRITE),
+    PUSH_DOWN_JOIN_ON_ASSERT_NUM_ROWS(RuleTypeClass.REWRITE),
 
     // expression of plan rewrite
     EXTRACT_IN_PREDICATE_FROM_OR(RuleTypeClass.REWRITE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownJoinOnAssertNumRows.java

@@ -0,0 +1,229 @@
+// 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.rules.Rule;
+import org.apache.doris.nereids.rules.RuleType;
+import org.apache.doris.nereids.trees.expressions.Alias;
+import org.apache.doris.nereids.trees.expressions.AssertNumRowsElement;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.NamedExpression;
+import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.plans.JoinType;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalAssertNumRows;
+import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
+
+import com.google.common.collect.Lists;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Push down join when one child is LogicalAssertNumRows.
+ * select * from T1 join T2 where T1.b=T2.b and T1.a > (select x from T3 ...)
+ *
+ * <pre>
+ * Case 1: Push to left child
+ * Before:
+ *     topJoin(T1.a > x)
+ *       |-- Project
+ *       |     `-- bottomJoin(T1.b = T2.b)
+ *       |           |-- Scan(T1)
+ *       |           `-- Scan(T2)
+ *       `-- LogicalAssertNumRows(output=(x, ...))
+ *
+ * After:
+ *     Project
+ *       |-- topJoin(T1.b = T2.b)
+ *               |-- bottomJoin(T1.a > x)
+ *                   |-- Scan(T1)
+ *                   `-- LogicalAssertNumRows(output=(x, ...))
+ *               `-- Scan(T2)
+ *
+ * Case 2: Push to right child
+ * Before:
+ *     topJoin(T2.a > x)
+ *       |-- Project
+ *       |     `-- bottomJoin(T1.b = T2.b)
+ *       |           |-- Scan(T1)
+ *       |           `-- Scan(T2)
+ *       `-- LogicalAssertNumRows(output=(x, ...))
+ *
+ * After:
+ *     Project
+ *       |-- topJoin(T1.b = T2.b)
+ *              |--Scan(T1)
+ *             `-- bottomJoin(T2.a > x)
+ *                   |-- Scan(T2)
+ *                   `-- LogicalAssertNumRows(output=(x, ...))
+ * </pre>
+ */
+public class PushDownJoinOnAssertNumRows extends OneRewriteRuleFactory {
+    @Override
+    public Rule build() {
+        return logicalJoin()
+                .when(topJoin -> pattenCheck(topJoin))
+                .then(topJoin -> pushDownAssertNumRowsJoin(topJoin))
+                .toRule(RuleType.PUSH_DOWN_JOIN_ON_ASSERT_NUM_ROWS);
+    }
+
+    private boolean pattenCheck(LogicalJoin topJoin) {
+        // 1. right is LogicalAssertNumRows or LogicalProject->LogicalAssertNumRows
+        // 2. left is join or project->join
+        // 3. only one join condition.
+        if (topJoin.getJoinType() != JoinType.INNER_JOIN && topJoin.getJoinType() != JoinType.CROSS_JOIN) {
+            return false;
+        }
+        LogicalJoin bottomJoin;
+        Plan left = topJoin.left();
+        Plan right = topJoin.right();
+        if (!isAssertOneRowEqOrProjectAssertOneRowEq(right)) {
+            return false;
+        }
+        if (left instanceof LogicalJoin) {
+            bottomJoin = (LogicalJoin) left;
+        } else if (left instanceof LogicalProject && left.child(0) instanceof LogicalJoin) {
+            bottomJoin = (LogicalJoin) left.child(0);
+        } else {
+            return false;
+        }
+
+        if (bottomJoin.getJoinType() != JoinType.INNER_JOIN && bottomJoin.getJoinType() != JoinType.CROSS_JOIN) {
+            return false;
+        }
+
+        if (joinOnAssertOneRowEq(bottomJoin)) {
+            return false;
+        }
+
+        if (topJoin.getHashJoinConjuncts().isEmpty()) {
+            return topJoin.getOtherJoinConjuncts().size() == 1;
+        }
+        return false;
+    }
+
+    private boolean isAssertOneRowEqOrProjectAssertOneRowEq(Plan plan) {
+        if (plan instanceof LogicalProject) {
+            plan = plan.child(0);
+        }
+        if (plan instanceof LogicalAssertNumRows) {
+            AssertNumRowsElement assertNumRowsElement = ((LogicalAssertNumRows) plan).getAssertNumRowsElement();
+            if (assertNumRowsElement.getAssertion() == AssertNumRowsElement.Assertion.EQ
+                    || assertNumRowsElement.getDesiredNumOfRows() == 1L) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    private boolean joinOnAssertOneRowEq(LogicalJoin join) {
+        return isAssertOneRowEqOrProjectAssertOneRowEq(join.right())
+                || isAssertOneRowEqOrProjectAssertOneRowEq(join.left());
+    }
+
+    private Plan pushDownAssertNumRowsJoin(LogicalJoin topJoin) {
+        Plan assertBranch = topJoin.right();
+        Expression condition = (Expression) topJoin.getOtherJoinConjuncts().get(0);
+        List<Alias> aliasUsedInConditionFromLeftProject = new ArrayList<>();
+        LogicalJoin<? extends Plan, ? extends Plan> bottomJoin;
+        if (topJoin.left() instanceof LogicalProject) {
+            LogicalProject<? extends Plan> leftProject = (LogicalProject<? extends Plan>) topJoin.left();
+            for (NamedExpression namedExpression : leftProject.getProjects()) {
+                if (namedExpression instanceof Alias && condition.getInputSlots().contains(namedExpression.toSlot())) {
+                    aliasUsedInConditionFromLeftProject.add((Alias) namedExpression);
+                }
+            }
+            condition = leftProject.pushDownExpressionPastProject(condition);
+            bottomJoin = (LogicalJoin<? extends Plan, ? extends Plan>) leftProject.child();
+        } else {
+            bottomJoin = (LogicalJoin<? extends Plan, ? extends Plan>) topJoin.left();
+        }
+        Plan bottomLeft = bottomJoin.left();
+        Plan bottomRight = bottomJoin.right();
+
+        List<Slot> conditionSlotsFromTopLeft = condition.getInputSlots().stream()
+                .filter(slot -> topJoin.left().getOutputSet().contains(slot))
+                .collect(Collectors.toList());
+        if (bottomLeft.getOutputSet().containsAll(conditionSlotsFromTopLeft)) {
+            // push to bottomLeft
+            Plan newBottomLeft;
+            if (aliasUsedInConditionFromLeftProject.isEmpty()) {
+                newBottomLeft = bottomLeft;
+            } else {
+                newBottomLeft = projectAliasOnPlan(aliasUsedInConditionFromLeftProject, bottomLeft);
+            }
+            LogicalJoin<? extends Plan, ? extends Plan> newBottomJoin = new LogicalJoin<>(
+                    topJoin.getJoinType(),
+                    topJoin.getHashJoinConjuncts(),
+                    topJoin.getOtherJoinConjuncts(),
+                    newBottomLeft,
+                    assertBranch,
+                    topJoin.getJoinReorderContext());
+            LogicalJoin<? extends Plan, ? extends Plan> newTopJoin = (LogicalJoin<? extends Plan, ? extends Plan>)
+                    bottomJoin.withChildren(newBottomJoin, bottomRight);
+            if (topJoin.left() instanceof LogicalProject) {
+                LogicalProject<? extends Plan> upperProject = projectAliasOnPlan(
+                        aliasUsedInConditionFromLeftProject, topJoin.left());
+                return upperProject.withChildren(newTopJoin);
+            } else {
+                return newTopJoin;
+            }
+        } else if (bottomRight.getOutputSet().containsAll(conditionSlotsFromTopLeft)) {
+            Plan newBottomRight;
+            if (aliasUsedInConditionFromLeftProject.isEmpty()) {
+                newBottomRight = bottomRight;
+            } else {
+                newBottomRight = projectAliasOnPlan(aliasUsedInConditionFromLeftProject, bottomRight);
+            }
+            LogicalJoin<? extends Plan, ? extends Plan> newBottomJoin = new LogicalJoin<>(
+                    topJoin.getJoinType(),
+                    topJoin.getHashJoinConjuncts(),
+                    topJoin.getOtherJoinConjuncts(),
+                    newBottomRight,
+                    assertBranch,
+                    topJoin.getJoinReorderContext());
+            LogicalJoin<? extends Plan, ? extends Plan> newTopJoin = (LogicalJoin<? extends Plan, ? extends Plan>)
+                    bottomJoin.withChildren(bottomLeft, newBottomJoin);
+            if (topJoin.left() instanceof LogicalProject) {
+                LogicalProject<? extends Plan> upperProject = projectAliasOnPlan(
+                        aliasUsedInConditionFromLeftProject, topJoin.left());
+                return upperProject.withChildren(newTopJoin);
+            } else {
+                return newTopJoin;
+            }
+        }
+        return null;
+    }
+
+    private LogicalProject<? extends Plan> projectAliasOnPlan(List<Alias> projections, Plan child) {
+        if (child instanceof LogicalProject) {
+            LogicalProject<? extends Plan> project = (LogicalProject<? extends Plan>) child;
+            List<NamedExpression> newProjections =
+                    Lists.newArrayList(project.getProjects());
+            newProjections.addAll(projections);
+            return project.withProjects(newProjections);
+        } else {
+            List<NamedExpression> newProjections = Lists.newArrayList(child.getOutput());
+            newProjections.addAll(projections);
+            return new LogicalProject<>(newProjections, child);
+        }
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalProject.java

@@ -305,4 +305,23 @@ public void computeFd(DataTrait.Builder builder) {
             builder.addDeps(expr.getInputSlots(), ImmutableSet.of(expr.toSlot()));
         }
     }
+
+    /**
+     *
+     * example:
+     * expression: x + 1
+     * project(a+b as x)
+     * then before project, the expression is a+b+1
+     *
+     */
+    public Expression pushDownExpressionPastProject(Expression expression) {
+        HashMap projectMap = new HashMap();
+        for (NamedExpression namedExpression : projects) {
+            if (namedExpression instanceof Alias) {
+                Alias alias = (Alias) namedExpression;
+                projectMap.put(alias.toSlot(), alias.child());
+            }
+        }
+        return ExpressionUtils.replace(expression, projectMap);
+    }
 }