1717
1818package org .apache .doris .nereids .rules .rewrite ;
1919
20- import org .apache .doris .nereids .annotation . DependsRules ;
20+ import org .apache .doris .nereids .jobs . JobContext ;
2121import org .apache .doris .nereids .properties .DataTrait ;
2222import org .apache .doris .nereids .properties .FuncDeps ;
23- import org .apache .doris .nereids .rules . Rule ;
24- import org .apache .doris .nereids .rules . RuleType ;
23+ import org .apache .doris .nereids .trees . expressions . Alias ;
24+ import org .apache .doris .nereids .trees . expressions . ExprId ;
2525import org .apache .doris .nereids .trees .expressions .Expression ;
2626import org .apache .doris .nereids .trees .expressions .NamedExpression ;
2727import org .apache .doris .nereids .trees .expressions .Slot ;
28+ import org .apache .doris .nereids .trees .expressions .functions .agg .AnyValue ;
2829import org .apache .doris .nereids .trees .plans .Plan ;
30+ import org .apache .doris .nereids .trees .plans .algebra .Aggregate ;
2931import org .apache .doris .nereids .trees .plans .logical .LogicalAggregate ;
32+ import org .apache .doris .nereids .trees .plans .logical .LogicalCTEConsumer ;
33+ import org .apache .doris .nereids .trees .plans .logical .LogicalFilter ;
34+ import org .apache .doris .nereids .trees .plans .logical .LogicalProject ;
35+ import org .apache .doris .nereids .trees .plans .visitor .CustomRewriter ;
36+ import org .apache .doris .nereids .trees .plans .visitor .DefaultPlanRewriter ;
3037
31- import com .google .common .collect .ImmutableList ;
38+ import com .google .common .collect .LinkedHashMultimap ;
39+ import com .google .common .collect .Multimap ;
3240
3341import java .util .ArrayList ;
3442import java .util .HashMap ;
3543import java .util .HashSet ;
44+ import java .util .LinkedHashMap ;
3645import java .util .List ;
3746import java .util .Map ;
3847import java .util .Map .Entry ;
3948import java .util .Set ;
4049
41-
4250/**
4351 * Eliminate group by key based on fd item information.
4452 * such as:
4553 * for a -> b, we can get:
4654 * group by a, b, c => group by a, c
55+ *
56+ * When a group-by key is FD-redundant but still needed in the output,
57+ * it is wrapped with any_value() and assigned a fresh ExprId.
58+ * Upper plan references are rewritten via ExprIdRewriter so that
59+ * all ancestor nodes see the new ExprIds.
4760 */
48- @ DependsRules ({EliminateGroupBy .class , ColumnPruning .class })
49- public class EliminateGroupByKey implements RewriteRuleFactory {
61+ public class EliminateGroupByKey extends DefaultPlanRewriter <Map <ExprId , ExprId >> implements CustomRewriter {
62+ private ExprIdRewriter exprIdReplacer ;
63+
64+ @ Override
65+ public Plan rewriteRoot (Plan plan , JobContext jobContext ) {
66+ if (!plan .containsType (Aggregate .class )) {
67+ return plan ;
68+ }
69+ Map <ExprId , ExprId > replaceMap = new HashMap <>();
70+ ExprIdRewriter .ReplaceRule replaceRule = new ExprIdRewriter .ReplaceRule (replaceMap , false );
71+ exprIdReplacer = new ExprIdRewriter (replaceRule , jobContext );
72+ return plan .accept (this , replaceMap );
73+ }
74+
75+ @ Override
76+ public Plan visit (Plan plan , Map <ExprId , ExprId > replaceMap ) {
77+ plan = visitChildren (this , plan , replaceMap );
78+ plan = exprIdReplacer .rewriteExpr (plan , replaceMap );
79+ return plan ;
80+ }
81+
82+ @ Override
83+ public Plan visitLogicalProject (LogicalProject <? extends Plan > proj , Map <ExprId , ExprId > replaceMap ) {
84+ proj = visitChildren (this , proj , replaceMap );
85+
86+ // Find the Aggregate child, possibly through a Filter
87+ Plan child = proj .child (0 );
88+ LogicalAggregate <? extends Plan > agg ;
89+ boolean hasFilter = child instanceof LogicalFilter ;
90+ if (hasFilter && child .child (0 ) instanceof LogicalAggregate ) {
91+ agg = (LogicalAggregate <? extends Plan >) child .child (0 );
92+ } else if (child instanceof LogicalAggregate ) {
93+ agg = (LogicalAggregate <? extends Plan >) child ;
94+ } else {
95+ return exprIdReplacer .rewriteExpr (proj , replaceMap );
96+ }
97+
98+ // Don't transform if source repeat is present
99+ if (agg .getSourceRepeat ().isPresent ()) {
100+ return exprIdReplacer .rewriteExpr (proj , replaceMap );
101+ }
102+
103+ // Rewrite proj and the filter (if present) through the replaceMap accumulated
104+ // by visitChildren, so that ExprId replacements from nested rewrites
105+ // (e.g. inner aggregates) are reflected in the required-output slot set.
106+ proj = (LogicalProject <? extends Plan >) exprIdReplacer .rewriteExpr (proj , replaceMap );
107+ if (hasFilter ) {
108+ child = exprIdReplacer .rewriteExpr (child , replaceMap );
109+ }
110+
111+ // Compute requireOutput: slots needed by the Project (and Filter, if present)
112+ Set <Slot > requireOutput = new HashSet <>(proj .getInputSlots ());
113+ if (hasFilter ) {
114+ requireOutput .addAll (child .getInputSlots ());
115+ }
116+
117+ // Transform the aggregate
118+ EliminateResult result = eliminateGroupByKeyWithMap (agg , requireOutput );
119+ if (!result .changed ) {
120+ return proj ;
121+ }
122+
123+ // Merge into the global replaceMap so that all ancestor nodes get rewritten
124+ replaceMap .putAll (result .replaceMap );
125+
126+ // Rebuild the child chain with the new aggregate,
127+ // and rewrite the Filter (if present) and Project expressions
128+ Plan newChild ;
129+ if (hasFilter ) {
130+ Plan updatedFilter = child .withChildren (result .newAgg );
131+ newChild = exprIdReplacer .rewriteExpr (updatedFilter , replaceMap );
132+ } else {
133+ newChild = result .newAgg ;
134+ }
135+ Plan newProj = exprIdReplacer .rewriteExpr (proj .withChildren (newChild ), replaceMap );
136+ return newProj ;
137+ }
50138
51139 @ Override
52- public List <Rule > buildRules () {
53- return ImmutableList .of (
54- RuleType .ELIMINATE_GROUP_BY_KEY .build (
55- logicalProject (logicalAggregate ().when (agg -> !agg .getSourceRepeat ().isPresent ()))
56- .then (proj -> {
57- LogicalAggregate <? extends Plan > agg = proj .child ();
58- LogicalAggregate <Plan > newAgg = eliminateGroupByKey (agg , proj .getInputSlots ());
59- if (newAgg == null ) {
60- return null ;
61- }
62- return proj .withChildren (newAgg );
63- })),
64- RuleType .ELIMINATE_FILTER_GROUP_BY_KEY .build (
65- logicalProject (logicalFilter (logicalAggregate ()
66- .when (agg -> !agg .getSourceRepeat ().isPresent ())))
67- .then (proj -> {
68- LogicalAggregate <? extends Plan > agg = proj .child ().child ();
69- Set <Slot > requireSlots = new HashSet <>(proj .getInputSlots ());
70- requireSlots .addAll (proj .child (0 ).getInputSlots ());
71- LogicalAggregate <Plan > newAgg = eliminateGroupByKey (agg , requireSlots );
72- if (newAgg == null ) {
73- return null ;
74- }
75- return proj .withChildren (proj .child ().withChildren (newAgg ));
76- })
77- )
78- );
140+ public Plan visitLogicalCTEConsumer (LogicalCTEConsumer cteConsumer , Map <ExprId , ExprId > replaceMap ) {
141+ // When a producer aggregate's output slot is wrapped with any_value(),
142+ // a fresh ExprId is recorded in replaceMap. The CTE consumer's producerToConsumerSlotMap
143+ // still references the old ExprId, so we must rebuild both maps with the new ExprIds.
144+ Map <Slot , Slot > newConsumerToProducer = new LinkedHashMap <>();
145+ Multimap <Slot , Slot > newProducerToConsumer = LinkedHashMultimap .create ();
146+ for (Slot producerSlot : cteConsumer .getConsumerToProducerOutputMap ().values ()) {
147+ ExprId newExprId = resolveExprIdChain (producerSlot .getExprId (), replaceMap );
148+ Slot effectiveProducerSlot = newExprId != null
149+ ? (Slot ) producerSlot .withExprId (newExprId )
150+ : producerSlot ;
151+ for (Slot consumerSlot : cteConsumer .getProducerToConsumerOutputMap ().get (producerSlot )) {
152+ newProducerToConsumer .put (effectiveProducerSlot , consumerSlot );
153+ newConsumerToProducer .put (consumerSlot , effectiveProducerSlot );
154+ }
155+ }
156+ return cteConsumer .withTwoMaps (newConsumerToProducer , newProducerToConsumer );
157+ }
158+
159+ /** Follow transitive ExprId chain to find the final replacement, or null if none. */
160+ private static ExprId resolveExprIdChain (ExprId exprId , Map <ExprId , ExprId > replaceMap ) {
161+ ExprId newId = replaceMap .get (exprId );
162+ if (newId == null ) {
163+ return null ;
164+ }
165+ ExprId lastId = newId ;
166+ while (true ) {
167+ ExprId next = replaceMap .get (lastId );
168+ if (next == null ) {
169+ return lastId ;
170+ }
171+ lastId = next ;
172+ }
173+ }
174+
175+ /** Result of eliminateGroupByKey: the new aggregate and a map of old->new ExprIds. */
176+ private static class EliminateResult {
177+ final LogicalAggregate <Plan > newAgg ;
178+ final Map <ExprId , ExprId > replaceMap ;
179+ final boolean changed ;
180+
181+ EliminateResult (LogicalAggregate <Plan > newAgg , Map <ExprId , ExprId > replaceMap , boolean changed ) {
182+ this .newAgg = newAgg ;
183+ this .replaceMap = replaceMap ;
184+ this .changed = changed ;
185+ }
79186 }
80187
81- LogicalAggregate < Plan > eliminateGroupByKey (LogicalAggregate <? extends Plan > agg , Set <Slot > requireOutput ) {
82- Set < Expression > removeExpression = findCanBeRemovedExpressions (agg , requireOutput ,
188+ EliminateResult eliminateGroupByKeyWithMap (LogicalAggregate <? extends Plan > agg , Set <Slot > requireOutput ) {
189+ FindResult result = findCanBeRemovedExpressionsInternal (agg , requireOutput ,
83190 agg .child ().getLogicalProperties ().getTrait ());
191+ Set <Expression > removeExpression = result .removeExpression ;
192+ Set <Expression > wrapWithAnyValue = result .wrapWithAnyValue ;
193+
84194 List <Expression > newGroupExpression = new ArrayList <>();
85195 for (Expression expression : agg .getGroupByExpressions ()) {
86- if (!removeExpression .contains (expression )) {
196+ if (!removeExpression .contains (expression )
197+ && !wrapWithAnyValue .contains (expression )) {
87198 newGroupExpression .add (expression );
88199 }
89200 }
90201 List <NamedExpression > newOutput = new ArrayList <>();
202+ Map <ExprId , ExprId > replaceMap = new HashMap <>();
203+ boolean changed = !removeExpression .isEmpty () || !wrapWithAnyValue .isEmpty ();
91204 for (NamedExpression expression : agg .getOutputExpressions ()) {
92- if (! removeExpression .contains (expression )) {
93- newOutput . add ( expression ) ;
205+ if (removeExpression .contains (expression )) {
206+ continue ;
94207 }
208+ if (wrapWithAnyValue .contains (expression )) {
209+ // expression is FD-redundant but needed in output: wrap with any_value
210+ // Use fresh ExprId (auto-generated by Alias) to avoid ExprId collision,
211+ // and record the mapping for rewriting upper plan references.
212+ Alias newAlias = new Alias (new AnyValue (expression .toSlot ()), expression .getName ());
213+ replaceMap .put (expression .getExprId (), newAlias .getExprId ());
214+ expression = newAlias ;
215+ }
216+ newOutput .add (expression );
95217 }
96- return agg .withGroupByAndOutput (newGroupExpression , newOutput );
218+ return new EliminateResult ( agg .withGroupByAndOutput (newGroupExpression , newOutput ), replaceMap , changed );
97219 }
98220
99221 /**
100- * return removeExpression
222+ * Return expressions that can be completely removed from both group-by and output.
223+ * Kept for backward compatibility with external callers (e.g. PushDownAggThroughJoinOnPkFk).
101224 */
102225 public static Set <Expression > findCanBeRemovedExpressions (LogicalAggregate <? extends Plan > agg ,
103226 Set <Slot > requireOutput , DataTrait dataTrait ) {
227+ FindResult result = findCanBeRemovedExpressionsInternal (agg , requireOutput , dataTrait );
228+ return new HashSet <>(result .removeExpression );
229+ }
230+
231+ /** Result of findCanBeRemovedExpressionsInternal: two sets of expressions. */
232+ private static class FindResult {
233+ final Set <Expression > removeExpression ; // remove from group-by and output
234+ final Set <Expression > wrapWithAnyValue ; // remove from group-by, wrap with ANY_VALUE in output
235+
236+ FindResult (Set <Expression > removeExpression , Set <Expression > wrapWithAnyValue ) {
237+ this .removeExpression = removeExpression ;
238+ this .wrapWithAnyValue = wrapWithAnyValue ;
239+ }
240+ }
241+
242+ private static FindResult findCanBeRemovedExpressionsInternal (LogicalAggregate <? extends Plan > agg ,
243+ Set <Slot > requireOutput , DataTrait dataTrait ) {
104244 Map <Expression , Set <Slot >> groupBySlots = new HashMap <>();
105245 Set <Slot > validSlots = new HashSet <>();
106246 for (Expression expression : agg .getGroupByExpressions ()) {
@@ -110,17 +250,24 @@ public static Set<Expression> findCanBeRemovedExpressions(LogicalAggregate<? ext
110250
111251 FuncDeps funcDeps = dataTrait .getAllValidFuncDeps (validSlots );
112252 if (funcDeps .isEmpty ()) {
113- return new HashSet <>();
253+ return new FindResult ( new HashSet <>(), new HashSet <>() );
114254 }
115255
116256 Set <Set <Slot >> minGroupBySlots = funcDeps .eliminateDeps (new HashSet <>(groupBySlots .values ()), requireOutput );
117257 Set <Expression > removeExpression = new HashSet <>();
258+ Set <Expression > wrapWithAnyValue = new HashSet <>();
118259 for (Entry <Expression , Set <Slot >> entry : groupBySlots .entrySet ()) {
119- if (!minGroupBySlots .contains (entry .getValue ())
120- && !requireOutput .containsAll (entry .getValue ())) {
121- removeExpression .add (entry .getKey ());
260+ if (!minGroupBySlots .contains (entry .getValue ())) {
261+ // FD redundant: can remove from group-by
262+ if (!requireOutput .containsAll (entry .getValue ())) {
263+ // Not needed in output either: remove completely
264+ removeExpression .add (entry .getKey ());
265+ } else {
266+ // Still needed in output: remove from group-by, wrap with ANY_VALUE in output
267+ wrapWithAnyValue .add (entry .getKey ());
268+ }
122269 }
123270 }
124- return removeExpression ;
271+ return new FindResult ( removeExpression , wrapWithAnyValue ) ;
125272 }
126273}
0 commit comments