|
29 | 29 | import org.apache.doris.nereids.trees.expressions.Slot; |
30 | 30 | import org.apache.doris.nereids.trees.expressions.SlotReference; |
31 | 31 | 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; |
32 | 34 | import org.apache.doris.nereids.trees.plans.Plan; |
33 | 35 | import org.apache.doris.nereids.trees.plans.algebra.Aggregate; |
34 | 36 | import org.apache.doris.nereids.trees.plans.physical.AbstractPhysicalPlan; |
@@ -64,12 +66,21 @@ public Plan visitPhysicalHashAggregate(PhysicalHashAggregate<? extends Plan> agg |
64 | 66 | * Shared CSE projection logic for PhysicalHashAggregate. |
65 | 67 | * Extracts common sub-expressions from |
66 | 68 | * 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> |
67 | 79 | */ |
68 | 80 | private <T extends AbstractPhysicalPlan & Aggregate<? extends Plan>> |
69 | 81 | Plan projectAggregateCse(T aggregate) { |
70 | 82 | // 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) { |
73 | 84 | return aggregate; |
74 | 85 | } |
75 | 86 |
|
@@ -161,6 +172,60 @@ Plan projectAggregateCse(T aggregate) { |
161 | 172 | project = project.withPhysicalPropertiesAndStats(projectPhysicalProperties, project.getStats()); |
162 | 173 | return (Plan) aggregate.withAggOutput(aggOutputReplaced) |
163 | 174 | .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); |
164 | 229 | } else { |
165 | 230 | List<NamedExpression> projections = new ArrayList<>(); |
166 | 231 | projections.addAll(inputSlots); |
|
0 commit comments