ESQL: Move most evaluators to BlockFactory (elastic#101491)

This replaces most of the calls to build a `Block.Builder` on `ExpressionEvaluator` subclasses with calls to methods on the `BlockFactory`. Mostly it changes the methods on the generated code.
not-napoleon · Oct 31, 2023 · 7609780 · 7609780
1 parent 1722c0a
commit 7609780
Show file tree

Hide file tree

Showing 194 changed files with 928 additions and 752 deletions.
diff --git a/.../compute/gen/src/main/java/org/elasticsearch/compute/gen/ConvertEvaluatorImplementer.java b/.../compute/gen/src/main/java/org/elasticsearch/compute/gen/ConvertEvaluatorImplementer.java
@@ -23,6 +23,7 @@
 import javax.lang.model.util.Elements;
 
 import static org.elasticsearch.compute.gen.Methods.appendMethod;
+import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
 import static org.elasticsearch.compute.gen.Methods.getMethod;
 import static org.elasticsearch.compute.gen.Types.ABSTRACT_CONVERT_FUNCTION_EVALUATOR;
 import static org.elasticsearch.compute.gen.Types.BLOCK;
@@ -33,6 +34,7 @@
 import static org.elasticsearch.compute.gen.Types.SOURCE;
 import static org.elasticsearch.compute.gen.Types.VECTOR;
 import static org.elasticsearch.compute.gen.Types.blockType;
+import static org.elasticsearch.compute.gen.Types.builderType;
 import static org.elasticsearch.compute.gen.Types.vectorType;
 
 public class ConvertEvaluatorImplementer {
@@ -131,35 +133,37 @@ private MethodSpec evalVector() {
             builder.nextControlFlow("catch (Exception e)");
             {
                 builder.addStatement("registerException(e)");
-                builder.addStatement("return Block.constantNullBlock(positionCount, driverContext.blockFactory())");
+                builder.addStatement("return driverContext.blockFactory().newConstantNullBlock(positionCount)");
             }
             builder.endControlFlow();
         }
         builder.endControlFlow();
 
-        ClassName returnBlockType = blockType(resultType);
-        builder.addStatement(
-            "$T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory())",
-            returnBlockType,
-            returnBlockType
+        ClassName resultBuilderType = builderType(blockType(resultType));
+        builder.beginControlFlow(
+            "try ($T builder = driverContext.blockFactory().$L(positionCount))",
+            resultBuilderType,
+            buildFromFactory(resultBuilderType)
         );
-        builder.beginControlFlow("for (int p = 0; p < positionCount; p++)");
         {
-            builder.beginControlFlow("try");
+            builder.beginControlFlow("for (int p = 0; p < positionCount; p++)");
             {
-                builder.addStatement("builder.$L($N)", appendMethod(resultType), evalValueCall("vector", "p", scratchPadName));
-            }
-            builder.nextControlFlow("catch (Exception e)");
-            {
-                builder.addStatement("registerException(e)");
-                builder.addStatement("builder.appendNull()");
+                builder.beginControlFlow("try");
+                {
+                    builder.addStatement("builder.$L($N)", appendMethod(resultType), evalValueCall("vector", "p", scratchPadName));
+                }
+                builder.nextControlFlow("catch (Exception e)");
+                {
+                    builder.addStatement("registerException(e)");
+                    builder.addStatement("builder.appendNull()");
+                }
+                builder.endControlFlow();
             }
             builder.endControlFlow();
+            builder.addStatement("return builder.build()");
         }
         builder.endControlFlow();
 
-        builder.addStatement("return builder.build()");
-
         return builder.build();
     }
 
@@ -170,11 +174,11 @@ private MethodSpec evalBlock() {
         TypeName blockType = blockType(argumentType);
         builder.addStatement("$T block = ($T) b", blockType, blockType);
         builder.addStatement("int positionCount = block.getPositionCount()");
-        TypeName resultBlockType = blockType(resultType);
+        TypeName resultBuilderType = builderType(blockType(resultType));
         builder.beginControlFlow(
-            "try ($T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory()))",
-            resultBlockType,
-            resultBlockType
+            "try ($T builder = driverContext.blockFactory().$L(positionCount))",
+            resultBuilderType,
+            buildFromFactory(resultBuilderType)
         );
         String scratchPadName = null;
         if (argumentType.equals(BYTES_REF)) {

diff --git a/...in/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/EvaluatorImplementer.java b/...in/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/EvaluatorImplementer.java
@@ -33,6 +33,7 @@
 import javax.lang.model.util.Elements;
 
 import static org.elasticsearch.compute.gen.Methods.appendMethod;
+import static org.elasticsearch.compute.gen.Methods.buildFromFactory;
 import static org.elasticsearch.compute.gen.Methods.getMethod;
 import static org.elasticsearch.compute.gen.Types.BLOCK_REF;
 import static org.elasticsearch.compute.gen.Types.BYTES_REF;
@@ -45,6 +46,7 @@
 import static org.elasticsearch.compute.gen.Types.SOURCE;
 import static org.elasticsearch.compute.gen.Types.WARNINGS;
 import static org.elasticsearch.compute.gen.Types.blockType;
+import static org.elasticsearch.compute.gen.Types.builderType;
 import static org.elasticsearch.compute.gen.Types.vectorType;
 
 public class EvaluatorImplementer {
@@ -158,11 +160,11 @@ private MethodSpec realEval(boolean blockStyle) {
                 builder.addParameter(a.dataType(blockStyle), a.paramName(blockStyle));
             }
         });
+        TypeName builderType = builderType(resultDataType);
         builder.beginControlFlow(
-            "try($T.Builder result = $T.$L(positionCount, driverContext.blockFactory()))",
-            resultDataType,
-            resultDataType,
-            resultDataType.simpleName().endsWith("Vector") ? "newVectorBuilder" : "newBlockBuilder"
+            "try($T result = driverContext.blockFactory().$L(positionCount))",
+            builderType,
+            buildFromFactory(builderType)
         );
         {
             processFunction.args.stream().forEach(a -> a.createScratch(builder));

diff --git a/x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/Methods.java b/x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/Methods.java
@@ -19,14 +19,28 @@
 import javax.lang.model.util.ElementFilter;
 
 import static org.elasticsearch.compute.gen.Types.BOOLEAN_BLOCK;
+import static org.elasticsearch.compute.gen.Types.BOOLEAN_BLOCK_BUILDER;
 import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR;
+import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR_BUILDER;
+import static org.elasticsearch.compute.gen.Types.BOOLEAN_VECTOR_FIXED_BUILDER;
 import static org.elasticsearch.compute.gen.Types.BYTES_REF_BLOCK;
+import static org.elasticsearch.compute.gen.Types.BYTES_REF_BLOCK_BUILDER;
+import static org.elasticsearch.compute.gen.Types.BYTES_REF_VECTOR_BUILDER;
 import static org.elasticsearch.compute.gen.Types.DOUBLE_BLOCK;
+import static org.elasticsearch.compute.gen.Types.DOUBLE_BLOCK_BUILDER;
 import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR;
+import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR_BUILDER;
+import static org.elasticsearch.compute.gen.Types.DOUBLE_VECTOR_FIXED_BUILDER;
 import static org.elasticsearch.compute.gen.Types.INT_BLOCK;
+import static org.elasticsearch.compute.gen.Types.INT_BLOCK_BUILDER;
 import static org.elasticsearch.compute.gen.Types.INT_VECTOR;
+import static org.elasticsearch.compute.gen.Types.INT_VECTOR_BUILDER;
+import static org.elasticsearch.compute.gen.Types.INT_VECTOR_FIXED_BUILDER;
 import static org.elasticsearch.compute.gen.Types.LONG_BLOCK;
+import static org.elasticsearch.compute.gen.Types.LONG_BLOCK_BUILDER;
 import static org.elasticsearch.compute.gen.Types.LONG_VECTOR;
+import static org.elasticsearch.compute.gen.Types.LONG_VECTOR_BUILDER;
+import static org.elasticsearch.compute.gen.Types.LONG_VECTOR_FIXED_BUILDER;
 
 /**
  * Finds declared methods for the code generator.
@@ -95,6 +109,56 @@ static String appendMethod(TypeName t) {
         throw new IllegalArgumentException("unknown append method for [" + t + "]");
     }
 
+    /**
+     * Returns the name of the method used to build {@code t} instances
+     * from a {@code BlockFactory}.
+     */
+    static String buildFromFactory(TypeName t) {
+        if (t.equals(BOOLEAN_BLOCK_BUILDER)) {
+            return "newBooleanBlockBuilder";
+        }
+        if (t.equals(BOOLEAN_VECTOR_FIXED_BUILDER)) {
+            return "newBooleanVectorFixedBuilder";
+        }
+        if (t.equals(BOOLEAN_VECTOR_BUILDER)) {
+            return "newBooleanVectorBuilder";
+        }
+        if (t.equals(BYTES_REF_BLOCK_BUILDER)) {
+            return "newBytesRefBlockBuilder";
+        }
+        if (t.equals(BYTES_REF_VECTOR_BUILDER)) {
+            return "newBytesRefVectorBuilder";
+        }
+        if (t.equals(INT_BLOCK_BUILDER)) {
+            return "newIntBlockBuilder";
+        }
+        if (t.equals(INT_VECTOR_FIXED_BUILDER)) {
+            return "newIntVectorFixedBuilder";
+        }
+        if (t.equals(INT_VECTOR_BUILDER)) {
+            return "newIntVectorBuilder";
+        }
+        if (t.equals(LONG_BLOCK_BUILDER)) {
+            return "newLongBlockBuilder";
+        }
+        if (t.equals(LONG_VECTOR_FIXED_BUILDER)) {
+            return "newLongVectorFixedBuilder";
+        }
+        if (t.equals(LONG_VECTOR_BUILDER)) {
+            return "newLongVectorBuilder";
+        }
+        if (t.equals(DOUBLE_BLOCK_BUILDER)) {
+            return "newDoubleBlockBuilder";
+        }
+        if (t.equals(DOUBLE_VECTOR_BUILDER)) {
+            return "newDoubleVectorBuilder";
+        }
+        if (t.equals(DOUBLE_VECTOR_FIXED_BUILDER)) {
+            return "newDoubleVectorFixedBuilder";
+        }
+        throw new IllegalArgumentException("unknown build method for [" + t + "]");
+    }
+
     /**
      * Returns the name of the method used to get {@code valueType} instances
      * from vectors or blocks.

diff --git a/.../esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/MvEvaluatorImplementer.java b/.../esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/MvEvaluatorImplementer.java
@@ -39,6 +39,8 @@
 import static org.elasticsearch.compute.gen.Types.SOURCE;
 import static org.elasticsearch.compute.gen.Types.WARNINGS;
 import static org.elasticsearch.compute.gen.Types.blockType;
+import static org.elasticsearch.compute.gen.Types.builderType;
+import static org.elasticsearch.compute.gen.Types.vectorFixedBuilderType;
 import static org.elasticsearch.compute.gen.Types.vectorType;
 
 public class MvEvaluatorImplementer {
@@ -196,28 +198,19 @@ private MethodSpec evalShell(
         builder.beginControlFlow("try (ref)");
         builder.addStatement("$T v = ($T) ref.block()", blockType, blockType);
         builder.addStatement("int positionCount = v.getPositionCount()");
+        TypeName builderType;
         if (nullable) {
-            TypeName resultBlockType = blockType(resultType);
-            builder.beginControlFlow(
-                "try ($T.Builder builder = $T.newBlockBuilder(positionCount, driverContext.blockFactory()))",
-                resultBlockType,
-                resultBlockType
-            );
+            builderType = builderType(blockType(resultType));
         } else if (resultType.equals(BYTES_REF)) {
-            TypeName resultVectorType = vectorType(resultType);
-            builder.beginControlFlow(
-                "try ($T.Builder builder = $T.newVectorBuilder(positionCount, driverContext.blockFactory()))",
-                resultVectorType,
-                resultVectorType
-            );
+            builderType = builderType(vectorType(resultType));
         } else {
-            TypeName resultVectorType = vectorType(resultType);
-            builder.beginControlFlow(
-                "try ($T.FixedBuilder builder = $T.newVectorFixedBuilder(positionCount, driverContext.blockFactory()))",
-                resultVectorType,
-                resultVectorType
-            );
+            builderType = vectorFixedBuilderType(resultType);
         }
+        builder.beginControlFlow(
+            "try ($T builder = driverContext.blockFactory().$L(positionCount))",
+            builderType,
+            Methods.buildFromFactory(builderType)
+        );
 
         if (false == workType.equals(fieldType) && workType.isPrimitive() == false) {
             builder.addStatement("$T work = new $T()", workType, workType);

diff --git a/x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/Types.java b/x-pack/plugin/esql/compute/gen/src/main/java/org/elasticsearch/compute/gen/Types.java
@@ -58,6 +58,17 @@ public class Types {
     static final ClassName LONG_VECTOR = ClassName.get(DATA_PACKAGE, "LongVector");
     static final ClassName DOUBLE_VECTOR = ClassName.get(DATA_PACKAGE, "DoubleVector");
 
+    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 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 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");
@@ -198,6 +209,56 @@ static ClassName vectorType(String elementType) {
         throw new IllegalArgumentException("unknown vector type for [" + elementType + "]");
     }
 
+    static ClassName builderType(TypeName resultType) {
+        if (resultType.equals(BOOLEAN_BLOCK)) {
+            return BOOLEAN_BLOCK_BUILDER;
+        }
+        if (resultType.equals(BOOLEAN_VECTOR)) {
+            return BOOLEAN_VECTOR_BUILDER;
+        }
+        if (resultType.equals(BYTES_REF_BLOCK)) {
+            return BYTES_REF_BLOCK_BUILDER;
+        }
+        if (resultType.equals(BYTES_REF_VECTOR)) {
+            return BYTES_REF_VECTOR_BUILDER;
+        }
+        if (resultType.equals(INT_BLOCK)) {
+            return INT_BLOCK_BUILDER;
+        }
+        if (resultType.equals(INT_VECTOR)) {
+            return INT_VECTOR_BUILDER;
+        }
+        if (resultType.equals(LONG_BLOCK)) {
+            return LONG_BLOCK_BUILDER;
+        }
+        if (resultType.equals(LONG_VECTOR)) {
+            return LONG_VECTOR_BUILDER;
+        }
+        if (resultType.equals(DOUBLE_BLOCK)) {
+            return DOUBLE_BLOCK_BUILDER;
+        }
+        if (resultType.equals(DOUBLE_VECTOR)) {
+            return DOUBLE_VECTOR_BUILDER;
+        }
+        throw new IllegalArgumentException("unknown builder type for [" + resultType + "]");
+    }
+
+    static ClassName vectorFixedBuilderType(TypeName elementType) {
+        if (elementType.equals(TypeName.BOOLEAN)) {
+            return BOOLEAN_VECTOR_FIXED_BUILDER;
+        }
+        if (elementType.equals(TypeName.INT)) {
+            return INT_VECTOR_FIXED_BUILDER;
+        }
+        if (elementType.equals(TypeName.LONG)) {
+            return LONG_VECTOR_FIXED_BUILDER;
+        }
+        if (elementType.equals(TypeName.DOUBLE)) {
+            return DOUBLE_VECTOR_FIXED_BUILDER;
+        }
+        throw new IllegalArgumentException("unknown vector fixed builder type for [" + elementType + "]");
+    }
+
     static ClassName arrayVectorType(TypeName elementType) {
         if (elementType.equals(TypeName.BOOLEAN)) {
             return BOOLEAN_ARRAY_VECTOR;

diff --git a/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java b/x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/BlockFactory.java
@@ -109,6 +109,9 @@ public BooleanBlock.Builder newBooleanBlockBuilder(int estimatedSize) {
         return new BooleanBlockBuilder(estimatedSize, this);
     }
 
+    /**
+     * Build a {@link BooleanVector.FixedBuilder} that never grows.
+     */
     public BooleanVector.FixedBuilder newBooleanVectorFixedBuilder(int size) {
         return new BooleanVectorFixedBuilder(size, this);
     }
@@ -172,6 +175,9 @@ public IntVector.Builder newIntVectorBuilder(int estimatedSize) {
         return new IntVectorBuilder(estimatedSize, this);
     }
 
+    /**
+     * Build a {@link IntVector.FixedBuilder} that never grows.
+     */
     public IntVector.FixedBuilder newIntVectorFixedBuilder(int size) {
         return new IntVectorFixedBuilder(size, this);
     }
@@ -236,7 +242,10 @@ public LongVector.Builder newLongVectorBuilder(int estimatedSize) {
         return new LongVectorBuilder(estimatedSize, this);
     }
 
-    LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
+    /**
+     * Build a {@link LongVector.FixedBuilder} that never grows.
+     */
+    public LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
         return new LongVectorFixedBuilder(size, this);
     }
 
@@ -286,7 +295,10 @@ public DoubleVector.Builder newDoubleVectorBuilder(int estimatedSize) {
         return new DoubleVectorBuilder(estimatedSize, this);
     }
 
-    DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
+    /**
+     * Build a {@link DoubleVector.FixedBuilder} that never grows.
+     */
+    public DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
         return new DoubleVectorFixedBuilder(size, this);
     }
 

diff --git a/...ck/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/MockBlockFactory.java b/...ck/plugin/esql/compute/src/test/java/org/elasticsearch/compute/data/MockBlockFactory.java
@@ -245,7 +245,8 @@ public LongVector.Builder newLongVectorBuilder(int estimatedSize) {
         return b;
     }
 
-    LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
+    @Override
+    public LongVector.FixedBuilder newLongVectorFixedBuilder(int size) {
         var b = super.newLongVectorFixedBuilder(size);
         track(b, trackDetail());
         return b;
@@ -286,7 +287,8 @@ public DoubleVector.Builder newDoubleVectorBuilder(int estimatedSize) {
         return b;
     }
 
-    DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
+    @Override
+    public DoubleVector.FixedBuilder newDoubleVectorFixedBuilder(int size) {
         var b = super.newDoubleVectorFixedBuilder(size);
         track(b, trackDetail());
         return b;