From 1d96c16c45cdcc1a1e90081344d6cfc2b8b68542 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Tue, 28 Oct 2025 17:21:01 -0700
Subject: [PATCH] do not deduplicate intermediate attributes

---
 .../AbstractPhysicalOperationProviders.java       | 13 ++++---------
 .../xpack/esql/planner/AggregateMapper.java       | 15 ++++-----------
 2 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AbstractPhysicalOperationProviders.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AbstractPhysicalOperationProviders.java
index 0702a6e3db642..fbb6520d5cde7 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AbstractPhysicalOperationProviders.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AbstractPhysicalOperationProviders.java
@@ -40,10 +40,8 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.function.Consumer;
 
@@ -250,7 +248,6 @@ private record AggFunctionSupplierContext(AggregatorFunctionSupplier supplier, L
     private static class IntermediateInputs {
         private final List<Attribute> inputAttributes;
         private int nextOffset;
-        private final Map<AggregateFunction, Integer> offsets = new HashMap<>();
 
         IntermediateInputs(AggregateExec aggregateExec) {
             inputAttributes = aggregateExec.child().output();
@@ -259,12 +256,10 @@ private static class IntermediateInputs {
 
         List<Attribute> nextInputAttributes(AggregateFunction af, boolean grouping) {
             int intermediateStateSize = AggregateMapper.intermediateStateDesc(af, grouping).size();
-            int offset = offsets.computeIfAbsent(af, unused -> {
-                int v = nextOffset;
-                nextOffset += intermediateStateSize;
-                return v;
-            });
-            return inputAttributes.subList(offset, offset + intermediateStateSize);
+            int endOffset = nextOffset + intermediateStateSize;
+            List<Attribute> attributes = inputAttributes.subList(nextOffset, endOffset);
+            nextOffset = endOffset;
+            return attributes;
         }
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java
index 3981b71f316b0..6c684fa121079 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java
@@ -13,7 +13,6 @@
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
@@ -23,9 +22,8 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
 
-import java.util.HashSet;
+import java.util.ArrayList;
 import java.util.List;
-import java.util.Set;
 import java.util.stream.Stream;
 
 /**
@@ -42,17 +40,12 @@ public static List<NamedExpression> mapGrouping(List<? extends NamedExpression>
     }
 
     private static List<NamedExpression> doMapping(List<? extends NamedExpression> aggregates, boolean grouping) {
-        Set<Expression> seen = new HashSet<>();
-        AttributeMap.Builder<NamedExpression> attrToExpressionsBuilder = AttributeMap.builder();
+        List<NamedExpression> namedExpressions = new ArrayList<>();
         for (NamedExpression agg : aggregates) {
             Expression inner = Alias.unwrap(agg);
-            if (seen.add(inner)) {
-                for (var ne : computeEntryForAgg(agg.name(), inner, grouping)) {
-                    attrToExpressionsBuilder.put(ne.toAttribute(), ne);
-                }
-            }
+            namedExpressions.addAll(computeEntryForAgg(agg.name(), inner, grouping));
         }
-        return attrToExpressionsBuilder.build().values().stream().toList();
+        return namedExpressions;
     }
 
     public static List<IntermediateStateDesc> intermediateStateDesc(AggregateFunction fn, boolean grouping) {