[improvement](nereids) Support CTE split for DISTINCT aggregate functions

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

@@ -41,7 +41,6 @@
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue;
 import org.apache.doris.nereids.trees.expressions.functions.generator.Unnest;
-import org.apache.doris.nereids.trees.expressions.literal.Literal;
 import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
@@ -183,7 +182,7 @@ private LogicalPlan normalizeAgg(LogicalAggregate<Plan> aggregate, Optional<Logi
  for (Expression arg : aggFunc.children()) {
  // should not push down literal under aggregate
  // e.g. group_concat(distinct xxx, ','), the ',' literal show stay in aggregate
- if (arg instanceof Literal) {
+ if (arg.isConstant()) {
  continue;
  }
  if (arg.containsType(SubqueryExpr.class, WindowExpression.class, Unnest.class,
@@ -201,7 +200,7 @@ private LogicalPlan normalizeAgg(LogicalAggregate<Plan> aggregate, Optional<Logi
  for (Expression arg : aggFunc.children()) {
  // should not push down literal under aggregate
  // e.g. group_concat(distinct xxx, ','), the ',' literal show stay in aggregate
- if (arg instanceof Literal) {
+ if (arg.isConstant()) {
  continue;
  }
 

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

@@ -247,6 +247,9 @@ private static Set<Expression> collectNeedToSlotArgsOfGroupingScalarFuncAndAggFu
  ImmutableSet.Builder<Expression> argumentsOfAggregateFunctionBuilder = ImmutableSet.builder();
  for (AggregateFunction function : aggregateFunctions) {
  for (Expression arg : function.getArguments()) {
+ if (arg.isConstant()) {
+ continue;
+ }
  if (arg instanceof OrderExpression) {
  argumentsOfAggregateFunctionBuilder.add(arg.child(0));
  } else {

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

@@ -119,19 +119,20 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, Distinct
  }
 
  private boolean shouldUseMultiDistinct(LogicalAggregate<? extends Plan> agg) {
- boolean mustUseCte = AggregateUtils.containsCountDistinctMultiExpr(agg);
+ boolean mustUseCte = AggregateUtils.containsNotSupportMultiDistinctFunction(agg);
  boolean mustUseMulti = agg.getSourceRepeat().isPresent();
  if (mustUseCte && mustUseMulti) {
  throw new AnalysisException(
- "Unsupported query: GROUPING SETS/ROLLUP/CUBE cannot be used with a combination of "
- + "multi-column COUNT(DISTINCT) and other COUNT(DISTINCT) expressions.\n\n"
+ "Unsupported query: GROUPING SETS/ROLLUP/CUBE cannot be used with multiple DISTINCT "
+ + "argument groups when any DISTINCT aggregate does not support multi-distinct execution.\n\n"
  + "Unsupported scenarios:\n"
+ + "• ARRAY_AGG(DISTINCT a) with ARRAY_AGG(DISTINCT b) + GROUPING\n"
  + "• COUNT(DISTINCT a, b) with COUNT(DISTINCT a) + GROUPING\n"
  + "• COUNT(DISTINCT a, b) with COUNT(DISTINCT a, c) + GROUPING\n\n"
  + "Supported scenarios:\n"
- + "• Single COUNT(DISTINCT a, b) + GROUPING\n"
- + "• Multiple COUNT(DISTINCT single_column) + "
- + "GROUPING (e.g., COUNT(DISTINCT a), COUNT(DISTINCT b))");
+ + "• A single DISTINCT argument group with GROUPING\n"
+ + "• Multiple DISTINCT argument groups when all DISTINCT aggregates support "
+ + "multi-distinct execution (e.g., COUNT(DISTINCT a), COUNT(DISTINCT b))");
  }
  if (mustUseCte) {
  return false;

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

@@ -42,6 +42,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -152,10 +153,11 @@ private static void collectDistinctAndNonDistinctFunctions(LogicalAggregate<? ex
  List<List<Alias>> aliases, List<Alias> otherAggFuncs) {
  Map<Set<Expression>, List<Alias>> distinctArgToAliasList = new LinkedHashMap<>();
  for (NamedExpression namedExpression : agg.getOutputExpressions()) {
- if (!(namedExpression instanceof Alias) || !(namedExpression.child(0) instanceof AggregateFunction)) {
+ if (!(namedExpression instanceof Alias) || !namedExpression.containsType(AggregateFunction.class)) {
  continue;
  }
- AggregateFunction aggFunc = (AggregateFunction) namedExpression.child(0);
+ Optional<Expression> op = namedExpression.collectFirst(e -> e instanceof AggregateFunction);
+ AggregateFunction aggFunc = (AggregateFunction) op.get();
  if (aggFunc.isDistinct()) {
  distinctArgToAliasList.computeIfAbsent(ImmutableSet.copyOf(aggFunc.getDistinctArguments()),
  k -> new ArrayList<>()).add((Alias) namedExpression);

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/CollectList.java

@@ -18,6 +18,7 @@
 package org.apache.doris.nereids.trees.expressions.functions.agg;
 
 import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
 import org.apache.doris.nereids.trees.expressions.literal.ArrayLiteral;
@@ -102,4 +103,18 @@ public List<FunctionSignature> getSignatures() {
  public Expression resultForEmptyInput() {
  return new ArrayLiteral(new ArrayList<>(), this.getDataType());
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
+
+ @Override
+ public void checkLegalityBeforeTypeCoercion() {
+ if (arity() == 2 && !getArgument(1).isConstant()) {
+ throw new AnalysisException(
+ "collect_list requires second parameter must be a constant: "
+ + this.toSql());
+ }
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/ExponentialMovingAverage.java

@@ -119,4 +119,10 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(1), getArgument(2)) : ImmutableList.of();
+ }
+
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/GroupBitmapXor.java

@@ -18,6 +18,7 @@
 package org.apache.doris.nereids.trees.expressions.functions.agg;
 
 import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
 import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
@@ -55,8 +56,7 @@ public GroupBitmapXor(boolean distinct, Expression arg) {
  }
 
  private GroupBitmapXor(boolean distinct, boolean alwaysNullable, Expression arg0) {
- /* distinct is meaningless to bitmap */
- super("group_bitmap_xor", false, alwaysNullable, arg0);
+ super("group_bitmap_xor", distinct, alwaysNullable, arg0);
  }
 
  /** constructor for withChildren and reuse signature */
@@ -92,4 +92,11 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public void checkLegalityBeforeTypeCoercion() {
+ if (distinct) {
+ throw new AnalysisException("group_bitmap_xor does not support DISTINCT");
+ }
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/Histogram.java

@@ -113,4 +113,9 @@ public List<FunctionSignature> getSignatures() {
  public Expression resultForEmptyInput() {
  return new VarcharLiteral("{\"num_buckets\":0,\"buckets\":[]}");
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/MapAgg.java

@@ -53,7 +53,7 @@ public MapAgg(Expression arg0, Expression arg1) {
  /**
  * constructor with 2 arguments.
  */
-private MapAgg(boolean distinct, Expression arg0, Expression arg1) {
+public MapAgg(boolean distinct, Expression arg0, Expression arg1) {
  super("map_agg_v1", distinct, arg0, arg1);
  }
 

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/MapAggV2.java

@@ -53,7 +53,7 @@ public MapAggV2(Expression arg0, Expression arg1) {
  /**
  * constructor with 2 arguments.
  */
-private MapAggV2(boolean distinct, Expression arg0, Expression arg1) {
+public MapAggV2(boolean distinct, Expression arg0, Expression arg1) {
  super("map_agg_v2", distinct, arg0, arg1);
  }
 

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/Percentile.java

@@ -107,4 +107,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/PercentileApprox.java

@@ -121,4 +121,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/PercentileApproxWeighted.java

@@ -125,4 +125,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0), getArgument(1)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/PercentileArray.java

@@ -111,4 +111,9 @@ public List<FunctionSignature> getSignatures() {
  public Expression resultForEmptyInput() {
  return new ArrayLiteral(new ArrayList<>(), this.getDataType());
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/PercentileReservoir.java

@@ -107,4 +107,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/TopN.java

@@ -134,4 +134,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/TopNArray.java

@@ -137,4 +137,9 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
  public List<FunctionSignature> getSignatures() {
  return SIGNATURES;
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0)) : ImmutableList.of();
+ }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/agg/TopNWeighted.java

@@ -219,4 +219,9 @@ public void checkLegalityAfterRewrite() {
  + this.toSql());
  }
  }
+
+ @Override
+ public List<Expression> getDistinctArguments() {
+ return distinct ? ImmutableList.of(getArgument(0), getArgument(1)) : ImmutableList.of();
+ }
 }

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

@@ -148,6 +148,17 @@ public static boolean containsCountDistinctMultiExpr(LogicalAggregate<? extends
  expr instanceof Count && ((Count) expr).isDistinct() && expr.arity() > 1);
  }
 
+ /** e.g. Aggregation with avg(distinct a)(not support multiDistinct) or count(distinct a,b) will return true*/
+ public static boolean containsNotSupportMultiDistinctFunction(LogicalAggregate<? extends Plan> aggregate) {
+ return ExpressionUtils.deapAnyMatch(aggregate.getOutputExpressions(), expr -> {
+ if (expr instanceof AggregateFunction && ((AggregateFunction) expr).isDistinct()) {
+ return !(expr instanceof SupportMultiDistinct)
+ || expr instanceof Count && ((Count) expr).isDistinct() && expr.arity() > 1;
+ }
+ return false;
+ });
+ }
+
  /** count agg function distinct group, up to 2*/
  public static int distinctArgumentGroupCountUpToTwo(Aggregate<? extends Plan> aggregate) {
  Set<Expression> distinctArgumentGroup = null;