tighten

Author: dnhatn

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

@@ -75,6 +75,9 @@ public EvaluatorImplementer(
         this.declarationType = (TypeElement) processFunction.getEnclosingElement();
         this.processFunction = new ProcessFunction(elements, types, processFunction, warnExceptions);
         this.estimateCost = estimateCost;
+        if (estimateCost <= 0) {
+            throw new IllegalArgumentException("estimateCost must be at least 1; got " + estimateCost);
+        }
         this.implementation = ClassName.get(
             elements.getPackageOf(declarationType).toString(),
             declarationType.getSimpleName() + extraName + "Evaluator"
@@ -197,6 +200,7 @@ private MethodSpec realEval(boolean blockStyle) {
             if (vectorize) {
                 realEvalWithVectorizedStyle(builder, resultDataType);
             } else {
+                builder.addStatement("int accumulatedCost = 0");
                 builder.beginControlFlow("position: for (int p = 0; p < positionCount; p++)");
                 {
                     if (blockStyle) {
@@ -249,8 +253,13 @@ private MethodSpec realEval(boolean blockStyle) {
                     if (processFunction.warnExceptions.isEmpty() == false) {
                         builder.beginControlFlow("try");
                     }
-
-                    builder.addStatement("driverContext.maybeCheckForEarlyTermination(" + estimateCost + ")");
+                    builder.addStatement("accumulatedCost += " + estimateCost);
+                    builder.beginControlFlow("if (accumulatedCost >= DriverContext.CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD)");
+                    {
+                        builder.addStatement("accumulatedCost = 0");
+                        builder.addStatement("driverContext.checkForEarlyTermination()");
+                    }
+                    builder.endControlFlow();
                     builder.addStatement(builtPattern, args.toArray());
 
                     if (processFunction.warnExceptions.isEmpty() == false) {
@@ -273,8 +282,8 @@ private MethodSpec realEval(boolean blockStyle) {
     }
 
     private void realEvalWithVectorizedStyle(MethodSpec.Builder builder, ClassName resultDataType) {
-        // generate the tight loop to allow vectorization
-        builder.addStatement("final int maxBatchSize = DriverContext.estimateBatchSizeForEarlyTermination(" + estimateCost + ")");
+        builder.addComment("generate a tight loop to allow vectorization");
+        builder.addStatement("int maxBatchSize = Math.max(DriverContext.CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD / $L, 1)", estimateCost);
         builder.beginControlFlow("for (int start = 0; start < positionCount; )");
         {
             builder.addStatement("int end = start + Math.min(positionCount - start, maxBatchSize)");

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/DriverContext.java

@@ -61,8 +61,8 @@ public class DriverContext {
     private final WarningsMode warningsMode;
 
     public static final int CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD = 2048;
+    private static final Runnable NO_OP = () -> {};
     private Runnable earlyTerminationChecker = () -> {};
-    private int accumulatedCostForEarlyTermination;
 
     public DriverContext(BigArrays bigArrays, BlockFactory blockFactory) {
         this(bigArrays, blockFactory, WarningsMode.COLLECT);
@@ -179,21 +179,10 @@ public void removeAsyncAction() {
         asyncActions.removeInstance();
     }
 
-    /**
-     * Accumulates the early termination cost and runs the early termination check if the accumulated cost passes the threshold.
-     */
-    public void maybeCheckForEarlyTermination(int estimateCost) {
-        accumulatedCostForEarlyTermination += estimateCost;
-        if (accumulatedCostForEarlyTermination >= CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD) {
-            checkForEarlyTermination();
-        }
-    }
-
     /**
      * Checks if the Driver associated with this DriverContext has been cancelled or early terminated.
      */
     public void checkForEarlyTermination() {
-        accumulatedCostForEarlyTermination = 0;
         earlyTerminationChecker.run();
     }
 
@@ -202,17 +191,10 @@ public void checkForEarlyTermination() {
      * This method should be called when associating this DriverContext with a driver.
      */
     public void initializeEarlyTerminationChecker(Runnable checker) {
-        this.earlyTerminationChecker = checker;
-    }
-
-    /**
-     * Returns the estimated batch size for early termination based on the given estimate cost for each item.
-     */
-    public static int estimateBatchSizeForEarlyTermination(int estimateCost) {
-        if (estimateCost <= 0) {
-            return CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD;
+        if (this.earlyTerminationChecker != NO_OP) {
+            throw new IllegalStateException("Early termination checker already initialized");
         }
-        return Math.max(CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD / estimateCost, 1);
+        this.earlyTerminationChecker = checker;
     }
 
     /**