ESQL: Remove AggregateMapper reflection, and delegate intermediate state to suppliers

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/AggregatorBenchmark.java

@@ -155,7 +155,7 @@ private static Operator operator(DriverContext driverContext, String grouping, S
 
         if (grouping.equals("none")) {
             return new AggregationOperator(
-                List.of(supplier(op, dataType, filter, 0).aggregatorFactory(AggregatorMode.SINGLE).apply(driverContext)),
+                List.of(supplier(op, dataType, filter).aggregatorFactory(AggregatorMode.SINGLE, List.of(0)).apply(driverContext)),
                 driverContext
             );
         }
@@ -182,33 +182,33 @@ private static Operator operator(DriverContext driverContext, String grouping, S
             default -> throw new IllegalArgumentException("unsupported grouping [" + grouping + "]");
         };
         return new HashAggregationOperator(
-            List.of(supplier(op, dataType, filter, groups.size()).groupingAggregatorFactory(AggregatorMode.SINGLE)),
+            List.of(supplier(op, dataType, filter).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(groups.size()))),
             () -> BlockHash.build(groups, driverContext.blockFactory(), 16 * 1024, false),
             driverContext
         );
     }
 
-    private static AggregatorFunctionSupplier supplier(String op, String dataType, String filter, int dataChannel) {
+    private static AggregatorFunctionSupplier supplier(String op, String dataType, String filter) {
         return filtered(switch (op) {
-            case COUNT -> CountAggregatorFunction.supplier(List.of(dataChannel));
+            case COUNT -> CountAggregatorFunction.supplier();
             case COUNT_DISTINCT -> switch (dataType) {
-                case LONGS -> new CountDistinctLongAggregatorFunctionSupplier(List.of(dataChannel), 3000);
-                case DOUBLES -> new CountDistinctDoubleAggregatorFunctionSupplier(List.of(dataChannel), 3000);
+                case LONGS -> new CountDistinctLongAggregatorFunctionSupplier(3000);
+                case DOUBLES -> new CountDistinctDoubleAggregatorFunctionSupplier(3000);
                 default -> throw new IllegalArgumentException("unsupported data type [" + dataType + "]");
             };
             case MAX -> switch (dataType) {
-                case LONGS -> new MaxLongAggregatorFunctionSupplier(List.of(dataChannel));
-                case DOUBLES -> new MaxDoubleAggregatorFunctionSupplier(List.of(dataChannel));
+                case LONGS -> new MaxLongAggregatorFunctionSupplier();
+                case DOUBLES -> new MaxDoubleAggregatorFunctionSupplier();
                 default -> throw new IllegalArgumentException("unsupported data type [" + dataType + "]");
             };
             case MIN -> switch (dataType) {
-                case LONGS -> new MinLongAggregatorFunctionSupplier(List.of(dataChannel));
-                case DOUBLES -> new MinDoubleAggregatorFunctionSupplier(List.of(dataChannel));
+                case LONGS -> new MinLongAggregatorFunctionSupplier();
+                case DOUBLES -> new MinDoubleAggregatorFunctionSupplier();
                 default -> throw new IllegalArgumentException("unsupported data type [" + dataType + "]");
             };
             case SUM -> switch (dataType) {
-                case LONGS -> new SumLongAggregatorFunctionSupplier(List.of(dataChannel));
-                case DOUBLES -> new SumDoubleAggregatorFunctionSupplier(List.of(dataChannel));
+                case LONGS -> new SumLongAggregatorFunctionSupplier();
+                case DOUBLES -> new SumDoubleAggregatorFunctionSupplier();
                 default -> throw new IllegalArgumentException("unsupported data type [" + dataType + "]");
             };
             default -> throw new IllegalArgumentException("unsupported op [" + op + "]");

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

@@ -31,6 +31,7 @@
 
 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;
 import static org.elasticsearch.compute.gen.Types.LIST_INTEGER;
 import static org.elasticsearch.compute.gen.Types.STRING;
 import static org.elasticsearch.compute.gen.Types.WARNINGS;
@@ -67,7 +68,6 @@ public AggregatorFunctionSupplierImplementer(
             createParameters.addAll(groupingAggregatorImplementer.createParameters());
         }
         this.createParameters = new ArrayList<>(createParameters);
-        this.createParameters.add(0, new Parameter(LIST_INTEGER, "channels"));
 
         this.implementation = ClassName.get(
             elements.getPackageOf(declarationType).toString(),
@@ -99,11 +99,9 @@ private TypeSpec type() {
         }
         createParameters.stream().forEach(p -> p.declareField(builder));
         builder.addMethod(ctor());
-        if (aggregatorImplementer != null) {
-            builder.addMethod(aggregator());
-        } else {
-            builder.addMethod(unsupportedNonGroupingAggregator());
-        }
+        builder.addMethod(nonGroupingIntermediateStateDesc());
+        builder.addMethod(groupingIntermediateStateDesc());
+        builder.addMethod(aggregator());
         builder.addMethod(groupingAggregator());
         builder.addMethod(describe());
         return builder.build();
@@ -123,19 +121,43 @@ private MethodSpec ctor() {
         return builder.build();
     }
 
-    private MethodSpec unsupportedNonGroupingAggregator() {
-        MethodSpec.Builder builder = MethodSpec.methodBuilder("aggregator")
-            .addParameter(DRIVER_CONTEXT, "driverContext")
-            .returns(Types.AGGREGATOR_FUNCTION);
+    private MethodSpec nonGroupingIntermediateStateDesc() {
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("nonGroupingIntermediateStateDesc");
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
-        builder.addStatement("throw new UnsupportedOperationException($S)", "non-grouping aggregator is not supported");
+        builder.returns(LIST_AGG_FUNC_DESC);
+
+        if (aggregatorImplementer == null) {
+            builder.addStatement("throw new UnsupportedOperationException($S)", "non-grouping aggregator is not supported");
+            return builder.build();
+        }
+
+        builder.addStatement("return $T.intermediateStateDesc()", aggregatorImplementer.implementation());
+
+        return builder.build();
+    }
+
+    private MethodSpec groupingIntermediateStateDesc() {
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("groupingIntermediateStateDesc");
+        builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
+        builder.returns(LIST_AGG_FUNC_DESC);
+
+        builder.addStatement("return $T.intermediateStateDesc()", groupingAggregatorImplementer.implementation());
+
         return builder.build();
     }
 
     private MethodSpec aggregator() {
         MethodSpec.Builder builder = MethodSpec.methodBuilder("aggregator");
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
         builder.addParameter(DRIVER_CONTEXT, "driverContext");
+        builder.addParameter(LIST_INTEGER, "channels");
+
+        if (aggregatorImplementer == null) {
+            builder.returns(Types.AGGREGATOR_FUNCTION);
+            builder.addStatement("throw new UnsupportedOperationException($S)", "non-grouping aggregator is not supported");
+            return builder.build();
+        }
+
         builder.returns(aggregatorImplementer.implementation());
 
         if (hasWarnings) {
@@ -162,6 +184,7 @@ private MethodSpec groupingAggregator() {
         MethodSpec.Builder builder = MethodSpec.methodBuilder("groupingAggregator");
         builder.addAnnotation(Override.class).addModifiers(Modifier.PUBLIC);
         builder.addParameter(DRIVER_CONTEXT, "driverContext");
+        builder.addParameter(LIST_INTEGER, "channels");
         builder.returns(groupingAggregatorImplementer.implementation());
 
         if (hasWarnings) {