ESQL: Ungrouped agg implementations for FIRST/LAST

x-pack/plugin/esql/compute/build.gradle

@@ -86,6 +86,7 @@ def addOccurrence(props, Occurrence) {
   newProps["Occurrence"] = Occurrence
   newProps["First"] = Occurrence == "First" ? "true" : ""
   newProps["Last"] = Occurrence == "Last" ? "true" : ""
+  newProps["occurrence"] = Occurrence.toLowerCase(Locale.ROOT)
   return newProps
 }
 
@@ -469,6 +470,27 @@ tasks.named('stringTemplates').configure {
     it.inputFile = stateInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/DoubleState.java"
   }
+
+  /*
+   * Generates pairwise states. We generate the ones that we need at the moment,
+   * but add more if you need more.
+   */
+  File twoStateInputFile = file("src/main/java/org/elasticsearch/compute/aggregation/X-2State.java.st")
+  [longProperties].forEach { v1 ->
+    [intProperties, longProperties, floatProperties, doubleProperties].forEach { v2 ->
+      {
+        var properties = [:]
+        v1.forEach { k, v -> properties["v1_" + k] = v}
+        v2.forEach { k, v -> properties["v2_" + k] = v}
+        template {
+          it.properties = properties
+          it.inputFile = twoStateInputFile
+          it.outputFile = "org/elasticsearch/compute/aggregation/${v1.Type}${v2.Type}State.java"
+        }
+      }
+    }
+  }
+
   File fallibleStateInputFile = new File("${projectDir}/src/main/java/org/elasticsearch/compute/aggregation/X-FallibleState.java.st")
   template {
     it.properties = booleanProperties

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

@@ -25,10 +25,15 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.util.Elements;
 
+import static org.elasticsearch.compute.gen.Methods.optionalStaticMethod;
+import static org.elasticsearch.compute.gen.Methods.requireArgs;
+import static org.elasticsearch.compute.gen.Methods.requireName;
+import static org.elasticsearch.compute.gen.Methods.requireType;
 import static org.elasticsearch.compute.gen.Types.AGGREGATOR_FUNCTION_SUPPLIER;
 import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
 import static org.elasticsearch.compute.gen.Types.LIST_AGG_FUNC_DESC;
@@ -210,17 +215,24 @@ private MethodSpec describe() {
         MethodSpec.Builder builder = MethodSpec.methodBuilder("describe").returns(String.class);
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
 
-        String name = declarationType.getSimpleName().toString();
-        name = name.replace("BytesRef", "Byte"); // The hack expects one word types so let's make BytesRef into Byte
-        String[] parts = name.split("(?=\\p{Upper})");
-        if (false == parts[parts.length - 1].equals("Aggregator") || parts.length < 3) {
-            throw new IllegalArgumentException("Can't generate description for " + declarationType.getSimpleName());
+        ExecutableElement describe = optionalStaticMethod(declarationType, requireType(STRING), requireName("describe"), requireArgs());
+        if (describe == null) {
+            String name = declarationType.getSimpleName().toString();
+            name = name.replace("BytesRef", "Byte"); // The hack expects one word types so let's make BytesRef into Byte
+            String[] parts = name.split("(?=\\p{Upper})");
+            if (false == parts[parts.length - 1].equals("Aggregator") || parts.length < 3) {
+                throw new IllegalArgumentException("Can't generate description for " + declarationType.getSimpleName());
+            }
+
+            String operation = Arrays.stream(parts, 0, parts.length - 2)
+                .map(s -> s.toLowerCase(Locale.ROOT))
+                .collect(Collectors.joining("_"));
+            String type = parts[parts.length - 2];
+
+            builder.addStatement("return $S", operation + " of " + type.toLowerCase(Locale.ROOT) + "s");
+        } else {
+            builder.addStatement("return $T.$L()", declarationType, "describe");
         }
-
-        String operation = Arrays.stream(parts, 0, parts.length - 2).map(s -> s.toLowerCase(Locale.ROOT)).collect(Collectors.joining("_"));
-        String type = parts[parts.length - 2];
-
-        builder.addStatement("return $S", operation + " of " + type.toLowerCase(Locale.ROOT) + "s");
         return builder.build();
     }
 }

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

@@ -198,6 +198,7 @@ private TypeSpec type() {
             builder.addMethod(addRawInputLoop(groupIdClass, true));
             builder.addMethod(addIntermediateInput(groupIdClass));
         }
+        builder.addMethod(maybeEnableGroupIdTracking());
         builder.addMethod(selectedMayContainUnseenGroups());
         builder.addMethod(evaluateIntermediate());
         builder.addMethod(evaluateFinal());
@@ -321,9 +322,11 @@ private MethodSpec prepareProcessRawInputPage() {
             builder.addStatement("$T $L = $L.asVector()", vectorType(p.type()), p.vectorName(), p.blockName());
             builder.beginControlFlow("if ($L == null)", p.vectorName());
             {
-                builder.beginControlFlow("if ($L.mayHaveNulls())", p.blockName());
-                builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
-                builder.endControlFlow();
+                builder.addStatement(
+                    "maybeEnableGroupIdTracking(seenGroupIds, "
+                        + aggParams.stream().map(AggregationParameter::blockName).collect(joining(", "))
+                        + ")"
+                );
                 returnAddInput(builder, false);
             }
             builder.endControlFlow();
@@ -351,6 +354,23 @@ private void returnAddInput(MethodSpec.Builder builder, boolean valuesAreVector)
         }
     }
 
+    private MethodSpec maybeEnableGroupIdTracking() {
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("maybeEnableGroupIdTracking");
+        builder.addModifiers(Modifier.PRIVATE).returns(TypeName.VOID);
+        builder.addParameter(SEEN_GROUP_IDS, "seenGroupIds");
+        for (AggregationParameter p : aggParams) {
+            builder.addParameter(blockType(p.type()), p.blockName());
+        }
+
+        for (AggregationParameter p : aggParams) {
+            builder.beginControlFlow("if ($L.mayHaveNulls())", p.blockName());
+            builder.addStatement("state.enableGroupIdTracking(seenGroupIds)");
+            builder.endControlFlow();
+        }
+
+        return builder.build();
+    }
+
     /**
      * Generate an {@code AddInput} implementation. That's a collection path optimized for the input data.
      */