pick (#32305) (#31858)

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

@@ -86,6 +86,8 @@
 import org.apache.doris.nereids.rules.rewrite.PruneOlapScanTablet;
 import org.apache.doris.nereids.rules.rewrite.PullUpCteAnchor;
 import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderApply;
+import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderLimit;
+import org.apache.doris.nereids.rules.rewrite.PullUpProjectUnderTopN;
 import org.apache.doris.nereids.rules.rewrite.PushConjunctsIntoEsScan;
 import org.apache.doris.nereids.rules.rewrite.PushConjunctsIntoJdbcScan;
 import org.apache.doris.nereids.rules.rewrite.PushConjunctsIntoOdbcScan;
@@ -321,6 +323,10 @@ public class Rewriter extends AbstractBatchJobExecutor {
                             new PushdownLimitDistinctThroughJoin(),
                             new PushdownTopNThroughWindow(),
                             new CreatePartitionTopNFromWindow()
+                    ),
+                    topDown(
+                            new PullUpProjectUnderTopN(),
+                            new PullUpProjectUnderLimit()
                     )
             ),
             // TODO: these rules should be implementation rules, and generate alternative physical plans.

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

@@ -260,6 +260,8 @@ public enum RuleType {
     PUSH_TOP_N_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),
     PUSH_TOP_N_THROUGH_PROJECT_WINDOW(RuleTypeClass.REWRITE),
     PUSH_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),
+    PULL_UP_PROJECT_UNDER_TOPN(RuleTypeClass.REWRITE),
+    PULL_UP_PROJECT_UNDER_LIMIT(RuleTypeClass.REWRITE),
     // limit distinct push down
     PUSH_LIMIT_DISTINCT_THROUGH_JOIN(RuleTypeClass.REWRITE),
     PUSH_LIMIT_DISTINCT_THROUGH_PROJECT_JOIN(RuleTypeClass.REWRITE),

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

@@ -0,0 +1,57 @@
+// 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.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.util.PlanUtils;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Pull up Project under Limit.
+ */
+public class PullUpProjectUnderLimit extends OneRewriteRuleFactory {
+    @Override
+    public Rule build() {
+        return logicalLimit(logicalProject(logicalJoin().when(j -> j.getJoinType().isLeftRightOuterOrCrossJoin()))
+                .whenNot(p -> p.isAllSlots()))
+                .then(limit -> {
+                    LogicalProject<LogicalJoin<Plan, Plan>> project = limit.child();
+                    Set<Slot> allUsedSlots = project.getProjects().stream().flatMap(ne -> ne.getInputSlots().stream())
+                            .collect(Collectors.toSet());
+                    Set<Slot> outputSet = project.child().getOutputSet();
+                    if (outputSet.size() == allUsedSlots.size()) {
+                        Preconditions.checkState(outputSet.equals(allUsedSlots));
+                        return project.withChildren(limit.withChildren(project.child()));
+                    } else {
+                        Plan columnProject = PlanUtils.projectOrSelf(ImmutableList.copyOf(allUsedSlots),
+                                project.child());
+                        return project.withChildren(limit.withChildren(columnProject));
+                    }
+                }).toRule(RuleType.PULL_UP_PROJECT_UNDER_LIMIT);
+    }
+}

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

@@ -0,0 +1,81 @@
+// 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.Expression;
+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.ExpressionUtils;
+import org.apache.doris.nereids.util.PlanUtils;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Pull up Project under TopN for PushDownTopNThroughJoin
+ */
+public class PullUpProjectUnderTopN extends OneRewriteRuleFactory {
+    @Override
+    public Rule build() {
+        return logicalTopN(
+                logicalProject(logicalJoin().when(j -> j.getJoinType().isLeftRightOuterOrCrossJoin()))
+                        .whenNot(p -> p.isAllSlots()))
+                .then(topN -> {
+                    LogicalProject<LogicalJoin<Plan, Plan>> project = topN.child();
+                    Set<Slot> outputSet = project.child().getOutputSet();
+
+                    Map<Slot, Expression> slotMap = ExpressionUtils.generateReplaceMap(project.getProjects());
+                    List<OrderKey> newOrderKeys = new ArrayList<>();
+                    for (OrderKey orderKey : topN.getOrderKeys()) {
+                        if (!(orderKey.getExpr() instanceof Slot)) {
+                            return null;
+                        }
+                        Expression expression = slotMap.get((Slot) orderKey.getExpr());
+                        if (expression instanceof Slot) {
+                            newOrderKeys.add(orderKey.withExpression(expression));
+                        } else {
+                            return null;
+                        }
+                    }
+
+                    Set<Slot> allUsedSlots = project.getProjects().stream().flatMap(ne -> ne.getInputSlots().stream())
+                            .collect(Collectors.toSet());
+                    LogicalTopN<Plan> newTopN = topN.withOrderKeys(newOrderKeys);
+                    if (outputSet.size() == allUsedSlots.size()) {
+                        Preconditions.checkState(outputSet.equals(allUsedSlots));
+                        return project.withChildren(newTopN.withChildren(project.child()));
+                    } else {
+                        Plan columnProject = PlanUtils.projectOrSelf(ImmutableList.copyOf(allUsedSlots),
+                                project.child());
+                        return project.withChildren(newTopN.withChildren(columnProject));
+                    }
+                }).toRule(RuleType.PULL_UP_PROJECT_UNDER_TOPN);
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/literal/DateLiteral.java

@@ -163,11 +163,10 @@ static String normalize(String s) {
         StringBuilder sb = new StringBuilder();
 
         int i = 0;
+        // date and time contains 6 number part at most, so we just need normal 6 number part
+        int partNumber = 0;
 
         // handle two digit year
-        if (!isPunctuation(s.charAt(2)) && !isPunctuation(s.charAt(4))) {
-            throw new AnalysisException("date/datetime literal [" + s + "] is invalid");
-        }
         if (isPunctuation(s.charAt(2))) {
             String yy = s.substring(0, 2);
             int year = Integer.parseInt(yy);
@@ -178,11 +177,10 @@ static String normalize(String s) {
             }
             sb.append(yy);
             i = 2;
+            partNumber += 1;
         }
 
         // normalize leading 0 for date and time
-        // date and time contains 6 number part at most, so we just need normal 6 number part
-        int partNumber = 0;
         while (i < s.length() && partNumber < 6) {
             char c = s.charAt(i);
             if (Character.isDigit(c)) {
@@ -193,39 +191,37 @@ static String normalize(String s) {
                 }
                 int len = j - i;
                 if (len == 4 || len == 2) {
-                    for (int k = i; k < j; k++) {
-                        sb.append(s.charAt(k));
-                    }
+                    sb.append(s, i, j);
                 } else if (len == 1) {
-                    sb.append('0').append(c);
+                    if (partNumber == 0) {
+                        sb.append("000").append(c);
+                    } else {
+                        sb.append('0').append(c);
+                    }
                 } else {
                     throw new AnalysisException("date/datetime literal [" + s + "] is invalid");
                 }
                 i = j;
                 partNumber += 1;
             } else if (isPunctuation(c) || c == ' ' || c == 'T') {
-                sb.append(c);
                 i += 1;
+                if (partNumber < 3 && isPunctuation(c)) {
+                    sb.append('-');
+                } else if (partNumber == 3) {
+                    while (i < s.length() && (isPunctuation(s.charAt(i)) || s.charAt(i) == ' ' || s.charAt(i) == 'T')) {
+                        i += 1;
+                    }
+                    sb.append(' ');
+                } else if (partNumber > 3 && isPunctuation(c)) {
+                    sb.append(':');
+                } else {
+                    throw new AnalysisException("date/datetime literal [" + s + "] is invalid");
+                }
             } else {
                 break;
             }
         }
 
-        // replace punctuation with '-'
-        replacePunctuation(s, sb, '-', 4);
-        replacePunctuation(s, sb, '-', 7);
-        // Replace punctuation with ' '
-        if (sb.length() > 10 && sb.charAt(10) != ' ') {
-            if (sb.charAt(10) == 'T') {
-                sb.setCharAt(10, ' ');
-            } else {
-                replacePunctuation(s, sb, ' ', 10);
-            }
-        }
-        // replace punctuation with ':'
-        replacePunctuation(s, sb, ':', 13);
-        replacePunctuation(s, sb, ':', 16);
-
         // add missing Minute Second in Time part
         if (sb.length() == 13) {
             sb.append(":00:00");

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

@@ -152,6 +152,10 @@ public final boolean isRightOuterJoin() {
         return this == RIGHT_OUTER_JOIN;
     }
 
+    public final boolean isLeftRightOuterOrCrossJoin() {
+        return this == LEFT_OUTER_JOIN || this == RIGHT_OUTER_JOIN || this == CROSS_JOIN;
+    }
+
     public final boolean isLeftSemiOrAntiJoin() {
         return this == LEFT_SEMI_JOIN || this == LEFT_ANTI_JOIN || this == NULL_AWARE_LEFT_ANTI_JOIN;
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/util/ExpressionUtils.java

@@ -21,6 +21,7 @@
 import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
 import org.apache.doris.nereids.rules.expression.rules.FoldConstantRule;
 import org.apache.doris.nereids.trees.TreeNode;
+import org.apache.doris.nereids.trees.expressions.Alias;
 import org.apache.doris.nereids.trees.expressions.And;
 import org.apache.doris.nereids.trees.expressions.Cast;
 import org.apache.doris.nereids.trees.expressions.ComparisonPredicate;
@@ -44,6 +45,7 @@
 import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableList.Builder;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
@@ -252,6 +254,34 @@ public static Optional<Slot> extractSlotOrCastOnSlot(Expression expr) {
         }
     }
 
+    /**
+     * Generate replaceMap Slot -> Expression from NamedExpression[Expression as name]
+     */
+    public static Map<Slot, Expression> generateReplaceMap(List<NamedExpression> namedExpressions) {
+        ImmutableMap.Builder<Slot, Expression> replaceMap = ImmutableMap.builderWithExpectedSize(
+                namedExpressions.size() * 2);
+        for (NamedExpression namedExpression : namedExpressions) {
+            if (namedExpression instanceof Alias) {
+                // Avoid cast to alias, retrieving the first child expression.
+                replaceMap.put(namedExpression.toSlot(), namedExpression.child(0));
+            }
+        }
+        return replaceMap.build();
+    }
+
+    /**
+     * replace NameExpression.
+     */
+    public static NamedExpression replaceNameExpression(NamedExpression expr,
+            Map<? extends Expression, ? extends Expression> replaceMap) {
+        Expression newExpr = replace(expr, replaceMap);
+        if (newExpr instanceof NamedExpression) {
+            return (NamedExpression) newExpr;
+        } else {
+            return new Alias(expr.getExprId(), newExpr, expr.getName());
+        }
+    }
+
     /**
      * Replace expression node in the expression tree by `replaceMap` in top-down manner.
      * For example.

fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/literal/DateTimeLiteralTest.java

@@ -33,6 +33,13 @@ void reject() {
         // });
     }
 
+    @Test
+    void mysqlStrangeCase() {
+        new DateTimeV2Literal("0-08-01 13:21:03");
+        new DateTimeV2Literal("0001-01-01: 00:01:01.001");
+        new DateTimeV2Literal("2021?01?01 00.00.00");
+    }
+
     @Test
     void testBasic() {
         Consumer<DateTimeV2Literal> assertFunc = (datetime) -> {