|
20 | 20 | import org.elasticsearch.xpack.esql.rule.Rule; |
21 | 21 |
|
22 | 22 | import java.util.ArrayList; |
| 23 | +import java.util.HashSet; |
23 | 24 | import java.util.List; |
| 25 | +import java.util.Set; |
| 26 | + |
| 27 | +import static org.elasticsearch.xpack.esql.optimizer.rules.logical.TemporaryNameUtils.locallyUniqueTemporaryName; |
24 | 28 |
|
25 | 29 | /** |
26 | 30 | * Replace aliasing evals (eval x=a) with a projection which can be further combined / simplified. |
@@ -53,47 +57,63 @@ public LogicalPlan apply(LogicalPlan logicalPlan) { |
53 | 57 | private LogicalPlan rule(Eval eval) { |
54 | 58 | LogicalPlan plan = eval; |
55 | 59 |
|
56 | | - // holds simple aliases such as b = a, c = b, d = c |
57 | | - AttributeMap.Builder<Expression> basicAliasesBuilder = AttributeMap.builder(); |
58 | | - // same as above but keeps the original expression |
59 | | - AttributeMap.Builder<NamedExpression> basicAliasSourcesBuilder = AttributeMap.builder(); |
| 60 | + // Mostly, holds simple aliases from the eval, such as b = a, c = b, d = c, so we can resolve them in subsequent eval fields |
| 61 | + AttributeMap.Builder<Expression> renamesToPropagate = AttributeMap.builder(); |
| 62 | + // the aliases for the final projection - mostly, same as above but holds the final aliases rather than the original attributes |
| 63 | + AttributeMap.Builder<NamedExpression> projectionAliases = AttributeMap.builder(); |
| 64 | + // The names of attributes that are required to perform the aliases in a subsequent projection - if the next eval field |
| 65 | + // shadows one of these names, the subsequent projection won't work, so we need to perform a temporary rename. |
| 66 | + Set<String> namesRequiredForProjectionAliases = new HashSet<>(); |
60 | 67 |
|
61 | | - List<Alias> keptFields = new ArrayList<>(); |
| 68 | + List<Alias> newEvalFields = new ArrayList<>(); |
62 | 69 |
|
63 | 70 | var fields = eval.fields(); |
64 | | - for (int i = 0, size = fields.size(); i < size; i++) { |
65 | | - Alias field = fields.get(i); |
| 71 | + for (Alias alias : fields) { |
| 72 | + // propagate all previous aliases into the current field |
| 73 | + Alias field = (Alias) alias.transformUp(e -> renamesToPropagate.build().resolve(e, e)); |
| 74 | + String name = field.name(); |
| 75 | + Attribute attribute = field.toAttribute(); |
66 | 76 | Expression child = field.child(); |
67 | | - var attribute = field.toAttribute(); |
68 | | - // put the aliases in a separate map to separate the underlying resolve from other aliases |
69 | | - if (child instanceof Attribute) { |
70 | | - basicAliasesBuilder.put(attribute, child); |
71 | | - basicAliasSourcesBuilder.put(attribute, field); |
| 77 | + |
| 78 | + if (child instanceof Attribute renamedAttribute) { |
| 79 | + // Basic renaming - let's do that in the subsequent projection |
| 80 | + renamesToPropagate.put(attribute, renamedAttribute); |
| 81 | + projectionAliases.put(attribute, field); |
| 82 | + namesRequiredForProjectionAliases.add(renamedAttribute.name()); |
| 83 | + } else if (namesRequiredForProjectionAliases.contains(name)) { |
| 84 | + // Not a basic renaming, needs to remain in the eval. |
| 85 | + // The field may shadow one of the attributes that we will need to correctly perform the subsequent projection. |
| 86 | + // So, rename it in the eval! |
| 87 | + |
| 88 | + Alias newField = new Alias(field.source(), locallyUniqueTemporaryName(name), child, null, true); |
| 89 | + Attribute newAttribute = newField.toAttribute(); |
| 90 | + Alias reRenamedField = new Alias(field.source(), name, newAttribute, field.id(), field.synthetic()); |
| 91 | + projectionAliases.put(attribute, reRenamedField); |
| 92 | + // the renaming also needs to be propagated to eval fields to the right |
| 93 | + renamesToPropagate.put(attribute, newAttribute); |
| 94 | + |
| 95 | + newEvalFields.add(newField); |
72 | 96 | } else { |
73 | | - // be lazy and start replacing name aliases only if needed |
74 | | - if (basicAliasesBuilder.build().size() > 0) { |
75 | | - // update the child through the field |
76 | | - field = (Alias) field.transformUp(e -> basicAliasesBuilder.build().resolve(e, e)); |
77 | | - } |
78 | | - keptFields.add(field); |
| 97 | + // still not a basic renaming, but no risk of shadowing |
| 98 | + newEvalFields.add(field); |
79 | 99 | } |
80 | 100 | } |
81 | 101 |
|
82 | 102 | // at least one alias encountered, move it into a project |
83 | | - if (basicAliasesBuilder.build().size() > 0) { |
| 103 | + if (renamesToPropagate.build().size() > 0) { |
84 | 104 | // preserve the eval output (takes care of shadowing and order) but replace the basic aliases |
85 | 105 | List<NamedExpression> projections = new ArrayList<>(eval.output()); |
86 | | - var basicAliasSources = basicAliasSourcesBuilder.build(); |
| 106 | + var projectionAliasesMap = projectionAliases.build(); |
87 | 107 | // replace the removed aliases with their initial definition - however use the output to preserve the shadowing |
88 | 108 | for (int i = projections.size() - 1; i >= 0; i--) { |
89 | 109 | NamedExpression project = projections.get(i); |
90 | | - projections.set(i, basicAliasSources.getOrDefault(project, project)); |
| 110 | + projections.set(i, projectionAliasesMap.getOrDefault(project, project)); |
91 | 111 | } |
92 | 112 |
|
93 | 113 | LogicalPlan child = eval.child(); |
94 | | - if (keptFields.size() > 0) { |
| 114 | + if (newEvalFields.size() > 0) { |
95 | 115 | // replace the eval with just the kept fields |
96 | | - child = new Eval(eval.source(), eval.child(), keptFields); |
| 116 | + child = new Eval(eval.source(), eval.child(), newEvalFields); |
97 | 117 | } |
98 | 118 | // put the projection in place |
99 | 119 | plan = new Project(eval.source(), child, projections); |
|
0 commit comments