apache · morrySnow · Jul 5, 2024 · Jun 5, 2024
@@ -21,8 +21,6 @@
 import org.apache.doris.catalog.Column;
 import org.apache.doris.common.AnalysisException;
 
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
 import com.google.gson.annotations.SerializedName;
 
 import java.util.List;
@@ -38,10 +36,6 @@ public enum MTMVPartitionType {
         SELF_MANAGE
     }
 
-    public static final ImmutableSet<String> MTMV_PARTITION_FUNCTIONS = new ImmutableSortedSet.Builder<String>(
-            String.CASE_INSENSITIVE_ORDER).add("date_trunc")
-            .build();
-
     @SerializedName("pt")
     private MTMVPartitionType partitionType;
     @SerializedName("rt")

@@ -104,9 +104,9 @@ public Optional<Pair<Id, Statistics>> getPlanStatistics(CascadesContext cascades
             return Optional.empty();
         }
         RelationId relationId = null;
-        List<LogicalOlapScan> logicalOlapScan = this.getScanPlan().collectFirst(LogicalOlapScan.class::isInstance);
-        if (!logicalOlapScan.isEmpty()) {
-            relationId = logicalOlapScan.get(0).getRelationId();
+        Optional<LogicalOlapScan> logicalOlapScan = this.getScanPlan().collectFirst(LogicalOlapScan.class::isInstance);
+        if (logicalOlapScan.isPresent()) {
+            relationId = logicalOlapScan.get().getRelationId();
         }
         return Optional.of(Pair.of(relationId, normalizeStatisticsColumnExpression(mtmvCache.getStatistics())));
     }

@@ -24,6 +24,8 @@
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.RollUpTrait;
 
+import com.google.common.collect.ImmutableList;
+
 import java.util.List;
 import java.util.Set;
 
@@ -63,7 +65,8 @@ public abstract Function doRollup(
     protected static List<Expression> extractArguments(Expression functionWithAny, Function actualFunction) {
         Set<Object> exprSetToRemove = functionWithAny.collectToSet(expr -> !(expr instanceof Any));
         return actualFunction.collectFirst(expr ->
-                exprSetToRemove.stream().noneMatch(exprToRemove -> exprToRemove.equals(expr)));
+                        exprSetToRemove.stream().noneMatch(exprToRemove -> exprToRemove.equals(expr)))
+                .map(expr -> ImmutableList.of((Expression) expr)).orElse(ImmutableList.of());
     }
 
     /**

@@ -23,6 +23,7 @@
 import org.apache.doris.nereids.rules.expression.rules.FoldConstantRule;
 import org.apache.doris.nereids.rules.expression.rules.InPredicateDedup;
 import org.apache.doris.nereids.rules.expression.rules.InPredicateToEqualToRule;
+import org.apache.doris.nereids.rules.expression.rules.MergeDateTrunc;
 import org.apache.doris.nereids.rules.expression.rules.NormalizeBinaryPredicatesRule;
 import org.apache.doris.nereids.rules.expression.rules.SimplifyArithmeticComparisonRule;
 import org.apache.doris.nereids.rules.expression.rules.SimplifyArithmeticRule;
@@ -53,6 +54,7 @@ public class ExpressionNormalization extends ExpressionRewrite {
                 DigitalMaskingConvert.INSTANCE,
                 SimplifyArithmeticComparisonRule.INSTANCE,
                 ConvertAggStateCast.INSTANCE,
+                MergeDateTrunc.INSTANCE,
                 CheckCast.INSTANCE
             )
     );

@@ -0,0 +1,78 @@
+// 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.expression.rules;
+
+import org.apache.doris.nereids.rules.expression.ExpressionPatternMatcher;
+import org.apache.doris.nereids.rules.expression.ExpressionPatternRuleFactory;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.DateTrunc;
+import org.apache.doris.nereids.trees.expressions.literal.Interval.TimeUnit;
+import org.apache.doris.nereids.trees.expressions.literal.Literal;
+import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * Rewrite rule to convert
+ * For example:
+ * date_trunc(date_trunc(data_slot, 'hour'), 'day') -> date_trunc(data_slot, 'day')
+ */
+public class MergeDateTrunc implements ExpressionPatternRuleFactory {
+
+    public static MergeDateTrunc INSTANCE = new MergeDateTrunc();
+    public static ImmutableSet<TimeUnit> UN_SUPPORT_TIME_UNIT = ImmutableSet.of(TimeUnit.WEEK, TimeUnit.QUARTER);
+
+    @Override
+    public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
+        return ImmutableList.of(
+                matchesTopType(DateTrunc.class)
+                        .when(dateTrunc -> dateTrunc.getArgument(0) instanceof DateTrunc)
+                        .then(MergeDateTrunc::rewrite)
+        );
+    }
+
+    private static Expression rewrite(DateTrunc dateTrunc) {
+
+        Expression parentTimeUnitArgument = dateTrunc.getArgument(1);
+        if (!(parentTimeUnitArgument instanceof Literal)) {
+            return dateTrunc;
+        }
+        Optional<TimeUnit> parentTimeUnit = TimeUnit.of(((Literal) parentTimeUnitArgument).getStringValue());
+        DateTrunc childDateTrunc = (DateTrunc) dateTrunc.getArgument(0);
+        Expression childTimeUnitArgument = childDateTrunc.getArgument(1);
+        if (!(childTimeUnitArgument instanceof Literal)) {
+            return dateTrunc;
+        }
+        Optional<TimeUnit> childTimeUnit = TimeUnit.of(((Literal) childTimeUnitArgument).getStringValue());
+        if (!parentTimeUnit.isPresent() || !childTimeUnit.isPresent()) {
+            return dateTrunc;
+        }
+        if (UN_SUPPORT_TIME_UNIT.contains(parentTimeUnit.get())
+                || UN_SUPPORT_TIME_UNIT.contains(childTimeUnit.get())) {
+            return dateTrunc;
+        }
+        if (parentTimeUnit.get().getLevel() < childTimeUnit.get().getLevel()) {
+            return dateTrunc;
+        }
+        return new DateTrunc(childDateTrunc.getArgument(0), new VarcharLiteral(parentTimeUnit.get().toString()));
+    }
+}
@@ -271,15 +271,15 @@ default <T> Set<T> collectToSet(Predicate<TreeNode<NODE_TYPE>> predicate) {
     /**
      * Collect the nodes that satisfied the predicate firstly.
      */
-    default <T> List<T> collectFirst(Predicate<TreeNode<NODE_TYPE>> predicate) {
+    default <T> Optional<T> collectFirst(Predicate<TreeNode<NODE_TYPE>> predicate) {
         List<TreeNode<NODE_TYPE>> result = new ArrayList<>();
         foreach(node -> {
             if (result.isEmpty() && predicate.test(node)) {
                 result.add(node);
             }
             return !result.isEmpty();
         });
-        return (List<T>) ImmutableList.copyOf(result);
+        return result.isEmpty() ? Optional.empty() : Optional.of((T) result.get(0));
     }
 
     /**

@@ -25,6 +25,9 @@
 import org.apache.doris.nereids.types.DateType;
 
 import com.google.common.collect.ImmutableList;
+import org.apache.commons.lang3.EnumUtils;
+
+import java.util.Optional;
 
 /**
  * Interval for timestamp calculation.
@@ -61,30 +64,46 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
      * Supported time unit.
      */
     public enum TimeUnit {
-        YEAR("YEAR", false),
-        MONTH("MONTH", false),
-        WEEK("WEEK", false),
-        DAY("DAY", false),
-        HOUR("HOUR", true),
-        MINUTE("MINUTE", true),
-        SECOND("SECOND", true);
+        YEAR("YEAR", false, 800),
+        MONTH("MONTH", false, 700),
+        QUARTER("QUARTER", false, 600),
+        WEEK("WEEK", false, 500),
+        DAY("DAY", false, 400),
+        HOUR("HOUR", true, 300),
+        MINUTE("MINUTE", true, 200),
+        SECOND("SECOND", true, 100);
 
         private final String description;
-
         private final boolean isDateTimeUnit;
+        /**
+         * Time unit level, second level is low, year level is high
+         */
+        private final int level;
 
-        TimeUnit(String description, boolean isDateTimeUnit) {
+        TimeUnit(String description, boolean isDateTimeUnit, int level) {
             this.description = description;
             this.isDateTimeUnit = isDateTimeUnit;
+            this.level = level;
         }
 
         public boolean isDateTimeUnit() {
             return isDateTimeUnit;
         }
 
+        public int getLevel() {
+            return level;
+        }
+
         @Override
         public String toString() {
             return description;
         }
+
+        /**
+         * Construct time unit by name
+         */
+        public static Optional<TimeUnit> of(String name) {
+            return Optional.ofNullable(EnumUtils.getEnumIgnoreCase(TimeUnit.class, name));
+        }
     }
 }