Skip to content

Commit be82a74

Browse files
fang-xing-esqlfang-xing-esql
committed
handle subquery with fork and full text function better
1 parent 1853615 commit be82a74

File tree

18 files changed

+514
-44
lines changed

18 files changed

+514
-44
lines changed

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,6 +1355,70 @@ public void testTopLevelFilterWithSubqueriesInFromCommand() throws IOException {
13551355
assertResultMap(result, matchesList().item(matchesMap().entry("name", "count(*)").entry("type", "long")), List.of(List.of(8)));
13561356
}
13571357

1358+
public void testNestedSubqueries() throws IOException {
1359+
assumeTrue("subqueries in from command", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
1360+
1361+
bulkLoadTestData(10);
1362+
1363+
ResponseException re = expectThrows(
1364+
ResponseException.class,
1365+
() -> runEsqlSync(
1366+
requestObjectBuilder().query(
1367+
format(
1368+
null,
1369+
"from {}, (from {}, (from {} | where integer > 1) | where integer < 8) | stats count(*)",
1370+
testIndexName(),
1371+
testIndexName(),
1372+
testIndexName()
1373+
)
1374+
)
1375+
)
1376+
);
1377+
String error = re.getMessage().replaceAll("\\\\\n\s+\\\\", "");
1378+
assertThat(error, containsString("VerificationException"));
1379+
assertThat(error, containsString("Nested subqueries are not supported"));
1380+
}
1381+
1382+
public void testSubqueryWithFork() throws IOException {
1383+
assumeTrue("subqueries in from command", EsqlCapabilities.Cap.SUBQUERY_IN_FROM_COMMAND.isEnabled());
1384+
1385+
bulkLoadTestData(10);
1386+
1387+
ResponseException re = expectThrows(
1388+
ResponseException.class,
1389+
() -> runEsqlSync(
1390+
requestObjectBuilder().query(
1391+
format(
1392+
null,
1393+
"from {}, (from {} | where integer > 1) | fork (where long > 2) (where ip == \"127.0.0.1\") | stats count(*)",
1394+
testIndexName(),
1395+
testIndexName()
1396+
)
1397+
)
1398+
)
1399+
);
1400+
String error = re.getMessage().replaceAll("\\\\\n\s+\\\\", "");
1401+
assertThat(error, containsString("VerificationException"));
1402+
assertThat(error, containsString("FORK after subquery is not supported"));
1403+
1404+
re = expectThrows(
1405+
ResponseException.class,
1406+
() -> runEsqlSync(
1407+
requestObjectBuilder().query(
1408+
format(
1409+
null,
1410+
"from {}, (from {} | where integer > 1 | fork (where long > 2) ( where ip == \"127.0.0.1\")) | stats count(*)",
1411+
testIndexName(),
1412+
testIndexName()
1413+
)
1414+
)
1415+
)
1416+
);
1417+
error = re.getMessage().replaceAll("\\\\\n\s+\\\\", "");
1418+
assertThat(error, containsString("VerificationException"));
1419+
assertThat(error, containsString("FORK inside subquery is not supported"));
1420+
}
1421+
13581422
private static String queryWithComplexFieldNames(int field) {
13591423
StringBuilder query = new StringBuilder();
13601424
query.append(" | keep ").append(randomAlphaOfLength(10)).append(1);

‎x-pack/plugin/esql/qa/testFixtures/src/main/resources/subquery.csv-spec‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -948,3 +948,21 @@ k8s-downsampled | 2024-05-09T23:30:00.000Z | 10.10.20.33 | prod | 21
948948
k8s-downsampled | 2024-05-09T23:30:00.000Z | 10.10.20.34 | staging | 821 | null
949949
k8s-downsampled | 2024-05-09T23:30:00.000Z | 10.10.20.34 | staging | 838 | null
950950
;
951+
952+
subqueryInFromWithFullTextFunctionInMainQueryAndSubquery
953+
required_capability: fork_v9
954+
required_capability: subquery_in_from_command
955+
required_capability: match_function
956+
957+
FROM books, (FROM books | WHERE author:"Faulkner" and ratings > 4)
958+
| WHERE match(title, "Faulkner")
959+
| KEEP book_no, title, author, ratings
960+
| SORT book_no
961+
;
962+
963+
book_no:keyword | title:text | author:text | ratings:double
964+
2713 | Collected Stories of William Faulkner | William Faulkner | 4.53000020980835
965+
2713 | Collected Stories of William Faulkner | William Faulkner | 4.53000020980835
966+
2883 | A Summer of Faulkner: As I Lay Dying/The Sound and the Fury/Light in August (Oprah's Book Club) | William Faulkner | 3.890000104904175
967+
5119 | William Faulkner | William Faulkner | 4.0
968+
;

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -865,7 +865,8 @@ private LogicalPlan resolveFork(Fork fork, AnalyzerContext context) {
865865
if (attrType.isCounter()) {
866866
attrType = attrType.noCounter();
867867
}
868-
return new Alias(source, attr.name(), new Literal(attr.source(), null, attrType));
868+
// use the current fork branch's source as the source of the alias, instead of the original FieldAttribute's source.
869+
return new Alias(source, attr.name(), new Literal(source, null, attrType));
869870
}).toList();
870871

871872
// add the missing columns

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/FullTextFunction.java‎

Lines changed: 42 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.index.query.QueryBuilder;
1717
import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
1818
import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware;
19+
import org.elasticsearch.xpack.esql.capabilities.PostOptimizationPlanVerificationAware;
1920
import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
2021
import org.elasticsearch.xpack.esql.capabilities.RewriteableAware;
2122
import org.elasticsearch.xpack.esql.capabilities.TranslationAware;
@@ -29,6 +30,7 @@
2930
import org.elasticsearch.xpack.esql.core.tree.Source;
3031
import org.elasticsearch.xpack.esql.core.type.DataType;
3132
import org.elasticsearch.xpack.esql.core.type.MultiTypeEsField;
33+
import org.elasticsearch.xpack.esql.core.util.Holder;
3234
import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
3335
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction;
3436
import org.elasticsearch.xpack.esql.expression.predicate.logical.BinaryLogic;
@@ -41,6 +43,7 @@
4143
import org.elasticsearch.xpack.esql.plan.logical.Limit;
4244
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
4345
import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
46+
import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
4447
import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders;
4548
import org.elasticsearch.xpack.esql.planner.TranslatorHandler;
4649
import org.elasticsearch.xpack.esql.querydsl.query.TranslationAwareExpressionQuery;
@@ -73,7 +76,8 @@ public abstract class FullTextFunction extends Function
7376
EvaluatorMapper,
7477
ExpressionScoreMapper,
7578
PostOptimizationVerificationAware,
76-
RewriteableAware {
79+
RewriteableAware,
80+
PostOptimizationPlanVerificationAware {
7781

7882
private final Expression query;
7983
private final QueryBuilder queryBuilder;
@@ -192,26 +196,30 @@ private static void checkFullTextQueryFunctions(LogicalPlan plan, Failures failu
192196
failures.add(fail(condition, "[SCORE] function can't be used in WHERE"));
193197
}
194198

195-
List.of(QueryString.class, Kql.class).forEach(functionClass -> {
196-
// Check for limitations of QSTR and KQL function.
199+
// Defer the check to children commands to post optimization plan verification,
200+
// after the candidate predicates are pushed down into subqueries
201+
if (hasSubqueryInChildrenPlans(f) == false) {
202+
List.of(QueryString.class, Kql.class).forEach(functionClass -> {
203+
// Check for limitations of QSTR and KQL function.
204+
checkCommandsBeforeExpression(
205+
plan,
206+
condition,
207+
functionClass,
208+
lp -> (lp instanceof Filter || lp instanceof OrderBy || lp instanceof EsRelation),
209+
fullTextFunction -> "[" + fullTextFunction.functionName() + "] " + fullTextFunction.functionType(),
210+
failures
211+
);
212+
});
213+
197214
checkCommandsBeforeExpression(
198215
plan,
199216
condition,
200-
functionClass,
201-
lp -> (lp instanceof Filter || lp instanceof OrderBy || lp instanceof EsRelation),
202-
fullTextFunction -> "[" + fullTextFunction.functionName() + "] " + fullTextFunction.functionType(),
217+
FullTextFunction.class,
218+
lp -> (lp instanceof Limit == false) && (lp instanceof Aggregate == false),
219+
m -> "[" + m.functionName() + "] " + m.functionType(),
203220
failures
204221
);
205-
});
206-
207-
checkCommandsBeforeExpression(
208-
plan,
209-
condition,
210-
FullTextFunction.class,
211-
lp -> (lp instanceof Limit == false) && (lp instanceof Aggregate == false),
212-
m -> "[" + m.functionName() + "] " + m.functionType(),
213-
failures
214-
);
222+
}
215223
checkFullTextFunctionsParents(condition, failures);
216224
} else if (plan instanceof Aggregate agg) {
217225
checkFullTextFunctionsInAggs(agg, failures);
@@ -343,7 +351,8 @@ public static void fieldVerifier(LogicalPlan plan, FullTextFunction function, Ex
343351
var fieldAttribute = fieldAsFieldAttribute(field);
344352
if (fieldAttribute == null) {
345353
plan.forEachExpression(function.getClass(), m -> {
346-
if (function.children().contains(field)) {
354+
// Defer the field check to post optimization plan verification if there are subqueries in the plan
355+
if (m.children().contains(field) && hasSubqueryInChildrenPlans(plan) == false) {
347356
failures.add(
348357
fail(
349358
field,
@@ -435,4 +444,20 @@ public static FieldAttribute fieldAsFieldAttribute(Expression field) {
435444
public void postOptimizationVerification(Failures failures) {
436445
resolveTypeQuery(query(), sourceText(), forPostOptimizationValidation(query(), failures));
437446
}
447+
448+
@Override
449+
public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
450+
// check plan again after predicates are pushed down into subqueries
451+
return FullTextFunction::checkFullTextQueryFunctions;
452+
}
453+
454+
private static boolean hasSubqueryInChildrenPlans(LogicalPlan plan) {
455+
Holder<Boolean> hasSubqueryInChildrenPlans = new Holder<>(false);
456+
plan.forEachDown(p -> {
457+
if (p instanceof UnionAll) {
458+
hasSubqueryInChildrenPlans.set(true);
459+
}
460+
});
461+
return hasSubqueryInChildrenPlans.get();
462+
}
438463
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Match.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,15 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
404404
};
405405
}
406406

407+
@Override
408+
public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
409+
// check plan again after predicates are pushed down into subqueries
410+
return (plan, failures) -> {
411+
super.postOptimizationPlanVerification().accept(plan, failures);
412+
fieldVerifier(plan, this, field, failures);
413+
};
414+
}
415+
407416
public Object queryAsObject() {
408417
Object queryAsObject = Foldables.queryAsObject(query(), sourceText());
409418

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/MatchPhrase.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,15 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
257257
};
258258
}
259259

260+
@Override
261+
public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
262+
// check plan again after predicates are pushed down into subqueries
263+
return (plan, failures) -> {
264+
super.postOptimizationPlanVerification().accept(plan, failures);
265+
fieldVerifier(plan, this, field, failures);
266+
};
267+
}
268+
260269
public Object queryAsObject() {
261270
Object queryAsObject = Foldables.queryAsObject(query(), sourceText());
262271

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/fulltext/Term.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,15 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
117117
};
118118
}
119119

120+
@Override
121+
public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
122+
// check plan again after predicates are pushed down into subqueries
123+
return (plan, failures) -> {
124+
super.postOptimizationPlanVerification().accept(plan, failures);
125+
fieldVerifier(plan, this, field, failures);
126+
};
127+
}
128+
120129
@Override
121130
public Expression replaceChildren(List<Expression> newChildren) {
122131
return new Term(source(), newChildren.get(0), newChildren.get(1), queryBuilder());

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/vector/Knn.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,15 @@ public void postOptimizationVerification(Failures failures) {
340340
}
341341
}
342342

343+
@Override
344+
public BiConsumer<LogicalPlan, Failures> postOptimizationPlanVerification() {
345+
// check plan again after predicates are pushed down into subqueries
346+
return (plan, failures) -> {
347+
super.postOptimizationPlanVerification().accept(plan, failures);
348+
fieldVerifier(plan, this, field, failures);
349+
};
350+
}
351+
343352
@Override
344353
public Expression replaceChildren(List<Expression> newChildren) {
345354
return new Knn(

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalVerifier.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ void checkPlanConsistency(LogicalPlan optimizedPlan, Failures failures, Failures
5454
if (ex instanceof PostOptimizationVerificationAware va) {
5555
va.postOptimizationVerification(failures);
5656
}
57+
if (ex instanceof PostOptimizationPlanVerificationAware vpa) {
58+
vpa.postOptimizationPlanVerification().accept(p, failures);
59+
}
5760
});
5861
}
5962
});

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
/**
3434
* Push down filters that can be evaluated by the UnionAll child to each child, below the added
3535
* {@code Limit} and {@code Subquery}, so that the filters can be pushed down further to the
36-
* data source when possible. Filters that cannot be pushed down remain above the UnionAll.
36+
* data source when possible. Filters that cannot be pushed down remain above the UnionAll.
3737
*
3838
* Also push down the {@code Limit} added by {@code AddImplicitForkLimit} below the
3939
* {@code Subquery}, so that the other rules related to {@code Limit} optimization can be applied.

0 commit comments

Comments
 (0)