ESQL: Fix treating all fields as MV in COUNT pushdown (#106720)

Fix a mistake in #106690 that accidentally prevented COUNT(field) from being pushed down in case field is single-valued.
Add test to avoid future regressions.

docs/changelog/106720.yaml

@@ -0,0 +1,5 @@
+pr: 106720
+summary: "ESQL: Fix treating all fields as MV in COUNT pushdown"
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

@@ -30,6 +30,7 @@ import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.esql.EsqlTestUtils;
 import org.junit.After;
 import org.junit.Before;
 
@@ -79,12 +80,8 @@ public abstract class RestEsqlTestCase extends ESRestTestCase {
     private static final String MAPPING_ALL_TYPES;
 
     static {
-        try (InputStream mappingPropertiesStream = RestEsqlTestCase.class.getResourceAsStream("/mapping-all-types.json")) {
-            String properties = new String(mappingPropertiesStream.readAllBytes(), StandardCharsets.UTF_8);
-            MAPPING_ALL_TYPES = "{\"mappings\": " + properties + "}";
-        } catch (IOException ex) {
-            throw new RuntimeException(ex);
-        }
+        String properties = EsqlTestUtils.loadUtf8TextFile("/mapping-all-types.json");
+        MAPPING_ALL_TYPES = "{\"mappings\": " + properties + "}";
     }
 
     private static final String DOCUMENT_TEMPLATE = """

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -34,6 +34,9 @@ import org.elasticsearch.xpack.ql.type.TypesTests;
 import org.elasticsearch.xpack.ql.util.StringUtils;
 import org.junit.Assert;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -146,6 +149,14 @@ public final class EsqlTestUtils {
         return TypesTests.loadMapping(EsqlDataTypeRegistry.INSTANCE, name, true);
     }
 
+    public static String loadUtf8TextFile(String name) {
+        try (InputStream textStream = EsqlTestUtils.class.getResourceAsStream(name)) {
+            return new String(textStream.readAllBytes(), StandardCharsets.UTF_8);
+        } catch (IOException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
+
     public static EnrichResolution emptyPolicyResolution() {
         return new EnrichResolution();
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/stats/SearchStats.java

@@ -198,8 +198,9 @@ public class SearchStats {
                 // fields are MV per default
                 var sv = new boolean[] { false };
                 for (SearchExecutionContext context : contexts) {
-                    MappedFieldType mappedType = context.isFieldMapped(field) ? null : context.getFieldType(field);
+                    MappedFieldType mappedType = context.isFieldMapped(field) ? context.getFieldType(field) : null;
                     if (mappedType != null) {
+                        sv[0] = true;
                         doWithContexts(r -> {
                             sv[0] &= detectSingleValue(r, mappedType, field);
                             return sv[0];

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

@@ -9,12 +9,16 @@ package org.elasticsearch.xpack.esql.optimizer;
 
 import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 
+import org.apache.lucene.search.IndexSearcher;
 import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.index.mapper.MapperServiceTestCase;
+import org.elasticsearch.index.mapper.ParsedDocument;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
-import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
 import org.elasticsearch.xpack.esql.EsqlTestUtils;
 import org.elasticsearch.xpack.esql.EsqlTestUtils.TestSearchStats;
@@ -53,8 +57,10 @@ import org.elasticsearch.xpack.ql.index.IndexResolution;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 import org.elasticsearch.xpack.ql.type.EsField;
+import org.elasticsearch.xpack.ql.util.Holder;
 import org.junit.Before;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
@@ -78,7 +84,7 @@ import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.nullValue;
 
 //@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE,org.elasticsearch.compute:TRACE", reason = "debug")
-public class LocalPhysicalPlanOptimizerTests extends ESTestCase {
+public class LocalPhysicalPlanOptimizerTests extends MapperServiceTestCase {
 
     private static final String PARAM_FORMATTING = "%1$s";
 
@@ -270,6 +276,70 @@ public class LocalPhysicalPlanOptimizerTests extends ESTestCase {
         assertThat(plan.anyMatch(EsQueryExec.class::isInstance), is(true));
     }
 
+    public void testCountPushdownForSvAndMvFields() throws IOException {
+        String properties = EsqlTestUtils.loadUtf8TextFile("/mapping-basic.json");
+        String mapping = "{\"mappings\": " + properties + "}";
+
+        String query = """
+            from test
+            | stats c = count(salary)
+            """;
+
+        PhysicalPlan plan;
+
+        List<List<String>> docsCasesWithoutPushdown = List.of(
+            // No pushdown yet in case of MVs
+            List.of("{ \"salary\" : [1,2] }"),
+            List.of("{ \"salary\" : [1,2] }", "{ \"salary\" : null}")
+        );
+        for (List<String> docs : docsCasesWithoutPushdown) {
+            plan = planWithMappingAndDocs(query, mapping, docs);
+            // No EsSatsQueryExec as leaf of the plan.
+            assertThat(plan.anyMatch(EsQueryExec.class::isInstance), is(true));
+        }
+
+        // Cases where we can push this down as a COUNT(*) since there are only SVs
+        List<List<String>> docsCasesWithPushdown = List.of(List.of(), List.of("{ \"salary\" : 1 }"), List.of("{ \"salary\": null }"));
+        for (List<String> docs : docsCasesWithPushdown) {
+            plan = planWithMappingAndDocs(query, mapping, docs);
+
+            Holder<EsStatsQueryExec> leaf = new Holder<>();
+            plan.forEachDown(p -> {
+                if (p instanceof EsStatsQueryExec s) {
+                    leaf.set(s);
+                }
+            });
+
+            String expectedStats = """
+                [Stat[name=salary, type=COUNT, query={
+                  "exists" : {
+                    "field" : "salary",
+                    "boost" : 1.0
+                  }
+                }]]""";
+            assertNotNull(leaf.get());
+            assertThat(leaf.get().stats().toString(), equalTo(expectedStats));
+        }
+    }
+
+    private PhysicalPlan planWithMappingAndDocs(String query, String mapping, List<String> docs) throws IOException {
+        MapperService mapperService = createMapperService(mapping);
+        List<ParsedDocument> parsedDocs = docs.stream().map(d -> mapperService.documentMapper().parse(source(d))).toList();
+
+        Holder<PhysicalPlan> plan = new Holder<>(null);
+        withLuceneIndex(mapperService, indexWriter -> {
+            for (ParsedDocument parsedDoc : parsedDocs) {
+                indexWriter.addDocument(parsedDoc.rootDoc());
+            }
+        }, directoryReader -> {
+            IndexSearcher searcher = newSearcher(directoryReader);
+            SearchExecutionContext ctx = createSearchExecutionContext(mapperService, searcher);
+            plan.set(plan(query, new SearchStats(List.of(ctx))));
+        });
+
+        return plan.get();
+    }
+
     // optimized doesn't know yet how to break down different multi count
     public void testCountMultipleFieldsWithFilter() {
         var plan = plan("""