Update tests

Prevent optimization across multiple fields
elastic · Oct 2, 2023 · 671f7ce · 671f7ce
1 parent 935e608
commit 671f7ce
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/...esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java b/...esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java
@@ -66,6 +66,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
@@ -329,6 +330,14 @@ protected PhysicalPlan rule(AggregateExec aggregateExec) {
             if (aggregateExec.child() instanceof EsQueryExec queryExec) {
                 var tuple = pushableStats(aggregateExec);
 
+                // for the moment support pushing count just for one field
+                List<Stat> stats = tuple.v2();
+                if (stats.size() > 1) {
+                    if (stats.stream().map(Stat::name).collect(Collectors.toSet()).size() > 1) {
+                        return aggregateExec;
+                    }
+                }
+
                 // TODO: handle case where some aggs cannot be pushed down by breaking the aggs into two sources (regular + stats) + union
                 // use the stats since the attributes are larger in size (due to seen)
                 if (tuple.v2().size() == aggregateExec.aggregates().size()) {

diff --git a/...esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java b/...esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizerTests.java
@@ -44,6 +44,7 @@
 import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec.FieldSort;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec;
+import org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec.Stat;
 import org.elasticsearch.xpack.esql.plan.physical.EstimatesRowSize;
 import org.elasticsearch.xpack.esql.plan.physical.EvalExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
@@ -1872,7 +1873,20 @@ public void testCountOneFieldWithFilter() {
             | where salary > 1000
             | stats c = count(salary)
             """));
-        assertThat(plan.anyMatch(EsQueryExec.class::isInstance), is(true));
+
+        var limit = as(plan, LimitExec.class);
+        var agg = as(limit.child(), AggregateExec.class);
+        assertThat(agg.getMode(), is(FINAL));
+        assertThat(Expressions.names(agg.aggregates()), contains("c"));
+        var exchange = as(agg.child(), ExchangeExec.class);
+        var esStatsQuery = as(exchange.child(), EsStatsQueryExec.class);
+        assertThat(esStatsQuery.limit(), is(nullValue()));
+        assertThat(Expressions.names(esStatsQuery.output()), contains("count", "seen"));
+        var stat = as(esStatsQuery.stats().get(0), Stat.class);
+        assertThat(stat.query(), is(QueryBuilders.existsQuery("salary")));
+        var expected = wrapWithSingleQuery(QueryBuilders.rangeQuery("salary").gt(1000), "salary");
+        assertThat(expected.toString(), is(esStatsQuery.query().toString()));
+
     }
 
     // optimized doesn't know yet how to push down count over field