ES|QL Add primitive float variants of all aggregators to the compute engine (#109781)

This commit adds primitive float variants of all ES|QL aggregators to the compute engine.

Much of the changes are mechanical and/or are generated, and retain similar type input/output params as other aggregators.

Co-authored-by: Iván Cea Fontenla <[email protected]>

Author: ChrisHegarty

docs/changelog/109781.yaml

@@ -0,0 +1,5 @@
+pr: 109781
+summary: ES|QL Add primitive float variants of all aggregators to the compute engine
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -410,6 +410,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  stateInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/LongState.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  stateInputFile
+    it.outputFile = "org/elasticsearch/compute/aggregation/FloatState.java"
+  }
   template {
     it.properties = doubleProperties
     it.inputFile =  stateInputFile
@@ -463,6 +468,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  arrayStateInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/DoubleArrayState.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  arrayStateInputFile
+    it.outputFile = "org/elasticsearch/compute/aggregation/FloatArrayState.java"
+  }
   File valuesAggregatorInputFile = new File("${projectDir}/src/main/java/org/elasticsearch/compute/aggregation/X-ValuesAggregator.java.st")
   template {
     it.properties = intProperties
@@ -474,6 +484,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  valuesAggregatorInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/ValuesLongAggregator.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  valuesAggregatorInputFile
+    it.outputFile = "org/elasticsearch/compute/aggregation/ValuesFloatAggregator.java"
+  }
   template {
     it.properties = doubleProperties
     it.inputFile =  valuesAggregatorInputFile
@@ -496,6 +511,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  rateAggregatorInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/RateLongAggregator.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  rateAggregatorInputFile
+    it.outputFile = "org/elasticsearch/compute/aggregation/RateFloatAggregator.java"
+  }
   template {
     it.properties = doubleProperties
     it.inputFile =  rateAggregatorInputFile
@@ -514,6 +534,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  topListAggregatorInputFile
     it.outputFile = "org/elasticsearch/compute/aggregation/TopListLongAggregator.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  topListAggregatorInputFile
+    it.outputFile = "org/elasticsearch/compute/aggregation/TopListFloatAggregator.java"
+  }
   template {
     it.properties = doubleProperties
     it.inputFile =  topListAggregatorInputFile
@@ -667,6 +692,11 @@ tasks.named('stringTemplates').configure {
     it.inputFile =  bucketedSortInputFile
     it.outputFile = "org/elasticsearch/compute/data/sort/LongBucketedSort.java"
   }
+  template {
+    it.properties = floatProperties
+    it.inputFile =  bucketedSortInputFile
+    it.outputFile = "org/elasticsearch/compute/data/sort/FloatBucketedSort.java"
+  }
   template {
     it.properties = doubleProperties
     it.inputFile =  bucketedSortInputFile

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

@@ -44,6 +44,8 @@
 import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR;
 import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
 import static org.elasticsearch.compute.gen.Types.ELEMENT_TYPE;
+import static org.elasticsearch.compute.gen.Types.FLOAT_BLOCK;
+import static org.elasticsearch.compute.gen.Types.FLOAT_VECTOR;
 import static org.elasticsearch.compute.gen.Types.INTERMEDIATE_STATE_DESC;
 import static org.elasticsearch.compute.gen.Types.INT_BLOCK;
 import static org.elasticsearch.compute.gen.Types.INT_VECTOR;
@@ -136,6 +138,8 @@ static String valueType(ExecutableElement init, ExecutableElement combine) {
         switch (initReturn) {
             case "double":
                 return "double";
+            case "float":
+                return "float";
             case "long":
                 return "long";
             case "int":
@@ -151,6 +155,7 @@ static ClassName valueBlockType(ExecutableElement init, ExecutableElement combin
         return switch (valueType(init, combine)) {
             case "boolean" -> BOOLEAN_BLOCK;
             case "double" -> DOUBLE_BLOCK;
+            case "float" -> FLOAT_BLOCK;
             case "long" -> LONG_BLOCK;
             case "int" -> INT_BLOCK;
             case "org.apache.lucene.util.BytesRef" -> BYTES_REF_BLOCK;
@@ -162,6 +167,7 @@ static ClassName valueVectorType(ExecutableElement init, ExecutableElement combi
         return switch (valueType(init, combine)) {
             case "boolean" -> BOOLEAN_VECTOR;
             case "double" -> DOUBLE_VECTOR;
+            case "float" -> FLOAT_VECTOR;
             case "long" -> LONG_VECTOR;
             case "int" -> INT_VECTOR;
             case "org.apache.lucene.util.BytesRef" -> BYTES_REF_VECTOR;
@@ -445,6 +451,8 @@ private String primitiveStateMethod() {
                 return "longValue";
             case "org.elasticsearch.compute.aggregation.DoubleState":
                 return "doubleValue";
+            case "org.elasticsearch.compute.aggregation.FloatState":
+                return "floatValue";
             default:
                 throw new IllegalArgumentException(
                     "don't know how to fetch primitive values from " + stateType + ". define combineIntermediate."
@@ -495,6 +503,9 @@ private void primitiveStateToResult(MethodSpec.Builder builder) {
             case "org.elasticsearch.compute.aggregation.DoubleState":
                 builder.addStatement("blocks[offset] = driverContext.blockFactory().newConstantDoubleBlockWith(state.doubleValue(), 1)");
                 return;
+            case "org.elasticsearch.compute.aggregation.FloatState":
+                builder.addStatement("blocks[offset] = driverContext.blockFactory().newConstantFloatBlockWith(state.floatValue(), 1)");
+                return;
             default:
                 throw new IllegalArgumentException("don't know how to convert state to result: " + stateType);
         }
@@ -521,7 +532,7 @@ private MethodSpec close() {
     private boolean hasPrimitiveState() {
         return switch (stateType.toString()) {
             case "org.elasticsearch.compute.aggregation.IntState", "org.elasticsearch.compute.aggregation.LongState",
-                "org.elasticsearch.compute.aggregation.DoubleState" -> true;
+                "org.elasticsearch.compute.aggregation.DoubleState", "org.elasticsearch.compute.aggregation.FloatState" -> true;
             default -> false;
         };
     }

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

@@ -585,7 +585,7 @@ private MethodSpec close() {
     private boolean hasPrimitiveState() {
         return switch (stateType.toString()) {
             case "org.elasticsearch.compute.aggregation.IntArrayState", "org.elasticsearch.compute.aggregation.LongArrayState",
-                "org.elasticsearch.compute.aggregation.DoubleArrayState" -> true;
+                "org.elasticsearch.compute.aggregation.DoubleArrayState", "org.elasticsearch.compute.aggregation.FloatArrayState" -> true;
             default -> false;
         };
     }

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

@@ -211,6 +211,7 @@ static String vectorAccessorName(String elementTypeName) {
             case "INT" -> "getInt";
             case "LONG" -> "getLong";
             case "DOUBLE" -> "getDouble";
+            case "FLOAT" -> "getFloat";
             case "BYTES_REF" -> "getBytesRef";
             default -> throw new IllegalArgumentException(
                 "don't know how to fetch primitive values from " + elementTypeName + ". define combineIntermediate."

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

@@ -41,12 +41,14 @@ public class Types {
     static final ClassName INT_BLOCK = ClassName.get(DATA_PACKAGE, "IntBlock");
     static final ClassName LONG_BLOCK = ClassName.get(DATA_PACKAGE, "LongBlock");
     static final ClassName DOUBLE_BLOCK = ClassName.get(DATA_PACKAGE, "DoubleBlock");
+    static final ClassName FLOAT_BLOCK = ClassName.get(DATA_PACKAGE, "FloatBlock");
 
     static final ClassName BOOLEAN_BLOCK_BUILDER = BOOLEAN_BLOCK.nestedClass("Builder");
     static final ClassName BYTES_REF_BLOCK_BUILDER = BYTES_REF_BLOCK.nestedClass("Builder");
     static final ClassName INT_BLOCK_BUILDER = INT_BLOCK.nestedClass("Builder");
     static final ClassName LONG_BLOCK_BUILDER = LONG_BLOCK.nestedClass("Builder");
     static final ClassName DOUBLE_BLOCK_BUILDER = DOUBLE_BLOCK.nestedClass("Builder");
+    static final ClassName FLOAT_BLOCK_BUILDER = FLOAT_BLOCK.nestedClass("Builder");
 
     static final ClassName ELEMENT_TYPE = ClassName.get(DATA_PACKAGE, "ElementType");
 
@@ -55,35 +57,41 @@ public class Types {
     static final ClassName INT_VECTOR = ClassName.get(DATA_PACKAGE, "IntVector");
     static final ClassName LONG_VECTOR = ClassName.get(DATA_PACKAGE, "LongVector");
     static final ClassName DOUBLE_VECTOR = ClassName.get(DATA_PACKAGE, "DoubleVector");
+    static final ClassName FLOAT_VECTOR = ClassName.get(DATA_PACKAGE, "FloatVector");
 
     static final ClassName BOOLEAN_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "BooleanVector", "Builder");
     static final ClassName BYTES_REF_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "BytesRefVector", "Builder");
     static final ClassName INT_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "IntVector", "Builder");
     static final ClassName LONG_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "LongVector", "Builder");
     static final ClassName DOUBLE_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "DoubleVector", "Builder");
+    static final ClassName FLOAT_VECTOR_BUILDER = ClassName.get(DATA_PACKAGE, "FloatVector", "Builder");
 
     static final ClassName BOOLEAN_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "BooleanVector", "FixedBuilder");
     static final ClassName INT_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "IntVector", "FixedBuilder");
     static final ClassName LONG_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "LongVector", "FixedBuilder");
     static final ClassName DOUBLE_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "DoubleVector", "FixedBuilder");
+    static final ClassName FLOAT_VECTOR_FIXED_BUILDER = ClassName.get(DATA_PACKAGE, "FloatVector", "FixedBuilder");
 
     static final ClassName BOOLEAN_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "BooleanArrayVector");
     static final ClassName BYTES_REF_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "BytesRefArrayVector");
     static final ClassName INT_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "IntArrayVector");
     static final ClassName LONG_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "LongArrayVector");
     static final ClassName DOUBLE_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "DoubleArrayVector");
+    static final ClassName FLOAT_ARRAY_VECTOR = ClassName.get(DATA_PACKAGE, "FloatArrayVector");
 
     static final ClassName BOOLEAN_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "BooleanArrayBlock");
     static final ClassName BYTES_REF_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "BytesRefArrayBlock");
     static final ClassName INT_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "IntArrayBlock");
     static final ClassName LONG_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "LongArrayBlock");
     static final ClassName DOUBLE_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "DoubleArrayBlock");
+    static final ClassName FLOAT_ARRAY_BLOCK = ClassName.get(DATA_PACKAGE, "FloatArrayBlock");
 
     static final ClassName BOOLEAN_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantBooleanVector");
     static final ClassName BYTES_REF_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantBytesRefVector");
     static final ClassName INT_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantIntVector");
     static final ClassName LONG_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantLongVector");
     static final ClassName DOUBLE_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantDoubleVector");
+    static final ClassName FLOAT_CONSTANT_VECTOR = ClassName.get(DATA_PACKAGE, "ConstantFloatVector");
 
     static final ClassName AGGREGATOR_FUNCTION = ClassName.get(AGGREGATION_PACKAGE, "AggregatorFunction");
     static final ClassName AGGREGATOR_FUNCTION_SUPPLIER = ClassName.get(AGGREGATION_PACKAGE, "AggregatorFunctionSupplier");
@@ -162,6 +170,9 @@ static ClassName blockType(String elementType) {
         if (elementType.equalsIgnoreCase(TypeName.DOUBLE.toString())) {
             return DOUBLE_BLOCK;
         }
+        if (elementType.equalsIgnoreCase(TypeName.FLOAT.toString())) {
+            return FLOAT_BLOCK;
+        }
         throw new IllegalArgumentException("unknown vector type for [" + elementType + "]");
     }
 
@@ -181,6 +192,9 @@ static ClassName vectorType(TypeName elementType) {
         if (elementType.equals(TypeName.DOUBLE)) {
             return DOUBLE_VECTOR;
         }
+        if (elementType.equals(TypeName.FLOAT)) {
+            return FLOAT_VECTOR;
+        }
         throw new IllegalArgumentException("unknown vector type for [" + elementType + "]");
     }
 
@@ -200,6 +214,9 @@ static ClassName vectorType(String elementType) {
         if (elementType.equalsIgnoreCase(TypeName.DOUBLE.toString())) {
             return DOUBLE_VECTOR;
         }
+        if (elementType.equalsIgnoreCase(TypeName.FLOAT.toString())) {
+            return FLOAT_VECTOR;
+        }
         throw new IllegalArgumentException("unknown vector type for [" + elementType + "]");
     }
 
@@ -234,6 +251,12 @@ static ClassName builderType(TypeName resultType) {
         if (resultType.equals(DOUBLE_VECTOR)) {
             return DOUBLE_VECTOR_BUILDER;
         }
+        if (resultType.equals(FLOAT_BLOCK)) {
+            return FLOAT_BLOCK_BUILDER;
+        }
+        if (resultType.equals(FLOAT_VECTOR)) {
+            return FLOAT_VECTOR_BUILDER;
+        }
         throw new IllegalArgumentException("unknown builder type for [" + resultType + "]");
     }
 
@@ -250,6 +273,9 @@ static ClassName vectorFixedBuilderType(TypeName elementType) {
         if (elementType.equals(TypeName.DOUBLE)) {
             return DOUBLE_VECTOR_FIXED_BUILDER;
         }
+        if (elementType.equals(TypeName.FLOAT)) {
+            return FLOAT_VECTOR_FIXED_BUILDER;
+        }
         throw new IllegalArgumentException("unknown vector fixed builder type for [" + elementType + "]");
     }