Skip to content

Commit f1cfc38

Browse files
authored
[fe](cse) Extract aggregate-argument CSE below distribute (#66815)
### What problem does this PR solve? Extract aggregate-argument CSE below distribute
1 parent 5870fb2 commit f1cfc38

3 files changed

Lines changed: 141 additions & 2 deletions

File tree

fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/ProjectAggregateExpressionsForCse.java

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
import org.apache.doris.nereids.trees.expressions.Slot;
3030
import org.apache.doris.nereids.trees.expressions.SlotReference;
3131
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
32+
import org.apache.doris.nereids.trees.plans.AggMode;
33+
import org.apache.doris.nereids.trees.plans.AggPhase;
3234
import org.apache.doris.nereids.trees.plans.Plan;
3335
import org.apache.doris.nereids.trees.plans.algebra.Aggregate;
3436
import org.apache.doris.nereids.trees.plans.physical.AbstractPhysicalPlan;
@@ -64,12 +66,21 @@ public Plan visitPhysicalHashAggregate(PhysicalHashAggregate<? extends Plan> agg
6466
* Shared CSE projection logic for PhysicalHashAggregate.
6567
* Extracts common sub-expressions from
6668
* aggregate function arguments into a project node beneath the aggregate.
69+
*
70+
* <p>For one-phase aggregates whose child is a PhysicalDistribute
71+
* (aggregate -> distribute -> scan), the CSE project is inserted below the
72+
* distribute so that the distribution-key slots stay intact and the exchange
73+
* only carries the (already pruned) aggregate input. The translator's bucketed
74+
* fusion (fusing one-phase aggregate + distribute into BucketedAggregationNode)
75+
* builds directly on the distribute's child, so the fused plan naturally
76+
* becomes BucketedAgg(sum(x), max(x)) -> Project(a+b AS x) -> scan and the
77+
* common aggregate argument is evaluated once per row instead of once per
78+
* aggregate function.</p>
6779
*/
6880
private <T extends AbstractPhysicalPlan & Aggregate<? extends Plan>>
6981
Plan projectAggregateCse(T aggregate) {
7082
// For multi-phase aggregates, only process the 1st phase.
71-
// Bucketed agg is always single-phase, but keep the same guard for safety.
72-
if (aggregate.child() instanceof PhysicalDistribute || aggregate.child() instanceof Aggregate) {
83+
if (aggregate.child() instanceof Aggregate) {
7384
return aggregate;
7485
}
7586

@@ -161,6 +172,60 @@ Plan projectAggregateCse(T aggregate) {
161172
project = project.withPhysicalPropertiesAndStats(projectPhysicalProperties, project.getStats());
162173
return (Plan) aggregate.withAggOutput(aggOutputReplaced)
163174
.withChildren(project);
175+
} else if (aggregate.child() instanceof PhysicalDistribute) {
176+
// One-phase (INPUT_TO_RESULT) aggregate over a distribute
177+
// (aggregate -> distribute -> scan): insert the CSE project between
178+
// the distribute and its child, instead of between the aggregate and
179+
// the distribute. This keeps the aggregate's child as a distribute
180+
// (so bucketed fusion and the property machinery still see the same
181+
// shape), and the project lands inside the scan
182+
// fragment, so the common aggregate argument is computed once per row
183+
// before the exchange. After bucketed fusion bypasses the distribute,
184+
// the executed plan is BucketedAgg(sum(x), max(x)) -> Project(a+b AS x)
185+
// -> scan.
186+
//
187+
// Only the one-phase shape reaches here with complex aggregate
188+
// arguments: two-phase GLOBAL aggregates (BUFFER_TO_RESULT) reference
189+
// the local phase's intermediate slots, so no CSE candidate is
190+
// extracted for them anyway. Guard explicitly anyway to keep the
191+
// intent clear and to stay safe if a future aggregate function
192+
// surfaces a non-slot argument on the GLOBAL phase.
193+
if (!(aggregate instanceof PhysicalHashAggregate)) {
194+
return aggregate;
195+
}
196+
PhysicalHashAggregate<? extends Plan> hashAggregate =
197+
(PhysicalHashAggregate<? extends Plan>) aggregate;
198+
if (hashAggregate.getAggPhase() != AggPhase.GLOBAL
199+
|| hashAggregate.getAggMode() != AggMode.INPUT_TO_RESULT) {
200+
return aggregate;
201+
}
202+
PhysicalDistribute<?> distribute = (PhysicalDistribute<?>) aggregate.child();
203+
List<NamedExpression> projections = new ArrayList<>();
204+
projections.addAll(inputSlots);
205+
projections.addAll(cseCandidates.values());
206+
List<Slot> projectOutput = new ImmutableList.Builder<Slot>()
207+
.addAll(inputSlots)
208+
.addAll(slotMap.values())
209+
.build();
210+
LogicalProperties projectLogicalProperties = new LogicalProperties(
211+
() -> projectOutput,
212+
() -> DataTrait.EMPTY_TRAIT
213+
);
214+
AbstractPhysicalPlan distributeChild = ((AbstractPhysicalPlan) distribute.child());
215+
PhysicalProperties projectPhysicalProperties = ChildOutputPropertyDeriver.computeProjectOutputProperties(
216+
projections, distributeChild.getPhysicalProperties());
217+
PhysicalProject<? extends Plan> project = new PhysicalProject<>(projections, Optional.empty(),
218+
projectLogicalProperties,
219+
projectPhysicalProperties,
220+
distributeChild.getStats(),
221+
distribute.child());
222+
// withChildren keeps the distribution spec and physical properties of the
223+
// distribute unchanged; its output now comes from the CSE project, which
224+
// still carries every distribution-key slot (the group-by slots are part
225+
// of inputSlots above).
226+
PhysicalDistribute<Plan> newDistribute = distribute.withChildren(ImmutableList.of(project));
227+
return (Plan) aggregate.withAggOutput(aggOutputReplaced)
228+
.withChildren(newDistribute);
164229
} else {
165230
List<NamedExpression> projections = new ArrayList<>();
166231
projections.addAll(inputSlots);
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
-- This file is automatically generated. You should know what you did if you want to edit this
2+
-- !one_phase_join_result --
3+
g1 33 19 33 19
4+
g2 22 15 22 15
5+
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
suite("cse_agg_distribute") {
19+
sql "SET enable_nereids_planner=true"
20+
sql "SET enable_fallback_to_original_planner=false"
21+
sql "SET runtime_filter_mode=OFF"
22+
23+
sql "DROP TABLE IF EXISTS cse_agg_distribute_tbl"
24+
sql """
25+
CREATE TABLE cse_agg_distribute_tbl (
26+
id int,
27+
grp varchar(20),
28+
a int,
29+
b int
30+
) DUPLICATE KEY(id)
31+
DISTRIBUTED BY HASH(id) BUCKETS 3
32+
PROPERTIES('replication_num' = '1')
33+
"""
34+
sql """ INSERT INTO cse_agg_distribute_tbl VALUES
35+
(1, 'g1', 1, 2),
36+
(2, 'g2', 3, 4),
37+
(3, 'g1', 5, 6),
38+
(4, 'g2', 7, 8),
39+
(5, 'g1', 9, 10)
40+
"""
41+
42+
// SUM(a+b) and MAX(a+b) share the same argument, so the aggregate-argument
43+
// CSE must extract "a+b" into a project node and make both functions
44+
// reference the extracted slot, instead of re-evaluating a+b per function.
45+
String query = "SELECT grp, SUM(a+b), MAX(a+b) FROM cse_agg_distribute_tbl GROUP BY grp"
46+
47+
// ---------------------------------------------------------------------
48+
// one-phase aggregate over a distribute (the aggregate is a join child,
49+
// so the distribute is required by the join): the CSE project must be
50+
// inserted below the distribute, keeping the distribution-key slots
51+
// intact. Both aggregates must reference the extracted slot (4
52+
// occurrences: SUM/MAX of each side).
53+
// ---------------------------------------------------------------------
54+
sql "set agg_phase=1"
55+
sql "set enable_bucketed_hash_agg=false"
56+
String joinQuery = """
57+
SELECT t1.grp, t1.s, t1.m, t2.s2, t2.m2 FROM
58+
(SELECT grp, SUM(a+b) s, MAX(a+b) m FROM cse_agg_distribute_tbl GROUP BY grp) t1
59+
JOIN (SELECT grp, SUM(a+b) s2, MAX(a+b) m2 FROM cse_agg_distribute_tbl GROUP BY grp) t2
60+
ON t1.grp = t2.grp
61+
"""
62+
explain {
63+
sql("${joinQuery}")
64+
contains("VEXCHANGE")
65+
contains("VSELECT")
66+
multiContains("cast(a as BIGINT) + cast(b as BIGINT))[#", 4)
67+
}
68+
order_qt_one_phase_join_result """${joinQuery} ORDER BY t1.grp"""
69+
}

0 commit comments

Comments
 (0)