showcase

Author: dnhatn

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

@@ -21,7 +21,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.function.BiConsumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -41,6 +40,7 @@
 import static org.elasticsearch.compute.gen.Types.BOOLEAN_BLOCK;
 import static org.elasticsearch.compute.gen.Types.BYTES_REF;
 import static org.elasticsearch.compute.gen.Types.BYTES_REF_BLOCK;
+import static org.elasticsearch.compute.gen.Types.BYTES_REF_VECTOR;
 import static org.elasticsearch.compute.gen.Types.DOUBLE_BLOCK;
 import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
 import static org.elasticsearch.compute.gen.Types.EXPRESSION_EVALUATOR;
@@ -257,13 +257,7 @@ private MethodSpec realEval(boolean blockStyle) {
                         builder.beginControlFlow("try");
                     }
                     if (executionCost > 0) {
-                        builder.addStatement("accumulatedCost += " + executionCost);
-                        builder.beginControlFlow("if (accumulatedCost >= DriverContext.CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD)");
-                        {
-                            builder.addStatement("accumulatedCost = 0");
-                            builder.addStatement("driverContext.checkForEarlyTermination()");
-                        }
-                        builder.endControlFlow();
+                        addEarlyTerminationCheck(builder, executionCost);
                     }
                     builder.addStatement(builtPattern, args.toArray());
 
@@ -287,49 +281,49 @@ private MethodSpec realEval(boolean blockStyle) {
     }
 
     private void realEvalWithVectorizedStyle(MethodSpec.Builder builder, ClassName resultDataType) {
-        boolean checkForEarlyTerminationPerRow = processFunction.args.stream()
-            .anyMatch(a -> a.dataType(false).getClass().getSimpleName().startsWith("BytesRef"));
-        BiConsumer<String, Object[]> innerLoop = (label, bounds) -> {
-            builder.beginControlFlow(label + "for (int p = " + bounds[0] + "; p < " + bounds[1] + "; p++)");
-            {
-                processFunction.args.forEach(a -> a.unpackValues(builder, false));
-                StringBuilder pattern = new StringBuilder();
-                List<Object> args = new ArrayList<>();
-                pattern.append("$T.$N(");
-                args.add(declarationType);
-                args.add(processFunction.function.getSimpleName());
-                processFunction.args.forEach(a -> {
-                    if (args.size() > 2) {
-                        pattern.append(", ");
-                    }
-                    a.buildInvocation(pattern, args, false);
-                });
-                pattern.append(")");
-                String builtPattern;
-                if (processFunction.builderArg == null) {
-                    builtPattern = "result.$L(p, " + pattern + ")";
-                    args.add(0, appendMethod(resultDataType));
-                } else {
-                    builtPattern = pattern.toString();
+        boolean checkEarlyTerminationPerRow = executionCost > 0
+            && processFunction.args.stream().anyMatch(a -> a.dataType(false).equals(BYTES_REF_VECTOR));
+        if (checkEarlyTerminationPerRow) {
+            builder.addStatement("int accumulatedCost = 0");
+        }
+        builder.beginControlFlow("position: for (int p = 0; p < positionCount; p++)");
+        {
+            processFunction.args.forEach(a -> a.unpackValues(builder, false));
+            StringBuilder pattern = new StringBuilder();
+            List<Object> args = new ArrayList<>();
+            pattern.append("$T.$N(");
+            args.add(declarationType);
+            args.add(processFunction.function.getSimpleName());
+            processFunction.args.forEach(a -> {
+                if (args.size() > 2) {
+                    pattern.append(", ");
                 }
-                builder.addStatement(builtPattern, args.toArray());
+                a.buildInvocation(pattern, args, false);
+            });
+            pattern.append(")");
+            String builtPattern;
+            if (processFunction.builderArg == null) {
+                builtPattern = "result.$L(p, " + pattern + ")";
+                args.add(0, appendMethod(resultDataType));
+            } else {
+                builtPattern = pattern.toString();
             }
-            builder.endControlFlow();
-        };
-        if (executionCost == 0) {
-            innerLoop.accept("position: ", new Object[] { 0, "positionCount" });
-        } else {
-            // generate tight loops to allow vectorization
-            builder.addStatement("final int batchSize = DriverContext.batchSizeForEarlyTermination($L)", executionCost);
-            builder.beginControlFlow("for (int start = 0; start < positionCount; )");
-            {
-                builder.addStatement("int end = start + Math.min(positionCount - start, batchSize)");
-                builder.addStatement("driverContext.checkForEarlyTermination()");
-                innerLoop.accept("", new Object[] { "start", "end" });
-                builder.addStatement("start = end");
+            if (checkEarlyTerminationPerRow) {
+                addEarlyTerminationCheck(builder, executionCost);
             }
-            builder.endControlFlow();
+            builder.addStatement(builtPattern, args.toArray());
         }
+        builder.endControlFlow();
+    }
+
+    static void addEarlyTerminationCheck(MethodSpec.Builder builder, int executionCost) {
+        builder.addStatement("accumulatedCost += $L", executionCost);
+        builder.beginControlFlow("if (accumulatedCost >= DriverContext.CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD)");
+        {
+            builder.addStatement("accumulatedCost = 0");
+            builder.addStatement("driverContext.checkForEarlyTermination()");
+        }
+        builder.endControlFlow();
     }
 
     private static void skipNull(MethodSpec.Builder builder, String value) {

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

@@ -139,10 +139,7 @@ public Driver(
                 DriverSleeps.empty()
             )
         );
-        driverContext.initializeEarlyTerminationChecker(() -> {
-            ensureNotCancelled();
-            checkForEarlyTermination();
-        });
+        driverContext.initializeEarlyTerminationChecker(this::checkForEarlyTerminationOrCancellation);
     }
 
     /**
@@ -252,7 +249,6 @@ public void abort(Exception reason, ActionListener<Void> listener) {
     }
 
     private IsBlockedResult runSingleLoopIteration() {
-        ensureNotCancelled();
         boolean movedPage = false;
 
         for (int i = 0; i < activeOperators.size() - 1; i++) {
@@ -265,6 +261,7 @@ private IsBlockedResult runSingleLoopIteration() {
             }
 
             if (op.isFinished() == false && nextOp.needsInput()) {
+                checkForEarlyTerminationOrCancellation();
                 Page page = op.getOutput();
                 if (page == null) {
                     // No result, just move to the next iteration
@@ -273,12 +270,19 @@ private IsBlockedResult runSingleLoopIteration() {
                     page.releaseBlocks();
                 } else {
                     // Non-empty result from the previous operation, move it to the next operation
+                    try {
+                        checkForEarlyTerminationOrCancellation();
+                    } catch (DriverEarlyTerminationException | TaskCancelledException e) {
+                        page.releaseBlocks();
+                        throw e;
+                    }
                     nextOp.addInput(page);
                     movedPage = true;
                 }
             }
 
             if (op.isFinished()) {
+                checkForEarlyTerminationOrCancellation();
                 nextOp.finish();
             }
         }
@@ -350,7 +354,8 @@ private static class DriverEarlyTerminationException extends RuntimeException {
 
     }
 
-    private void checkForEarlyTermination() throws DriverEarlyTerminationException {
+    private void checkForEarlyTerminationOrCancellation() throws DriverEarlyTerminationException, TaskCancelledException {
+        ensureNotCancelled();
         // 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++) {

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

@@ -194,17 +194,6 @@ public void initializeEarlyTerminationChecker(Runnable checker) {
         this.earlyTerminationChecker = checker;
     }
 
-    public static int batchSizeForEarlyTermination(int executionCost) {
-        if (executionCost <= 0) {
-            throw new IllegalStateException("executionCost must be positive");
-        }
-        final int size = (DriverContext.CHECK_FOR_EARLY_TERMINATION_COST_THRESHOLD + executionCost - 1) / executionCost;
-        if (size <= 32) {
-            return 32;
-        }
-        return Integer.highestOneBit(size - 1) << 1;
-    }
-
     /**
      * Evaluators should use this function to decide their warning behavior.
      * @return an appropriate {@link WarningsMode}

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/DriverContextTests.java

@@ -24,11 +24,9 @@
 import org.junit.Before;
 
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
@@ -158,23 +156,6 @@ public void testWaitForAsyncActions() {
         Releasables.closeExpectNoException(driverContext.getSnapshot());
     }
 
-    public void testBatchSizeForEarlyTermination() {
-        Map<Integer, Integer> tests = new HashMap<>();
-        tests.put(1, 2048);
-        tests.put(between(2, 3), 1024);
-        tests.put(between(4, 7), 512);
-        tests.put(between(8, 15), 256);
-        tests.put(between(16, 31), 128);
-        tests.put(between(32, 63), 64);
-        tests.put(between(64, 127), 32);
-        tests.put(atLeast(128), 32);
-        for (Map.Entry<Integer, Integer> e : tests.entrySet()) {
-            Integer cost = e.getKey();
-            Integer batchSize = e.getValue();
-            assertThat("cost=" + cost, DriverContext.batchSizeForEarlyTermination(cost), equalTo(batchSize));
-        }
-    }
-
     static TestDriver newTestDriver(int unused) {
         var driverContext = new AssertingDriverContext();
         return new TestDriver(driverContext, randomInt(128), driverContext.bigArrays());