Check for early termination

Author: dnhatn

x-pack/plugin/esql/compute/ann/src/main/java/org/elasticsearch/compute/ann/Evaluator.java

@@ -42,4 +42,10 @@
      * into a warning and turn into a null value.
      */
     Class<? extends Exception>[] warnExceptions() default {};
+
+    /**
+     * The estimated cost of evaluating one row of this evaluator. The Driver periodically checks for cancellation or early termination
+     * after it has accumulated the cost over the threshold (2028). Increase this estimate for expensive evaluators.
+     */
+    int estimateCost() default 1;
 }

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

@@ -61,18 +61,20 @@ public class EvaluatorImplementer {
     private final TypeElement declarationType;
     private final ProcessFunction processFunction;
     private final ClassName implementation;
+    private final int estimateCost;
     private final boolean processOutputsMultivalued;
 
     public EvaluatorImplementer(
         Elements elements,
         javax.lang.model.util.Types types,
         ExecutableElement processFunction,
         String extraName,
+        int estimateCost,
         List<TypeMirror> warnExceptions
     ) {
         this.declarationType = (TypeElement) processFunction.getEnclosingElement();
         this.processFunction = new ProcessFunction(elements, types, processFunction, warnExceptions);
-
+        this.estimateCost = estimateCost;
         this.implementation = ClassName.get(
             elements.getPackageOf(declarationType).toString(),
             declarationType.getSimpleName() + extraName + "Evaluator"
@@ -175,7 +177,7 @@ private MethodSpec realEval(boolean blockStyle) {
         boolean vectorize = false;
         if (blockStyle == false && processFunction.warnExceptions.isEmpty() && processOutputsMultivalued == false) {
             ClassName type = processFunction.resultDataType(false);
-            vectorize = type.simpleName().startsWith("BytesRef") == false;
+            vectorize = type.simpleName().startsWith("BytesRef") == false && processFunction.warnExceptions.isEmpty();
         }
 
         TypeName builderType = vectorize ? vectorFixedBuilderType(elementType(resultDataType)) : builderType(resultDataType);
@@ -192,78 +194,119 @@ private MethodSpec realEval(boolean blockStyle) {
             });
 
             processFunction.args.stream().forEach(a -> a.createScratch(builder));
-
-            builder.beginControlFlow("position: for (int p = 0; p < positionCount; p++)");
-            {
-                if (blockStyle) {
-                    if (processOutputsMultivalued == false) {
-                        processFunction.args.stream().forEach(a -> a.skipNull(builder));
-                    } else {
-                        builder.addStatement("boolean allBlocksAreNulls = true");
-                        // allow block type inputs to be null
-                        processFunction.args.stream().forEach(a -> {
-                            if (a instanceof StandardProcessFunctionArg as) {
-                                as.skipNull(builder);
-                            } else if (a instanceof BlockProcessFunctionArg ab) {
-                                builder.beginControlFlow("if (!$N.isNull(p))", ab.paramName(blockStyle));
-                                {
-                                    builder.addStatement("allBlocksAreNulls = false");
+            if (vectorize) {
+                realEvalWithVectorizedStyle(builder, resultDataType);
+            } else {
+                builder.beginControlFlow("position: for (int p = 0; p < positionCount; p++)");
+                {
+                    if (blockStyle) {
+                        if (processOutputsMultivalued == false) {
+                            processFunction.args.stream().forEach(a -> a.skipNull(builder));
+                        } else {
+                            builder.addStatement("boolean allBlocksAreNulls = true");
+                            // allow block type inputs to be null
+                            processFunction.args.stream().forEach(a -> {
+                                if (a instanceof StandardProcessFunctionArg as) {
+                                    as.skipNull(builder);
+                                } else if (a instanceof BlockProcessFunctionArg ab) {
+                                    builder.beginControlFlow("if (!$N.isNull(p))", ab.paramName(blockStyle));
+                                    {
+                                        builder.addStatement("allBlocksAreNulls = false");
+                                    }
+                                    builder.endControlFlow();
                                 }
-                                builder.endControlFlow();
-                            }
-                        });
+                            });
 
-                        builder.beginControlFlow("if (allBlocksAreNulls)");
-                        {
-                            builder.addStatement("result.appendNull()");
-                            builder.addStatement("continue position");
+                            builder.beginControlFlow("if (allBlocksAreNulls)");
+                            {
+                                builder.addStatement("result.appendNull()");
+                                builder.addStatement("continue position");
+                            }
+                            builder.endControlFlow();
                         }
+                    }
+                    processFunction.args.stream().forEach(a -> a.unpackValues(builder, blockStyle));
+
+                    StringBuilder pattern = new StringBuilder();
+                    List<Object> args = new ArrayList<>();
+                    pattern.append(processOutputsMultivalued ? "$T.$N(result, p, " : "$T.$N(");
+                    args.add(declarationType);
+                    args.add(processFunction.function.getSimpleName());
+                    processFunction.args.stream().forEach(a -> {
+                        if (args.size() > 2) {
+                            pattern.append(", ");
+                        }
+                        a.buildInvocation(pattern, args, blockStyle);
+                    });
+                    pattern.append(")");
+                    String builtPattern;
+                    if (processFunction.builderArg == null) {
+                        builtPattern = "result.$L(" + pattern + ")";
+                        args.add(0, appendMethod(resultDataType));
+                    } else {
+                        builtPattern = pattern.toString();
+                    }
+                    if (processFunction.warnExceptions.isEmpty() == false) {
+                        builder.beginControlFlow("try");
+                    }
+
+                    builder.addStatement("driverContext.maybeCheckForEarlyTermination(" + estimateCost + ")");
+                    builder.addStatement(builtPattern, args.toArray());
+
+                    if (processFunction.warnExceptions.isEmpty() == false) {
+                        String catchPattern = "catch ("
+                            + processFunction.warnExceptions.stream().map(m -> "$T").collect(Collectors.joining(" | "))
+                            + " e)";
+                        builder.nextControlFlow(catchPattern, processFunction.warnExceptions.stream().map(m -> TypeName.get(m)).toArray());
+                        builder.addStatement("warnings().registerException(e)");
+                        builder.addStatement("result.appendNull()");
                         builder.endControlFlow();
                     }
                 }
-                processFunction.args.stream().forEach(a -> a.unpackValues(builder, blockStyle));
+                builder.endControlFlow();
+            }
+            builder.addStatement("return result.build()");
+        }
+        builder.endControlFlow();
 
+        return builder.build();
+    }
+
+    private void realEvalWithVectorizedStyle(MethodSpec.Builder builder, ClassName resultDataType) {
+        // generate the tight loop to allow vectorization
+        builder.addStatement("final int maxBatchSize = DriverContext.estimateBatchSizeForEarlyTermination(" + estimateCost + ")");
+        builder.beginControlFlow("for (int start = 0; start < positionCount; )");
+        {
+            builder.addStatement("int end = start + Math.min(positionCount - start, maxBatchSize)");
+            builder.addStatement("driverContext.checkForEarlyTermination()");
+            builder.beginControlFlow("for (int p = start; p < end; p++)");
+            {
+                processFunction.args.forEach(a -> a.unpackValues(builder, false));
                 StringBuilder pattern = new StringBuilder();
                 List<Object> args = new ArrayList<>();
-                pattern.append(processOutputsMultivalued ? "$T.$N(result, p, " : "$T.$N(");
+                pattern.append("$T.$N(");
                 args.add(declarationType);
                 args.add(processFunction.function.getSimpleName());
-                processFunction.args.stream().forEach(a -> {
+                processFunction.args.forEach(a -> {
                     if (args.size() > 2) {
                         pattern.append(", ");
                     }
-                    a.buildInvocation(pattern, args, blockStyle);
+                    a.buildInvocation(pattern, args, false);
                 });
                 pattern.append(")");
                 String builtPattern;
                 if (processFunction.builderArg == null) {
-                    builtPattern = vectorize ? "result.$L(p, " + pattern + ")" : "result.$L(" + pattern + ")";
+                    builtPattern = "result.$L(p, " + pattern + ")";
                     args.add(0, appendMethod(resultDataType));
                 } else {
                     builtPattern = pattern.toString();
                 }
-                if (processFunction.warnExceptions.isEmpty() == false) {
-                    builder.beginControlFlow("try");
-                }
-
                 builder.addStatement(builtPattern, args.toArray());
-
-                if (processFunction.warnExceptions.isEmpty() == false) {
-                    String catchPattern = "catch ("
-                        + processFunction.warnExceptions.stream().map(m -> "$T").collect(Collectors.joining(" | "))
-                        + " e)";
-                    builder.nextControlFlow(catchPattern, processFunction.warnExceptions.stream().map(m -> TypeName.get(m)).toArray());
-                    builder.addStatement("warnings().registerException(e)");
-                    builder.addStatement("result.appendNull()");
-                    builder.endControlFlow();
-                }
             }
             builder.endControlFlow();
-            builder.addStatement("return result.build()");
+            builder.addStatement("start = end");
         }
         builder.endControlFlow();
-
-        return builder.build();
     }
 
     private static void skipNull(MethodSpec.Builder builder, String value) {

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

@@ -82,6 +82,7 @@ public boolean process(Set<? extends TypeElement> set, RoundEnvironment roundEnv
                                 env.getTypeUtils(),
                                 (ExecutableElement) evaluatorMethod,
                                 evaluatorAnn.extraName(),
+                                evaluatorAnn.estimateCost(),
                                 warnExceptionsTypes
                             ).sourceFile(),
                             env

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

@@ -139,6 +139,10 @@ public Driver(
                 DriverSleeps.empty()
             )
         );
+        driverContext.initializeEarlyTerminationChecker(() -> {
+            ensureNotCancelled();
+            checkForEarlyTermination();
+        });
     }
 
     /**
@@ -186,7 +190,13 @@ SubscribableListener<Void> run(TimeValue maxTime, int maxIterations, LongSupplie
         long nextStatus = startTime + statusNanos;
         int iter = 0;
         while (true) {
-            IsBlockedResult isBlocked = runSingleLoopIteration();
+            final IsBlockedResult isBlocked;
+            try {
+                isBlocked = runSingleLoopIteration();
+            } catch (DriverEarlyTerminationException ignored) {
+                closeEarlyFinishedOperators();
+                continue;
+            }
             iter++;
             if (isBlocked.listener().isDone() == false) {
                 updateStatus(nowSupplier.getAsLong() - startTime, iter, DriverStatus.Status.ASYNC, isBlocked.reason());
@@ -273,6 +283,20 @@ private IsBlockedResult runSingleLoopIteration() {
             }
         }
 
+        closeEarlyFinishedOperators();
+
+        if (movedPage == false) {
+            return oneOf(
+                activeOperators.stream()
+                    .map(Operator::isBlocked)
+                    .filter(laf -> laf.listener().isDone() == false)
+                    .collect(Collectors.toList())
+            );
+        }
+        return Operator.NOT_BLOCKED;
+    }
+
+    private void closeEarlyFinishedOperators() {
         for (int index = activeOperators.size() - 1; index >= 0; index--) {
             if (activeOperators.get(index).isFinished()) {
                 /*
@@ -298,16 +322,6 @@ private IsBlockedResult runSingleLoopIteration() {
                 break;
             }
         }
-
-        if (movedPage == false) {
-            return oneOf(
-                activeOperators.stream()
-                    .map(Operator::isBlocked)
-                    .filter(laf -> laf.listener().isDone() == false)
-                    .collect(Collectors.toList())
-            );
-        }
-        return Operator.NOT_BLOCKED;
     }
 
     public void cancel(String reason) {
@@ -332,6 +346,22 @@ private void ensureNotCancelled() {
         }
     }
 
+    private static class DriverEarlyTerminationException extends RuntimeException {
+
+    }
+
+    private void checkForEarlyTermination() throws DriverEarlyTerminationException {
+        // If the last operation is finished, then we can discard all operations in the driver
+        if (activeOperators.size() >= 2 && activeOperators.getLast().isFinished()) {
+            for (int i = 0; i < activeOperators.size() - 1; i++) {
+                Operator op = activeOperators.get(i);
+                if (op.isFinished() == false) {
+                    throw new DriverEarlyTerminationException();
+                }
+            }
+        }
+    }
+
     public static void start(
         ThreadContext threadContext,
         Executor executor,

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

@@ -60,6 +60,12 @@ public class DriverContext {
 
     private final WarningsMode warningsMode;
 
+    private static final Runnable NO_OP = () -> {};
+
+    public static final int CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD = 2048;
+    private Runnable earlyTerminationChecker = NO_OP;
+    private int accumulatedCostForEarlyTermination;
+
     public DriverContext(BigArrays bigArrays, BlockFactory blockFactory) {
         this(bigArrays, blockFactory, WarningsMode.COLLECT);
     }
@@ -175,6 +181,44 @@ 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() {
+        assert earlyTerminationChecker != NO_OP : "cancellation or early termination checker not initialized";
+        accumulatedCostForEarlyTermination = 0;
+        earlyTerminationChecker.run();
+    }
+
+    /**
+     * Initializes the early termination or cancellation checker for this DriverContext.
+     * This method should be called when associating this DriverContext with a driver.
+     */
+    public void initializeEarlyTerminationChecker(Runnable checker) {
+        assert earlyTerminationChecker == NO_OP : "cancellation or early termination checker already initialized";
+        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;
+        }
+        return Math.max(CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD / estimateCost, 1);
+    }
+
     /**
      * Evaluators should use this function to decide their warning behavior.
      * @return an appropriate {@link WarningsMode}