diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/PushDownAggregation.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/PushDownAggregation.java index 07aba94b2318a0..3c260ce0254463 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/PushDownAggregation.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/PushDownAggregation.java @@ -119,6 +119,12 @@ public Plan visitLogicalAggregate(LogicalAggregate agg, JobConte boolean hasDecomposedAggIf = false; boolean containsNullToNonNull = false; Map> 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 decomposedIfBodies = Maps.newHashMap(); + Map decomposedTrueParts = Maps.newHashMap(); + Map decomposedFalseParts = Maps.newHashMap(); Set allAggFunctions = agg.getAggregateFunctions(); boolean hasDistinctAgg = allAggFunctions.stream().anyMatch(AggregateFunction::isDistinct); if (hasDistinctAgg) { @@ -171,11 +177,14 @@ public Plan visitLogicalAggregate(LogicalAggregate 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())); @@ -241,19 +250,36 @@ public Plan visitLogicalAggregate(LogicalAggregate 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 replaceMap = new HashMap<>(); - List 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 = @@ -275,6 +301,17 @@ public Plan visitLogicalAggregate(LogicalAggregate 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> getCommonDistinctKeys(Set aggFunctions) { Set commonDistinctKeys = null; for (AggregateFunction aggFunction : aggFunctions) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriterTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriterTest.java index c22cdad81ce220..04d820f55a7ef5 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriterTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/eageraggregation/EagerAggRewriterTest.java @@ -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; @@ -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 { @@ -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"); @@ -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 > `: 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 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