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 (2048). 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,23 @@ 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;
+        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"
@@ -175,7 +180,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 +197,125 @@ 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.addStatement("int accumulatedCost = 0");
+                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("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) {
+                        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) {
+        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)");
+            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) {
+                // rerun to prune early terminated operators
+                continue;
+            }
             iter++;
             if (isBlocked.listener().isDone() == false) {
                 updateStatus(nowSupplier.getAsLong() - startTime, iter, DriverStatus.Status.ASYNC, isBlocked.reason());
@@ -243,6 +253,33 @@ public void abort(Exception reason, ActionListener<Void> listener) {
 
     private IsBlockedResult runSingleLoopIteration() {
         ensureNotCancelled();
+
+        for (int index = activeOperators.size() - 1; index >= 0; index--) {
+            if (activeOperators.get(index).isFinished()) {
+                /*
+                 * Close and remove this operator and all source operators in the
+                 * most paranoid possible way. Closing operators shouldn't throw,
+                 * but if it does, this will make sure we don't try to close any
+                 * that succeed twice.
+                 */
+                List<Operator> finishedOperators = this.activeOperators.subList(0, index + 1);
+                Iterator<Operator> itr = finishedOperators.iterator();
+                while (itr.hasNext()) {
+                    Operator op = itr.next();
+                    statusOfCompletedOperators.add(new DriverStatus.OperatorStatus(op.toString(), op.status()));
+                    op.close();
+                    itr.remove();
+                }
+
+                // Finish the next operator, which is now the first operator.
+                if (activeOperators.isEmpty() == false) {
+                    Operator newRootOperator = activeOperators.get(0);
+                    newRootOperator.finish();
+                }
+                break;
+            }
+        }
+
         boolean movedPage = false;
 
         for (int i = 0; i < activeOperators.size() - 1; i++) {
@@ -273,32 +310,6 @@ private IsBlockedResult runSingleLoopIteration() {
             }
         }
 
-        for (int index = activeOperators.size() - 1; index >= 0; index--) {
-            if (activeOperators.get(index).isFinished()) {
-                /*
-                 * Close and remove this operator and all source operators in the
-                 * most paranoid possible way. Closing operators shouldn't throw,
-                 * but if it does, this will make sure we don't try to close any
-                 * that succeed twice.
-                 */
-                List<Operator> finishedOperators = this.activeOperators.subList(0, index + 1);
-                Iterator<Operator> itr = finishedOperators.iterator();
-                while (itr.hasNext()) {
-                    Operator op = itr.next();
-                    statusOfCompletedOperators.add(new DriverStatus.OperatorStatus(op.toString(), op.status()));
-                    op.close();
-                    itr.remove();
-                }
-
-                // Finish the next operator, which is now the first operator.
-                if (activeOperators.isEmpty() == false) {
-                    Operator newRootOperator = activeOperators.get(0);
-                    newRootOperator.finish();
-                }
-                break;
-            }
-        }
-
         if (movedPage == false) {
             return oneOf(
                 activeOperators.stream()
@@ -332,6 +343,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,10 @@ 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 = NO_OP;
+
     public DriverContext(BigArrays bigArrays, BlockFactory blockFactory) {
         this(bigArrays, blockFactory, WarningsMode.COLLECT);
     }
@@ -175,6 +179,24 @@ public void removeAsyncAction() {
         asyncActions.removeInstance();
     }
 
+    /**
+     * Checks if the Driver associated with this DriverContext has been cancelled or early terminated.
+     */
+    public void checkForEarlyTermination() {
+        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) {
+        if (this.earlyTerminationChecker != NO_OP) {
+            throw new IllegalStateException("Early termination checker already initialized");
+        }
+        this.earlyTerminationChecker = checker;
+    }
+
     /**
      * Evaluators should use this function to decide their warning behavior.
      * @return an appropriate {@link WarningsMode}