[Refactor] refactor multi distinct rewrite strategy (StarRocks#40549)

Signed-off-by: packy92 <[email protected]>
xuzifu666 · Feb 4, 2024 · bfa9c72 · bfa9c72
1 parent 28b477f
commit bfa9c72
Show file tree

Hide file tree

Showing 20 changed files with 629 additions and 365 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java
@@ -60,6 +60,7 @@
 import com.starrocks.sql.optimizer.rule.transformation.PushLimitAndFilterToCTEProduceRule;
 import com.starrocks.sql.optimizer.rule.transformation.RemoveAggregationFromAggTable;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteGroupingSetsByCTERule;
+import com.starrocks.sql.optimizer.rule.transformation.RewriteMultiDistinctRule;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteSimpleAggToMetaScanRule;
 import com.starrocks.sql.optimizer.rule.transformation.SeparateProjectRule;
 import com.starrocks.sql.optimizer.rule.transformation.SkewJoinOptimizeRule;
@@ -427,7 +428,6 @@ private OptExpression logicalRuleRewrite(
             deriveLogicalProperty(tree);
         }
 
-        ruleRewriteIterative(tree, rootTaskContext, RuleSetType.MULTI_DISTINCT_REWRITE);
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PUSH_DOWN_PREDICATE);
 
         // No heavy metadata operation before external table partition prune
@@ -436,6 +436,7 @@ private OptExpression logicalRuleRewrite(
         // rewrite before SplitScanORToUnionRule
         ruleRewriteOnlyOnce(tree, rootTaskContext, new SplitDatePredicateRule());
         ruleRewriteOnlyOnce(tree, rootTaskContext, RuleSetType.PARTITION_PRUNE);
+        ruleRewriteIterative(tree, rootTaskContext, new RewriteMultiDistinctRule());
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PRUNE_EMPTY_OPERATOR);
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PRUNE_PROJECT);
 

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java
@@ -52,8 +52,7 @@
 import com.starrocks.sql.optimizer.rewrite.ReplaceColumnRefRewriter;
 import com.starrocks.sql.optimizer.rewrite.ScalarOperatorRewriter;
 import com.starrocks.sql.optimizer.statistics.ColumnStatistic;
-import org.apache.commons.collections.CollectionUtils;
-import org.apache.commons.collections.MapUtils;
+import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.collections4.SetUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -715,17 +714,26 @@ public static boolean mustGenerateMultiStageAggregate(Operator inputOp, Operator
             aggs = ((PhysicalHashAggregateOperator) inputOp).getAggregations();
         }
 
-        if (MapUtils.isEmpty(aggs)) {
-            return false;
-        } else {
-            // Must do multiple stage aggregate when aggregate distinct function has array type
-            // Must generate three, four phase aggregate for distinct aggregate with multi columns
-            return aggs.values().stream().anyMatch(callOperator -> callOperator.isDistinct()
-                    && (callOperator.getChildren().size() > 1 ||
-                    callOperator.getChildren().stream().anyMatch(c -> c.getType().isComplexType())));
+        for (CallOperator callOperator : aggs.values()) {
+            if (callOperator.isDistinct()) {
+                String fnName = callOperator.getFnName();
+                List<ScalarOperator> children = callOperator.getChildren();
+                if (children.size() > 1 || children.stream().anyMatch(c -> c.getType().isComplexType())) {
+                    return true;
+                }
+                if (FunctionSet.GROUP_CONCAT.equalsIgnoreCase(fnName) || FunctionSet.AVG.equalsIgnoreCase(fnName)) {
+                    return true;
+                } else if (FunctionSet.ARRAY_AGG.equalsIgnoreCase(fnName))  {
+                    if (children.size() > 1 || children.get(0).getType().isDecimalOfAnyVersion()) {
+                        return true;
+                    }
+                }
+            }
         }
+        return false;
     }
 
+    // without distinct function, the common distinctCols is an empty list.
     public static Optional<List<ColumnRefOperator>> extractCommonDistinctCols(Collection<CallOperator> aggCallOperators) {
         Set<ColumnRefOperator> distinctChildren = Sets.newHashSet();
         for (CallOperator callOperator : aggCallOperators) {

diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/scalar/ScalarOperatorUtil.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/scalar/ScalarOperatorUtil.java
@@ -23,7 +23,6 @@
 import com.starrocks.catalog.ScalarType;
 import com.starrocks.catalog.Type;
 import com.starrocks.server.GlobalStateMgr;
-import com.starrocks.sql.analyzer.DecimalV3FunctionAnalyzer;
 import com.starrocks.sql.optimizer.rewrite.ScalarOperatorRewriter;
 
 import static com.starrocks.catalog.Function.CompareMode.IS_IDENTICAL;
@@ -76,14 +75,4 @@ public static Function findSumFn(Type[] argTypes) {
         }
         return newFn;
     }
-
-    public static CallOperator buildMultiSumDistinct(CallOperator oldFunctionCall) {
-        Function multiDistinctSum = DecimalV3FunctionAnalyzer.convertSumToMultiDistinctSum(
-                oldFunctionCall.getFunction(), oldFunctionCall.getChild(0).getType());
-        ScalarOperatorRewriter scalarOpRewriter = new ScalarOperatorRewriter();
-        return (CallOperator) scalarOpRewriter.rewrite(
-                new CallOperator(
-                        FunctionSet.MULTI_DISTINCT_SUM, multiDistinctSum.getReturnType(),
-                        oldFunctionCall.getChildren(), multiDistinctSum), DEFAULT_TYPE_CAST_RULE);
-    }
 }
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/RuleSet.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/RuleSet.java
@@ -139,8 +139,6 @@
 import com.starrocks.sql.optimizer.rule.transformation.RewriteCountIfFunction;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteDuplicateAggregateFnRule;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteHllCountDistinctRule;
-import com.starrocks.sql.optimizer.rule.transformation.RewriteMultiDistinctByCTERule;
-import com.starrocks.sql.optimizer.rule.transformation.RewriteMultiDistinctRule;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteSimpleAggToMetaScanRule;
 import com.starrocks.sql.optimizer.rule.transformation.RewriteSumByAssociativeRule;
 import com.starrocks.sql.optimizer.rule.transformation.ScalarApply2AnalyticRule;
@@ -361,11 +359,6 @@ public class RuleSet {
                 new RewriteCountIfFunction()
         ));
 
-        REWRITE_RULES.put(RuleSetType.MULTI_DISTINCT_REWRITE, ImmutableList.of(
-                new RewriteMultiDistinctByCTERule(),
-                new RewriteMultiDistinctRule()
-        ));
-
         REWRITE_RULES.put(RuleSetType.PRUNE_PROJECT, ImmutableList.of(
                 new PruneProjectRule(),
                 new PruneProjectEmptyRule(),

diff --git a/...mation/RewriteMultiDistinctByCTERule.java → ...formation/MultiDistinctByCTERewriter.java b/...mation/RewriteMultiDistinctByCTERule.java → ...formation/MultiDistinctByCTERewriter.java
@@ -23,7 +23,6 @@
 import com.starrocks.catalog.Function;
 import com.starrocks.catalog.FunctionSet;
 import com.starrocks.catalog.ScalarType;
-import com.starrocks.qe.ConnectContext;
 import com.starrocks.server.GlobalStateMgr;
 import com.starrocks.sql.analyzer.DecimalV3FunctionAnalyzer;
 import com.starrocks.sql.optimizer.ExpressionContext;
@@ -34,15 +33,13 @@
 import com.starrocks.sql.optimizer.base.ColumnRefSet;
 import com.starrocks.sql.optimizer.operator.Operator;
 import com.starrocks.sql.optimizer.operator.OperatorBuilderFactory;
-import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.logical.LogicalAggregationOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalCTEAnchorOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalCTEConsumeOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalCTEProduceOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalFilterOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalJoinOperator;
 import com.starrocks.sql.optimizer.operator.logical.LogicalProjectOperator;
-import com.starrocks.sql.optimizer.operator.pattern.Pattern;
 import com.starrocks.sql.optimizer.operator.scalar.BinaryPredicateOperator;
 import com.starrocks.sql.optimizer.operator.scalar.CallOperator;
 import com.starrocks.sql.optimizer.operator.scalar.CastOperator;
@@ -51,7 +48,6 @@
 import com.starrocks.sql.optimizer.rewrite.ReplaceColumnRefRewriter;
 import com.starrocks.sql.optimizer.rewrite.ScalarOperatorRewriter;
 import com.starrocks.sql.optimizer.rewrite.scalar.ImplicitCastRule;
-import com.starrocks.sql.optimizer.rule.RuleType;
 
 import java.util.LinkedList;
 import java.util.List;
@@ -123,48 +119,11 @@
  *
  *
  */
-public class RewriteMultiDistinctByCTERule extends TransformationRule {
-    private final ScalarOperatorRewriter scalarRewriter = new ScalarOperatorRewriter();
-
-    public RewriteMultiDistinctByCTERule() {
-        super(RuleType.TF_REWRITE_MULTI_DISTINCT_BY_CTE,
-                Pattern.create(OperatorType.LOGICAL_AGGR).addChildren(Pattern.create(
-                        OperatorType.PATTERN_LEAF)));
-    }
-
-    @Override
-    public boolean check(OptExpression input, OptimizerContext context) {
-        // check cte is disabled or hasNoGroup false
-        LogicalAggregationOperator agg = (LogicalAggregationOperator) input.getOp();
-        List<CallOperator> distinctAggOperatorList = agg.getAggregations().values().stream()
-                .filter(CallOperator::isDistinct).collect(Collectors.toList());
-        boolean hasMultiColumns = distinctAggOperatorList.stream().anyMatch(f -> f.getDistinctChildren().size() > 1);
-        if (hasMultiColumns && distinctAggOperatorList.size() > 1) {
-            return true;
-        }
-
-        if (!context.getSessionVariable().isCboCteReuse()) {
-            return false;
-        }
-
-        if (agg.hasSkew() && distinctAggOperatorList.size() > 1 && !agg.getGroupingKeys().isEmpty()) {
-            return true;
-        }
-
-        if (agg.hasLimit() && !ConnectContext.get().getSessionVariable().isPreferCTERewrite()) {
-            return false;
-        }
-
-        if (!hasMultiColumns && agg.getGroupingKeys().size() > 1) {
-            return false;
-        }
+public class MultiDistinctByCTERewriter {
 
-        return distinctAggOperatorList.size() > 1 || agg.getAggregations().values().stream()
-                .anyMatch(call -> call.isDistinct() && call.getFnName().equals(FunctionSet.AVG));
-    }
+    private final ScalarOperatorRewriter scalarRewriter = new ScalarOperatorRewriter();
 
-    @Override
-    public List<OptExpression> transform(OptExpression input, OptimizerContext context) {
+    public List<OptExpression> transformImpl(OptExpression input, OptimizerContext context) {
         ColumnRefFactory columnRefFactory = context.getColumnRefFactory();
         // define cteId
         int cteId = context.getCteContext().getNextCteId();
@@ -453,4 +412,5 @@ private LogicalCTEConsumeOperator buildCteConsume(OptExpression cteProduce, Colu
 
         return new LogicalCTEConsumeOperator(cteId, consumeOutputMap);
     }
+
 }