|
|
@@ -112,9 +112,11 @@ import org.elasticsearch.xpack.ql.type.EsField;
|
|
|
import org.junit.BeforeClass;
|
|
|
|
|
|
import java.lang.reflect.Constructor;
|
|
|
+import java.util.Arrays;
|
|
|
import java.util.List;
|
|
|
import java.util.Map;
|
|
|
import java.util.Set;
|
|
|
+import java.util.function.Function;
|
|
|
|
|
|
import static java.util.Arrays.asList;
|
|
|
import static java.util.Collections.emptyList;
|
|
|
@@ -173,6 +175,19 @@ public class LogicalPlanOptimizerTests extends ESTestCase {
|
|
|
private static Analyzer analyzerAirports;
|
|
|
private static EnrichResolution enrichResolution;
|
|
|
|
|
|
+ private static class SubstitutionOnlyOptimizer extends LogicalPlanOptimizer {
|
|
|
+ static SubstitutionOnlyOptimizer INSTANCE = new SubstitutionOnlyOptimizer(new LogicalOptimizerContext(EsqlTestUtils.TEST_CFG));
|
|
|
+
|
|
|
+ SubstitutionOnlyOptimizer(LogicalOptimizerContext optimizerContext) {
|
|
|
+ super(optimizerContext);
|
|
|
+ }
|
|
|
+
|
|
|
+ @Override
|
|
|
+ protected List<Batch<LogicalPlan>> batches() {
|
|
|
+ return List.of(substitutions());
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
@BeforeClass
|
|
|
public static void init() {
|
|
|
parser = new EsqlParser();
|
|
|
@@ -3272,6 +3287,177 @@ public class LogicalPlanOptimizerTests extends ESTestCase {
|
|
|
assertThat(Expressions.attribute(fields.get(1)), is(Expressions.attribute(sum_argument)));
|
|
|
}
|
|
|
|
|
|
+ /**
|
|
|
+ * Expects after running the {@link LogicalPlanOptimizer#substitutions()}:
|
|
|
+ *
|
|
|
+ * Limit[1000[INTEGER]]
|
|
|
+ * \_EsqlProject[[s{r}#3, s_expr{r}#5, s_null{r}#7, w{r}#10]]
|
|
|
+ * \_Project[[s{r}#3, s_expr{r}#5, s_null{r}#7, w{r}#10]]
|
|
|
+ * \_Eval[[MVSUM([1, 2][INTEGER]) * $$COUNT$s$0{r}#25 AS s, MVSUM(314.0[DOUBLE] / 100[INTEGER]) * $$COUNT$s$0{r}#25 AS s
|
|
|
+ * _expr, MVSUM(null[NULL]) * $$COUNT$s$0{r}#25 AS s_null]]
|
|
|
+ * \_Aggregate[[w{r}#10],[COUNT(*[KEYWORD]) AS $$COUNT$s$0, w{r}#10]]
|
|
|
+ * \_Eval[[emp_no{f}#15 % 2[INTEGER] AS w]]
|
|
|
+ * \_EsRelation[test][_meta_field{f}#21, emp_no{f}#15, first_name{f}#16, ..]
|
|
|
+ */
|
|
|
+ public void testSumOfLiteral() {
|
|
|
+ var plan = plan("""
|
|
|
+ from test
|
|
|
+ | stats s = sum([1,2]),
|
|
|
+ s_expr = sum(314.0/100),
|
|
|
+ s_null = sum(null)
|
|
|
+ by w = emp_no % 2
|
|
|
+ | keep s, s_expr, s_null, w
|
|
|
+ """, SubstitutionOnlyOptimizer.INSTANCE);
|
|
|
+
|
|
|
+ var limit = as(plan, Limit.class);
|
|
|
+ var esqlProject = as(limit.child(), EsqlProject.class);
|
|
|
+ var project = as(esqlProject.child(), Project.class);
|
|
|
+ var eval = as(project.child(), Eval.class);
|
|
|
+ var agg = as(eval.child(), Aggregate.class);
|
|
|
+
|
|
|
+ var exprs = eval.fields();
|
|
|
+ // s = count(*) * 3
|
|
|
+ var s = as(exprs.get(0), Alias.class);
|
|
|
+ assertThat(s.name(), equalTo("s"));
|
|
|
+ var mul = as(s.child(), Mul.class);
|
|
|
+ var mvSum = as(mul.left(), MvSum.class);
|
|
|
+ assertThat(mvSum.fold(), equalTo(3));
|
|
|
+ var count = as(mul.right(), ReferenceAttribute.class);
|
|
|
+ assertThat(count.name(), equalTo("$$COUNT$s$0"));
|
|
|
+
|
|
|
+ // s_expr = count(*) * 3.14
|
|
|
+ var s_expr = as(exprs.get(1), Alias.class);
|
|
|
+ assertThat(s_expr.name(), equalTo("s_expr"));
|
|
|
+ var mul_expr = as(s_expr.child(), Mul.class);
|
|
|
+ var mvSum_expr = as(mul_expr.left(), MvSum.class);
|
|
|
+ assertThat(mvSum_expr.fold(), equalTo(3.14));
|
|
|
+ var count_expr = as(mul_expr.right(), ReferenceAttribute.class);
|
|
|
+ assertThat(count_expr.name(), equalTo("$$COUNT$s$0"));
|
|
|
+
|
|
|
+ // s_null = null
|
|
|
+ var s_null = as(exprs.get(2), Alias.class);
|
|
|
+ assertThat(s_null.name(), equalTo("s_null"));
|
|
|
+ var mul_null = as(s_null.child(), Mul.class);
|
|
|
+ var mvSum_null = as(mul_null.left(), MvSum.class);
|
|
|
+ assertThat(mvSum_null.field(), equalTo(NULL));
|
|
|
+ var count_null = as(mul_null.right(), ReferenceAttribute.class);
|
|
|
+ assertThat(count_null.name(), equalTo("$$COUNT$s$0"));
|
|
|
+
|
|
|
+ var count_agg = as(Alias.unwrap(agg.aggregates().get(0)), Count.class);
|
|
|
+ assertThat(count_agg.children().get(0), instanceOf(Literal.class));
|
|
|
+ var w = as(Alias.unwrap(agg.groupings().get(0)), ReferenceAttribute.class);
|
|
|
+ assertThat(w.name(), equalTo("w"));
|
|
|
+ }
|
|
|
+
|
|
|
+ private record AggOfLiteralTestCase(
|
|
|
+ String aggFunctionName,
|
|
|
+ Class<? extends Expression> substitution,
|
|
|
+ Function<int[], Object> aggMultiValue
|
|
|
+ ) {};
|
|
|
+
|
|
|
+ private static List<AggOfLiteralTestCase> AGG_OF_CONST_CASES = List.of(
|
|
|
+ new AggOfLiteralTestCase("avg", MvAvg.class, ints -> ((double) Arrays.stream(ints).sum()) / ints.length),
|
|
|
+ new AggOfLiteralTestCase("min", MvMin.class, ints -> Arrays.stream(ints).min().getAsInt()),
|
|
|
+ new AggOfLiteralTestCase("max", MvMax.class, ints -> Arrays.stream(ints).max().getAsInt())
|
|
|
+ );
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Aggs of literals in case that the agg can be simply replaced by a corresponding mv-function;
|
|
|
+ * e.g. avg([1,2,3]) which is equivalent to mv_avg([1,2,3]).
|
|
|
+ *
|
|
|
+ * Expects after running the {@link LogicalPlanOptimizer#substitutions()}:
|
|
|
+ *
|
|
|
+ * Limit[1000[INTEGER]]
|
|
|
+ * \_EsqlProject[[s{r}#3, s_expr{r}#5, s_null{r}#7]]
|
|
|
+ * \_Project[[s{r}#3, s_expr{r}#5, s_null{r}#7]]
|
|
|
+ * \_Eval[[MVAVG([1, 2][INTEGER]) AS s, MVAVG(314.0[DOUBLE] / 100[INTEGER]) AS s_expr, MVAVG(null[NULL]) AS s_null]]
|
|
|
+ * \_LocalRelation[[{e}#21],[ConstantNullBlock[positions=1]]]
|
|
|
+ */
|
|
|
+ public void testAggOfLiteral() {
|
|
|
+ for (AggOfLiteralTestCase testCase : AGG_OF_CONST_CASES) {
|
|
|
+ String query = LoggerMessageFormat.format(null, """
|
|
|
+ from test
|
|
|
+ | stats s = {}([1,2]),
|
|
|
+ s_expr = {}(314.0/100),
|
|
|
+ s_null = {}(null)
|
|
|
+ | keep s, s_expr, s_null
|
|
|
+ """, testCase.aggFunctionName, testCase.aggFunctionName, testCase.aggFunctionName);
|
|
|
+
|
|
|
+ var plan = plan(query, SubstitutionOnlyOptimizer.INSTANCE);
|
|
|
+
|
|
|
+ var limit = as(plan, Limit.class);
|
|
|
+ var esqlProject = as(limit.child(), EsqlProject.class);
|
|
|
+ var project = as(esqlProject.child(), Project.class);
|
|
|
+ var eval = as(project.child(), Eval.class);
|
|
|
+ var singleRowRelation = as(eval.child(), LocalRelation.class);
|
|
|
+ var singleRow = singleRowRelation.supplier().get();
|
|
|
+ assertThat(singleRow.length, equalTo(1));
|
|
|
+ assertThat(singleRow[0].getPositionCount(), equalTo(1));
|
|
|
+
|
|
|
+ var exprs = eval.fields();
|
|
|
+ var s = as(exprs.get(0), Alias.class);
|
|
|
+ assertThat(s.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s.child().fold(), equalTo(testCase.aggMultiValue.apply(new int[] { 1, 2 })));
|
|
|
+ var s_expr = as(exprs.get(1), Alias.class);
|
|
|
+ assertThat(s_expr.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s_expr.child().fold(), equalTo(3.14));
|
|
|
+ var s_null = as(exprs.get(2), Alias.class);
|
|
|
+ assertThat(s_null.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s_null.child().fold(), equalTo(null));
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ /**
|
|
|
+ * Like {@link LogicalPlanOptimizerTests#testAggOfLiteral()} but with a grouping key.
|
|
|
+ *
|
|
|
+ * Expects after running the {@link LogicalPlanOptimizer#substitutions()}:
|
|
|
+ *
|
|
|
+ * Limit[1000[INTEGER]]
|
|
|
+ * \_EsqlProject[[s{r}#3, s_expr{r}#5, s_null{r}#7, emp_no{f}#13]]
|
|
|
+ * \_Project[[s{r}#3, s_expr{r}#5, s_null{r}#7, emp_no{f}#13]]
|
|
|
+ * \_Eval[[MVAVG([1, 2][INTEGER]) AS s, MVAVG(314.0[DOUBLE] / 100[INTEGER]) AS s_expr, MVAVG(null[NULL]) AS s_null]]
|
|
|
+ * \_Aggregate[[emp_no{f}#13],[emp_no{f}#13]]
|
|
|
+ * \_EsRelation[test][_meta_field{f}#19, emp_no{f}#13, first_name{f}#14, ..]
|
|
|
+ */
|
|
|
+ public void testAggOfLiteralGrouped() {
|
|
|
+ for (AggOfLiteralTestCase testCase : AGG_OF_CONST_CASES) {
|
|
|
+ String query = LoggerMessageFormat.format(null, """
|
|
|
+ from test
|
|
|
+ | stats s = {}([1,2]),
|
|
|
+ s_expr = {}(314.0/100),
|
|
|
+ s_null = {}(null)
|
|
|
+ by emp_no
|
|
|
+ | keep s, s_expr, s_null, emp_no
|
|
|
+ """, testCase.aggFunctionName, testCase.aggFunctionName, testCase.aggFunctionName);
|
|
|
+
|
|
|
+ var plan = plan(query, SubstitutionOnlyOptimizer.INSTANCE);
|
|
|
+
|
|
|
+ var limit = as(plan, Limit.class);
|
|
|
+ var esqlProject = as(limit.child(), EsqlProject.class);
|
|
|
+ var project = as(esqlProject.child(), Project.class);
|
|
|
+ var eval = as(project.child(), Eval.class);
|
|
|
+ var agg = as(eval.child(), Aggregate.class);
|
|
|
+ assertThat(agg.child(), instanceOf(EsRelation.class));
|
|
|
+
|
|
|
+ // Assert exprs
|
|
|
+ var exprs = eval.fields();
|
|
|
+
|
|
|
+ var s = as(exprs.get(0), Alias.class);
|
|
|
+ assertThat(s.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s.child().fold(), equalTo(testCase.aggMultiValue.apply(new int[] { 1, 2 })));
|
|
|
+ var s_expr = as(exprs.get(1), Alias.class);
|
|
|
+ assertThat(s_expr.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s_expr.child().fold(), equalTo(3.14));
|
|
|
+ var s_null = as(exprs.get(2), Alias.class);
|
|
|
+ assertThat(s_null.child(), instanceOf(testCase.substitution));
|
|
|
+ assertThat(s_null.child().fold(), equalTo(null));
|
|
|
+
|
|
|
+ // Assert that the aggregate only does the grouping by emp_no
|
|
|
+ assertThat(Expressions.names(agg.groupings()), contains("emp_no"));
|
|
|
+ assertThat(agg.aggregates().size(), equalTo(1));
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
public void testEmptyMappingIndex() {
|
|
|
EsIndex empty = new EsIndex("empty_test", emptyMap(), emptySet());
|
|
|
IndexResolution getIndexResultAirports = IndexResolution.valid(empty);
|
|
|
@@ -3455,9 +3641,13 @@ public class LogicalPlanOptimizerTests extends ESTestCase {
|
|
|
}
|
|
|
|
|
|
private LogicalPlan plan(String query) {
|
|
|
+ return plan(query, logicalOptimizer);
|
|
|
+ }
|
|
|
+
|
|
|
+ private LogicalPlan plan(String query, LogicalPlanOptimizer optimizer) {
|
|
|
var analyzed = analyzer.analyze(parser.createStatement(query));
|
|
|
// System.out.println(analyzed);
|
|
|
- var optimized = logicalOptimizer.optimize(analyzed);
|
|
|
+ var optimized = optimizer.optimize(analyzed);
|
|
|
// System.out.println(optimized);
|
|
|
return optimized;
|
|
|
}
|
Alexander Spies