handle methods in grouping aggregator

Author: idegtiarenko

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

@@ -36,6 +36,10 @@
 import static org.elasticsearch.compute.gen.AggregatorImplementer.capitalize;
 import static org.elasticsearch.compute.gen.Methods.findMethod;
 import static org.elasticsearch.compute.gen.Methods.findRequiredMethod;
+import static org.elasticsearch.compute.gen.Methods.requireAnyArgs;
+import static org.elasticsearch.compute.gen.Methods.requireName;
+import static org.elasticsearch.compute.gen.Methods.requirePrimitiveOrImplements;
+import static org.elasticsearch.compute.gen.Methods.requireStaticMethod;
 import static org.elasticsearch.compute.gen.Methods.vectorAccessorName;
 import static org.elasticsearch.compute.gen.Types.BIG_ARRAYS;
 import static org.elasticsearch.compute.gen.Types.BLOCK_ARRAY;
@@ -73,7 +77,6 @@ public class GroupingAggregatorImplementer {
     private final ExecutableElement combine;
     private final ExecutableElement combineStates;
     private final ExecutableElement evaluateFinal;
-    private final ExecutableElement combineIntermediate;
     private final List<Parameter> createParameters;
     private final ClassName implementation;
     private final List<AggregatorImplementer.IntermediateStateDesc> intermediateState;
@@ -92,9 +95,16 @@ public GroupingAggregatorImplementer(
         this.declarationType = declarationType;
         this.warnExceptions = warnExceptions;
 
-        this.init = findRequiredMethod(declarationType, new String[] { "init", "initGrouping" }, e -> true);
+        this.init = requireStaticMethod(
+            declarationType,
+            // This should be more restrictive and require org.elasticsearch.compute.aggregation.AggregatorState
+            requirePrimitiveOrImplements(elements, Types.RELEASABLE),
+            requireName("init", "initGrouping"),
+            requireAnyArgs("<arbitrary init arguments>")
+        );
         this.aggState = AggregationState.create(elements, init.getReturnType(), warnExceptions.isEmpty() == false, true);
 
+        // TODO optional timestamp
         this.combine = findRequiredMethod(declarationType, new String[] { "combine" }, e -> {
             if (e.getParameters().size() == 0) {
                 return false;
@@ -106,7 +116,6 @@ public GroupingAggregatorImplementer(
         this.aggParam = AggregationParameter.create(combine.getParameters().get(combine.getParameters().size() - 1).asType());
 
         this.combineStates = findMethod(declarationType, "combineStates");
-        this.combineIntermediate = findMethod(declarationType, "combineIntermediate");
         this.evaluateFinal = findMethod(declarationType, "evaluateFinal");
         this.createParameters = init.getParameters()
             .stream()
@@ -557,6 +566,7 @@ private MethodSpec addIntermediateInput() {
                 });
                 builder.endControlFlow();
             } else {
+                // TODO combineIntermediate with optional block parameter
                 builder.addStatement("$T.combineIntermediate(state, groupId, " + intermediateStateRowAccess() + ")", declarationType);
             }
             builder.endControlFlow();

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

@@ -22,6 +22,7 @@
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
+import javax.lang.model.type.TypeKind;
 import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.ElementFilter;
 import javax.lang.model.util.Elements;
@@ -149,11 +150,13 @@ private static boolean isImplementing(Elements elements, TypeName type, TypeName
     }
 
     private static Stream<TypeName> allInterfacesOf(Elements elements, TypeName type) {
-        return elements.getTypeElement(type.toString())
-            .getInterfaces()
-            .stream()
-            .map(TypeName::get)
-            .flatMap(anInterface -> Stream.concat(Stream.of(anInterface), allInterfacesOf(elements, anInterface)));
+        var typeElement = elements.getTypeElement(type.toString());
+        var superType = Stream.of(typeElement.getSuperclass()).filter(sType -> sType.getKind() != TypeKind.NONE).map(TypeName::get);
+        var interfaces = typeElement.getInterfaces().stream().map(TypeName::get);
+        return Stream.concat(
+            superType.flatMap(sType -> allInterfacesOf(elements, sType)),
+            interfaces.flatMap(anInterface -> Stream.concat(Stream.of(anInterface), allInterfacesOf(elements, anInterface)))
+        );
     }
 
     private static Stream<TypeElement> typeAndSuperType(TypeElement declarationType) {