Avoid dropping aggregate groupings in local plans

docs/changelog/129370.yaml

@@ -0,0 +1,6 @@
+pr: 129370
+summary: Avoid dropping aggregate groupings in local plans
+area: ES|QL
+type: bug
+issues:
+ - 128054

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java

@@ -1680,6 +1680,36 @@ public void testQueryOnEmptyDataIndex() {
         }
     }
 
+    public void testGroupingStatsOnMissingFields() {
+        assertAcked(client().admin().indices().prepareCreate("test_missing_fields").setMapping("data", "type=long"));
+        long oneValue = between(1, 1000);
+        indexDoc("missing_field_index", "1", "data", oneValue);
+        refresh("missing_field_index");
+        QueryPragmas pragmas = randomPragmas();
+        pragmas = new QueryPragmas(
+            Settings.builder().put(pragmas.getSettings()).put(QueryPragmas.MAX_CONCURRENT_SHARDS_PER_NODE.getKey(), 1).build()
+        );
+        try (var r = run("FROM missing_field_index,test | STATS s = sum(data) BY color, tag | SORT color", pragmas)) {
+            var rows = getValuesList(r);
+            assertThat(rows, hasSize(4));
+            for (List<Object> row : rows) {
+                assertThat(row, hasSize(3));
+            }
+            assertThat(rows.get(0).get(0), equalTo(20L));
+            assertThat(rows.get(0).get(1), equalTo("blue"));
+            assertNull(rows.get(0).get(2));
+            assertThat(rows.get(1).get(0), equalTo(10L));
+            assertThat(rows.get(1).get(1), equalTo("green"));
+            assertNull(rows.get(1).get(2));
+            assertThat(rows.get(2).get(0), equalTo(30L));
+            assertThat(rows.get(2).get(1), equalTo("red"));
+            assertNull(rows.get(2).get(2));
+            assertThat(rows.get(3).get(0), equalTo(oneValue));
+            assertNull(rows.get(3).get(1));
+            assertNull(rows.get(3).get(2));
+        }
+    }
+
     private void assertEmptyIndexQueries(String from) {
         try (EsqlQueryResponse resp = run(from + "METADATA _source | KEEP _source | LIMIT 1")) {
             assertFalse(resp.values().hasNext());
@@ -1817,6 +1847,8 @@ private void createAndPopulateIndex(String indexName, Settings additionalSetting
                     "time",
                     "type=long",
                     "color",
+                    "type=keyword",
+                    "tag",
                     "type=keyword"
                 )
         );

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/RuleUtils.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules;
 
-import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -16,13 +15,11 @@
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
-import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.function.Predicate;
 
 public final class RuleUtils {
@@ -41,37 +38,20 @@ public static Tuple<List<Alias>, List<NamedExpression>> aliasedNulls(
         List<Attribute> outputAttributes,
         Predicate<Attribute> shouldBeReplaced
     ) {
-        Map<DataType, Alias> nullLiterals = Maps.newLinkedHashMapWithExpectedSize(DataType.types().size());
         List<NamedExpression> newProjections = new ArrayList<>(outputAttributes.size());
+        List<Alias> nullLiterals = new ArrayList<>(outputAttributes.size());
         for (Attribute attr : outputAttributes) {
             NamedExpression projection;
             if (shouldBeReplaced.test(attr)) {
-                DataType dt = attr.dataType();
-                Alias nullAlias = nullLiterals.get(dt);
-                // save the first field as null (per datatype)
-                if (nullAlias == null) {
-                    // Keep the same id so downstream query plans don't need updating
-                    // NOTE: THIS IS BRITTLE AND CAN LEAD TO BUGS.
-                    // In case some optimizer rule or so inserts a plan node that requires the field BEFORE the Eval that we're adding
-                    // on top of the EsRelation, this can trigger a field extraction in the physical optimizer phase, causing wrong
-                    // layouts due to a duplicate name id.
-                    // If someone reaches here AGAIN when debugging e.g. ClassCastExceptions NPEs from wrong layouts, we should probably
-                    // give up on this approach and instead insert EvalExecs in InsertFieldExtraction.
-                    Alias alias = new Alias(attr.source(), attr.name(), Literal.of(attr, null), attr.id());
-                    nullLiterals.put(dt, alias);
-                    projection = alias.toAttribute();
-                }
-                // otherwise point to it since this avoids creating field copies
-                else {
-                    projection = new Alias(attr.source(), attr.name(), nullAlias.toAttribute(), attr.id());
-                }
+                Alias alias = new Alias(attr.source(), attr.name(), Literal.of(attr, null), attr.id());
+                nullLiterals.add(alias);
+                projection = alias.toAttribute();
             } else {
                 projection = attr;
             }
             newProjections.add(projection);
         }
-
-        return new Tuple<>(new ArrayList<>(nullLiterals.values()), newProjections);
+        return new Tuple<>(nullLiterals, newProjections);
     }
 
     /**

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

@@ -744,6 +744,14 @@ public void testUnionTypesInferNonNullAggConstraint() {
         assertEquals("integer_long_field", unionTypeField.fieldName().string());
     }
 
+    public void testGroupingByMissingFields() {
+        var plan = plan("FROM test | STATS AVG(salary) BY first_name, last_name ");
+        var testStats = statsForMissingField("first_name", "last_name");
+        var localPlan = localPlan(plan, testStats);
+        Aggregate aggregate = (Aggregate) localPlan.collectFirstChildren(p -> p instanceof Aggregate).getFirst();
+        assertThat(aggregate.groupings(), hasSize(2));
+    }
+
     private IsNotNull isNotNull(Expression field) {
         return new IsNotNull(EMPTY, field);
     }

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

@@ -45,7 +45,6 @@
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
-import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
@@ -133,6 +132,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
 import static org.elasticsearch.xpack.esql.plan.physical.EsStatsQueryExec.StatsType;
 import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
@@ -925,21 +925,12 @@ public void testMissingFieldsDoNotGetExtracted() {
                 "salary"
             )
         );
-        // emp_no
-        assertThat(projections.get(1), instanceOf(ReferenceAttribute.class));
-        // first_name
-        assertThat(projections.get(2), instanceOf(ReferenceAttribute.class));
-
-        // last_name --> first_name
-        var nullAlias = Alias.unwrap(projections.get(8));
-        assertThat(Expressions.name(nullAlias), is("first_name"));
-        // salary --> emp_no
-        nullAlias = Alias.unwrap(projections.get(10));
-        assertThat(Expressions.name(nullAlias), is("emp_no"));
-        // check field extraction is skipped and that evaled fields are not extracted anymore
         var field = as(project.child(), FieldExtractExec.class);
         var fields = field.attributesToExtract();
         assertThat(Expressions.names(fields), contains("_meta_field", "gender", "hire_date", "job", "job.raw", "languages", "long_noidx"));
+        var eval = as(field.child(), EvalExec.class);
+        List<String> nullFields = Expressions.names(eval.fields());
+        assertThat(nullFields, containsInAnyOrder("first_name", "last_name", "emp_no", "salary"));
     }
 
     /*