Skip to content

[fix](nereids) eager aggregation: never substitute a literal inside the If condition #69582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
raghav-reglobe wants to merge 1 commit into
base: master
Choose a base branch
from
+98 −7
Open

[fix](nereids) eager aggregation: never substitute a literal inside the If condition #69582

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 44 additions & 7 deletions ...ain/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/PushDownAggregation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,12 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, JobConte
boolean hasDecomposedAggIf = false;
boolean containsNullToNonNull = false;
Map<NamedExpression, List<AggregateFunction>> aggFunctionsForOutputExpressions = Maps.newHashMap();
// A decomposed Sum/Max/Min(If(cond, a, b)) is rebuilt at the top as func(If(cond, pushedA, pushedB)).
// The rebuild must target the two VALUE positions only: a literal that also appears inside `cond`
// (sum(if(flag > 0, 1, 0)) -> the 0) must never be swapped for the pushed BIGINT slot.
Map<AggregateFunction, If> decomposedIfBodies = Maps.newHashMap();
Map<AggregateFunction, AggregateFunction> decomposedTrueParts = Maps.newHashMap();
Map<AggregateFunction, AggregateFunction> decomposedFalseParts = Maps.newHashMap();
Set<AggregateFunction> allAggFunctions = agg.getAggregateFunctions();
boolean hasDistinctAgg = allAggFunctions.stream().anyMatch(AggregateFunction::isDistinct);
if (hasDistinctAgg) {
Expand Down Expand Up @@ -171,11 +177,14 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, JobConte
AggregateFunction aggTrue = (AggregateFunction) aggFunction.withChildren(body.getTrueValue());
aggFunctions.add(aggTrue);
funcs.add(aggTrue);
decomposedIfBodies.put(aggFunction, body);
decomposedTrueParts.put(aggFunction, aggTrue);
if (!(body.getFalseValue() instanceof NullLiteral)) {
AggregateFunction aggFalse =
(AggregateFunction) aggFunction.withChildren(body.getFalseValue());
aggFunctions.add(aggFalse);
funcs.add(aggFalse);
decomposedFalseParts.put(aggFunction, aggFalse);
}
groupKeys.addAll(body.getCondition().getInputSlots()
.stream().map(slot -> (SlotReference) slot).collect(Collectors.toList()));
Expand Down Expand Up @@ -241,19 +250,36 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, JobConte
// for min(A), replaceMap: A->minA
// for sum(A), replaceMap: A->sumA
// for count(A), replaceMap: count(A)->sum(countA), because count needs rollup to sum
// Keys are whole aggregate functions: ExpressionUtils.replace walks top-down and
// short-circuits at a match, so nothing inside an If condition (or beside the
// function) is ever substituted. Keying on func.child(0) replaced EVERY equal
// expression in the output - for sum(if(flag > 0, 1, 0)) the condition's 0 became
// the pushed BIGINT sum slot: a BE type error on TINYINT/BOOLEAN columns.
Map<Expression, Expression> replaceMap = new HashMap<>();
List<AggregateFunction> relatedAggFunc = aggFunctionsForOutputExpressions.get(ne);
for (AggregateFunction func : relatedAggFunc) {
Alias pushedAlias = pushDownContext.getAliasMap().get(func);
ExprId pushId = pushedAlias.getExprId();
if (!state.hasAggFuncOutput(pushId)) {
for (Object obj : ne.collect(AggregateFunction.class::isInstance)) {
AggregateFunction func = (AggregateFunction) obj;
If body = decomposedIfBodies.get(func);
if (body != null) {
Expression trueValue = pushedSlotOf(
decomposedTrueParts.get(func), pushDownContext, state);
AggregateFunction falsePart = decomposedFalseParts.get(func);
Expression falseValue = falsePart == null ? body.getFalseValue()
: pushedSlotOf(falsePart, pushDownContext, state);
if (trueValue == null || falseValue == null) {
return agg;
}
replaceMap.put(func, func.withChildren(
new If(body.getCondition(), trueValue, falseValue)));
continue;
}
Expression value = pushedSlotOf(func, pushDownContext, state);
if (value == null) {
continue;
}
Expression value = state.getPushedAggFuncSlot(pushId);
if (func instanceof Count) {
replaceMap.put(func, new Sum0(value));
} else if (func.arity() > 0) {
replaceMap.put(func.child(0), value);
replaceMap.put(func, func.withChildren(value));
}
}
NamedExpression replaceAliasExpr =
Expand All @@ -275,6 +301,17 @@ public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, JobConte
return agg;
}

/** The slot the pushed-down copy of {@code func} produces below the join, or null when it has none. */
private static Expression pushedSlotOf(AggregateFunction func, PushDownAggContext pushDownContext,
BilateralState state) {
Alias pushedAlias = func == null ? null : pushDownContext.getAliasMap().get(func);
if (pushedAlias == null) {
return null;
}
ExprId pushId = pushedAlias.getExprId();
return state.hasAggFuncOutput(pushId) ? state.getPushedAggFuncSlot(pushId) : null;
}

private Optional<Set<SlotReference>> getCommonDistinctKeys(Set<AggregateFunction> aggFunctions) {
Set<SlotReference> commonDistinctKeys = null;
for (AggregateFunction aggFunction : aggFunctions) {
Expand Down
54 changes: 54 additions & 0 deletions ...st/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,13 @@
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.rules.exploration.join.JoinReorderContext;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.ComparisonPredicate;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
import org.apache.doris.nereids.trees.expressions.functions.scalar.If;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.plans.JoinType;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
Expand All @@ -40,7 +44,9 @@
import org.junit.jupiter.api.Test;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

class EagerAggRewriterTest extends TestWithFeService implements MemoPatternMatchSupported {
Expand Down Expand Up @@ -69,6 +75,15 @@ protected void runBeforeAll() throws Exception {
+ ")\n"
+ "DUPLICATE KEY(id3)\n"
+ "DISTRIBUTED BY HASH(id3) BUCKETS 10\n"
+ "PROPERTIES (\"replication_num\" = \"1\")\n",
"CREATE TABLE IF NOT EXISTS t4 (\n"
+ " id4 int not null,\n"
+ " flag boolean,\n"
+ " tiny tinyint,\n"
+ " qty bigint\n"
+ ")\n"
+ "DUPLICATE KEY(id4)\n"
+ "DISTRIBUTED BY HASH(id4) BUCKETS 10\n"
+ "PROPERTIES (\"replication_num\" = \"1\")\n"
);
connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION");
Expand Down Expand Up @@ -112,6 +127,45 @@ void testNotPushAggCaseWhenToNullableSideOfOuterJoin() {
}
}

@Test
void testDecomposedIfKeepsConditionLiteral() {
// sum(if(flag > 0, 1, 0)): the condition's literal 0 is the same expression as the false-branch value.
// The eager-aggregation rebuild must substitute the two branch VALUES with the pushed slots and leave the
// comparison operand alone. Before the fix the plan read `flag > <pushed bigint slot>`: a BE
// "not same type in function gt" on BOOLEAN/TINYINT columns (the only columns whose comparison literal
// stays equal to the branch literal after type coercion), and for max/min the constant slot by luck.
connectContext.getSessionVariable().setEagerAggregationMode(1);
connectContext.getSessionVariable().setDisableJoinReorder(true);
try {
String sql = "select sum(if(t4.flag > 0, 1, 0)), max(if(t4.tiny = 1, 1, 0)),"
+ " sum(if(t4.qty > 1, 1, 0)), t2.id2 from t4 join t2 on t4.id4 = t2.id2 group by t2.id2";
PlanChecker checker = PlanChecker.from(connectContext).analyze(sql).rewrite();
// the rewrite still fires: the decomposed parts sit below the join
checker.matches(logicalAggregate(logicalProject(logicalJoin(logicalAggregate(), any()))));
LogicalAggregate<?> top = (LogicalAggregate<?>) checker.getPlan()
.collectFirst(LogicalAggregate.class::isInstance).get();
List<If> ifs = new ArrayList<>();
for (NamedExpression ne : top.getOutputExpressions()) {
for (Object obj : ne.collect(If.class::isInstance)) {
ifs.add((If) obj);
}
}
Assertions.assertEquals(3, ifs.size(), top.treeString());
for (If body : ifs) {
// both branch values became the pushed slots ...
Assertions.assertTrue(body.getTrueValue() instanceof SlotReference, body.toSql());
Assertions.assertTrue(body.getFalseValue() instanceof SlotReference, body.toSql());
// ... and the condition still compares against its literal, never against a pushed slot
ComparisonPredicate cmp = (ComparisonPredicate) body.getCondition();
Assertions.assertTrue(cmp.child(1) instanceof Literal, cmp.toSql());
Assertions.assertFalse(cmp.child(1).anyMatch(SlotReference.class::isInstance), cmp.toSql());
}
} finally {
connectContext.getSessionVariable().setEagerAggregationMode(0);
connectContext.getSessionVariable().setDisableJoinReorder(false);
}
}

@Test
void testPushDownCount() {
// Test count pushdown: count(a) should be pushed down and
Expand Down
Loading