|
18 | 18 | import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute; |
19 | 19 | import org.elasticsearch.xpack.esql.expression.function.grouping.GroupingFunction; |
20 | 20 | import org.elasticsearch.xpack.esql.plan.logical.Aggregate; |
| 21 | +import org.elasticsearch.xpack.esql.plan.logical.Eval; |
21 | 22 | import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; |
22 | 23 | import org.elasticsearch.xpack.esql.plan.logical.Project; |
23 | 24 | import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan; |
24 | 25 |
|
25 | 26 | import java.util.ArrayList; |
| 27 | +import java.util.HashSet; |
26 | 28 | import java.util.LinkedHashSet; |
27 | 29 | import java.util.List; |
| 30 | +import java.util.Set; |
28 | 31 |
|
29 | 32 | public final class CombineProjections extends OptimizerRules.OptimizerRule<UnaryPlan> { |
| 33 | + // don't drop groupings from a local plan, as the layout has already been agreed upon |
| 34 | + private final boolean local; |
30 | 35 |
|
31 | | - public CombineProjections() { |
| 36 | + public CombineProjections(boolean local) { |
32 | 37 | super(OptimizerRules.TransformDirection.UP); |
| 38 | + this.local = local; |
33 | 39 | } |
34 | 40 |
|
35 | 41 | @Override |
@@ -60,27 +66,91 @@ protected LogicalPlan rule(UnaryPlan plan) { |
60 | 66 | return plan; |
61 | 67 | } |
62 | 68 |
|
63 | | - // Agg with underlying Project (group by on sub-queries) |
64 | | - if (plan instanceof Aggregate a) { |
65 | | - if (child instanceof Project p) { |
66 | | - var groupings = a.groupings(); |
67 | | - List<NamedExpression> groupingAttrs = new ArrayList<>(a.groupings().size()); |
68 | | - for (Expression grouping : groupings) { |
69 | | - if (grouping instanceof Attribute attribute) { |
70 | | - groupingAttrs.add(attribute); |
71 | | - } else if (grouping instanceof Alias as && as.child() instanceof GroupingFunction.NonEvaluatableGroupingFunction) { |
72 | | - groupingAttrs.add(as); |
| 69 | + if (plan instanceof Aggregate a && child instanceof Project p) { |
| 70 | + var groupings = a.groupings(); |
| 71 | + |
| 72 | + // sanity checks |
| 73 | + for (Expression grouping : groupings) { |
| 74 | + if ((grouping instanceof Attribute |
| 75 | + || grouping instanceof Alias as && as.child() instanceof GroupingFunction.NonEvaluatableGroupingFunction) == false) { |
| 76 | + // After applying ReplaceAggregateNestedExpressionWithEval, |
| 77 | + // evaluatable groupings can only contain attributes. |
| 78 | + throw new EsqlIllegalArgumentException("Expected an attribute or grouping function, got {}", grouping); |
| 79 | + } |
| 80 | + } |
| 81 | + assert groupings.size() <= 1 |
| 82 | + || groupings.stream() |
| 83 | + .anyMatch(group -> group.anyMatch(expr -> expr instanceof GroupingFunction.NonEvaluatableGroupingFunction)) == false |
| 84 | + : "CombineProjections only tested with a single CATEGORIZE with no additional groups"; |
| 85 | + |
| 86 | + // Collect the alias map for resolving the source (f1 = 1, f2 = f1, etc..) |
| 87 | + AttributeMap.Builder<Attribute> aliasesBuilder = AttributeMap.builder(); |
| 88 | + for (NamedExpression ne : p.projections()) { |
| 89 | + // Record the aliases. |
| 90 | + // Projections are just aliases for attributes, so casting is safe. |
| 91 | + aliasesBuilder.put(ne.toAttribute(), (Attribute) Alias.unwrap(ne)); |
| 92 | + } |
| 93 | + var aliases = aliasesBuilder.build(); |
| 94 | + |
| 95 | + // Propagate any renames from the lower projection into the upper groupings. |
| 96 | + List<Expression> resolvedGroupings = new ArrayList<>(); |
| 97 | + for (Expression grouping : groupings) { |
| 98 | + Expression transformed = grouping.transformUp(Attribute.class, as -> aliases.resolve(as, as)); |
| 99 | + resolvedGroupings.add(transformed); |
| 100 | + } |
| 101 | + |
| 102 | + // This can lead to duplicates in the groupings: e.g. |
| 103 | + // | EVAL x = y | STATS ... BY x, y |
| 104 | + if (local) { |
| 105 | + // On the data node, the groupings must be preserved because they affect the physical output (see |
| 106 | + // AbstractPhysicalOperationProviders#intermediateAttributes). |
| 107 | + // In case that propagating the lower projection leads to duplicates in the resolved groupings, we'll leave an Eval in place |
| 108 | + // of the original projection to create new attributes for the duplicate groups. |
| 109 | + Set<Expression> seenResolvedGroupings = new HashSet<>(resolvedGroupings.size()); |
| 110 | + List<Expression> newGroupings = new ArrayList<>(); |
| 111 | + List<Alias> aliasesAgainstDuplication = new ArrayList<>(); |
| 112 | + |
| 113 | + for (int i = 0; i < groupings.size(); i++) { |
| 114 | + Expression resolvedGrouping = resolvedGroupings.get(i); |
| 115 | + if (seenResolvedGroupings.add(resolvedGrouping)) { |
| 116 | + newGroupings.add(resolvedGrouping); |
73 | 117 | } else { |
74 | | - // After applying ReplaceAggregateNestedExpressionWithEval, |
75 | | - // evaluatable groupings can only contain attributes. |
76 | | - throw new EsqlIllegalArgumentException("Expected an Attribute, got {}", grouping); |
| 118 | + // resolving the renames leads to a duplicate here - we need to alias the underlying attribute this refers to. |
| 119 | + // should really only be 1 attribute, anyway, but going via .references() includes the case of a |
| 120 | + // GroupingFunction.NonEvaluatableGroupingFunction. |
| 121 | + Attribute coreAttribute = resolvedGrouping.references().iterator().next(); |
| 122 | + |
| 123 | + Alias renameAgainstDuplication = new Alias( |
| 124 | + coreAttribute.source(), |
| 125 | + TemporaryNameUtils.locallyUniqueTemporaryName(coreAttribute.name()), |
| 126 | + coreAttribute |
| 127 | + ); |
| 128 | + aliasesAgainstDuplication.add(renameAgainstDuplication); |
| 129 | + |
| 130 | + // propagate the new alias into the new grouping |
| 131 | + AttributeMap.Builder<Attribute> resolverBuilder = AttributeMap.builder(); |
| 132 | + resolverBuilder.put(coreAttribute, renameAgainstDuplication.toAttribute()); |
| 133 | + AttributeMap<Attribute> resolver = resolverBuilder.build(); |
| 134 | + |
| 135 | + newGroupings.add(resolvedGrouping.transformUp(Attribute.class, attr -> resolver.resolve(attr, attr))); |
77 | 136 | } |
78 | 137 | } |
79 | | - plan = a.with( |
80 | | - p.child(), |
81 | | - combineUpperGroupingsAndLowerProjections(groupingAttrs, p.projections()), |
82 | | - combineProjections(a.aggregates(), p.projections()) |
83 | | - ); |
| 138 | + |
| 139 | + LogicalPlan newChild = aliasesAgainstDuplication.isEmpty() |
| 140 | + ? p.child() |
| 141 | + : new Eval(p.source(), p.child(), aliasesAgainstDuplication); |
| 142 | + plan = a.with(newChild, newGroupings, combineProjections(a.aggregates(), p.projections())); |
| 143 | + } else { |
| 144 | + // On the coordinator, we can just discard the duplicates. |
| 145 | + // All substitutions happen before; groupings must be attributes at this point except for non-evaluatable groupings which |
| 146 | + // will be an alias like `c = CATEGORIZE(attribute)`. |
| 147 | + // Due to such aliases, we can't use an AttributeSet to deduplicate. But we can use a regular set to deduplicate based on |
| 148 | + // regular equality (i.e. based on names) instead of name ids. |
| 149 | + // TODO: The deduplication based on simple equality will be insufficient in case of multiple non-evaluatable groupings, e.g. |
| 150 | + // for `| EVAL x = y | STATS ... BY CATEGORIZE(x), CATEGORIZE(y)`. That will require semantic equality instead. Also |
| 151 | + // applies in the local case below. |
| 152 | + List<Expression> newGroupings = new ArrayList<>(new LinkedHashSet<>(resolvedGroupings)); |
| 153 | + plan = a.with(p.child(), newGroupings, combineProjections(a.aggregates(), p.projections())); |
84 | 154 | } |
85 | 155 | } |
86 | 156 |
|
@@ -143,39 +213,6 @@ private static List<NamedExpression> combineProjections(List<? extends NamedExpr |
143 | 213 | return replaced; |
144 | 214 | } |
145 | 215 |
|
146 | | - private static List<Expression> combineUpperGroupingsAndLowerProjections( |
147 | | - List<? extends NamedExpression> upperGroupings, |
148 | | - List<? extends NamedExpression> lowerProjections |
149 | | - ) { |
150 | | - assert upperGroupings.size() <= 1 |
151 | | - || upperGroupings.stream() |
152 | | - .anyMatch(group -> group.anyMatch(expr -> expr instanceof GroupingFunction.NonEvaluatableGroupingFunction)) == false |
153 | | - : "CombineProjections only tested with a single CATEGORIZE with no additional groups"; |
154 | | - // Collect the alias map for resolving the source (f1 = 1, f2 = f1, etc..) |
155 | | - AttributeMap.Builder<Attribute> aliasesBuilder = AttributeMap.builder(); |
156 | | - for (NamedExpression ne : lowerProjections) { |
157 | | - // Record the aliases. |
158 | | - // Projections are just aliases for attributes, so casting is safe. |
159 | | - aliasesBuilder.put(ne.toAttribute(), (Attribute) Alias.unwrap(ne)); |
160 | | - } |
161 | | - var aliases = aliasesBuilder.build(); |
162 | | - |
163 | | - // Propagate any renames from the lower projection into the upper groupings. |
164 | | - // This can lead to duplicates: e.g. |
165 | | - // | EVAL x = y | STATS ... BY x, y |
166 | | - // All substitutions happen before; groupings must be attributes at this point except for non-evaluatable groupings which will be |
167 | | - // an alias like `c = CATEGORIZE(attribute)`. |
168 | | - // Therefore, it is correct to deduplicate based on simple equality (based on names) instead of name ids (Set vs. AttributeSet). |
169 | | - // TODO: The deduplication based on simple equality will be insufficient in case of multiple non-evaluatable groupings, e.g. for |
170 | | - // `| EVAL x = y | STATS ... BY CATEGORIZE(x), CATEGORIZE(y)`. That will require semantic equality instead. |
171 | | - LinkedHashSet<NamedExpression> resolvedGroupings = new LinkedHashSet<>(); |
172 | | - for (NamedExpression ne : upperGroupings) { |
173 | | - NamedExpression transformed = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.resolve(a, a)); |
174 | | - resolvedGroupings.add(transformed); |
175 | | - } |
176 | | - return new ArrayList<>(resolvedGroupings); |
177 | | - } |
178 | | - |
179 | 216 | /** |
180 | 217 | * Replace grouping alias previously contained in the aggregations that might have been projected away. |
181 | 218 | */ |
|
0 commit comments