Skip to content

Commit 2145a40

Browse files
smalyshevsmalyshev
committed
moar tests
1 parent 8c0c305 commit 2145a40

File tree

4 files changed

+292
-1
lines changed

4 files changed

+292
-1
lines changed

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTestUtils.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,15 @@ public static EnrichResolution defaultEnrichResolution() {
194194
"languages_idx",
195195
"mapping-languages.json"
196196
);
197+
loadEnrichPolicyResolution(
198+
enrichResolution,
199+
Enrich.Mode.REMOTE,
200+
MATCH_TYPE,
201+
"languages_remote",
202+
"language_code",
203+
"languages_idx",
204+
"mapping-languages.json"
205+
);
197206
return enrichResolution;
198207
}
199208

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/AbstractLogicalPlanOptimizerTests.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.elasticsearch.xpack.esql.index.EsIndex;
2020
import org.elasticsearch.xpack.esql.index.IndexResolution;
2121
import org.elasticsearch.xpack.esql.parser.EsqlParser;
22+
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
2223
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
2324
import org.junit.BeforeClass;
2425

@@ -27,6 +28,7 @@
2728
import java.util.Set;
2829

2930
import static java.util.Collections.emptyMap;
31+
import static org.elasticsearch.xpack.core.enrich.EnrichPolicy.MATCH_TYPE;
3032
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_VERIFIER;
3133
import static org.elasticsearch.xpack.esql.EsqlTestUtils.emptyInferenceResolution;
3234
import static org.elasticsearch.xpack.esql.EsqlTestUtils.loadMapping;
@@ -76,6 +78,24 @@ public static void init() {
7678
logicalOptimizer = new LogicalPlanOptimizer(logicalOptimizerCtx);
7779
enrichResolution = new EnrichResolution();
7880
AnalyzerTestUtils.loadEnrichPolicyResolution(enrichResolution, "languages_idx", "id", "languages_idx", "mapping-languages.json");
81+
AnalyzerTestUtils.loadEnrichPolicyResolution(
82+
enrichResolution,
83+
Enrich.Mode.REMOTE,
84+
MATCH_TYPE,
85+
"languages_remote",
86+
"id",
87+
"languages_idx",
88+
"mapping-languages.json"
89+
);
90+
AnalyzerTestUtils.loadEnrichPolicyResolution(
91+
enrichResolution,
92+
Enrich.Mode.COORDINATOR,
93+
MATCH_TYPE,
94+
"languages_coordinator",
95+
"id",
96+
"languages_idx",
97+
"mapping-languages.json"
98+
);
7999

80100
// Most tests used data from the test index, so we load it here, and use it in the plan() function.
81101
mapping = loadMapping("mapping-basic.json");

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

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
9+
10+
import org.elasticsearch.xpack.esql.VerificationException;
11+
import org.elasticsearch.xpack.esql.expression.Foldables;
12+
import org.elasticsearch.xpack.esql.optimizer.AbstractLogicalPlanOptimizerTests;
13+
import org.elasticsearch.xpack.esql.plan.logical.Dissect;
14+
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
15+
import org.elasticsearch.xpack.esql.plan.logical.Eval;
16+
import org.elasticsearch.xpack.esql.plan.logical.Limit;
17+
import org.elasticsearch.xpack.esql.plan.logical.Project;
18+
19+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
20+
import static org.hamcrest.Matchers.containsString;
21+
import static org.hamcrest.Matchers.equalTo;
22+
import static org.hamcrest.Matchers.is;
23+
import static org.hamcrest.Matchers.not;
24+
25+
public class HoistRemoteEnrichLimitTests extends AbstractLogicalPlanOptimizerTests {
26+
27+
/**
28+
* <pre>
29+
* Limit[10[INTEGER],true,false]
30+
* \_Enrich[REMOTE,languages_remote[KEYWORD],id{r}#4,{"match":{"indices":[],"match_field":"id",
31+
* "enrich_fields":["language_code","language_name"]}},{=languages_idx},[language_code{r}#20, language_name{r}#21]]
32+
* \_Eval[[emp_no{f}#6 AS id#4]]
33+
* \_Limit[10[INTEGER],false,true]
34+
* \_EsRelation[test][_meta_field{f}#12, emp_no{f}#6, first_name{f}#7, ...]
35+
* </pre>
36+
*/
37+
public void testLimitWithinRemoteEnrich() {
38+
var plan = plan(randomFrom("""
39+
from test
40+
| EVAL id = emp_no
41+
| LIMIT 10
42+
| ENRICH _remote:languages_remote
43+
""", """
44+
from test
45+
| LIMIT 10
46+
| EVAL id = emp_no
47+
| ENRICH _remote:languages_remote
48+
""")); // it should be the same in any order
49+
50+
var limit = as(plan, Limit.class);
51+
assertTrue(limit.duplicated());
52+
assertFalse(limit.local());
53+
var enrich = as(limit.child(), Enrich.class);
54+
assertThat(enrich.mode(), is(Enrich.Mode.REMOTE));
55+
var eval = as(enrich.child(), Eval.class);
56+
var innerLimit = as(eval.child(), Limit.class);
57+
assertFalse(innerLimit.duplicated());
58+
assertTrue(innerLimit.local());
59+
}
60+
61+
/**
62+
* <pre>
63+
* Project[[salary{f}#19 AS wage#10, emp_no{f}#14 AS id#4, first_name{r}#11, language_code{r}#28, language_name{r}#29]]
64+
* \_Limit[5[INTEGER],true,false]
65+
* \_Enrich[REMOTE,languages_remote[KEYWORD],emp_no{f}#14,{"match":{"indices":[],"match_field":"id","enrich_fields":["languag
66+
* e_code","language_name"]}},{=languages_idx},[language_code{r}#28, language_name{r}#29]]
67+
* \_Dissect[first_name{f}#15,Parser[pattern=%{first_name}s, appendSeparator=, parser=org.elasticsearch.dissect.DissectParse
68+
* r@7d5c7931],[first_name{r}#11]]
69+
* \_Limit[5[INTEGER],false,true]
70+
* \_EsRelation[test][_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, ..]
71+
* </pre>
72+
*/
73+
public void testManyLimitsWithinRemoteEnrich() {
74+
var plan = plan("""
75+
from test
76+
| LIMIT 10
77+
| EVAL id = emp_no
78+
| KEEP first_name, salary, id
79+
| RENAME salary AS wage
80+
| DISSECT first_name "%{first_name}s"
81+
| LIMIT 5
82+
| ENRICH _remote:languages_remote
83+
""");
84+
85+
var project = as(plan, Project.class);
86+
var limit = as(project.child(), Limit.class);
87+
assertTrue(limit.duplicated());
88+
assertFalse(limit.local());
89+
assertThat(Foldables.limitValue(limit.limit(), limit.sourceText()), equalTo(5));
90+
var enrich = as(limit.child(), Enrich.class);
91+
assertThat(enrich.mode(), is(Enrich.Mode.REMOTE));
92+
var dissect = as(enrich.child(), Dissect.class);
93+
var innerLimit = as(dissect.child(), Limit.class);
94+
assertFalse(innerLimit.duplicated());
95+
assertTrue(innerLimit.local());
96+
assertThat(Foldables.limitValue(innerLimit.limit(), innerLimit.sourceText()), equalTo(5));
97+
}
98+
99+
/**
100+
* Project[[first_name{f}#14, salary{f}#18 AS wage#11, emp_no{f}#13 AS id#4, language_code{r}#32, language_name{r}#33]]
101+
* \_Limit[5[INTEGER],true,false]
102+
* \_Enrich[REMOTE,languages_remote[KEYWORD],emp_no{f}#13,{"match":{"indices":[],"match_field":"id","enrich_fields":["languag
103+
* e_code","language_name"]}},{=languages_idx},[language_code{r}#32, language_name{r}#33]]
104+
* \_Limit[5[INTEGER],true,true]
105+
* \_Enrich[REMOTE,languages_remote[KEYWORD],emp_no{f}#13,{"match":{"indices":[],"match_field":"id","enrich_fields":["languag
106+
* e_code","language_name"]}},{=languages_idx},[language_code{r}#27, language_name{r}#28]]
107+
* \_Limit[5[INTEGER],false,true]
108+
* \_EsRelation[test][_meta_field{f}#19, emp_no{f}#13, first_name{f}#14, ..]
109+
*/
110+
public void testLimitsWithinRemoteEnrichTwice() {
111+
var plan = plan("""
112+
from test
113+
| LIMIT 10
114+
| EVAL id = emp_no
115+
| KEEP first_name, salary, id
116+
| ENRICH _remote:languages_remote
117+
| RENAME salary AS wage
118+
| LIMIT 5
119+
| ENRICH _remote:languages_remote
120+
""");
121+
var project = as(plan, Project.class);
122+
var limit = as(project.child(), Limit.class);
123+
assertTrue(limit.duplicated());
124+
assertFalse(limit.local());
125+
assertThat(Foldables.limitValue(limit.limit(), limit.sourceText()), equalTo(5));
126+
var enrich = as(limit.child(), Enrich.class);
127+
assertThat(enrich.mode(), is(Enrich.Mode.REMOTE));
128+
var innerLimit = as(enrich.child(), Limit.class);
129+
assertTrue(innerLimit.duplicated());
130+
assertTrue(innerLimit.local());
131+
assertThat(Foldables.limitValue(innerLimit.limit(), innerLimit.sourceText()), equalTo(5));
132+
var secondEnrich = as(innerLimit.child(), Enrich.class);
133+
assertThat(secondEnrich.mode(), is(Enrich.Mode.REMOTE));
134+
var innermostLimit = as(secondEnrich.child(), Limit.class);
135+
assertFalse(innermostLimit.duplicated());
136+
assertTrue(innermostLimit.local());
137+
assertThat(Foldables.limitValue(innermostLimit.limit(), innermostLimit.sourceText()), equalTo(5));
138+
}
139+
140+
// These cases do not get hoisting, and it's ok
141+
public void testLimitWithinOtherEnrich() {
142+
String enrichPolicy = randomFrom("languages_idx", "_any:languages_idx", "_coordinator:languages_coordinator");
143+
var plan = plan(String.format("""
144+
from test
145+
| EVAL id = emp_no
146+
| LIMIT 10
147+
| ENRICH %s
148+
""", enrichPolicy));
149+
// Here ENRICH is on top - no hoisting happens
150+
var enrich = as(plan, Enrich.class);
151+
assertThat(enrich.mode(), not(is(Enrich.Mode.REMOTE)));
152+
}
153+
154+
// Non-cardinality preserving commands after limit
155+
public void testFilterLimitThenEnrich() {
156+
// Hoisting does not happen, so the verifier fails since LIMIT is before remote ENRICH
157+
failPlan("""
158+
from test
159+
| EVAL id = emp_no
160+
| LIMIT 10
161+
| WHERE first_name != "john"
162+
| ENRICH _remote:languages_remote
163+
""", "ENRICH with remote policy can't be executed after [LIMIT 10]");
164+
}
165+
166+
public void testMvExpandLimitThenEnrich() {
167+
// Hoisting does not happen, so the verifier fails since LIMIT is before remote ENRICH
168+
failPlan("""
169+
from test
170+
| EVAL id = emp_no
171+
| LIMIT 10
172+
| MV_EXPAND languages
173+
| ENRICH _remote:languages_remote
174+
""", "MV_EXPAND after LIMIT is incompatible with remote ENRICH");
175+
}
176+
177+
// Other cases where hoisting does not happen:
178+
// - ExecutesOn.COORDINATOR - this fails the verifier
179+
// - PipelineBreaker - all relevant ones are also ExecutesOn.COORDINATOR
180+
181+
private void failPlan(String esql, String reason) {
182+
var e = expectThrows(VerificationException.class, () -> plan(esql));
183+
assertThat(e.getMessage(), containsString(reason));
184+
}
185+
}

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

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,23 @@
1616
import org.elasticsearch.xpack.esql.expression.Order;
1717
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToInteger;
1818
import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
19+
import org.elasticsearch.xpack.esql.plan.logical.Enrich;
1920
import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
2021
import org.elasticsearch.xpack.esql.plan.logical.Eval;
2122
import org.elasticsearch.xpack.esql.plan.logical.Filter;
2223
import org.elasticsearch.xpack.esql.plan.logical.Limit;
2324
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
25+
import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
2426
import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
2527
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
2628
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
2729
import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
30+
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
31+
import org.elasticsearch.xpack.esql.plan.logical.join.JoinConfig;
32+
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
2833

2934
import java.util.List;
35+
import java.util.Map;
3036
import java.util.function.BiConsumer;
3137
import java.util.function.BiFunction;
3238

@@ -42,7 +48,7 @@
4248

4349
public class PushDownAndCombineLimitsTests extends ESTestCase {
4450

45-
private static class PushDownLimitTestCase<PlanType extends UnaryPlan> {
51+
private static class PushDownLimitTestCase<PlanType extends LogicalPlan> {
4652
private final Class<PlanType> clazz;
4753
private final BiFunction<LogicalPlan, Attribute, PlanType> planBuilder;
4854
private final BiConsumer<PlanType, PlanType> planChecker;
@@ -102,6 +108,25 @@ public void checkOptimizedPlan(LogicalPlan basePlan, LogicalPlan optimizedPlan)
102108
assertEquals(basePlan.rerankFields(), optimizedPlan.rerankFields());
103109
assertEquals(basePlan.scoreAttribute(), optimizedPlan.scoreAttribute());
104110
}
111+
),
112+
new PushDownLimitTestCase<>(
113+
Enrich.class,
114+
(plan, attr) -> new Enrich(
115+
EMPTY,
116+
plan,
117+
randomFrom(Enrich.Mode.ANY, Enrich.Mode.COORDINATOR),
118+
randomLiteral(KEYWORD),
119+
attr,
120+
null,
121+
Map.of(),
122+
List.of()
123+
),
124+
(basePlan, optimizedPlan) -> {
125+
assertEquals(basePlan.source(), optimizedPlan.source());
126+
assertEquals(basePlan.mode(), optimizedPlan.mode());
127+
assertEquals(basePlan.policyName(), optimizedPlan.policyName());
128+
assertEquals(basePlan.matchField(), optimizedPlan.matchField());
129+
}
105130
)
106131
);
107132

@@ -174,6 +199,58 @@ public void testNonPushableLimit() {
174199
}
175200
}
176201

202+
private static final List<PushDownLimitTestCase<? extends LogicalPlan>> DUPLICATING_TEST_CASES = List.of(
203+
new PushDownLimitTestCase<>(
204+
Enrich.class,
205+
(plan, attr) -> new Enrich(EMPTY, plan, Enrich.Mode.REMOTE, randomLiteral(KEYWORD), attr, null, Map.of(), List.of()),
206+
(basePlan, optimizedPlan) -> {
207+
assertEquals(basePlan.source(), optimizedPlan.source());
208+
assertEquals(basePlan.mode(), optimizedPlan.mode());
209+
assertEquals(basePlan.policyName(), optimizedPlan.policyName());
210+
assertEquals(basePlan.matchField(), optimizedPlan.matchField());
211+
var limit = as(optimizedPlan.child(), Limit.class);
212+
assertTrue(limit.local());
213+
assertFalse(limit.duplicated());
214+
}
215+
),
216+
new PushDownLimitTestCase<>(MvExpand.class, (plan, attr) -> new MvExpand(EMPTY, plan, attr, attr), (basePlan, optimizedPlan) -> {
217+
assertEquals(basePlan.source(), optimizedPlan.source());
218+
assertEquals(basePlan.expanded(), optimizedPlan.expanded());
219+
var limit = as(optimizedPlan.child(), Limit.class);
220+
assertFalse(limit.local());
221+
assertFalse(limit.duplicated());
222+
}),
223+
new PushDownLimitTestCase<>(
224+
Join.class,
225+
(plan, attr) -> new Join(EMPTY, plan, plan, new JoinConfig(JoinTypes.LEFT, List.of(), List.of(), attr)),
226+
(basePlan, optimizedPlan) -> {
227+
assertEquals(basePlan.source(), optimizedPlan.source());
228+
var limit = as(optimizedPlan.left(), Limit.class);
229+
assertFalse(limit.local());
230+
assertFalse(limit.duplicated());
231+
}
232+
)
233+
234+
);
235+
236+
public void testPushableLimitDuplicate() {
237+
FieldAttribute a = getFieldAttribute("a");
238+
FieldAttribute b = getFieldAttribute("b");
239+
EsRelation relation = relation().withAttributes(List.of(a, b));
240+
241+
for (PushDownLimitTestCase<? extends LogicalPlan> duplicatingTestCase : DUPLICATING_TEST_CASES) {
242+
int precedingLimitValue = randomIntBetween(1, 10_000);
243+
Limit precedingLimit = new Limit(EMPTY, new Literal(EMPTY, precedingLimitValue, INTEGER), relation);
244+
LogicalPlan duplicatingLimitTestPlan = duplicatingTestCase.buildPlan(precedingLimit, a);
245+
int upperLimitValue = randomIntBetween(1, 10_000);
246+
Limit upperLimit = new Limit(EMPTY, new Literal(EMPTY, upperLimitValue, INTEGER), duplicatingLimitTestPlan);
247+
Limit optimizedPlan = as(optimizePlan(upperLimit), Limit.class);
248+
duplicatingTestCase.checkOptimizedPlan(duplicatingLimitTestPlan, optimizedPlan.child());
249+
assertTrue(optimizedPlan.duplicated());
250+
assertFalse(optimizedPlan.local());
251+
}
252+
}
253+
177254
private LogicalPlan optimizePlan(LogicalPlan plan) {
178255
return new PushDownAndCombineLimits().apply(plan, unboundLogicalOptimizerContext());
179256
}

0 commit comments

Comments
 (0)