[ESQL] Fix TopN grouping aggregates bug with non-qualifying intermediate states

x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/GroupingAggregatorImplementer.java

@@ -203,10 +203,10 @@ private TypeSpec type() {
         for (ClassName groupIdClass : GROUP_IDS_CLASSES) {
             builder.addMethod(addRawInputLoop(groupIdClass, blockType(aggParam.type())));
             builder.addMethod(addRawInputLoop(groupIdClass, vectorType(aggParam.type())));
+            builder.addMethod(addIntermediateInput(groupIdClass));
         }
-        builder.addMethod(selectedMayContainUnseenGroups());
-        builder.addMethod(addIntermediateInput());
         builder.addMethod(addIntermediateRowInput());
+        builder.addMethod(selectedMayContainUnseenGroups());
         builder.addMethod(evaluateIntermediate());
         builder.addMethod(evaluateFinal());
         builder.addMethod(toStringMethod());
@@ -580,11 +580,12 @@ private MethodSpec selectedMayContainUnseenGroups() {
         return builder.build();
     }
 
-    private MethodSpec addIntermediateInput() {
+    private MethodSpec addIntermediateInput(TypeName groupsType) {
+        boolean groupsIsBlock = groupsType.toString().endsWith("Block");
         MethodSpec.Builder builder = MethodSpec.methodBuilder("addIntermediateInput");
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
         builder.addParameter(TypeName.INT, "positionOffset");
-        builder.addParameter(INT_VECTOR, "groups");
+        builder.addParameter(groupsType, "groups");
         builder.addParameter(PAGE, "page");
 
         builder.addStatement("state.enableGroupIdTracking(new $T.Empty())", SEEN_GROUP_IDS);
@@ -610,7 +611,18 @@ private MethodSpec addIntermediateInput() {
         }
         builder.beginControlFlow("for (int groupPosition = 0; groupPosition < groups.getPositionCount(); groupPosition++)");
         {
-            builder.addStatement("int groupId = groups.getInt(groupPosition)");
+            if (groupsIsBlock) {
+                builder.beginControlFlow("if (groups.isNull(groupPosition))");
+                builder.addStatement("continue");
+                builder.endControlFlow();
+                builder.addStatement("int groupStart = groups.getFirstValueIndex(groupPosition)");
+                builder.addStatement("int groupEnd = groupStart + groups.getValueCount(groupPosition)");
+                builder.beginControlFlow("for (int g = groupStart; g < groupEnd; g++)");
+                builder.addStatement("int groupId = groups.getInt(g)");
+            } else {
+                builder.addStatement("int groupId = groups.getInt(groupPosition)");
+            }
+
             if (aggState.declaredType().isPrimitive()) {
                 if (warnExceptions.isEmpty()) {
                     assert intermediateState.size() == 2;
@@ -661,6 +673,9 @@ private MethodSpec addIntermediateInput() {
                     declarationType
                 );
             }
+            if (groupsIsBlock) {
+                builder.endControlFlow();
+            }
             builder.endControlFlow();
         }
         return builder.build();