ESQL: Fix count optimization with pushable union types

docs/changelog/127225.yaml

@@ -0,0 +1,6 @@
+pr: 127225
+summary: Fix count optimization with pushable union types
+area: ES|QL
+type: bug
+issues:
+ - 127200

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

@@ -393,6 +393,30 @@ count:long | @timestamp:date
          0 | 2023-10-23T13:50:00.000Z
 ;
 
+multiIndexIpStatsPushableCount
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_str
+| STATS count=count(client_ip::ip)
+;
+
+count:long
+14
+;
+
+multiIndexIpStatsPushableCountEval
+required_capability: union_types
+required_capability: fix_count_pushdown_for_union_types
+
+FROM sample_data, sample_data_str
+| EVAL client_ip = client_ip::ip
+| STATS count=count(client_ip)
+;
+
+count:long
+14
+;
 
 multiIndexIpStringStatsInline2
 required_capability: union_types

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1038,7 +1038,13 @@ public enum Cap {
         /**
          * Support for the SAMPLE command
          */
-        SAMPLE(Build.current().isSnapshot());
+        SAMPLE(Build.current().isSnapshot()),
+
+        /**
+         * When pushing down {@code STATS count(field::type)} for a union type field, we wrongly used a synthetic attribute name in the
+         * query instead of the actual field name. This led to 0 counts instead of the correct result.
+         */
+        FIX_COUNT_PUSHDOWN_FOR_UNION_TYPES;
 
         private final boolean enabled;
 

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

@@ -66,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan, LocalLogicalOptimizerContext localLog
 
     private LogicalPlan missingToNull(LogicalPlan plan, Predicate<FieldAttribute> shouldBeRetained) {
         if (plan instanceof EsRelation relation) {
-            // Remove missing fields from the EsRelation because this is not where we will obtain them from; replace them by an Eval right
-            // after, instead. This allows us to safely re-use the attribute ids of the corresponding FieldAttributes.
+            // For any missing field, place an Eval right after the EsRelation to assign null values to that attribute (using the same name
+            // id!), thus avoiding that InsertFieldExtrations inserts a field extraction later.
             // This means that an EsRelation[field1, field2, field3] where field1 and field 3 are missing will be replaced by
             // Project[field1, field2, field3] <- keeps the ordering intact
             // \_Eval[field1 = null, field3 = null]
-            // \_EsRelation[field2]
+            // \_EsRelation[field1, field2, field3]
             List<Attribute> relationOutput = relation.output();
             Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
             List<NamedExpression> newProjections = new ArrayList<>(relationOutput.size());

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

@@ -94,9 +94,9 @@ private Tuple<List<Attribute>, List<EsStatsQueryExec.Stat>> pushableStats(
                             // check if regular field
                             else {
                                 if (target instanceof FieldAttribute fa) {
-                                    var fName = fa.name();
+                                    var fName = fa.fieldName();
                                     if (context.searchStats().isSingleValue(fName)) {
-                                        fieldName = fa.name();
+                                        fieldName = fName;
                                         query = QueryBuilders.existsQuery(fieldName);
                                     }
                                 }