Remove reflection from AggregateMapper

Author: ivancea

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/AggregateMapper.java

@@ -10,7 +10,6 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.compute.aggregation.IntermediateStateDesc;
 import org.elasticsearch.compute.data.ElementType;
-import org.elasticsearch.core.Tuple;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -30,10 +29,8 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Max;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.MedianAbsoluteDeviation;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Min;
-import org.elasticsearch.xpack.esql.expression.function.aggregate.NumericAggregate;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Percentile;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate;
-import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialAggregateFunction;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialCentroid;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.SpatialExtent;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.StdDev;
@@ -42,14 +39,8 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Top;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Values;
 
-import java.lang.invoke.MethodHandle;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.MethodType;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -67,9 +58,6 @@
  */
 final class AggregateMapper {
 
-    private static final List<String> NUMERIC = List.of("Int", "Long", "Double");
-    private static final List<String> SPATIAL_EXTRA_CONFIGS = List.of("SourceValues", "DocValues");
-
     /** List of all mappable ESQL agg functions (excludes surrogates like AVG = SUM/COUNT). */
     private static final List<? extends Class<? extends Function>> AGG_FUNCTIONS = List.of(
         Count.class,
@@ -91,18 +79,6 @@ final class AggregateMapper {
         ToPartial.class
     );
 
-    /** Record of agg Class, type, and grouping (or non-grouping). */
-    private record AggDef(Class<?> aggClazz, String type, String extra, boolean grouping) {
-        public AggDef withoutExtra() {
-            return new AggDef(aggClazz, type, "", grouping);
-        }
-    }
-
-    /** Map of AggDef types to intermediate named expressions. */
-    private static final Map<AggDef, List<IntermediateStateDesc>> MAPPER = AGG_FUNCTIONS.stream()
-        .flatMap(AggregateMapper::aggDefs)
-        .collect(Collectors.toUnmodifiableMap(aggDef -> aggDef, AggregateMapper::lookupIntermediateState));
-
     /** Cache of aggregates to intermediate expressions. */
     private final HashMap<Expression, List<NamedExpression>> cache = new HashMap<>();
 
@@ -144,127 +120,21 @@ private static List<NamedExpression> computeEntryForAgg(String aggAlias, Express
     }
 
     private static List<NamedExpression> entryForAgg(String aggAlias, AggregateFunction aggregateFunction, boolean grouping) {
-        var aggDef = new AggDef(
-            aggregateFunction.getClass(),
-            dataTypeToString(aggregateFunction.field().dataType(), aggregateFunction.getClass()),
-            aggregateFunction instanceof SpatialAggregateFunction ? "SourceValues" : "",
-            grouping
-        );
-        var is = getNonNull(aggDef);
-        return isToNE(is, aggAlias).toList();
-    }
-
-    /** Gets the agg from the mapper - wrapper around map::get for more informative failure.*/
-    private static List<IntermediateStateDesc> getNonNull(AggDef aggDef) {
-        var l = MAPPER.getOrDefault(aggDef, MAPPER.get(aggDef.withoutExtra()));
-        if (l == null) {
-            throw new EsqlIllegalArgumentException("Cannot find intermediate state for: " + aggDef);
-        }
-        return l;
-    }
-
-    private static Stream<AggDef> aggDefs(Class<?> clazz) {
-        List<String> types;
-        List<String> extraConfigs = List.of("");
-        if (NumericAggregate.class.isAssignableFrom(clazz)) {
-            types = NUMERIC;
-        } else if (Max.class.isAssignableFrom(clazz) || Min.class.isAssignableFrom(clazz)) {
-            types = List.of("Boolean", "Int", "Long", "Double", "Ip", "BytesRef");
-        } else if (clazz == Count.class) {
-            types = List.of(""); // no extra type distinction
-        } else if (clazz == SpatialCentroid.class) {
-            types = List.of("GeoPoint", "CartesianPoint");
-            extraConfigs = SPATIAL_EXTRA_CONFIGS;
-        } else if (clazz == SpatialExtent.class) {
-            types = List.of("GeoPoint", "CartesianPoint", "GeoShape", "CartesianShape");
-            extraConfigs = SPATIAL_EXTRA_CONFIGS;
-        } else if (Values.class.isAssignableFrom(clazz)) {
-            // TODO can't we figure this out from the function itself?
-            types = List.of("Int", "Long", "Double", "Boolean", "BytesRef");
-        } else if (Top.class.isAssignableFrom(clazz)) {
-            types = List.of("Boolean", "Int", "Long", "Double", "Ip", "BytesRef");
-        } else if (Rate.class.isAssignableFrom(clazz) || StdDev.class.isAssignableFrom(clazz)) {
-            types = List.of("Int", "Long", "Double");
-        } else if (FromPartial.class.isAssignableFrom(clazz) || ToPartial.class.isAssignableFrom(clazz)) {
-            types = List.of(""); // no type
-        } else if (CountDistinct.class.isAssignableFrom(clazz)) {
-            types = Stream.concat(NUMERIC.stream(), Stream.of("Boolean", "BytesRef")).toList();
+        List<IntermediateStateDesc> intermediateState;
+        if (aggregateFunction instanceof ToAggregator toAggregator) {
+            var supplier = toAggregator.supplier();
+            intermediateState = grouping ? supplier.groupingIntermediateStateDesc() : supplier.nonGroupingIntermediateStateDesc();
         } else {
-            assert false : "unknown aggregate type " + clazz;
-            throw new IllegalArgumentException("unknown aggregate type " + clazz);
-        }
-        return combinations(types, extraConfigs).flatMap(typeAndExtraConfig -> {
-            var type = typeAndExtraConfig.v1();
-            var extra = typeAndExtraConfig.v2();
-
-            if (clazz.isAssignableFrom(Rate.class)) {
-                // rate doesn't support non-grouping aggregations
-                return Stream.of(new AggDef(clazz, type, extra, true));
-            } else if (Objects.equals(type, "AggregateMetricDouble")) {
-                // TODO: support grouping aggregations for aggregate metric double
-                return Stream.of(new AggDef(clazz, type, extra, false));
-            } else {
-                return Stream.of(new AggDef(clazz, type, extra, true), new AggDef(clazz, type, extra, false));
-            }
-        });
-    }
-
-    private static Stream<Tuple<String, String>> combinations(List<String> types, List<String> extraConfigs) {
-        return types.stream().flatMap(type -> extraConfigs.stream().map(config -> new Tuple<>(type, config)));
-    }
-
-    /** Retrieves the intermediate state description for a given class, type, and grouping. */
-    private static List<IntermediateStateDesc> lookupIntermediateState(AggDef aggDef) {
-        try {
-            return (List<IntermediateStateDesc>) lookup(aggDef.aggClazz(), aggDef.type(), aggDef.extra(), aggDef.grouping()).invokeExact();
-        } catch (Throwable t) {
-            // invokeExact forces us to handle any Throwable thrown by lookup.
-            throw new EsqlIllegalArgumentException(t);
-        }
-    }
-
-    /** Looks up the intermediate state method for a given class, type, and grouping. */
-    private static MethodHandle lookup(Class<?> clazz, String type, String extra, boolean grouping) {
-        try {
-            return lookupRetry(clazz, type, extra, grouping);
-        } catch (IllegalAccessException | NoSuchMethodException | ClassNotFoundException e) {
-            throw new EsqlIllegalArgumentException(e);
+            throw new EsqlIllegalArgumentException("Aggregate has no defined intermediate state: " + aggregateFunction);
         }
-    }
-
-    private static MethodHandle lookupRetry(Class<?> clazz, String type, String extra, boolean grouping) throws IllegalAccessException,
-        NoSuchMethodException, ClassNotFoundException {
-        try {
-            return MethodHandles.lookup()
-                .findStatic(
-                    Class.forName(determineAggName(clazz, type, extra, grouping)),
-                    "intermediateStateDesc",
-                    MethodType.methodType(List.class)
-                );
-        } catch (NoSuchMethodException ignore) {
-            // Retry without the extra information.
-            return MethodHandles.lookup()
-                .findStatic(
-                    Class.forName(determineAggName(clazz, type, "", grouping)),
-                    "intermediateStateDesc",
-                    MethodType.methodType(List.class)
-                );
-        }
-    }
-
-    /** Determines the engines agg class name, for the given class, type, and grouping. */
-    private static String determineAggName(Class<?> clazz, String type, String extra, boolean grouping) {
-        return "org.elasticsearch.compute.aggregation."
-            + (clazz.getSimpleName().startsWith("Spatial") ? "spatial." : "")
-            + clazz.getSimpleName()
-            + type
-            + extra
-            + (grouping ? "Grouping" : "")
-            + "AggregatorFunction";
+        return intermediateStateToNamedExpressions(intermediateState, aggAlias).toList();
     }
 
     /** Maps intermediate state description to named expressions.  */
-    private static Stream<NamedExpression> isToNE(List<IntermediateStateDesc> intermediateStateDescs, String aggAlias) {
+    private static Stream<NamedExpression> intermediateStateToNamedExpressions(
+        List<IntermediateStateDesc> intermediateStateDescs,
+        String aggAlias
+    ) {
         return intermediateStateDescs.stream().map(is -> {
             final DataType dataType;
             if (Strings.isEmpty(is.dataType())) {
@@ -288,34 +158,4 @@ private static DataType toDataType(ElementType elementType) {
             case FLOAT, NULL, DOC, COMPOSITE, UNKNOWN -> throw new EsqlIllegalArgumentException("unsupported agg type: " + elementType);
         };
     }
-
-    /** Returns the string representation for the data type. This reflects the engine's aggs naming structure. */
-    private static String dataTypeToString(DataType type, Class<?> aggClass) {
-        if (aggClass == Count.class) {
-            return "";  // no type distinction
-        }
-        if (aggClass == ToPartial.class || aggClass == FromPartial.class) {
-            return "";
-        }
-        if ((aggClass == Max.class || aggClass == Min.class || aggClass == Top.class) && type.equals(DataType.IP)) {
-            return "Ip";
-        }
-
-        return switch (type) {
-            case BOOLEAN -> "Boolean";
-            case INTEGER, COUNTER_INTEGER -> "Int";
-            case LONG, DATETIME, COUNTER_LONG, DATE_NANOS -> "Long";
-            case DOUBLE, COUNTER_DOUBLE -> "Double";
-            case KEYWORD, IP, VERSION, TEXT, SEMANTIC_TEXT -> "BytesRef";
-            case GEO_POINT -> "GeoPoint";
-            case CARTESIAN_POINT -> "CartesianPoint";
-            case GEO_SHAPE -> "GeoShape";
-            case CARTESIAN_SHAPE -> "CartesianShape";
-            case AGGREGATE_METRIC_DOUBLE -> "AggregateMetricDouble";
-            case UNSUPPORTED, NULL, UNSIGNED_LONG, SHORT, BYTE, FLOAT, HALF_FLOAT, SCALED_FLOAT, OBJECT, SOURCE, DATE_PERIOD, TIME_DURATION,
-                DOC_DATA_TYPE, TSID_DATA_TYPE, PARTIAL_AGG -> throw new EsqlIllegalArgumentException(
-                    "illegal agg type: " + type.typeName()
-                );
-        };
-    }
 }