From db7b8a1a20fae53f450d4156bd0e6c893c10503a Mon Sep 17 00:00:00 2001 From: 924060929 Date: Tue, 18 Aug 2026 19:55:10 +0800 Subject: [PATCH] branch-4.2: [fix](local shuffle) Require hash input for distinct finalize agg without group keys #66570 Cherry-picked from #66570 --- .../apache/doris/planner/AggregationNode.java | 75 +++-- .../planner/LocalShuffleNodeCoverageTest.java | 295 ++++++++++++++++++ .../doris/qe/LocalExchangePlannerTest.java | 106 +++++++ .../test_local_shuffle_rqg_bugs.out | 6 + .../test_local_shuffle_rqg_bugs.groovy | 49 +++ 5 files changed, 511 insertions(+), 20 deletions(-) create mode 100644 regression-test/data/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.out diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/AggregationNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/AggregationNode.java index 8bd8a87b69899d..50818d00fa49eb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/AggregationNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/AggregationNode.java @@ -302,6 +302,8 @@ public Pair enforceAndDeriveLocalExchange( // PR #62438: when false, non-finalize agg falls back to BE base class. boolean enableLeBeforeAgg = sessionVariable.enableLocalExchangeBeforeAgg; boolean hasKeys = !aggInfo.getGroupingExprs().isEmpty(); + boolean selfOrInheritedShuffled = translatorContext.hasShuffleForCorrectnessAncestor(this) + || requiresShuffleForCorrectness(); // Each branch mirrors the corresponding BE operator's required_data_distribution() // check order 1:1. The helper baseClassRequire() expands BE's base class behavior. @@ -355,7 +357,16 @@ public Pair enforceAndDeriveLocalExchange( // early return also catches FIRST_MERGE, dropping the HASH requirement and // causing wrong-result (e.g. PASSTHROUGH over serial child breaks the // group-by-key invariant — DORIS-25413). - if (!hasKeys) { + if (!hasPartitionRequirement(selfOrInheritedShuffled)) { + // No effective partition key (no group keys, and no child distribute + // exprs set for a DISTINCT / followed-by-shuffle agg): the input + // distribution is irrelevant. A finalize agg with an effective key + // emits per-instance scalar values (sum0(multi_distinct_count(...)) + // above) that the parent sums, so same-key rows must stay in a single + // instance — this mirrors BE's `_partition_exprs` exactly, and keeps + // a directly called multi_distinct_count(col) (no distribute exprs) + // on the no-requirement path instead of collapsing it onto a zero-key + // HASH exchange. requireChild = needsFinalize ? LocalExchangeTypeRequire.noRequire() : baseClassRequire(connectContext); @@ -368,13 +379,16 @@ public Pair enforceAndDeriveLocalExchange( // FIRST_MERGE (correctness) or finalize+colocate → HASH. requireChild = parentRequire.autoRequireHash(); } else if (hasPartitionExprs(parentRequire)) { - // FE-only heuristic: finalize non-colocate with parent hash requirement - // → inherit parent's specific hash type. + // finalize non-colocate with a parent hash requirement → inherit the + // parent's specific hash type. requireChild = parentRequire.autoRequireHash(); } else { - // FE-only heuristic: finalize non-colocate without parent hash → skip - // LE (child Exchange already provides hash distribution). - requireChild = LocalExchangeTypeRequire.noRequire(); + // finalize non-colocate without a parent hash requirement: the input + // must still be key-aligned (group/distinct key), so require HASH + // explicitly instead of trusting the child's distribution. When the + // child already provides hash distribution, satisfy() passes and no + // LE is inserted, so this is safe and free in the common case. + requireChild = LocalExchangeTypeRequire.requireHash(); } } @@ -391,6 +405,34 @@ private LocalExchangeTypeRequire baseClassRequire(ConnectContext connectContext) : LocalExchangeTypeRequire.noRequire(); } + /** + * Whether this agg needs key-aligned (hash-partitioned) input from its child. + * Mirrors BE AggSinkOperatorX::update_operator's `_partition_exprs` exactly: + * non-empty grouping exprs, or the child distribute exprs when the plan set + * them for a DISTINCT (or followed-by-shuffle) agg. The test is on the + * *effective* key, not the function name: a directly called + * multi_distinct_count(col) has neither distribute exprs nor grouping exprs, + * so it stays on the no-requirement path — a zero-key HASH exchange would + * collapse the whole input onto one task per BE. A finalize agg with an + * effective key emits per-instance scalar values (the + * sum0(multi_distinct_count(...)) above) that the parent sums, so same-key + * rows must stay in a single instance. + */ + private boolean hasPartitionRequirement(boolean followedByShuffled) { + return !getLocalExchangeDistributeExprs(0, followedByShuffled).isEmpty(); + } + + private boolean hasDistinctAggregate() { + // Multi-distinct aggregates are detected by function name. Nereids rewrites + // count/sum/group_concat(distinct ...) into dedicated MultiDistinct* functions + // constructed with distinct=false, so by this legacy FunctionCallExpr layer + // isDistinct() is already false and the function name is the only signal. + return aggInfo.getAggregateExprs().stream() + .map(FunctionCallExpr::getFnName) + .map(name -> name.getFunction()) + .anyMatch(name -> name.startsWith("multi_distinct_")); + } + @Override protected List getSemanticPartitionExprs() { return aggInfo.getGroupingExprs(); @@ -406,18 +448,7 @@ protected List getLocalExchangeDistributeExprs(int childIndex, boolean fol // chain scatters same-group rows across N instances, leaving partial_preagg essentially a // no-op and breaking row-arrival order at downstream merge-finalize (e.g. group_concat). List childDist = getChildDistributeExprList(childIndex); - // Multi-distinct aggregates are detected by function name. Nereids rewrites - // count/sum(distinct ...) into dedicated MultiDistinct* functions constructed with - // distinct=false and a "multi_distinct_" name, so by this legacy FunctionCallExpr layer - // isDistinct() is already false and the function name is the only remaining signal — - // there is no structural flag to test here. - boolean hasDistinct = aggInfo.getAggregateExprs().stream() - .map(FunctionCallExpr::getFnName) - .filter(name -> name != null) - .map(name -> name.getFunction()) - .filter(name -> name != null) - .anyMatch(name -> name.startsWith("multi_distinct_")); - if (childDist != null && !childDist.isEmpty() && (followedByShuffled || hasDistinct)) { + if (childDist != null && !childDist.isEmpty() && (followedByShuffled || hasDistinctAggregate())) { return childDist; } return Lists.newArrayList(aggInfo.getGroupingExprs()); @@ -426,13 +457,17 @@ protected List getLocalExchangeDistributeExprs(int childIndex, boolean fol @Override public boolean requiresShuffleForCorrectness() { // Mirrors BE's AggSinkOperatorX::is_shuffled_operator() exactly: - // finalize agg with group keys needs hash-distributed input for correctness. + // finalize agg with partition exprs (group keys, or child distribute + // exprs set for a DISTINCT aggregate) needs hash-distributed input for + // correctness. The effective-key test is the node's own requirement + // (followedByShuffled=false); inherited shuffle state is added by the + // caller via selfOrInheritedShuffled. // GLOBAL dedup (!needsFinalize) is intentionally NOT included here — if a // GLOBAL dedup exists, a finalize agg always sits above it (e.g. DISTINCT_GLOBAL // above DISTINCT_LOCAL/GLOBAL_DEDUP), and the finalize agg propagates the flag // down via inheritedShuffled. A solo finalize agg satisfies hash distribution // through its own child requirement. - return needsFinalize && !aggInfo.getGroupingExprs().isEmpty(); + return needsFinalize && !getLocalExchangeDistributeExprs(0, false).isEmpty(); } private boolean canUseDistinctStreamingAgg(SessionVariable sessionVariable) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/planner/LocalShuffleNodeCoverageTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/LocalShuffleNodeCoverageTest.java index 02e68ba6b59def..59c26166c3f51b 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/LocalShuffleNodeCoverageTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/LocalShuffleNodeCoverageTest.java @@ -17,9 +17,11 @@ package org.apache.doris.planner; +import org.apache.doris.analysis.AggregateInfo; import org.apache.doris.analysis.AssertNumRowsElement; import org.apache.doris.analysis.BinaryPredicate; import org.apache.doris.analysis.Expr; +import org.apache.doris.analysis.FunctionCallExpr; import org.apache.doris.analysis.GroupingInfo; import org.apache.doris.analysis.JoinOperator; import org.apache.doris.analysis.OrderByElement; @@ -28,6 +30,7 @@ import org.apache.doris.analysis.SortInfo; import org.apache.doris.analysis.TupleDescriptor; import org.apache.doris.analysis.TupleId; +import org.apache.doris.catalog.FunctionName; import org.apache.doris.common.Pair; import org.apache.doris.common.UserException; import org.apache.doris.nereids.glue.translator.PlanTranslatorContext; @@ -35,6 +38,8 @@ import org.apache.doris.nereids.trees.plans.WindowFuncType; import org.apache.doris.planner.LocalExchangeNode.LocalExchangeType; import org.apache.doris.planner.LocalExchangeNode.LocalExchangeTypeRequire; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.SessionVariable; import org.apache.doris.thrift.TExplainLevel; import org.apache.doris.thrift.TPartitionType; import org.apache.doris.thrift.TPlanNode; @@ -855,6 +860,296 @@ public void testExchangeNodeBranches() { Assertions.assertEquals(LocalExchangeType.NOOP, noopOutput.second); } + @Test + public void testAggregationNodeDistinctFinalizeRequiresHash() { + // count(distinct k) without group-by: the finalize merge agg emits per-instance + // scalar values that the parent sums (sum0(multi_distinct_count(...)) above), so + // the input must be hash-partitioned by the distinct key. Pre-fix this agg got + // NoRequire and a PASSTHROUGH local exchange below scattered same-key rows across + // instances → the parent double-counted (result = correct × task count). + for (String fn : new String[] {"multi_distinct_count", "multi_distinct_sum", + "multi_distinct_group_concat"}) { + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction(fn)), /* groupByExprs */ true, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass(), + fn + " finalize agg must require hash input"); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + } + + @Test + public void testAggregationNodeDistinctFinalizeWithParentHashRequirement() { + // A parent that already requires hash must not change the agg's own hash demand. + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), /* groupByExprs */ true, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.requireHash()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass()); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testAggregationNodeDirectMultiDistinctNoKeyStaysNoRequire() { + // A directly called scalar multi_distinct_count(col) has isDistinct=false and + // no child distribute exprs (SplitAggWithoutDistinct builds a LOCAL aggregate + // without partition exprs). It must NOT be given a HASH requirement — a + // zero-key HASH exchange would collapse the whole input onto one task per BE. + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), + /* groupByExprs */ true, /* merge */ true, /* needsFinalize */ true, + LocalExchangeType.NOOP, null); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.NoRequire.class, agg.child.lastRequire.getClass(), + "direct multi_distinct with no effective key must stay NoRequire"); + Assertions.assertEquals(LocalExchangeType.NOOP, output.second); + Assertions.assertSame(agg.child, agg.node.getChild(0)); + } + + @Test + public void testAggregationNodeNoPartitionNonFinalizeBaseClassRequire() { + // COUNT(*)-style non-finalize (LOCAL) agg: no partition requirement, so + // the non-finalize arm of the first branch falls back to base class + // behavior (NOOP for a non-serial child). The agg exprs are non-empty + // (a plain count function) so the AggSink branch is exercised rather + // than DistinctStreamingAgg. + AggContext agg = buildAggContext( + Collections.singletonList(plainAggregateFunction("count")), /* groupByExprs */ true, + /* merge */ false, /* needsFinalize */ false, LocalExchangeType.NOOP, null); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.NoRequire.class, agg.child.lastRequire.getClass()); + Assertions.assertEquals(LocalExchangeType.NOOP, output.second); + Assertions.assertSame(agg.child, agg.node.getChild(0)); + } + + @Test + public void testAggregationNodeNoPartitionFinalizeStaysNoRequire() { + // COUNT(*)-style agg (no group keys, no DISTINCT aggregates) genuinely has no + // partition requirement: the input distribution is irrelevant. + AggContext agg = buildAggContext(Collections.singletonList(plainAggregateFunction("count")), /* groupByExprs */ true, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, null); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.NoRequire.class, agg.child.lastRequire.getClass()); + Assertions.assertEquals(LocalExchangeType.NOOP, output.second); + Assertions.assertSame(agg.child, agg.node.getChild(0)); + } + + @Test + public void testAggregationNodeDistinctLocalPhaseDefaultLeRequiresHash() { + // LOCAL (FIRST/SECOND, non-merge, non-finalize) phase of a distinct agg with the + // default enable_local_exchange_before_agg=true: BE requires HASH here + // (partition_exprs non-empty), so the FE must mirror that. + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), /* groupByExprs */ true, + /* merge */ false, /* needsFinalize */ false, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass(), + "LOCAL distinct phase with default LE requires hash"); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testAggregationNodeDistinctLocalPhaseWithLeDisabledStaysNoRequire() { + // LOCAL distinct phase + enable_local_exchange_before_agg=false → base class + // behavior (NOOP for a non-serial child): user explicitly opted out of pre-agg LE. + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), /* groupByExprs */ true, + /* merge */ false, /* needsFinalize */ false, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.enableLocalExchangeBeforeAgg = false; + Mockito.when(agg.connectContext.getSessionVariable()).thenReturn(sessionVariable); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.NoRequire.class, agg.child.lastRequire.getClass(), + "LOCAL distinct phase with LE disabled keeps no alignment requirement"); + Assertions.assertEquals(LocalExchangeType.NOOP, output.second); + Assertions.assertSame(agg.child, agg.node.getChild(0)); + } + + @Test + public void testAggregationNodeDistinctFirstMergeRequiresHash() { + // FIRST_MERGE (correctness-required) keeps the hash demand even when the + // user opts out of pre-agg local exchanges (enable_local_exchange_before_agg + // = false): removing the !isMerge() exemption must not weaken it. + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), /* groupByExprs */ true, + /* merge */ true, /* needsFinalize */ false, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.enableLocalExchangeBeforeAgg = false; + Mockito.when(agg.connectContext.getSessionVariable()).thenReturn(sessionVariable); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass(), + "FIRST_MERGE must keep the hash demand with enable_local_exchange_before_agg=false"); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testAggregationNodeGroupByFinalizeRequiresHash() { + // GROUP BY finalize agg requires hash input; when the parent has no hash + // requirement the semantic partition exprs (group keys) drive the decision. + AggContext agg = buildAggContext(Collections.singletonList(plainAggregateFunction("count")), /* groupByExprs */ false, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, null); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass()); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testAggregationNodeGroupByLocalPhaseWithLeDisabledStaysNoRequire() { + // GROUP BY local phase + enable_local_exchange_before_agg=false → base class + // behavior (NOOP for a non-serial child): user explicitly opted out of pre-agg LE. + // aggExprs is non-empty so the AggSink branch is exercised (an empty aggExprs + // would route through DistinctStreamingAgg with its own hash logic). + AggContext agg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), /* groupByExprs */ false, + /* merge */ false, /* needsFinalize */ false, LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + SessionVariable sessionVariable = new SessionVariable(); + sessionVariable.enableLocalExchangeBeforeAgg = false; + Mockito.when(agg.connectContext.getSessionVariable()).thenReturn(sessionVariable); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.NoRequire.class, agg.child.lastRequire.getClass()); + Assertions.assertEquals(LocalExchangeType.NOOP, output.second); + Assertions.assertSame(agg.child, agg.node.getChild(0)); + } + + @Test + public void testAggregationNodeRequiresShuffleForCorrectness() { + // Mirrors BE is_shuffled_operator(): finalize agg with partition exprs + // (group keys or DISTINCT aggregates) needs hash-distributed input. + AggContext distinctAgg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), + /* groupByExprs */ true, /* merge */ true, /* needsFinalize */ true, + LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + Assertions.assertTrue(distinctAgg.node.requiresShuffleForCorrectness(), + "distinct finalize agg must require shuffle for correctness"); + + AggContext noPartitionAgg = buildAggContext(Collections.singletonList(plainAggregateFunction("count")), /* groupByExprs */ true, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, null); + Assertions.assertFalse(noPartitionAgg.node.requiresShuffleForCorrectness(), + "COUNT(*) finalize agg has no partition requirement"); + + AggContext groupByAgg = buildAggContext(Collections.singletonList(plainAggregateFunction("count")), /* groupByExprs */ false, + /* merge */ true, /* needsFinalize */ true, LocalExchangeType.NOOP, null); + Assertions.assertTrue(groupByAgg.node.requiresShuffleForCorrectness(), + "GROUP BY finalize agg must require shuffle for correctness"); + + AggContext localAgg = buildAggContext(Collections.singletonList(multiDistinctFunction("multi_distinct_count")), + /* groupByExprs */ true, /* merge */ false, /* needsFinalize */ false, + LocalExchangeType.NOOP, KEYED_DISTRIBUTE_EXPRS); + Assertions.assertFalse(localAgg.node.requiresShuffleForCorrectness(), + "non-finalize agg does not require shuffle for correctness"); + } + + @Test + public void testAggregationNodeInheritedShuffleUsesChildDistributeExprs() { + // An intermediate agg that inherits a shuffle-for-correctness ancestor (e.g. + // DISTINCT_GLOBAL/FIRST_MERGE chain above a Union) keeps the child distribute + // exprs as its hash key even though the agg itself has no DISTINCT functions. + // The grouping key is deliberately different from the child distribution key: + // dropping the inherited state or selecting the grouping key must fail this test. + Expr groupingExpr = Mockito.mock(Expr.class, "groupingExpr"); + Expr childDistributeExpr = Mockito.mock(Expr.class, "childDistributeExpr"); + List childDistributeExprs = Collections.singletonList(childDistributeExpr); + AggContext agg = buildAggContext( + Collections.singletonList(plainAggregateFunction("count")), + Collections.singletonList(groupingExpr), /* merge */ true, + /* needsFinalize */ false, LocalExchangeType.NOOP, childDistributeExprs); + Mockito.when(agg.ctx.hasShuffleForCorrectnessAncestor(agg.node)).thenReturn(true); + Pair output = agg.node.enforceAndDeriveLocalExchange( + agg.ctx, null, LocalExchangeTypeRequire.noRequire()); + Assertions.assertEquals(LocalExchangeNode.RequireHash.class, agg.child.lastRequire.getClass(), + "inherited shuffle ancestor must keep the hash demand"); + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, output.second); + assertChildLocalExchangeType(agg.node, 0, LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + LocalExchangeNode exchangeNode = (LocalExchangeNode) agg.node.getChild(0); + Assertions.assertEquals(childDistributeExprs, exchangeNode.getDistributeExprLists(), + "inherited intermediate agg must hash by the child's distribution key"); + Assertions.assertNotEquals(Collections.singletonList(groupingExpr), exchangeNode.getDistributeExprLists(), + "the grouping key must not replace the inherited child distribution key"); + } + + /** A non-empty child distribute expr list, as fragment planning sets for a keyed DISTINCT agg. */ + private static final List KEYED_DISTRIBUTE_EXPRS = + Collections.singletonList(Mockito.mock(Expr.class)); + + private static class AggContext { + final AggregationNode node; + final PlanTranslatorContext ctx; + final TrackingPlanNode child; + final ConnectContext connectContext; + + AggContext(AggregationNode node, PlanTranslatorContext ctx, TrackingPlanNode child, + ConnectContext connectContext) { + this.node = node; + this.ctx = ctx; + this.child = child; + this.connectContext = connectContext; + } + } + + /** + * noGroupByExprs == true → no group keys (mirrors the scalar COUNT(DISTINCT)); + * distributeExprs != null → the plan set child distribute exprs for this agg + * (as fragment planning does for a DISTINCT agg), which is what makes + * hasPartitionRequirement() true for a keyed agg. + */ + private static AggContext buildAggContext(List aggExprs, boolean noGroupByExprs, + boolean merge, boolean needsFinalize, LocalExchangeType childProvided, + List distributeExprs) { + List groupingExprs = noGroupByExprs + ? Collections.emptyList() : Collections.singletonList(Mockito.mock(Expr.class)); + return buildAggContext(aggExprs, groupingExprs, merge, needsFinalize, + childProvided, distributeExprs); + } + + private static AggContext buildAggContext(List aggExprs, List groupingExprs, + boolean merge, boolean needsFinalize, LocalExchangeType childProvided, + List distributeExprs) { + PlanTranslatorContext ctx = Mockito.mock(PlanTranslatorContext.class); + ConnectContext connectContext = Mockito.mock(ConnectContext.class); + Mockito.when(connectContext.getSessionVariable()).thenReturn(new SessionVariable()); + Mockito.when(ctx.getConnectContext()).thenReturn(connectContext); + + AggregateInfo aggInfo = Mockito.mock(AggregateInfo.class); + Mockito.when(aggInfo.getOutputTupleId()).thenReturn(new TupleId(NEXT_ID.getAndIncrement())); + Mockito.when(aggInfo.getGroupingExprs()).thenReturn(new ArrayList<>(groupingExprs)); + Mockito.when(aggInfo.getAggregateExprs()).thenReturn(new ArrayList<>(aggExprs)); + Mockito.when(aggInfo.isMerge()).thenReturn(merge); + + TrackingPlanNode child = new TrackingPlanNode(nextPlanNodeId(), childProvided); + AggregationNode agg = new AggregationNode(nextPlanNodeId(), child, aggInfo); + if (distributeExprs != null) { + agg.setChildrenDistributeExprLists(Collections.singletonList(distributeExprs)); + } + if (!needsFinalize) { + agg.unsetNeedsFinalize(); + } + return new AggContext(agg, ctx, child, connectContext); + } + + private static FunctionCallExpr plainAggregateFunction(String functionName) { + FunctionCallExpr fce = Mockito.mock(FunctionCallExpr.class); + FunctionName fnName = Mockito.mock(FunctionName.class); + Mockito.when(fnName.getFunction()).thenReturn(functionName); + Mockito.when(fce.getFnName()).thenReturn(fnName); + return fce; + } + + private static FunctionCallExpr multiDistinctFunction(String functionName) { + FunctionCallExpr fce = Mockito.mock(FunctionCallExpr.class); + FunctionName fnName = Mockito.mock(FunctionName.class); + Mockito.when(fnName.getFunction()).thenReturn(functionName); + Mockito.when(fce.getFnName()).thenReturn(fnName); + return fce; + } + private static PlanNodeId nextPlanNodeId() { return new PlanNodeId(NEXT_ID.getAndIncrement()); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/LocalExchangePlannerTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/LocalExchangePlannerTest.java index 275df40cffaa34..d691028666f3c4 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/LocalExchangePlannerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/LocalExchangePlannerTest.java @@ -22,6 +22,7 @@ import org.apache.doris.common.UserException; import org.apache.doris.nereids.NereidsPlanner; import org.apache.doris.planner.AddLocalExchange; +import org.apache.doris.planner.AggregationNode; import org.apache.doris.planner.LocalExchangeNode; import org.apache.doris.planner.LocalExchangeNode.LocalExchangeType; import org.apache.doris.planner.PlanFragment; @@ -77,6 +78,11 @@ protected void setupLocalShuffleSession(java.util.function.Consumer { + sv.enableBroadcastJoinForcePassthrough = true; + sv.aggPhase = 1; + }); + assertFinalizeDistinctAggChildHashKeyedBy("select count(distinct a.k2) from test.t1 a " + + "left join [shuffle] test.t2 b on a.k2 = b.k2 " + + "left join [broadcast] test.t2 c on b.k1 = c.k1", + "k2"); + + } + + @Test + public void testCountDistinctNoGroupByWithoutForcePassthroughNoRedundantLe() throws Exception { + // Same multi_distinct shape but without broadcast-join force-passthrough: the + // shuffle join's probe output is already hash-partitioned by k2, which satisfies + // the finalize agg's hash demand — so no LOCAL_HASH local exchange may appear. + // The explicit aggPhase=1 + force-passthrough=false (reset in setup) pins the + // exact shape this test means to verify. + setupLocalShuffleSession(sv -> { + sv.enableBroadcastJoinForcePassthrough = false; + sv.aggPhase = 1; + }); + assertNoLocalExchangeOfType("select count(distinct a.k2) from test.t1 a " + + "left join [shuffle] test.t2 b on a.k2 = b.k2 " + + "left join [broadcast] test.t2 c on b.k1 = c.k1", + LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testDirectMultiDistinctNoKeyHasNoHashLe() throws Exception { + // A directly called scalar multi_distinct_count(k2) has no distribute exprs + // and no group keys: a zero-key LOCAL_HASH exchange would collapse the whole + // input onto one task per BE. The plan must not contain any LOCAL_HASH. + setupLocalShuffleSession(sv -> sv.aggPhase = 1); + assertNoLocalExchangeOfType("select multi_distinct_count(k2) from test.t1", + LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + @Test + public void testCountStarNoGroupByHasNoHashLe() throws Exception { + // COUNT(*) has no partition requirement: no LOCAL_HASH local exchange may appear + // anywhere in the plan (the two-phase agg only gets the PASSTHROUGH fan-out of + // the pooling scan). + setupLocalShuffleSession(null); + assertNoLocalExchangeOfType("select count(*) from test.t1", + LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE); + } + + /** + * Assert that every finalize DISTINCT agg (multi_distinct_* output) has a + * LOCAL_EXECUTION_HASH_SHUFFLE local exchange directly beneath it, keyed by + * {@code keyName}. This pins the agg-to-exchange edge and its partition + * expressions — a flatten-to-enum check could not distinguish a keyless or + * wrong-key HASH exchange. + */ + protected void assertFinalizeDistinctAggChildHashKeyedBy(String sql, String keyName) throws Exception { + StmtExecutor executor = executeNereidsSql("explain distributed plan " + sql); + NereidsPlanner planner = (NereidsPlanner) executor.planner(); + List finalizeAggs = new ArrayList<>(); + for (PlanFragment fragment : planner.getFragments()) { + collectFinalizeDistinctAggs(fragment.getPlanRoot(), finalizeAggs); + } + Assertions.assertFalse(finalizeAggs.isEmpty(), "no finalize DISTINCT agg found in plan"); + for (AggregationNode agg : finalizeAggs) { + PlanNode child = agg.getChild(0); + Assertions.assertTrue(child instanceof LocalExchangeNode, + "expected LocalExchangeNode directly below finalize DISTINCT agg, got: " + child); + LocalExchangeNode le = (LocalExchangeNode) child; + Assertions.assertEquals(LocalExchangeType.LOCAL_EXECUTION_HASH_SHUFFLE, le.getExchangeType(), + "hash LE below finalize DISTINCT agg must be LOCAL_EXECUTION_HASH_SHUFFLE"); + Assertions.assertFalse(le.getDistributeExprLists().isEmpty(), + "hash LE below finalize DISTINCT agg must be keyed"); + Assertions.assertTrue(le.getDistributeExprLists().stream() + .anyMatch(e -> e.toString().contains(keyName)), + "hash LE must be keyed by " + keyName + ", actual: " + le.getDistributeExprLists()); + } + } + + private void collectFinalizeDistinctAggs(PlanNode node, List found) { + // "output: multi_distinct_count(...)" pins the merge/finalize DISTINCT agg; + // the sum0(multi_distinct_count(...)) parent above it must not match. + if (node instanceof AggregationNode && node.getNodeExplainString("", TExplainLevel.NORMAL) + .contains("output: multi_distinct_count")) { + found.add((AggregationNode) node); + } + for (PlanNode child : node.getChildren()) { + collectFinalizeDistinctAggs(child, found); + } + } + @Test public void testBroadcastJoinPoolingShapeDsl() throws Exception { // doc rule "HashJoin / BROADCAST / 池化": diff --git a/regression-test/data/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.out b/regression-test/data/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.out new file mode 100644 index 00000000000000..507ebcdcc1241c --- /dev/null +++ b/regression-test/data/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.out @@ -0,0 +1,6 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !bug26_be_native -- +1 + +-- !bug26_fe_planned -- +1 diff --git a/regression-test/suites/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.groovy b/regression-test/suites/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.groovy index ca3fe027c47fe7..6a9ae2c8fd38e8 100644 --- a/regression-test/suites/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.groovy +++ b/regression-test/suites/nereids_p0/local_shuffle/test_local_shuffle_rqg_bugs.groovy @@ -1612,5 +1612,54 @@ suite("test_local_shuffle_rqg_bugs") { assertTrue(false, "Bug 25: COLOCATE+NLJ CROSS probe: ${t.message}") } + + // ============================================================ + // Bug 26: scalar count(distinct) over shuffle+broadcast joins returns + // correct-value × task-count when agg_phase=1 + broadcast-join + // force-passthrough with the FE local-shuffle planner. + // Root cause (FE-planned): AggregationNode handed NoRequire to a finalize + // merge agg with no group keys but DISTINCT aggregates; the PASSTHROUGH + // local exchange below the broadcast-join probe scattered same-key rows, + // and sum0(multi_distinct_count(...)) summed the overlapping per-instance + // values. Fixed by keying the hash requirement on the effective partition + // exprs (mirrors BE `_partition_exprs`). + // ============================================================ + try { + logger.info("Bug 26: count(distinct) under agg_phase=1 + broadcast force-passthrough") + sql "DROP TABLE IF EXISTS rqg_local_shuffle_distinct_t1" + sql "DROP TABLE IF EXISTS rqg_local_shuffle_distinct_t2" + sql """CREATE TABLE rqg_local_shuffle_distinct_t1 (pk INT NOT NULL, k2 INT NOT NULL) + ENGINE=OLAP DUPLICATE KEY(pk) DISTRIBUTED BY HASH(pk) BUCKETS 5 + PROPERTIES ("replication_num"="1")""" + sql """CREATE TABLE rqg_local_shuffle_distinct_t2 (pk INT NOT NULL, k2 INT NOT NULL, other INT NOT NULL) + ENGINE=OLAP DUPLICATE KEY(pk) DISTRIBUTED BY HASH(pk) BUCKETS 5 + PROPERTIES ("replication_num"="1")""" + // Two rows sharing the same distinct key. batch_size=1 with 4 local tasks + // forces the PASSTHROUGH exchange to send separate blocks to different + // channels, so the pre-fix plan counts the shared key once per task. + sql "INSERT INTO rqg_local_shuffle_distinct_t1 VALUES (1, 5), (2, 5)" + sql "INSERT INTO rqg_local_shuffle_distinct_t2 VALUES (1, 5, 10), (2, 5, 20)" + + def distinctJoinQuery = { vars -> """ + SELECT /*+SET_VAR(${vars})*/ + count(distinct t1.k2) AS cnt_distinct + FROM rqg_local_shuffle_distinct_t1 t1 + LEFT JOIN [shuffle] rqg_local_shuffle_distinct_t2 t2 ON t1.k2 = t2.k2 + LEFT JOIN [broadcast] rqg_local_shuffle_distinct_t2 t3 ON t2.pk = t3.pk + """ } + def distinctJoinVariables = "enable_sql_cache=false, agg_phase=1, " + + "enable_broadcast_join_force_passthrough=true, parallel_pipeline_task_num=4, batch_size=1" + // Pin both implementations to the mathematically correct result (1). Using + // one implementation as the other's oracle would let a shared bug pass. + order_qt_bug26_be_native(distinctJoinQuery( + "${distinctJoinVariables}, enable_local_shuffle_planner=false")) + order_qt_bug26_fe_planned(distinctJoinQuery( + "${distinctJoinVariables}, enable_local_shuffle_planner=true")) + logger.info("Bug 26: PASSED") + } catch (Throwable t) { + logger.error("Bug 26 FAILED: ${t.message}") + assertTrue(false, "Bug 26: ${t.message}") + } + logger.info("=== All RQG bug reproduction tests completed ===") }