Skip to content

Commit 21b7bde

Browse files
alex-spiesalex-spies
committed
Add first edge case optimizer tests
1 parent ac13587 commit 21b7bde

File tree

1 file changed

+150
-30
lines changed

1 file changed

+150
-30
lines changed

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownJoinPastProjectTests.java‎

Lines changed: 150 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -8,34 +8,45 @@
88
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
99

1010
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
11+
import org.elasticsearch.xpack.esql.core.expression.Alias;
1112
import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
13+
import org.elasticsearch.xpack.esql.core.expression.Expression;
14+
import org.elasticsearch.xpack.esql.core.expression.Expressions;
15+
import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
16+
import org.elasticsearch.xpack.esql.core.expression.Literal;
17+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
18+
import org.elasticsearch.xpack.esql.core.tree.Source;
19+
import org.elasticsearch.xpack.esql.core.type.DataType;
20+
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
1221
import org.elasticsearch.xpack.esql.optimizer.AbstractLogicalPlanOptimizerTests;
1322
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
23+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
1424
import org.elasticsearch.xpack.esql.plan.logical.Project;
1525
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
26+
import org.elasticsearch.xpack.esql.plan.logical.join.JoinConfig;
27+
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
1628

1729
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
1830
import static org.elasticsearch.xpack.esql.EsqlTestUtils.asLimit;
31+
import static org.hamcrest.Matchers.contains;
32+
import static org.hamcrest.Matchers.startsWith;
1933

2034
public class PushDownJoinPastProjectTests extends AbstractLogicalPlanOptimizerTests {
35+
2136
/**
2237
* Expects
23-
*
38+
* <p>
2439
* Project[[languages{f}#16, emp_no{f}#13, languages{f}#16 AS language_code#6, language_name{f}#27]]
2540
* \_Limit[1000[INTEGER],true]
26-
* \_Join[LEFT,[languages{f}#16],[languages{f}#16],[language_code{f}#26]]
27-
* |_Limit[1000[INTEGER],true]
28-
* | \_Join[LEFT,[languages{f}#16],[languages{f}#16],[language_code{f}#24]]
29-
* | |_Limit[1000[INTEGER],false]
30-
* | | \_EsRelation[test][_meta_field{f}#19, emp_no{f}#13, first_name{f}#14, ..]
31-
* | \_EsRelation[languages_lookup][LOOKUP][language_code{f}#24]
32-
* \_EsRelation[languages_lookup][LOOKUP][language_code{f}#26, language_name{f}#27]
41+
* \_Join[LEFT,[languages{f}#16],[languages{f}#16],[language_code{f}#26]]
42+
* |_Limit[1000[INTEGER],true]
43+
* | \_Join[LEFT,[languages{f}#16],[languages{f}#16],[language_code{f}#24]]
44+
* | |_Limit[1000[INTEGER],false]
45+
* | | \_EsRelation[test][_meta_field{f}#19, emp_no{f}#13, first_name{f}#14, ..]
46+
* | \_EsRelation[languages_lookup][LOOKUP][language_code{f}#24]
47+
* \_EsRelation[languages_lookup][LOOKUP][language_code{f}#26, language_name{f}#27]
3348
*/
34-
public void testMultipleLookupProject() {
35-
// TODO a test case where pushing down past the RENAME would shadow
36-
// analogous to
37-
// Project[[x{f}#1, y{f}#2 as z, $$y{r}#3 as y]]
38-
// \_Eval[[2 * x{f}#1 as $$y]]
49+
public void testMultipleLookups() {
3950
assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());
4051

4152
String query = """
@@ -57,29 +68,138 @@ public void testMultipleLookupProject() {
5768
var lookupRel1 = as(join1.right(), EsRelation.class);
5869
var limit1 = asLimit(join1.left(), 1000, true);
5970

60-
AttributeSet lookupIndexFields1 = lookupRel1.outputSet();
61-
var rightKeys1 = join1.config().rightFields();
62-
var leftKeys1 = join1.config().leftFields();
63-
// Left join key should be updated to use an attribute from the main index directly
64-
assertTrue(leftKeys1.size() == 1 && leftKeys1.get(0).name() == "languages");
65-
assertTrue(rightKeys1.size() == 1 && rightKeys1.get(0).name() == "language_code");
66-
assertTrue(lookupIndexFields1.contains(rightKeys1.get(0)));
67-
6871
var join2 = as(limit1.child(), Join.class);
6972
var lookupRel2 = as(join2.right(), EsRelation.class);
7073
var limit2 = asLimit(join2.left(), 1000, false);
7174

72-
AttributeSet lookupIndexFields2 = lookupRel2.outputSet();
73-
var rightKeys2 = join2.config().rightFields();
74-
var leftKeys2 = join2.config().leftFields();
75+
var mainRel = as(limit2.child(), EsRelation.class);
76+
7577
// Left join key should be updated to use an attribute from the main index directly
76-
assertTrue(leftKeys2.size() == 1 && leftKeys2.get(0).name() == "languages");
77-
assertTrue(rightKeys2.size() == 1 && rightKeys2.get(0).name() == "language_code");
78-
assertTrue(lookupIndexFields2.contains(rightKeys2.get(0)));
78+
assertLeftJoinConfig(join1.config(), "languages", mainRel.outputSet(), "language_code", lookupRel1.outputSet());
79+
assertLeftJoinConfig(join2.config(), "languages", mainRel.outputSet(), "language_code", lookupRel2.outputSet());
80+
}
81+
82+
/**
83+
* Expects
84+
* <p>
85+
* Project[[languages{f}#14 AS language_code#4, language_name{f}#23]]
86+
* \_Limit[1000[INTEGER],true]
87+
* \_Join[LEFT,[languages{f}#14],[languages{f}#14],[language_code{f}#22]]
88+
* |_Limit[1000[INTEGER],false]
89+
* | \_EsRelation[test][_meta_field{f}#17, emp_no{f}#11, first_name{f}#12, ..]
90+
* \_EsRelation[languages_lookup][LOOKUP][language_code{f}#22, language_name{f}#23]
91+
*/
92+
public void testShadowingBeforePushdown() {
93+
assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());
94+
95+
String query = """
96+
FROM test
97+
| RENAME languages AS language_code, last_name AS language_name
98+
| KEEP language_code, language_name
99+
| LOOKUP JOIN languages_lookup ON language_code
100+
""";
79101

102+
var plan = optimizedPlan(query);
103+
104+
var project = as(plan, Project.class);
105+
var limit1 = asLimit(project.child(), 1000, true);
106+
var join = as(limit1.child(), Join.class);
107+
var lookupRel = as(join.right(), EsRelation.class);
108+
var limit2 = asLimit(join.left(), 1000, false);
80109
var mainRel = as(limit2.child(), EsRelation.class);
81-
AttributeSet mainFields = mainRel.outputSet();
82-
assertTrue(mainFields.contains(leftKeys1.get(0)));
83-
assertTrue(mainFields.contains(leftKeys2.get(0)));
110+
111+
// Left join key should be updated to use an attribute from the main index directly
112+
assertLeftJoinConfig(join.config(), "languages", mainRel.outputSet(), "language_code", lookupRel.outputSet());
113+
}
114+
115+
/**
116+
* Expects
117+
* <p>
118+
* Project[[languages{f}#17 AS language_code#9, $$language_name$temp_name$27{r$}#28 AS foo#12, language_name{f}#26]]
119+
* \_Limit[1000[INTEGER],true]
120+
* \_Join[LEFT,[languages{f}#17],[languages{f}#17],[language_code{f}#25]]
121+
* |_Eval[[salary{f}#19 * 2[INTEGER] AS language_name#4, language_name{r}#4 AS $$language_name$temp_name$27#28]]
122+
* | \_Limit[1000[INTEGER],false]
123+
* | \_EsRelation[test][_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, ..]
124+
* \_EsRelation[languages_lookup][LOOKUP][language_code{f}#25, language_name{f}#26]
125+
*/
126+
public void testShadowingAfterPushdown() {
127+
assumeTrue("Requires LOOKUP JOIN", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());
128+
129+
String query = """
130+
FROM test
131+
| EVAL language_name = 2*salary
132+
| KEEP languages, language_name
133+
| RENAME languages AS language_code, language_name AS foo
134+
| LOOKUP JOIN languages_lookup ON language_code
135+
""";
136+
137+
var plan = optimizedPlan(query);
138+
139+
var project = as(plan, Project.class);
140+
var limit1 = asLimit(project.child(), 1000, true);
141+
var join = as(limit1.child(), Join.class);
142+
var lookupRel = as(join.right(), EsRelation.class);
143+
144+
var eval = as(join.left(), Eval.class);
145+
var limit2 = asLimit(eval.child(), 1000, false);
146+
var mainRel = as(limit2.child(), EsRelation.class);
147+
148+
var projections = project.projections();
149+
assertThat(Expressions.names(projections), contains("language_code", "foo", "language_name"));
150+
151+
var languages = unwrapAlias(projections.get(0), FieldAttribute.class);
152+
assertEquals("languages", languages.fieldName().string());
153+
assertTrue(mainRel.outputSet().contains(languages));
154+
155+
var tempName = unwrapAlias(projections.get(1), ReferenceAttribute.class);
156+
assertThat(tempName.name(), startsWith("$$language_name$temp_name$"));
157+
assertTrue(eval.outputSet().contains(tempName));
158+
159+
var languageName = as(projections.get(2), FieldAttribute.class);
160+
assertTrue(lookupRel.outputSet().contains(languageName));
161+
162+
var evalExprs = eval.fields();
163+
assertThat(Expressions.names(evalExprs), contains("language_name", tempName.name()));
164+
var originalLanguageName = unwrapAlias(evalExprs.get(1), ReferenceAttribute.class);
165+
assertEquals("language_name", originalLanguageName.name());
166+
assertTrue(originalLanguageName.semanticEquals(as(evalExprs.get(0), Alias.class).toAttribute()));
167+
168+
var mul = unwrapAlias(evalExprs.get(0), Mul.class);
169+
assertEquals(new Literal(Source.EMPTY, 2, DataType.INTEGER), mul.right());
170+
var salary = as(mul.left(), FieldAttribute.class);
171+
assertEquals("salary", salary.fieldName().string());
172+
assertTrue(mainRel.outputSet().contains(salary));
173+
174+
assertLeftJoinConfig(join.config(), "languages", mainRel.outputSet(), "language_code", lookupRel.outputSet());
175+
}
176+
177+
private static void assertLeftJoinConfig(
178+
JoinConfig config,
179+
String expectedLeftFieldName,
180+
AttributeSet leftSourceAttributes,
181+
String expectedRightFieldName,
182+
AttributeSet rightSourceAttributes
183+
) {
184+
assertSame(config.type(), JoinTypes.LEFT);
185+
186+
var leftKeys = config.leftFields();
187+
var rightKeys = config.rightFields();
188+
189+
assertEquals(1, leftKeys.size());
190+
var leftKey = leftKeys.get(0);
191+
assertEquals(expectedLeftFieldName, leftKey.name());
192+
assertTrue(leftSourceAttributes.contains(leftKey));
193+
194+
assertEquals(1, rightKeys.size());
195+
var rightKey = rightKeys.get(0);
196+
assertEquals(expectedRightFieldName, rightKey.name());
197+
assertTrue(rightSourceAttributes.contains(rightKey));
198+
}
199+
200+
private static <T> T unwrapAlias(Expression alias, Class<T> type) {
201+
var child = as(alias, Alias.class).child();
202+
203+
return as(child, type);
84204
}
85205
}

0 commit comments

Comments
 (0)