[Fix](nereids) fix NormalizeAgg, change the upper project projections rewrite logic (#36161) (#36622)

cherry-pick #36161 to branch-2.1

NormalizeAggregate rewrite logic has a bug, for sql like this:

SELECT
	CASE
		1 WHEN CAST( NULL AS SIGNED ) THEN NULL
		WHEN COUNT( DISTINCT CAST( NULL AS SIGNED ) ) THEN NULL
		ELSE null
	END ;

This is the plan after NormalizeAggregate, the LogicalAggregate only
output `count(DISTINCT cast(NULL as SIGNED))`#3, do not output cast(NULL
as SIGNED)#2, but the upper project use cast(NULL as SIGNED)#2, so Doris
report error "cast(NULL as SIGNED) not in aggregate's output".

LogicalResultSink[29] ( outputExprs=[__case_when_0#1] ) +--LogicalProject[26] ( distinct=false, projects=[CASE WHEN (1 = cast(NULL as SIGNED)#2) THEN NULL WHEN (1 = count(DISTINCT cast(NULL as SIGNED))#3) THEN NULL ELSE NULL END AS `CASE WHEN (1 = cast(NULL as SIGNED)) THEN NULL WHEN (1 = count(DISTINCT cast(NULL as SIGNED))) THEN NULL ELSE NULL END`#1], excepts=[] )
   +--LogicalAggregate[25] ( groupByExpr=[], outputExpr=[count(DISTINCT cast(NULL as SIGNED)#2) AS `count(DISTINCT cast(NULL as SIGNED))`#3], hasRepeat=false )
      +--LogicalProject[24] ( distinct=false, projects=[cast(NULL as SIGNED) AS `cast(NULL as SIGNED)`#2], excepts=[] )
         +--LogicalOneRowRelation ( projects=[0 AS `0`#0] )

The problem is that the cast(NULL as SIGNED)#2 should not outputted by
LogicalAggregate, cast(NULL as SIGNED) should be computed in
LogicalProject.
This pr change the upper project projections rewrite logic:
aggregateOutputs is rewritten and become the upper-level LogicalProject
projections. During the rewriting process, the expressions inside the
agg function can be rewritten with expressions in aggregate function
arguments and group by expressions, but the ones outside the agg
function can only be rewritten with group by expressions.

---------

Co-authored-by: moailing <moailing@selectdb.com>
This commit is contained in:
feiniaofeiafei
2024-06-27 12:17:18 +08:00
committed by GitHub
parent 23cf494b48
commit a8e9c89dc6
3 changed files with 77 additions and 9 deletions

View File

@ -190,10 +190,14 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
// 1. group by exprs
// 2. trivialAgg children
// 3. trivialAgg input slots
Set<Expression> allPushDownExprs =
Sets.union(groupingByExprs, Sets.union(needPushSelf, needPushInputSlots));
NormalizeToSlotContext bottomSlotContext =
NormalizeToSlotContext.buildContext(existsAlias, allPushDownExprs);
// We need to distinguish between expressions in aggregate function arguments and group by expressions.
NormalizeToSlotContext groupByExprContext = NormalizeToSlotContext.buildContext(existsAlias, groupingByExprs);
Set<Alias> existsAliasAndGroupByAlias = getExistsAlias(existsAlias, groupByExprContext.getNormalizeToSlotMap());
Set<Expression> argsOfAggFuncNeedPushDown = Sets.union(needPushSelf, needPushInputSlots);
NormalizeToSlotContext argsOfAggFuncNeedPushDownContext = NormalizeToSlotContext
.buildContext(existsAliasAndGroupByAlias, argsOfAggFuncNeedPushDown);
NormalizeToSlotContext bottomSlotContext = argsOfAggFuncNeedPushDownContext.mergeContext(groupByExprContext);
Set<NamedExpression> pushedGroupByExprs =
bottomSlotContext.pushDownToNamedExpression(groupingByExprs);
Set<NamedExpression> pushedTrivialAggChildren =
@ -256,8 +260,12 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
aggregate.withNormalized(normalizedGroupExprs, normalizedAggOutput, bottomPlan);
// create upper projects by normalize all output exprs in old LogicalAggregate
// In aggregateOutput, the expressions inside the agg function can be rewritten
// with expressions in aggregate function arguments and group by expressions,
// but the ones outside the agg function can only be rewritten with group by expressions.
// After the above two rewrites are completed, use aggregate output agg functions to rewrite.
List<NamedExpression> upperProjects = normalizeOutput(aggregateOutput,
bottomSlotContext, normalizedAggFuncsToSlotContext);
groupByExprContext, argsOfAggFuncNeedPushDownContext, normalizedAggFuncsToSlotContext);
// create a parent project node
LogicalProject<Plan> project = new LogicalProject<>(upperProjects, newAggregate);
@ -302,11 +310,18 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
}
private List<NamedExpression> normalizeOutput(List<NamedExpression> aggregateOutput,
NormalizeToSlotContext groupByToSlotContext, NormalizeToSlotContext normalizedAggFuncsToSlotContext) {
NormalizeToSlotContext groupByToSlotContext, NormalizeToSlotContext argsOfAggFuncNeedPushDownContext,
NormalizeToSlotContext normalizedAggFuncsToSlotContext) {
// build upper project, use two context to do pop up, because agg output maybe contain two part:
// group by keys and agg expressions
List<NamedExpression> upperProjects = groupByToSlotContext
.normalizeToUseSlotRefWithoutWindowFunction(aggregateOutput);
// group by keys and agg expressions
List<NamedExpression> upperProjects = new ArrayList<>();
for (Expression expr : aggregateOutput) {
Expression rewrittenExpr = expr.rewriteDownShortCircuit(
e -> normalizeAggFuncChildren(
argsOfAggFuncNeedPushDownContext, e));
upperProjects.add((NamedExpression) rewrittenExpr);
}
upperProjects = groupByToSlotContext.normalizeToUseSlotRefWithoutWindowFunction(upperProjects);
upperProjects = normalizedAggFuncsToSlotContext.normalizeToUseSlotRefWithoutWindowFunction(upperProjects);
Builder<NamedExpression> builder = new ImmutableList.Builder<>();
@ -338,4 +353,28 @@ public class NormalizeAggregate implements RewriteRuleFactory, NormalizeToSlot {
slots.addAll(ExpressionUtils.getInputSlotSet(expressions));
return slots;
}
private Set<Alias> getExistsAlias(Set<Alias> originAliases,
Map<Expression, NormalizeToSlotTriplet> groupingExprMap) {
Set<Alias> existsAlias = Sets.newHashSet();
existsAlias.addAll(originAliases);
for (NormalizeToSlotTriplet triplet : groupingExprMap.values()) {
if (triplet.pushedExpr instanceof Alias) {
Alias alias = (Alias) triplet.pushedExpr;
existsAlias.add(alias);
}
}
return existsAlias;
}
private Expression normalizeAggFuncChildren(NormalizeToSlotContext context, Expression expr) {
if (expr instanceof AggregateFunction) {
AggregateFunction function = (AggregateFunction) expr;
List<Expression> normalizedRealExpressions = context.normalizeToUseSlotRef(function.getArguments());
function = function.withChildren(normalizedRealExpressions);
return function;
} else {
return expr;
}
}
}