Skip to content

Commit c308c5a

Browse files
jan-elasticjan-elastic
committed
better verification errors + tests
1 parent 77562c8 commit c308c5a

File tree

3 files changed

+120
-17
lines changed

3 files changed

+120
-17
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/approximate/Approximate.java‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,8 @@
1111
import org.apache.logging.log4j.Logger;
1212
import org.elasticsearch.action.ActionListener;
1313
import org.elasticsearch.compute.data.LongBlock;
14-
import org.elasticsearch.xpack.esql.core.InvalidArgumentException;
14+
import org.elasticsearch.xpack.esql.VerificationException;
15+
import org.elasticsearch.xpack.esql.common.Failure;
1516
import org.elasticsearch.xpack.esql.core.expression.Alias;
1617
import org.elasticsearch.xpack.esql.core.expression.Expression;
1718
import org.elasticsearch.xpack.esql.core.expression.Literal;
@@ -26,17 +27,17 @@
2627
import org.elasticsearch.xpack.esql.plan.logical.Filter;
2728
import org.elasticsearch.xpack.esql.plan.logical.Fork;
2829
import org.elasticsearch.xpack.esql.plan.logical.Grok;
30+
import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
2931
import org.elasticsearch.xpack.esql.plan.logical.Keep;
3032
import org.elasticsearch.xpack.esql.plan.logical.LeafPlan;
3133
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
3234
import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
3335
import org.elasticsearch.xpack.esql.plan.logical.Rename;
3436
import org.elasticsearch.xpack.esql.plan.logical.Sample;
35-
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
37+
import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
3638
import org.elasticsearch.xpack.esql.session.Result;
3739

3840
import java.util.List;
39-
import java.util.Locale;
4041
import java.util.Set;
4142

4243
/**
@@ -99,7 +100,7 @@ public interface LogicalPlanRunner {
99100
/**
100101
* Commands that cannot be used anywhere in an approximated query.
101102
*/
102-
private static final Set<Class<? extends LogicalPlan>> INCOMPATIBLE_COMMANDS = Set.of(Fork.class, Join.class);
103+
private static final Set<Class<? extends LogicalPlan>> INCOMPATIBLE_COMMANDS = Set.of(Fork.class, InlineStats.class, LookupJoin.class);
103104

104105
// TODO: find a good default value, or alternative ways of setting it
105106
private static final int SAMPLE_ROW_COUNT = 10000;
@@ -131,31 +132,29 @@ private boolean verifyPlan() {
131132
throw new IllegalStateException("Expected pre-optimized plan");
132133
}
133134
if (logicalPlan.anyMatch(plan -> plan instanceof Aggregate) == false) {
134-
throw new InvalidArgumentException("query without [STATS] command cannot be approximated");
135+
throw new VerificationException(
136+
List.of(Failure.fail(logicalPlan.collectLeaves().getFirst(), "query without [STATS] cannot be approximated"))
137+
);
135138
}
136139
logicalPlan.forEachUp(plan -> {
137140
if (INCOMPATIBLE_COMMANDS.contains(plan.getClass())) {
138-
throw new InvalidArgumentException(
139-
"query with [" + plan.nodeName().toUpperCase(Locale.ROOT) + "] command cannot be approximated"
141+
throw new VerificationException(
142+
List.of(Failure.fail(plan, "query with [" + plan.sourceText() + "] cannot be approximated"))
140143
);
141144
}
142145
});
143146

144147
Holder<Boolean> encounteredStats = new Holder<>(false);
145148
Holder<Boolean> hasFilters = new Holder<>(false);
146149
logicalPlan.transformUp(plan -> {
147-
if (INCOMPATIBLE_COMMANDS.contains(plan.getClass())) {
148-
throw new InvalidArgumentException(
149-
"query with [" + plan.nodeName().toUpperCase(Locale.ROOT) + "] command cannot be approximated"
150-
);
151-
} else if (plan instanceof LeafPlan) {
150+
if (plan instanceof LeafPlan) {
152151
encounteredStats.set(false);
153152
} else if (encounteredStats.get() == false) {
154153
if (plan instanceof Aggregate) {
155154
encounteredStats.set(true);
156155
} else if (SWAPPABLE_COMMANDS.contains(plan.getClass()) == false && FILTER_COMMANDS.contains(plan.getClass()) == false) {
157-
throw new InvalidArgumentException(
158-
"query with [" + plan.nodeName().toUpperCase(Locale.ROOT) + "] before [STATS] command cannot be approximated"
156+
throw new VerificationException(
157+
List.of(Failure.fail(plan, "query with [" + plan.sourceText() + "] before [STATS] cannot be approximated"))
159158
);
160159
} else if (FILTER_COMMANDS.contains(plan.getClass())) {
161160
hasFilters.set(true);
@@ -178,7 +177,7 @@ private LogicalPlan sourceCountPlan() {
178177
Source.EMPTY,
179178
plan,
180179
List.of(),
181-
List.of(new Alias(Source.EMPTY, "approximate-count", new Count(Source.EMPTY, Literal.keyword(Source.EMPTY, "*"))))
180+
List.of(new Alias(Source.EMPTY, ".approximate-count", new Count(Source.EMPTY, Literal.keyword(Source.EMPTY, "*"))))
182181
);
183182
} else {
184183
plan = plan.children().getFirst();
@@ -227,7 +226,7 @@ private LogicalPlan countPlan(double sampleProbability) {
227226
Source.EMPTY,
228227
sample,
229228
List.of(),
230-
List.of(new Alias(Source.EMPTY, "approximate-count", new Count(Source.EMPTY, Literal.keyword(Source.EMPTY, "*"))))
229+
List.of(new Alias(Source.EMPTY, ".approximate-count", new Count(Source.EMPTY, Literal.keyword(Source.EMPTY, "*"))))
231230
);
232231
}
233232
} else {

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,12 @@ public void execute(EsqlQueryRequest request, EsqlExecutionInfo executionInfo, P
181181
analyzedPlan(parsed, executionInfo, request.filter(), new EsqlCCSUtils.CssPartialErrorsActionListener(executionInfo, listener) {
182182
@Override
183183
public void onResponse(LogicalPlan analyzedPlan) {
184-
SubscribableListener.<LogicalPlan>newForked(l -> preOptimizedPlan(analyzedPlan, l))
184+
SubscribableListener.<LogicalPlan>newForked(l -> preOptimizedPlan(analyzedPlan, l)).<LogicalPlan>andThen((l, p) -> {
185+
if (request.approximate()) {
186+
new Approximate(p); // to verify whether the pre-optimized plan is suitable for approximation
187+
}
188+
l.onResponse(p);
189+
})
185190
.<LogicalPlan>andThen((l, p) -> preMapper.preMapper(optimizedPlan(p), l))
186191
.<Result>andThen((l, p) -> executeOptimizedPlan(request, executionInfo, planRunner, p, l))
187192
.addListener(listener);

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/approximate/ApproximateTests.java‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
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.approximate;
9+
10+
import org.apache.lucene.util.SetOnce;
11+
import org.elasticsearch.action.ActionListener;
12+
import org.elasticsearch.test.ESTestCase;
13+
import org.elasticsearch.xpack.esql.VerificationException;
14+
import org.elasticsearch.xpack.esql.core.expression.FoldContext;
15+
import org.elasticsearch.xpack.esql.optimizer.LogicalPlanPreOptimizer;
16+
import org.elasticsearch.xpack.esql.optimizer.LogicalPreOptimizerContext;
17+
import org.elasticsearch.xpack.esql.parser.EsqlParser;
18+
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
19+
import org.hamcrest.Matcher;
20+
21+
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_CFG;
22+
import static org.hamcrest.Matchers.equalTo;
23+
24+
public class ApproximateTests extends ESTestCase {
25+
26+
private static final EsqlParser parser = new EsqlParser();
27+
private static final LogicalPlanPreOptimizer logicalPlanPreOptimizer = new LogicalPlanPreOptimizer(
28+
new LogicalPreOptimizerContext(FoldContext.small())
29+
);
30+
31+
public void testVerify_validQuery() throws Exception {
32+
verifyQuery("FROM index | EVAL x=1 | WHERE y<0.1 | SORT z | SAMPLE 0.1 | STATS COUNT() BY y | SORT z | MV_EXPAND x");
33+
}
34+
35+
public void testVerify_noStats() {
36+
assertError("FROM index | EVAL x = 1 | SORT y", equalTo("line 1:1: query without [STATS] cannot be approximated"));
37+
}
38+
39+
public void testVerify_incompatibleCommand() {
40+
assertError(
41+
"FROM index | FORK (EVAL x=1) (EVAL y=1) | STATS COUNT()",
42+
equalTo("line 1:14: query with [FORK (EVAL x=1) (EVAL y=1)] cannot be approximated")
43+
);
44+
assertError(
45+
"FROM index | STATS COUNT() | FORK (EVAL x=1) (EVAL y=1)",
46+
equalTo("line 1:30: query with [FORK (EVAL x=1) (EVAL y=1)] cannot be approximated")
47+
);
48+
49+
assertError(
50+
"FROM index | INLINESTATS COUNT() | STATS COUNT()",
51+
equalTo("line 1:14: query with [INLINESTATS COUNT()] cannot be approximated")
52+
);
53+
assertError(
54+
"FROM index | STATS COUNT() | INLINESTATS COUNT()",
55+
equalTo("line 1:30: query with [INLINESTATS COUNT()] cannot be approximated")
56+
);
57+
58+
assertError(
59+
"FROM index | LOOKUP JOIN index2 ON id | STATS COUNT()",
60+
equalTo("line 1:14: query with [LOOKUP JOIN index2 ON id] cannot be approximated")
61+
);
62+
assertError(
63+
"FROM index | STATS COUNT() | LOOKUP JOIN index2 ON id",
64+
equalTo("line 1:30: query with [LOOKUP JOIN index2 ON id] cannot be approximated")
65+
);
66+
}
67+
68+
public void testVerify_nonSwappableCommand() {
69+
assertError(
70+
"FROM index | MV_EXPAND x | STATS COUNT()",
71+
equalTo("line 1:14: query with [MV_EXPAND x] before [STATS] cannot be approximated")
72+
);
73+
assertError(
74+
"FROM index | CHANGE_POINT x ON t | STATS COUNT()",
75+
equalTo("line 1:14: query with [CHANGE_POINT x ON t] before [STATS] cannot be approximated")
76+
);
77+
assertError(
78+
"FROM index | LIMIT 100 | STATS COUNT()",
79+
equalTo("line 1:14: query with [LIMIT 100] before [STATS] cannot be approximated")
80+
);
81+
}
82+
83+
private void assertError(String esql, Matcher<String> matcher) {
84+
Exception e = assertThrows(VerificationException.class, () -> verifyQuery(esql));
85+
assertThat(e.getMessage().substring("Found 1 problem\n".length()), matcher);
86+
}
87+
88+
private void verifyQuery(String esql) throws Exception {
89+
SetOnce<LogicalPlan> resultHolder = new SetOnce<>();
90+
SetOnce<Exception> exceptionHolder = new SetOnce<>();
91+
LogicalPlan plan = parser.createStatement(esql, TEST_CFG);
92+
plan.setAnalyzed();
93+
logicalPlanPreOptimizer.preOptimize(plan, ActionListener.wrap(resultHolder::set, exceptionHolder::set));
94+
if (exceptionHolder.get() != null) {
95+
throw exceptionHolder.get();
96+
}
97+
new Approximate(resultHolder.get());
98+
}
99+
}

0 commit comments

Comments
 (0)