Skip to content

Commit f37c93e

Browse files
author
Andras Palinkas
authored
SQL: Resolve attributes recursively for improved subquery support (#69765) (#70325)
Previously we did not resolve the attributes recursively which meant that if a field or expression was re-aliased multiple times (through multiple levels of subqueries), the aliases were only resolved one level down. This led to failed query translation because `ReferenceAttribute`s were pointing to non-existing attributes during query translation. For example the query ```sql SELECT i AS j FROM ( SELECT int AS i FROM test) ORDER BY j ``` failed during translation because the `OrderBy` resolved the `j` ReferenceAttribute to another `i` ReferenceAttribute that was later removed by an Optimization: ``` OrderBy[[Order[j{r}#4,ASC,LAST]]] ! OrderBy[[Order[i{r}#2,ASC,LAST]]] \_Project[[j]] = \_Project[[j]] \_Project[[i]] ! \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] ! ``` By resolving the `Attributes` recursively both `j{r}` and `i{r}` will resolve to `test.int{f}` above: ``` OrderBy[[Order[test.int{f}#22,ASC,LAST]]] = OrderBy[[Order[test.int{f}#22,ASC,LAST]]] \_Project[[j]] = \_Project[[j]] \_Project[[i]] ! \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] \_EsRelation[test][date{f}#6, some{f}#7, some.string{f}#8, some.string..] ! ``` The scope of recursive resolution depends on how the `AttributeMap` is constructed and populated. Fixes #67237
1 parent a1ba75e commit f37c93e

File tree

10 files changed

+211
-36
lines changed

10 files changed

+211
-36
lines changed

‎x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeMap.java‎

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
*/
77
package org.elasticsearch.xpack.ql.expression;
88

9+
import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
10+
911
import java.util.AbstractSet;
1012
import java.util.Collection;
1113
import java.util.Iterator;
@@ -148,6 +150,8 @@ public static final <E> AttributeMap<E> emptyAttributeMap() {
148150
return EMPTY;
149151
}
150152

153+
private static final Object NOT_FOUND = new Object();
154+
151155
private final Map<AttributeWrapper, E> delegate;
152156
private Set<Attribute> keySet = null;
153157
private Collection<E> values = null;
@@ -238,14 +242,14 @@ public boolean isEmpty() {
238242
@Override
239243
public boolean containsKey(Object key) {
240244
if (key instanceof NamedExpression) {
241-
return delegate.keySet().contains(new AttributeWrapper(((NamedExpression) key).toAttribute()));
245+
return delegate.containsKey(new AttributeWrapper(((NamedExpression) key).toAttribute()));
242246
}
243247
return false;
244248
}
245249

246250
@Override
247251
public boolean containsValue(Object value) {
248-
return delegate.values().contains(value);
252+
return delegate.containsValue(value);
249253
}
250254

251255
@Override
@@ -258,10 +262,32 @@ public E get(Object key) {
258262

259263
@Override
260264
public E getOrDefault(Object key, E defaultValue) {
261-
E e;
262-
return (((e = get(key)) != null) || containsKey(key))
263-
? e
264-
: defaultValue;
265+
if (key instanceof NamedExpression) {
266+
return delegate.getOrDefault(new AttributeWrapper(((NamedExpression) key).toAttribute()), defaultValue);
267+
}
268+
return defaultValue;
269+
}
270+
271+
public E resolve(Object key) {
272+
return resolve(key, null);
273+
}
274+
275+
public E resolve(Object key, E defaultValue) {
276+
E value = defaultValue;
277+
E candidate = null;
278+
int allowedLookups = 1000;
279+
while ((candidate = get(key)) != null || containsKey(key)) {
280+
// instead of circling around, return
281+
if (candidate == key) {
282+
return candidate;
283+
}
284+
if (--allowedLookups == 0) {
285+
throw new QlIllegalArgumentException("Potential cycle detected");
286+
}
287+
key = candidate;
288+
value = candidate;
289+
}
290+
return value;
265291
}
266292

267293
@Override

‎x-pack/plugin/ql/src/test/java/org/elasticsearch/xpack/ql/expression/AttributeMapTests.java‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
package org.elasticsearch.xpack.ql.expression;
88

99
import org.elasticsearch.test.ESTestCase;
10+
import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
1011
import org.elasticsearch.xpack.ql.tree.Source;
1112
import org.elasticsearch.xpack.ql.type.DataTypes;
1213

@@ -18,6 +19,8 @@
1819
import java.util.Set;
1920

2021
import static java.util.stream.Collectors.toList;
22+
import static org.elasticsearch.xpack.ql.TestUtils.fieldAttribute;
23+
import static org.elasticsearch.xpack.ql.TestUtils.of;
2124
import static org.hamcrest.Matchers.arrayContaining;
2225
import static org.hamcrest.Matchers.arrayWithSize;
2326
import static org.hamcrest.Matchers.contains;
@@ -62,6 +65,64 @@ public void testAttributeMapWithSameAliasesCanResolveAttributes() {
6265
assertTrue(newAttributeMap.get(param2.toAttribute()) == param2.child());
6366
}
6467

68+
public void testResolve() {
69+
AttributeMap.Builder<Object> builder = AttributeMap.builder();
70+
Attribute one = a("one");
71+
Attribute two = fieldAttribute("two", DataTypes.INTEGER);
72+
Attribute three = fieldAttribute("three", DataTypes.INTEGER);
73+
Alias threeAlias = new Alias(Source.EMPTY, "three_alias", three);
74+
Alias threeAliasAlias = new Alias(Source.EMPTY, "three_alias_alias", threeAlias);
75+
builder.put(one, of("one"));
76+
builder.put(two, "two");
77+
builder.put(three, of("three"));
78+
builder.put(threeAlias.toAttribute(), threeAlias.child());
79+
builder.put(threeAliasAlias.toAttribute(), threeAliasAlias.child());
80+
AttributeMap<Object> map = builder.build();
81+
82+
assertEquals(of("one"), map.resolve(one));
83+
assertEquals("two", map.resolve(two));
84+
assertEquals(of("three"), map.resolve(three));
85+
assertEquals(of("three"), map.resolve(threeAlias));
86+
assertEquals(of("three"), map.resolve(threeAliasAlias));
87+
assertEquals(of("three"), map.resolve(threeAliasAlias, threeAlias));
88+
Attribute four = a("four");
89+
assertEquals("not found", map.resolve(four, "not found"));
90+
assertNull(map.resolve(four));
91+
assertEquals(four, map.resolve(four, four));
92+
}
93+
94+
public void testResolveOneHopCycle() {
95+
AttributeMap.Builder<Object> builder = AttributeMap.builder();
96+
Attribute a = fieldAttribute("a", DataTypes.INTEGER);
97+
Attribute b = fieldAttribute("b", DataTypes.INTEGER);
98+
builder.put(a, a);
99+
builder.put(b, a);
100+
AttributeMap<Object> map = builder.build();
101+
102+
assertEquals(a, map.resolve(a, "default"));
103+
assertEquals(a, map.resolve(b, "default"));
104+
assertEquals("default", map.resolve("non-existing-key", "default"));
105+
}
106+
107+
public void testResolveMultiHopCycle() {
108+
AttributeMap.Builder<Object> builder = AttributeMap.builder();
109+
Attribute a = fieldAttribute("a", DataTypes.INTEGER);
110+
Attribute b = fieldAttribute("b", DataTypes.INTEGER);
111+
Attribute c = fieldAttribute("c", DataTypes.INTEGER);
112+
Attribute d = fieldAttribute("d", DataTypes.INTEGER);
113+
builder.put(a, b);
114+
builder.put(b, c);
115+
builder.put(c, d);
116+
builder.put(d, a);
117+
AttributeMap<Object> map = builder.build();
118+
119+
// note: multi hop cycles should not happen, unless we have a
120+
// bug in the code that populates the AttributeMaps
121+
expectThrows(QlIllegalArgumentException.class, () -> {
122+
assertEquals(a, map.resolve(a, c));
123+
});
124+
}
125+
65126
private Alias createIntParameterAlias(int index, int value) {
66127
Source source = new Source(1, index * 5, "?");
67128
Literal literal = new Literal(source, value, DataTypes.INTEGER);

‎x-pack/plugin/sql/qa/server/src/main/resources/subselect.sql-spec‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,34 @@ basicGroupBy
1616
SELECT gender FROM (SELECT first_name AS f, last_name, gender FROM test_emp) GROUP BY gender ORDER BY gender ASC;
1717
basicGroupByAlias
1818
SELECT g FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY g ORDER BY g ASC;
19-
// @AwaitsFix(bugUrl = "follow-up to https://github.com/elastic/elasticsearch/pull/67216")
20-
basicGroupByWithFilterByAlias-Ignore
19+
basicGroupByWithFilterByAlias
2120
SELECT g FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) WHERE g IS NOT NULL GROUP BY g ORDER BY g ASC;
21+
basicGroupByRealiased
22+
SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY g ORDER BY g DESC NULLS last;
23+
basicGroupByRealiasedTwice
24+
SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) GROUP BY h ORDER BY h DESC NULLS last;
25+
basicOrderByRealiasedField
26+
SELECT g AS h FROM (SELECT first_name AS f, last_name, gender AS g FROM test_emp) ORDER BY g DESC NULLS first;
27+
28+
groupAndOrderByRealiasedExpression
29+
SELECT emp_group AS e, min_high_salary AS s
30+
FROM (
31+
SELECT emp_no % 2 AS emp_group, MIN(salary) AS min_high_salary
32+
FROM test_emp
33+
WHERE salary > 50000
34+
GROUP BY emp_group
35+
)
36+
ORDER BY e DESC;
37+
38+
multiLevelSelectStar
39+
SELECT * FROM (SELECT * FROM ( SELECT * FROM test_emp ));
40+
41+
multiLevelSelectStarWithAlias
42+
SELECT * FROM (SELECT * FROM ( SELECT * FROM test_emp ) b) c;
2243

44+
// AwaitsFix: https://github.com/elastic/elasticsearch/issues/69758
45+
filterAfterGroupBy-Ignore
46+
SELECT s2 AS s3 FROM (SELECT s AS s2 FROM ( SELECT salary AS s FROM test_emp) GROUP BY s2) WHERE s2 < 5 ORDER BY s3 DESC NULLS last;
2347

2448
countAndComplexCondition
2549
SELECT COUNT(*) as c FROM (SELECT * FROM test_emp WHERE gender IS NOT NULL) WHERE ABS(salary) > 0 GROUP BY gender ORDER BY gender;

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -638,12 +638,12 @@ protected LogicalPlan doRule(LogicalPlan plan) {
638638
AttributeMap.Builder<Expression> builder = AttributeMap.builder();
639639
// collect aliases
640640
child.forEachUp(p -> p.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child())));
641-
final Map<Attribute, Expression> collectRefs = builder.build();
641+
final AttributeMap<Expression> collectRefs = builder.build();
642642

643643
referencesStream = referencesStream.filter(r -> {
644644
for (Attribute attr : child.outputSet()) {
645645
if (attr instanceof ReferenceAttribute) {
646-
Expression source = collectRefs.getOrDefault(attr, attr);
646+
Expression source = collectRefs.resolve(attr, attr);
647647
// found a match, no need to resolve it further
648648
// so filter it out
649649
if (source.equals(r.child())) {

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java‎

Lines changed: 17 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -316,10 +316,11 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
316316
Map<Expression, Node<?>> missing = new LinkedHashMap<>();
317317

318318
o.order().forEach(oe -> {
319-
Expression e = oe.child();
319+
final Expression e = oe.child();
320+
final Expression resolvedE = attributeRefs.resolve(e, e);
320321

321322
// aggregates are allowed
322-
if (Functions.isAggregate(attributeRefs.getOrDefault(e, e))) {
323+
if (Functions.isAggregate(resolvedE)) {
323324
return;
324325
}
325326

@@ -340,8 +341,12 @@ private static boolean checkGroupByOrder(LogicalPlan p, Set<Failure> localFailur
340341
// e.g.: if "GROUP BY f2(f1(field))" you can "ORDER BY f4(f3(f2(f1(field))))"
341342
//
342343
// Also, make sure to compare attributes directly
343-
if (e.anyMatch(expression -> Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
344-
g -> expression.semanticEquals(expression instanceof Attribute ? Expressions.attribute(g) : g)))) {
344+
if (resolvedE.anyMatch(expression -> Expressions.anyMatch(groupingAndMatchingAggregatesAliases,
345+
g -> {
346+
Expression resolvedG = attributeRefs.resolve(g, g);
347+
resolvedG = expression instanceof Attribute ? Expressions.attribute(resolvedG) : resolvedG;
348+
return expression.semanticEquals(resolvedG);
349+
}))) {
345350
return;
346351
}
347352

@@ -406,7 +411,7 @@ private static boolean checkGroupByHavingHasOnlyAggs(Expression e, Set<Expressio
406411

407412
// resolve FunctionAttribute to backing functions
408413
if (e instanceof ReferenceAttribute) {
409-
e = attributeRefs.get(e);
414+
e = attributeRefs.resolve(e);
410415
}
411416

412417
// scalar functions can be a binary tree
@@ -484,7 +489,7 @@ private static boolean onlyRawFields(Iterable<? extends Expression> expressions,
484489

485490
expressions.forEach(e -> e.forEachDown(c -> {
486491
if (c instanceof ReferenceAttribute) {
487-
c = attributeRefs.getOrDefault(c, c);
492+
c = attributeRefs.resolve(c, c);
488493
}
489494
if (c instanceof Function) {
490495
localFailures.add(fail(c, "No functions allowed (yet); encountered [{}]", c.sourceText()));
@@ -579,7 +584,7 @@ private static boolean checkGroupMatch(Expression e, Node<?> source, List<Expres
579584

580585
// resolve FunctionAttribute to backing functions
581586
if (e instanceof ReferenceAttribute) {
582-
e = attributeRefs.get(e);
587+
e = attributeRefs.resolve(e);
583588
}
584589

585590
// scalar functions can be a binary tree
@@ -668,7 +673,7 @@ private static void checkFilterOnAggs(LogicalPlan p, Set<Failure> localFailures,
668673
LogicalPlan filterChild = filter.child();
669674
if (filterChild instanceof Aggregate == false) {
670675
filter.condition().forEachDown(Expression.class, e -> {
671-
if (Functions.isAggregate(attributeRefs.getOrDefault(e, e))) {
676+
if (Functions.isAggregate(attributeRefs.resolve(e, e))) {
672677
if (filterChild instanceof Project) {
673678
filter.condition().forEachDown(FieldAttribute.class,
674679
f -> localFailures.add(fail(e, "[{}] field must appear in the GROUP BY clause or in an aggregate function",
@@ -690,7 +695,7 @@ private static void checkFilterOnGrouping(LogicalPlan p, Set<Failure> localFailu
690695
if (p instanceof Filter) {
691696
Filter filter = (Filter) p;
692697
filter.condition().forEachDown(Expression.class, e -> {
693-
if (Functions.isGrouping(attributeRefs.getOrDefault(e, e))) {
698+
if (Functions.isGrouping(attributeRefs.resolve(e, e))) {
694699
localFailures
695700
.add(fail(e, "Cannot filter on grouping function [{}], use its argument instead", Expressions.name(e)));
696701
}
@@ -717,7 +722,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
717722
}
718723
};
719724
Consumer<Expression> checkForNested = e ->
720-
attributeRefs.getOrDefault(e, e).forEachUp(FieldAttribute.class, matchNested);
725+
attributeRefs.resolve(e, e).forEachUp(FieldAttribute.class, matchNested);
721726
Consumer<ScalarFunction> checkForNestedInFunction = f -> f.arguments().forEach(
722727
arg -> arg.forEachUp(FieldAttribute.class, matchNested));
723728

@@ -739,7 +744,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
739744

740745
// check in where (scalars not allowed)
741746
p.forEachDown(Filter.class, f -> f.condition().forEachUp(e ->
742-
attributeRefs.getOrDefault(e, e).forEachUp(ScalarFunction.class, sf -> {
747+
attributeRefs.resolve(e, e).forEachUp(ScalarFunction.class, sf -> {
743748
if (sf instanceof BinaryComparison == false &&
744749
sf instanceof IsNull == false &&
745750
sf instanceof IsNotNull == false &&
@@ -758,7 +763,7 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
758763

759764
// check in order by (scalars not allowed)
760765
p.forEachDown(OrderBy.class, ob -> ob.order().forEach(o -> o.forEachUp(e ->
761-
attributeRefs.getOrDefault(e, e).forEachUp(ScalarFunction.class, checkForNestedInFunction)
766+
attributeRefs.resolve(e, e).forEachUp(ScalarFunction.class, checkForNestedInFunction)
762767
)));
763768
if (nested.isEmpty() == false) {
764769
localFailures.add(

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -222,8 +222,8 @@ public LogicalPlan apply(LogicalPlan plan) {
222222
AttributeMap.Builder<Expression> builder = AttributeMap.builder();
223223
// collect aliases
224224
plan.forEachExpressionUp(Alias.class, a -> builder.put(a.toAttribute(), a.child()));
225-
final Map<Attribute, Expression> collectRefs = builder.build();
226-
java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefs.getOrDefault(r, r);
225+
final AttributeMap<Expression> collectRefs = builder.build();
226+
java.util.function.Function<ReferenceAttribute, Expression> replaceReference = r -> collectRefs.resolve(r, r);
227227

228228
plan = plan.transformUp(p -> {
229229
// non attribute defining plans get their references removed
@@ -300,7 +300,7 @@ static class PruneOrderByNestedFields extends OptimizerRule<Project> {
300300
private void findNested(Expression exp, AttributeMap<Function> functions, Consumer<FieldAttribute> onFind) {
301301
exp.forEachUp(e -> {
302302
if (e instanceof ReferenceAttribute) {
303-
Function f = functions.get(e);
303+
Function f = functions.resolve(e);
304304
if (f != null) {
305305
findNested(f, functions, onFind);
306306
}
@@ -578,7 +578,7 @@ private List<NamedExpression> combineProjections(List<? extends NamedExpression>
578578
// replace any matching attribute with a lower alias (if there's a match)
579579
// but clean-up non-top aliases at the end
580580
for (NamedExpression ne : upper) {
581-
NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.getOrDefault(a, a));
581+
NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.resolve(a, a));
582582
replaced.add((NamedExpression) CleanAliases.trimNonTopLevelAliases(replacedExp));
583583
}
584584
return replaced;

‎x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryFolder.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -600,10 +600,13 @@ else if (target.foldable()) {
600600
else {
601601
GroupByKey matchingGroup = null;
602602
if (groupingContext != null) {
603+
target = queryC.aliases().resolve(target, target);
604+
id = Expressions.id(target);
603605
matchingGroup = groupingContext.groupFor(target);
604606
Check.notNull(matchingGroup, "Cannot find group [{}]", Expressions.name(ne));
605607

606-
queryC = queryC.addColumn(new GroupByRef(matchingGroup.id(), null, isDateBased(ne.dataType())), id);
608+
queryC = queryC.addColumn(
609+
new GroupByRef(matchingGroup.id(), null, isDateBased(ne.dataType())), id);
607610
}
608611
// fallback
609612
else {
@@ -707,7 +710,7 @@ protected PhysicalPlan rule(OrderExec plan) {
707710

708711
// if it's a reference, get the target expression
709712
if (orderExpression instanceof ReferenceAttribute) {
710-
orderExpression = qContainer.aliases().get(orderExpression);
713+
orderExpression = qContainer.aliases().resolve(orderExpression);
711714
}
712715
String lookup = Expressions.id(orderExpression);
713716
GroupByKey group = qContainer.findGroupForAgg(lookup);

0 commit comments

Comments
 (0)