feiniaofeiafei commented on code in PR #63690:
URL: https://github.com/apache/doris/pull/63690#discussion_r3402095408
##########
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriter.java:
##########
@@ -534,21 +656,382 @@ public Plan visitLogicalRelation(LogicalRelation
relation, PushDownAggContext co
}
private Plan genAggregate(Plan child, PushDownAggContext context) {
- if (context.isValid() && checkStats(child, context)) {
+ if (isPushDisabledByVariable(context)) {
+ context.getBilateralState().registerNoCountSlot(child);
+ return child;
+ }
+ if (checkStats(child, context) || isPushEnabledByVariable(context)) {
List<NamedExpression> aggOutputExpressions = new ArrayList<>();
for (AggregateFunction func : context.getAggFunctions()) {
aggOutputExpressions.add(context.getAliasMap().get(func));
}
+ Alias countStarAlias = null;
+ boolean countStarAlreadyProjected = false;
+ Count countStar = new Count();
+ if (context.getAliasMap().containsKey(countStar)) {
+ countStarAlias = context.getAliasMap().get(countStar);
+ countStarAlreadyProjected = true;
+ } else {
+ countStarAlias = new Alias(countStar,
+ "cnt" +
context.getCascadesContext().getStatementContext().generateColumnName());
+ }
aggOutputExpressions.addAll(context.getGroupKeys());
+ if (countStarAlias != null && !countStarAlreadyProjected) {
+ aggOutputExpressions.add(countStarAlias);
+ }
LogicalAggregate genAgg = new
LogicalAggregate(context.getGroupKeys(), aggOutputExpressions, child);
NormalizeAggregate normalizeAggregate = new NormalizeAggregate();
- return normalizeAggregate.normalizeAgg(genAgg, Optional.empty(),
+ Plan normalized = normalizeAggregate.normalizeAgg(genAgg,
Optional.empty(),
context.getCascadesContext());
+
+ for (AggregateFunction func : context.getAggFunctions()) {
+ Alias a = context.getAliasMap().get(func);
+
context.getBilateralState().registerPushedAggFuncSlot(a.getExprId(),
a.toSlot());
+ }
+
+ if (countStarAlias != null) {
+ context.getBilateralState().registerCountSlot(normalized,
countStarAlias.toSlot());
+ } else {
+ context.getBilateralState().registerNoCountSlot(normalized);
+ }
+ return normalized;
} else {
+ context.getBilateralState().registerNoCountSlot(child);
return child;
}
}
+ // Build the canonical project above a rewritten join after
eager-aggregation pushdown.
+ // Responsibilities:
+ // 1. Restore the outputs expected by the parent rollup. If a join side
has a childContext, materialize
+ // that side's aggregate current values and group keys; otherwise
forward the original join outputs.
+ // 2. For inner joins, recover join multiplicity by multiplying
non-MIN/MAX aggregate current values by
+ // the opposite side's count slot when that side contributes rows to
the parent aggregate.
+ // 3. Append and register a synthetic join-count slot `cnt` (logical jcnt)
for upper-level rollup.
+ //
+ // The examples below are schematic. The real project may keep extra
forwarded slots such as join keys.
+ //
+ // Inner join + sum, single-side rewrite:
+ // Before:
+ // agg(sum(t1.a), sum(t2.a), gby t2.k)
+ // -> inner join(k = k)
+ // -> scan(t1)
+ // -> scan(t2)
+ // After:
+ // agg(sum(s1), sum(s2), gby t2.k)
+ // -> project(s1, t2.a * cnt1 as s2, t2.k, cnt1 as cnt)
+ // -> inner join(k = k)
+ // -> agg(sum(t1.a) as s1, count(*) as cnt1, gby k)
+ // -> scan(t1)
+ // -> scan(t2)
+ //
+ // Inner join + sum, bilateral rewrite:
+ // Before:
+ // agg(sum(t1.a), sum(t2.a), gby t2.k)
+ // -> inner join(k = k)
+ // -> scan(t1)
+ // -> scan(t2)
+ // After:
+ // agg(sum(s1'), sum(s2'), gby t2.k)
+ // -> project(s1 * cnt2 as s1', s2 * cnt1 as s2', t2.k, cnt1 * cnt2
as cnt)
+ // -> inner join(k = k)
+ // -> agg(sum(t1.a) as s1, count(*) as cnt1, gby k)
+ // -> scan(t1)
+ // -> agg(sum(t2.a) as s2, count(*) as cnt2, gby k)
+ // -> scan(t2)
+ //
+ // Inner join + count(col), single-side rewrite:
+ // Before:
+ // agg(count(t1.a), count(t2.a), gby t2.k)
+ // -> inner join(k = k)
+ // -> scan(t1)
+ // -> scan(t2)
+ // After:
+ // agg(sum0(c1), sum0(c2), gby t2.k)
+ // -> project(c1, if(t2.a is null, 0, 1) * cnt1 as c2, t2.k, cnt1 as
cnt)
+ // -> inner join(k = k)
+ // -> agg(count(t1.a) as c1, count(*) as cnt1, gby k)
+ // -> scan(t1)
+ // -> scan(t2)
+ //
+ // Inner join + count(col), bilateral rewrite:
+ // Before:
+ // agg(count(t1.a), count(t2.a), gby t2.k)
+ // -> inner join(k = k)
+ // -> scan(t1)
+ // -> scan(t2)
+ // After:
+ // agg(sum0(c1'), sum0(c2'), gby t2.k)
+ // -> project(c1 * cnt2 as c1', c2 * cnt1 as c2', t2.k, cnt1 * cnt2
as cnt)
+ // -> inner join(k = k)
+ // -> agg(count(t1.a) as c1, count(*) as cnt1, gby k)
+ // -> scan(t1)
+ // -> agg(count(t2.a) as c2, count(*) as cnt2, gby k)
+ // -> scan(t2)
+ // For count(*), the current row value is 1 instead of if(col is null,
0, 1).
+ //
+ // Semi/anti join:
+ // The project does not multiply by the opposite-side count
+ //
+ // Outer join:
+ // Aggregate outputs are not multiplied by the opposite-side count
either; only `cnt` changes:
+ // left outer join with left push -> project(s1, t2.k, cnt1 as cnt)
+ // right outer join with left push -> project(s1, t2.k, nvl(cnt1, 1)
as cnt)
+ private Plan buildCanonicalJoinProject(LogicalJoin<? extends Plan, ?
extends Plan> join, PushDownAggContext context,
+ Optional<PushDownAggContext> leftChildContext,
Optional<PushDownAggContext> rightChildContext,
+ Optional<Slot> leftCountSlot, Optional<Slot> rightCountSlot) {
+ List<NamedExpression> projections = new ArrayList<>();
+ Set<ExprId> outputIds = new HashSet<>();
+ boolean remainLeft = join.getJoinType().isRemainLeftJoin();
+ boolean remainRight = join.getJoinType().isRemainRightJoin();
+ boolean shouldAdjustLeft =
shouldUseJoinOppositeCntAdjustAggOutput(join, leftChildContext, rightCountSlot);
+ boolean shouldAdjustRight =
shouldUseJoinOppositeCntAdjustAggOutput(join, rightChildContext, leftCountSlot);
+
+ if (remainLeft) {
+ appendJoinSideOutputs(projections, outputIds, join.left(),
leftChildContext, context,
+ rightCountSlot, shouldAdjustLeft);
+ }
+ if (remainRight) {
+ appendJoinSideOutputs(projections, outputIds, join.right(),
rightChildContext, context,
+ leftCountSlot, shouldAdjustRight);
+ }
+
+ Optional<Expression> joinCount = computeJoinCount(join,
leftChildContext, rightChildContext,
+ leftCountSlot, rightCountSlot);
+ Optional<Slot> projectedCountSlot = Optional.empty();
+ if (joinCount.isPresent()) {
+ Alias countAlias = new Alias(joinCount.get(),
+ "joinCnt" +
context.getCascadesContext().getStatementContext().generateColumnName());
+ projections.add(countAlias);
+ projectedCountSlot = Optional.of(countAlias.toSlot());
+ }
+ LogicalProject<Plan> project = new LogicalProject<>(projections, join);
+ if (projectedCountSlot.isPresent()) {
+ context.getBilateralState().registerCountSlot(project,
+ (Slot) project.getOutput().get(project.getOutput().size()
- 1));
+ } else {
+ context.getBilateralState().registerNoCountSlot(project);
+ }
+ return project;
+ }
+
+ private void appendJoinSideOutputs(List<NamedExpression> projections,
Set<ExprId> outputIds, Plan originalSide,
+ Optional<PushDownAggContext> childContext, PushDownAggContext
parentContext,
+ Optional<Slot> oppositeCountSlot, boolean shouldAdjustOutput) {
+ if (childContext.isPresent()) {
+ for (AggregateFunction aggFunc :
childContext.get().getAggFunctions()) {
+ NamedExpression aggOutput = shouldAdjustOutput
+ ? adjustAggOutputUseOppositeCountOnJoin(aggFunc,
parentContext, oppositeCountSlot)
+ : buildAggOutputWithoutJoinAdjustment(aggFunc,
parentContext);
+ appendProjectionIfAbsent(projections, outputIds, aggOutput);
+ }
+ for (SlotReference groupKey : childContext.get().getGroupKeys()) {
+ appendProjectionIfAbsent(projections, outputIds, groupKey);
+ }
+ } else {
+ for (Slot slot : originalSide.getOutput()) {
+ appendProjectionIfAbsent(projections, outputIds, slot);
+ }
+ }
+ }
+
+ private void appendProjectionIfAbsent(List<NamedExpression> projections,
Set<ExprId> outputIds,
+ NamedExpression expression) {
+ if (outputIds.add(expression.getExprId())) {
+ projections.add(expression);
+ }
+ }
+
+ private boolean shouldUseJoinOppositeCntAdjustAggOutput(LogicalJoin<?
extends Plan, ? extends Plan> join,
+ Optional<PushDownAggContext> childContext, Optional<Slot>
oppositeCountSlot) {
+ return join.getJoinType().isInnerOrCrossJoin() &&
childContext.isPresent() && oppositeCountSlot.isPresent();
+ }
+
+ private Optional<Expression> computeJoinCount(LogicalJoin<? extends Plan,
? extends Plan> join,
+ Optional<PushDownAggContext> leftChildContext,
Optional<PushDownAggContext> rightChildContext,
+ Optional<Slot> leftCountSlot, Optional<Slot> rightCountSlot) {
+ JoinType joinType = join.getJoinType();
+ if (joinType.isInnerJoin()) {
+ if (leftCountSlot.isPresent() && rightCountSlot.isPresent()) {
+ return Optional.of(ExpressionUtils.rebuildSignature(
+ new Multiply(leftCountSlot.get(),
rightCountSlot.get())));
+ } else if (leftCountSlot.isPresent()) {
+ return Optional.of(leftCountSlot.get());
+ } else if (rightCountSlot.isPresent()) {
+ return Optional.of(rightCountSlot.get());
+ }
+ return Optional.empty();
+ }
+ if (joinType.isLeftOuterJoin()) {
+ if (leftChildContext.isPresent()) {
+ return leftCountSlot.map(cnt -> (Expression) cnt);
+ }
+ if (rightChildContext.isPresent()) {
+ return rightCountSlot.map(cnt -> (Expression)
ExpressionUtils.rebuildSignature(
+ new Nvl(cnt, BigIntLiteral.of(1))));
+ }
Review Comment:
I think the explicit Cast is not necessary here. The Nvl expression is now
built through TypeCoercionUtils.processBoundFunction(), so its normal
BoundFunction type coercion path will choose the common signature for cnt and
the literal and insert casts if needed.
For the current count(*) type, cnt and BigIntLiteral.of(1) are both BIGINT,
so coercion is identity. If the count return type changes in the future,
processBoundFunction() should still handle the type alignment through Nvl's
custom signature/type coercion path.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
