Fix TopN row size estimate (#119476)

Currently, we only account for attributes after the TopN operation when 
estimating its row size. This results in invalid estimates for the final
TopN on the coordinator and the node-level reduction TopN. This change
corrects the estimates and removes the workaround, paving the way for
enabling node-level reduction in ES|QL.

Closes #106956

Author: dnhatn

docs/changelog/119476.yaml

@@ -0,0 +1,6 @@
+pr: 119476
+summary: Fix TopN row size estimate
+area: ES|QL
+type: bug
+issues:
+ - 106956

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EstimatesRowSize.java

@@ -103,9 +103,6 @@ public String toString() {
 
     static int estimateSize(DataType dataType) {
         ElementType elementType = PlannerUtils.toElementType(dataType);
-        if (elementType == ElementType.DOC) {
-            throw new EsqlIllegalArgumentException("can't load a [doc] with field extraction");
-        }
         if (elementType == ElementType.UNKNOWN) {
             throw new EsqlIllegalArgumentException("[unknown] can't be the result of field extraction");
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/TopNExec.java

@@ -10,9 +10,11 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.Order;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 
@@ -95,6 +97,9 @@ public Integer estimatedRowSize() {
 
     @Override
     public PhysicalPlan estimateRowSize(State state) {
+        final List<Attribute> output = output();
+        final boolean needsSortedDocIds = output.stream().anyMatch(a -> a.dataType() == DataType.DOC_DATA_TYPE);
+        state.add(needsSortedDocIds, output);
         int size = state.consumeAllFields(true);
         return Objects.equals(this.estimatedRowSize, size) ? this : new TopNExec(source(), child(), order, limit, size);
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -345,6 +345,8 @@ private PhysicalOperation planExchangeSource(ExchangeSourceExec exchangeSource,
     }
 
     private PhysicalOperation planTopN(TopNExec topNExec, LocalExecutionPlannerContext context) {
+        final Integer rowSize = topNExec.estimatedRowSize();
+        assert rowSize != null && rowSize > 0 : "estimated row size [" + rowSize + "] wasn't set";
         PhysicalOperation source = plan(topNExec.child(), context);
 
         ElementType[] elementTypes = new ElementType[source.layout.numberOfChannels()];
@@ -385,24 +387,8 @@ private PhysicalOperation planTopN(TopNExec topNExec, LocalExecutionPlannerConte
         } else {
             throw new EsqlIllegalArgumentException("limit only supported with literal values");
         }
-
-        // TODO Replace page size with passing estimatedRowSize down
-        /*
-         * The 2000 below is a hack to account for incoming size and to make
-         * sure the estimated row size is never 0 which'd cause a divide by 0.
-         * But we should replace this with passing the estimate into the real
-         * topn and letting it actually measure the size of rows it produces.
-         * That'll be more accurate. And we don't have a path for estimating
-         * incoming rows. And we don't need one because we can estimate.
-         */
         return source.with(
-            new TopNOperatorFactory(
-                limit,
-                asList(elementTypes),
-                asList(encoders),
-                orders,
-                context.pageSize(2000 + topNExec.estimatedRowSize())
-            ),
+            new TopNOperatorFactory(limit, asList(elementTypes), asList(encoders), orders, context.pageSize(rowSize)),
             source.layout
         );
     }

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

@@ -143,6 +143,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -920,16 +921,16 @@ public void testQueryWithNull() {
 
         var optimized = optimizedPlan(plan);
         var topN = as(optimized, TopNExec.class);
-        // no fields are added after the top n - so 0 here
-        assertThat(topN.estimatedRowSize(), equalTo(0));
+        // all fields + nullsum are loaded in the final TopN
+        assertThat(topN.estimatedRowSize(), equalTo(allFieldRowSize + Integer.BYTES));
 
         var exchange = asRemoteExchange(topN.child());
         var project = as(exchange.child(), ProjectExec.class);
         var extract = as(project.child(), FieldExtractExec.class);
         var eval = as(extract.child(), EvalExec.class);
         var source = source(eval.child());
         // All fields loaded
-        assertThat(source.estimatedRowSize(), equalTo(allFieldRowSize + 3 * Integer.BYTES + Long.BYTES));
+        assertThat(source.estimatedRowSize(), equalTo(allFieldRowSize + 3 * Integer.BYTES + 2 * Integer.BYTES));
     }
 
     public void testPushAndInequalitiesFilter() {
@@ -1141,8 +1142,8 @@ public void testExtractorForEvalWithoutProject() throws Exception {
         var project = as(exchange.child(), ProjectExec.class);
         var extract = as(project.child(), FieldExtractExec.class);
         var topNLocal = as(extract.child(), TopNExec.class);
-        // two extra ints for forwards and backwards map
-        assertThat(topNLocal.estimatedRowSize(), equalTo(allFieldRowSize + Integer.BYTES * 2));
+        // all fields plus nullsum and shards, segments, docs and two extra ints for forwards and backwards map
+        assertThat(topNLocal.estimatedRowSize(), equalTo(allFieldRowSize + Integer.BYTES + Integer.BYTES * 2 + Integer.BYTES * 3));
 
         var eval = as(topNLocal.child(), EvalExec.class);
         var source = source(eval.child());
@@ -7650,6 +7651,31 @@ public void testScoreTopN() {
         assertTrue(esRelation.output().stream().anyMatch(a -> a.name().equals(MetadataAttribute.SCORE) && a instanceof MetadataAttribute));
     }
 
+    public void testReductionPlanForTopN() {
+        int limit = between(1, 100);
+        var plan = physicalPlan(String.format(Locale.ROOT, """
+            FROM test
+            | sort emp_no
+            | LIMIT %d
+            """, limit));
+        Tuple<PhysicalPlan, PhysicalPlan> plans = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(plan, config);
+        PhysicalPlan reduction = PlannerUtils.reductionPlan(plans.v2());
+        TopNExec reductionTopN = as(reduction, TopNExec.class);
+        assertThat(reductionTopN.estimatedRowSize(), equalTo(allFieldRowSize));
+        assertThat(reductionTopN.limit().fold(), equalTo(limit));
+    }
+
+    public void testReductionPlanForAggs() {
+        var plan = physicalPlan("""
+            FROM test
+            | stats x = sum(salary) BY first_name
+            """);
+        Tuple<PhysicalPlan, PhysicalPlan> plans = PlannerUtils.breakPlanBetweenCoordinatorAndDataNode(plan, config);
+        PhysicalPlan reduction = PlannerUtils.reductionPlan(plans.v2());
+        AggregateExec reductionAggs = as(reduction, AggregateExec.class);
+        assertThat(reductionAggs.estimatedRowSize(), equalTo(58)); // double and keyword
+    }
+
     @SuppressWarnings("SameParameterValue")
     private static void assertFilterCondition(
         Filter filter,