handle methods in grouping aggregator

Author: idegtiarenko

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

@@ -22,9 +22,10 @@
 
 import java.util.Arrays;
 import java.util.List;
-import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
@@ -34,9 +35,8 @@
 
 import static java.util.stream.Collectors.joining;
 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.requireAnyType;
 import static org.elasticsearch.compute.gen.Methods.requireArgs;
 import static org.elasticsearch.compute.gen.Methods.requireName;
 import static org.elasticsearch.compute.gen.Methods.requirePrimitiveOrImplements;
@@ -79,7 +79,6 @@ public class GroupingAggregatorImplementer {
     private final List<TypeMirror> warnExceptions;
     private final ExecutableElement init;
     private final ExecutableElement combine;
-    private final ExecutableElement combineStates;
     private final List<Parameter> createParameters;
     private final ClassName implementation;
     private final List<AggregatorImplementer.IntermediateStateDesc> intermediateState;
@@ -100,25 +99,21 @@ public GroupingAggregatorImplementer(
 
         this.init = requireStaticMethod(
             declarationType,
-            // This should be more restrictive and require org.elasticsearch.compute.aggregation.AggregatorState
+            // This should be more restrictive and require org.elasticsearch.compute.aggregation.GroupingAggregatorState
             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;
-            }
-            TypeName firstParamType = TypeName.get(e.getParameters().get(0).asType());
-            return Objects.equals(firstParamType.toString(), aggState.declaredType().toString());
-        });
-        // TODO support multiple parameters
+        this.combine = requireStaticMethod(
+            declarationType,
+            aggState.declaredType().isPrimitive() ? requireType(aggState.declaredType()) : requireVoidType(),
+            requireName("combine"),
+            combineArgs(aggState, includeTimestampVector)
+        );
         this.aggParam = AggregationParameter.create(combine.getParameters().get(combine.getParameters().size() - 1).asType());
 
-        this.combineStates = findMethod(declarationType, "combineStates");
         this.createParameters = init.getParameters()
             .stream()
             .map(Parameter::from)
@@ -136,6 +131,25 @@ public GroupingAggregatorImplementer(
         this.includeTimestampVector = includeTimestampVector;
     }
 
+    private static Methods.ArgumentMatcher combineArgs(AggregationState aggState, boolean includeTimestampVector) {
+        if (aggState.declaredType().isPrimitive()) {
+            return requireArgs(requireType(aggState.declaredType()), requireAnyType("<aggregation input column type>"));
+        } else if (includeTimestampVector) {
+            return requireArgs(
+                requireType(aggState.declaredType()),
+                requireType(TypeName.INT),
+                requireType(TypeName.LONG), // @timestamp
+                requireAnyType("<aggregation input column type>")
+            );
+        } else {
+            return requireArgs(
+                requireType(aggState.declaredType()),
+                requireType(TypeName.INT),
+                requireAnyType("<aggregation input column type>")
+            );
+        }
+    }
+
     public ClassName implementation() {
         return implementation;
     }
@@ -568,22 +582,33 @@ private MethodSpec addIntermediateInput() {
                 });
                 builder.endControlFlow();
             } else {
-                // TODO combineIntermediate with optional block parameter
-                builder.addStatement("$T.combineIntermediate(state, groupId, " + intermediateStateRowAccess() + ")", declarationType);
+                var stateHasBlock = intermediateState.stream().anyMatch(AggregatorImplementer.IntermediateStateDesc::block);
+                requireStaticMethod(
+                    declarationType,
+                    requireVoidType(),
+                    requireName("combineIntermediate"),
+                    requireArgs(
+                        Stream.of(
+                            Stream.of(aggState.declaredType(), TypeName.INT), // aggState and groupId
+                            intermediateState.stream().map(AggregatorImplementer.IntermediateStateDesc::combineArgType),
+                            Stream.of(TypeName.INT).filter(p -> stateHasBlock) // position
+                        ).flatMap(Function.identity()).map(Methods::requireType).toArray(Methods.TypeMatcher[]::new)
+                    )
+                );
+
+                builder.addStatement(
+                    "$T.combineIntermediate(state, groupId, "
+                        + intermediateState.stream().map(desc -> desc.access("groupPosition + positionOffset")).collect(joining(", "))
+                        + (stateHasBlock ? ", groupPosition + positionOffset" : "")
+                        + ")",
+                    declarationType
+                );
             }
             builder.endControlFlow();
         }
         return builder.build();
     }
 
-    String intermediateStateRowAccess() {
-        String rowAccess = intermediateState.stream().map(desc -> desc.access("groupPosition + positionOffset")).collect(joining(", "));
-        if (intermediateState.stream().anyMatch(AggregatorImplementer.IntermediateStateDesc::block)) {
-            rowAccess += ", groupPosition + positionOffset";
-        }
-        return rowAccess;
-    }
-
     private MethodSpec addIntermediateRowInput() {
         MethodSpec.Builder builder = MethodSpec.methodBuilder("addIntermediateRowInput");
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);

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

@@ -9,21 +9,18 @@
 
 import com.squareup.javapoet.TypeName;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
-import javax.lang.model.element.Element;
 import javax.lang.model.element.ExecutableElement;
 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;
 
@@ -168,32 +165,6 @@ private static Stream<TypeElement> typeAndSuperType(TypeElement declarationType)
         }
     }
 
-    static ExecutableElement findRequiredMethod(TypeElement declarationType, String[] names, Predicate<ExecutableElement> filter) {
-        ExecutableElement result = findMethod(names, filter, declarationType, superClassOf(declarationType));
-        if (result == null) {
-            if (names.length == 1) {
-                throw new IllegalArgumentException(declarationType + "#" + names[0] + " is required");
-            }
-            throw new IllegalArgumentException("one of " + declarationType + "#" + Arrays.toString(names) + " is required");
-        }
-        return result;
-    }
-
-    static ExecutableElement findMethod(TypeElement declarationType, String name) {
-        return findMethod(new String[] { name }, e -> true, declarationType, superClassOf(declarationType));
-    }
-
-    private static TypeElement superClassOf(TypeElement declarationType) {
-        TypeMirror superclass = declarationType.getSuperclass();
-        if (superclass instanceof DeclaredType declaredType) {
-            Element superclassElement = declaredType.asElement();
-            if (superclassElement instanceof TypeElement) {
-                return (TypeElement) superclassElement;
-            }
-        }
-        return null;
-    }
-
     static ExecutableElement findMethod(TypeElement declarationType, String[] names, Predicate<ExecutableElement> filter) {
         return findMethod(names, filter, declarationType);
     }