Skip to content

Commit 58e617f

Browse files
jan-elasticjan-elastic
committed
rewire
1 parent 7f118b1 commit 58e617f

File tree

3 files changed

+42
-34
lines changed

3 files changed

+42
-34
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,16 +226,18 @@ public interface LogicalPlanRunner {
226226

227227
private final LogicalPlan logicalPlan;
228228
private final boolean preservesRows;
229+
private final LogicalPlanRunner runner;
229230

230-
public Approximate(LogicalPlan logicalPlan) {
231+
public Approximate(LogicalPlan logicalPlan, LogicalPlanRunner logicalPlanRunner) {
231232
this.logicalPlan = logicalPlan;
232233
this.preservesRows = verifyPlan(logicalPlan);
234+
this.runner = logicalPlanRunner;
233235
}
234236

235237
/**
236238
* Computes approximate results for the logical plan.
237239
*/
238-
public void approximate(LogicalPlanRunner runner, ActionListener<Result> listener) {
240+
public void approximate(ActionListener<Result> listener) {
239241
runner.run(sourceCountPlan(), sourceCountListener(runner, listener));
240242
}
241243

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -337,15 +337,13 @@ private void executeSubPlans(
337337
ActionListener.runAfter(listener, executionInfo::finishSubPlans)
338338
);
339339
} else if (request.approximate()) {
340-
new Approximate(optimizedPlan).approximate(
341-
(p, l) -> runner.run(
342-
logicalPlanToPhysicalPlan(optimizedPlan(p, logicalPlanOptimizer), request, physicalPlanOptimizer),
343-
configuration,
344-
foldContext,
345-
l
346-
),
347-
listener
340+
Approximate.LogicalPlanRunner logicalPlanRunner = (p, l) -> runner.run(
341+
logicalPlanToPhysicalPlan(optimizedPlan(p, logicalPlanOptimizer), request, physicalPlanOptimizer),
342+
configuration,
343+
foldContext,
344+
l
348345
);
346+
new Approximate(optimizedPlan, logicalPlanRunner).approximate(listener);
349347
} else {
350348
PhysicalPlan physicalPlan = logicalPlanToPhysicalPlan(optimizedPlan, request, physicalPlanOptimizer);
351349
// execute main plan

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

Lines changed: 32 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
import java.util.List;
4545
import java.util.function.Predicate;
4646

47-
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_CFG;
4847
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
4948
import static org.hamcrest.CoreMatchers.allOf;
5049
import static org.hamcrest.CoreMatchers.not;
@@ -110,11 +109,11 @@ protected List<String> filteredWarnings() {
110109
}
111110

112111
public void testVerify_validQuery() throws Exception {
113-
createApproximate("FROM test | WHERE emp_no<99 | SORT last_name | MV_EXPAND salary | STATS COUNT() BY gender");
114-
createApproximate("FROM test | CHANGE_POINT salary ON emp_no | EVAL x=1 | DROP emp_no | STATS SUM(salary) BY x");
115-
createApproximate("FROM test | LIMIT 1000 | KEEP gender, emp_no | RENAME gender AS whatever | STATS MEDIAN(emp_no)");
116-
createApproximate("FROM test | EVAL blah=1 | GROK last_name \"%{IP:x}\" | SAMPLE 0.1 | STATS a=COUNT() | LIMIT 100 | SORT a");
117-
createApproximate("ROW i=[1,2,3] | EVAL x=TO_STRING(i) | DISSECT x \"%{x}\" | STATS i=10*POW(PERCENTILE(i, 0.5), 2) | LIMIT 10");
112+
verify("FROM test | WHERE emp_no<99 | SORT last_name | MV_EXPAND salary | STATS COUNT() BY gender");
113+
verify("FROM test | CHANGE_POINT salary ON emp_no | EVAL x=1 | DROP emp_no | STATS SUM(salary) BY x");
114+
verify("FROM test | LIMIT 1000 | KEEP gender, emp_no | RENAME gender AS whatever | STATS MEDIAN(emp_no)");
115+
verify("FROM test | EVAL blah=1 | GROK last_name \"%{IP:x}\" | SAMPLE 0.1 | STATS a=COUNT() | LIMIT 100 | SORT a");
116+
verify("ROW i=[1,2,3] | EVAL x=TO_STRING(i) | DISSECT x \"%{x}\" | STATS i=10*POW(PERCENTILE(i, 0.5), 2) | LIMIT 10");
118117
}
119118

120119
public void testVerify_exactlyOneStats() {
@@ -176,9 +175,9 @@ public void testVerify_incompatibleAggregation() {
176175
}
177176

178177
public void testCountPlan_largeDataNoFilters() throws Exception {
179-
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)");
180178
TestRunner runner = new TestRunner(1_000_000_000, 1_000_000_000);
181-
approximate.approximate(runner, TestRunner.resultCloser);
179+
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)", runner);
180+
approximate.approximate(TestRunner.resultCloser);
182181
// One pass is needed to get the number of rows, and approximation is executed immediately
183182
// after that with the correct sample probability.
184183
assertThat(runner.invocations, hasSize(2));
@@ -187,9 +186,9 @@ public void testCountPlan_largeDataNoFilters() throws Exception {
187186
}
188187

189188
public void testCountPlan_smallDataNoFilters() throws Exception {
190-
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)");
191189
TestRunner runner = new TestRunner(1_000, 1_000);
192-
approximate.approximate(runner, TestRunner.resultCloser);
190+
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)", runner);
191+
approximate.approximate(TestRunner.resultCloser);
193192
// One pass is needed to get the number of rows, and the original query is executed
194193
// immediately after that without sampling.
195194
assertThat(runner.invocations, hasSize(2));
@@ -198,9 +197,9 @@ public void testCountPlan_smallDataNoFilters() throws Exception {
198197
}
199198

200199
public void testCountPlan_largeDataAfterFiltering() throws Exception {
201-
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)");
202200
TestRunner runner = new TestRunner(1_000_000_000_000L, 1_000_000_000);
203-
approximate.approximate(runner, TestRunner.resultCloser);
201+
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)", runner);
202+
approximate.approximate(TestRunner.resultCloser);
204203
// One pass is needed to get the number of rows, then a few passes to get a good sample
205204
// probability, and finally approximation is executed.
206205
assertThat(runner.invocations, hasSize(4));
@@ -211,9 +210,9 @@ public void testCountPlan_largeDataAfterFiltering() throws Exception {
211210
}
212211

213212
public void testCountPlan_smallDataAfterFiltering() throws Exception {
214-
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)");
215213
TestRunner runner = new TestRunner(1_000_000_000_000_000_000L, 100);
216-
approximate.approximate(runner, TestRunner.resultCloser);
214+
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)", runner);
215+
approximate.approximate(TestRunner.resultCloser);
217216
// One pass is needed to get the number of rows, then a few passes to get a good sample
218217
// probability, and finally the original query is executed without sampling.
219218
assertThat(runner.invocations, hasSize(5));
@@ -225,9 +224,9 @@ public void testCountPlan_smallDataAfterFiltering() throws Exception {
225224
}
226225

227226
public void testCountPlan_smallDataBeforeFiltering() throws Exception {
228-
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)");
229227
TestRunner runner = new TestRunner(1_000, 10);
230-
approximate.approximate(runner, TestRunner.resultCloser);
228+
Approximate approximate = createApproximate("FROM test | WHERE emp_no < 1 | STATS SUM(emp_no)", runner);
229+
approximate.approximate(TestRunner.resultCloser);
231230
// One pass is needed to get the number of rows, and the original query is executed
232231
// immediately after that without sampling.
233232
assertThat(runner.invocations, hasSize(2));
@@ -236,9 +235,9 @@ public void testCountPlan_smallDataBeforeFiltering() throws Exception {
236235
}
237236

238237
public void testApproximatePlan_createsConfidenceInterval() throws Exception {
239-
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)");
240238
TestRunner runner = new TestRunner(1_000_000_000, 1_000_000_000);
241-
approximate.approximate(runner, TestRunner.resultCloser);
239+
Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)", runner);
240+
approximate.approximate(TestRunner.resultCloser);
242241
// One pass is needed to get the number of rows, and approximation is executed immediately
243242
// after that with the correct sample probability.
244243
assertThat(runner.invocations, hasSize(2));
@@ -249,11 +248,12 @@ public void testApproximatePlan_createsConfidenceInterval() throws Exception {
249248
}
250249

251250
public void testApproximatePlan_dependentConfidenceIntervals() throws Exception {
251+
TestRunner runner = new TestRunner(1_000_000_000, 1_000_000_000);
252252
Approximate approximate = createApproximate(
253-
"FROM test | STATS x=COUNT() | EVAL a=x*x, b=7, c=TO_STRING(x), d=MV_APPEND(x, 1::LONG), e=a+POW(b, 2)"
253+
"FROM test | STATS x=COUNT() | EVAL a=x*x, b=7, c=TO_STRING(x), d=MV_APPEND(x, 1::LONG), e=a+POW(b, 2)",
254+
runner
254255
);
255-
TestRunner runner = new TestRunner(1_000_000_000, 1_000_000_000);
256-
approximate.approximate(runner, TestRunner.resultCloser);
256+
approximate.approximate(TestRunner.resultCloser);
257257
// One pass is needed to get the number of rows, and approximation is executed immediately
258258
// after that with the correct sample probability.
259259
assertThat(runner.invocations, hasSize(2));
@@ -303,11 +303,19 @@ public void describeTo(Description description) {
303303
}
304304

305305
private void assertError(String esql, Matcher<String> matcher) {
306-
Exception e = assertThrows(VerificationException.class, () -> createApproximate(esql));
306+
Exception e = assertThrows(VerificationException.class, () -> verify(esql));
307307
assertThat(e.getMessage().substring("Found 1 problem\n".length()), matcher);
308308
}
309309

310-
private Approximate createApproximate(String query) throws Exception {
310+
private void verify(String query) throws Exception {
311+
Approximate.verifyPlan(getLogicalPlan(query));
312+
}
313+
314+
private Approximate createApproximate(String query, Approximate.LogicalPlanRunner runner) throws Exception {
315+
return new Approximate(getLogicalPlan(query), runner);
316+
}
317+
318+
private LogicalPlan getLogicalPlan(String query) throws Exception {
311319
SetOnce<LogicalPlan> resultHolder = new SetOnce<>();
312320
SetOnce<Exception> exceptionHolder = new SetOnce<>();
313321
LogicalPlan plan = parser.createStatement(query, new QueryParams());
@@ -317,6 +325,6 @@ private Approximate createApproximate(String query) throws Exception {
317325
if (exceptionHolder.get() != null) {
318326
throw exceptionHolder.get();
319327
}
320-
return new Approximate(resultHolder.get());
328+
return resultHolder.get();
321329
}
322330
}

0 commit comments

Comments
 (0)