ESQL: fix COUNT filter pushdown (#117503) (#117652)

* ESQL: fix COUNT filter pushdown (#117503)

If `COUNT` agg has a filter applied, this must also be push down to source. This currently does not happen, but this issue is masked currently by two factors:
* a logical optimisation, `ExtractAggregateCommonFilter` that extracts the filter out of the STATS entirely (and pushes it to source then from a `WHERE`);
* the phisical plan optimisation implementing the  push down, `PushStatsToSource`, currently only applies if there's just one agg function to push down.

However, this fix needs to be applied since:
* it's still present in versions prior to `ExtractAggregateCommonFilter` introduction;
* the defect might resurface when the restriction in `PushStatsToSource` is lifted.

Fixes #115522.

(cherry picked from commit 560e0c5)

* 8.16 adaptation

Author: bpintea

docs/changelog/117503.yaml

@@ -0,0 +1,6 @@
+pr: 117503
+summary: Fix COUNT filter pushdown
+area: ES|QL
+type: bug
+issues:
+ - 115522

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

@@ -2537,6 +2537,57 @@ c2:l |c2_f:l |m2:i |m2_f:i |c:l
 1    |1      |5    |5      |21
 ;
 
+simpleCountOnFieldWithFilteringAndNoGrouping
+required_capability: per_agg_filtering
+from employees
+| stats c1 = count(emp_no) where emp_no < 10042
+;
+
+c1:long
+41
+;
+
+simpleCountOnFieldWithFilteringOnDifferentFieldAndNoGrouping
+required_capability: per_agg_filtering
+from employees
+| stats c1 = count(hire_date) where emp_no < 10042
+;
+
+c1:long
+41
+;
+
+simpleCountOnStarWithFilteringAndNoGrouping
+required_capability: per_agg_filtering
+from employees
+| stats c1 = count(*) where emp_no < 10042
+;
+
+c1:long
+41
+;
+
+simpleCountWithFilteringAndNoGroupingOnFieldWithNulls
+required_capability: per_agg_filtering
+from employees
+| stats c1 = count(birth_date) where emp_no <= 10050
+;
+
+c1:long
+40
+;
+
+
+simpleCountWithFilteringAndNoGroupingOnFieldWithMultivalues
+required_capability: per_agg_filtering
+from employees
+| stats c1 = count(job_positions) where emp_no <= 10003
+;
+
+c1:long
+3
+;
+
 filterIsAlwaysTrue
 required_capability: per_agg_filtering
 FROM employees

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushStatsToSource.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.util.Queries;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Count;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext;
@@ -25,12 +26,15 @@
 import org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.planner.AbstractPhysicalOperationProviders;
+import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 
 import java.util.ArrayList;
 import java.util.List;
 
+import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
+import static org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushFiltersToSource.canPushToSource;
 import static org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec.StatsType.COUNT;
 
 /**
@@ -98,6 +102,13 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                                 }
                             }
                             if (fieldName != null) {
+                                if (count.hasFilter()) {
+                                    if (canPushToSource(count.filter(), fa -> false) == false) {
+                                        return null; // can't push down
+                                    }
+                                    var countFilter = PlannerUtils.TRANSLATOR_HANDLER.asQuery(count.filter());
+                                    query = Queries.combine(Queries.Clause.MUST, asList(countFilter.asBuilder(), query));
+                                }
                                 return new EsStatsQueryExec.Stat(fieldName, COUNT, query);
                             }
                         }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java

@@ -67,6 +67,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Function;
 
 import static java.util.Arrays.asList;
 import static org.elasticsearch.compute.aggregation.AggregatorMode.FINAL;
@@ -373,6 +374,54 @@ public void testMultiCountAllWithFilter() {
         assertThat(plan.anyMatch(EsQueryExec.class::isInstance), is(true));
     }
 
+    @SuppressWarnings("unchecked")
+    public void testSingleCountWithStatsFilter() {
+        var analyzer = makeAnalyzer("mapping-default.json", new EnrichResolution());
+        var plannerOptimizer = new TestPlannerOptimizer(config, analyzer);
+        var plan = plannerOptimizer.plan("""
+            from test
+            | stats c = count(hire_date) where emp_no < 10042
+            """, IS_SV_STATS);
+
+        var limit = as(plan, LimitExec.class);
+        var agg = as(limit.child(), AggregateExec.class);
+        assertThat(agg.getMode(), is(FINAL));
+        var exchange = as(agg.child(), ExchangeExec.class);
+        var esStatsQuery = as(exchange.child(), EsStatsQueryExec.class);
+        assertThat(esStatsQuery.stats().size(), is(1));
+
+        Function<String, String> compact = s -> s.replaceAll("\\s+", "");
+        assertThat(compact.apply(esStatsQuery.stats().get(0).query().toString()), is(compact.apply("""
+            {
+                "bool": {
+                    "must": [
+                        {
+                            "esql_single_value": {
+                                "field": "emp_no",
+                                "next": {
+                                    "range": {
+                                        "emp_no": {
+                                            "lt": 10042,
+                                            "boost": 1.0
+                                        }
+                                    }
+                                },
+                                "source": "emp_no < 10042@2:36"
+                            }
+                        },
+                        {
+                            "exists": {
+                                "field": "hire_date",
+                                "boost": 1.0
+                            }
+                        }
+                    ],
+                    "boost": 1.0
+                }
+            }
+            """)));
+    }
+
     /**
      * Expecting
      * LimitExec[1000[INTEGER]]