ESQL: Speed up CASE for some parameters

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/EvalBenchmark.java

@@ -24,13 +24,15 @@
 import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.predicate.regex.RLikePattern;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
 import org.elasticsearch.xpack.esql.evaluator.EvalMapper;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Abs;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin;
@@ -52,6 +54,7 @@
 
 import java.time.Duration;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
@@ -90,6 +93,8 @@ public class EvalBenchmark {
             "abs",
             "add",
             "add_double",
+            "case_1_eager",
+            "case_1_lazy",
             "date_trunc",
             "equal_to_const",
             "long_equal_to_long",
@@ -124,6 +129,18 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                     layout(doubleField)
                 ).get(driverContext);
             }
+            case "case_1_eager", "case_1_lazy" -> {
+                FieldAttribute f1 = longField();
+                FieldAttribute f2 = longField();
+                Expression condition = new Equals(Source.EMPTY, f1, new Literal(Source.EMPTY, 1L, DataType.LONG));
+                Expression lhs = f1;
+                Expression rhs = f2;
+                if (operation.endsWith("lazy")) {
+                    lhs = new Add(Source.EMPTY, lhs, new Literal(Source.EMPTY, 1L, DataType.LONG));
+                    rhs = new Add(Source.EMPTY, rhs, new Literal(Source.EMPTY, 1L, DataType.LONG));
+                }
+                yield EvalMapper.toEvaluator(new Case(Source.EMPTY, condition, List.of(lhs, rhs)), layout(f1, f2)).get(driverContext);
+            }
             case "date_trunc" -> {
                 FieldAttribute timestamp = new FieldAttribute(
                     Source.EMPTY,
@@ -215,6 +232,28 @@ private static void checkExpected(String operation, Page actual) {
                     }
                 }
             }
+            case "case_1_eager" -> {
+                LongVector f1 = actual.<LongBlock>getBlock(0).asVector();
+                LongVector f2 = actual.<LongBlock>getBlock(1).asVector();
+                LongVector result = actual.<LongBlock>getBlock(2).asVector();
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    long expected = f1.getLong(i) == 1 ? f1.getLong(i) : f2.getLong(i);
+                    if (result.getLong(i) != expected) {
+                        throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
+                    }
+                }
+            }
+            case "case_1_lazy" -> {
+                LongVector f1 = actual.<LongBlock>getBlock(0).asVector();
+                LongVector f2 = actual.<LongBlock>getBlock(1).asVector();
+                LongVector result = actual.<LongBlock>getBlock(2).asVector();
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    long expected = 1 + (f1.getLong(i) == 1 ? f1.getLong(i) : f2.getLong(i));
+                    if (result.getLong(i) != expected) {
+                        throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + result.getLong(i) + "]");
+                    }
+                }
+            }
             case "date_trunc" -> {
                 LongVector v = actual.<LongBlock>getBlock(1).asVector();
                 long oneDay = TimeValue.timeValueHours(24).millis();
@@ -279,6 +318,15 @@ private static Page page(String operation) {
                 }
                 yield new Page(builder.build());
             }
+            case "case_1_eager", "case_1_lazy" -> {
+                var f1 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
+                var f2 = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    f1.appendLong(i);
+                    f2.appendLong(-i);
+                }
+                yield new Page(f1.build(), f2.build());
+            }
             case "long_equal_to_long" -> {
                 var lhs = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);
                 var rhs = blockFactory.newLongBlockBuilder(BLOCK_LENGTH);

docs/changelog/112295.yaml

@@ -0,0 +1,5 @@
+pr: 112295
+summary: "ESQL: Speed up CASE for some parameters"
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -63,6 +63,16 @@ public interface ExpressionEvaluator extends Releasable {
         /** A Factory for creating ExpressionEvaluators. */
         interface Factory {
             ExpressionEvaluator get(DriverContext context);
+
+            /**
+             * {@code true} if it is safe and fast to evaluate this expression non-lazily
+             * in contexts like {@code CASE}. This defaults to {@code false}, but expressions
+             * that evaluate quickly and can not produce warnings may override this to
+             * {@code true} to get a significant speed-up in {@code CASE}-like operations.
+             */
+            default boolean safeToEvalInLazy() {
+                return false;
+            }
         }
 
         /**

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/evaluator/EvalMapper.java

@@ -176,6 +176,11 @@ public ExpressionEvaluator get(DriverContext driverContext) {
                 public String toString() {
                     return "Attribute[channel=" + channel + "]";
                 }
+
+                @Override
+                public boolean safeToEvalInLazy() {
+                    return true;
+                }
             }
             return new AttributeFactory(layout.get(attr.id()).channel());
         }
@@ -209,6 +214,11 @@ public ExpressionEvaluator get(DriverContext driverContext) {
                 public String toString() {
                     return "LiteralsEvaluator[lit=" + lit + "]";
                 }
+
+                @Override
+                public boolean safeToEvalInLazy() {
+                    return true;
+                }
             }
             return new LiteralsEvaluatorFactory(lit);
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BooleanBlock;
+import org.elasticsearch.compute.data.BooleanVector;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
@@ -293,6 +294,12 @@ public ExpressionEvaluator.Factory toEvaluator(Function<Expression, ExpressionEv
         return new ExpressionEvaluator.Factory() {
             @Override
             public ExpressionEvaluator get(DriverContext context) {
+                if (conditionsFactories.size() == 1
+                    && conditionsFactories.get(0).value.safeToEvalInLazy()
+                    && elseValueFactory.safeToEvalInLazy()) {
+
+                    return new EagerEvaluator(context, conditionsFactories.get(0).apply(context), elseValueFactory.get(context));
+                }
                 return new CaseEvaluator(
                     context,
                     resultType,
@@ -314,7 +321,7 @@ public String toString() {
         };
     }
 
-    record ConditionEvaluatorSupplier(ExpressionEvaluator.Factory condition, ExpressionEvaluator.Factory value)
+    private record ConditionEvaluatorSupplier(ExpressionEvaluator.Factory condition, ExpressionEvaluator.Factory value)
         implements
             Function<DriverContext, ConditionEvaluator> {
         @Override
@@ -328,7 +335,9 @@ public String toString() {
         }
     }
 
-    record ConditionEvaluator(EvalOperator.ExpressionEvaluator condition, EvalOperator.ExpressionEvaluator value) implements Releasable {
+    private record ConditionEvaluator(EvalOperator.ExpressionEvaluator condition, EvalOperator.ExpressionEvaluator value)
+        implements
+            Releasable {
         @Override
         public void close() {
             Releasables.closeExpectNoException(condition, value);
@@ -393,4 +402,39 @@ public String toString() {
             return "CaseEvaluator[resultType=" + resultType + ", conditions=" + conditions + ", elseVal=" + elseVal + ']';
         }
     }
+
+    private record EagerEvaluator(DriverContext driverContext, ConditionEvaluator condition, EvalOperator.ExpressionEvaluator elseVal)
+        implements
+            EvalOperator.ExpressionEvaluator {
+        @Override
+        public Block eval(Page page) {
+            try (
+                BooleanVector lhsOrRhs = ((BooleanBlock) condition.condition.eval(page)).asVector();
+                // NOCOMMIT replace above with toMask
+                Block lhs = condition.value.eval(page);
+                Block rhs = elseVal.eval(page)
+            ) {
+                try (Block.Builder builder = lhs.elementType().newBlockBuilder(lhs.getTotalValueCount(), driverContext.blockFactory())) {
+                    for (int p = 0; p < lhs.getPositionCount(); p++) {
+                        if (lhsOrRhs.getBoolean(p)) {
+                            builder.copyFrom(lhs, p, p + 1);
+                        } else {
+                            builder.copyFrom(rhs, p, p + 1);
+                        }
+                    }
+                    return builder.build();
+                }
+            }
+        }
+
+        @Override
+        public void close() {
+            Releasables.closeExpectNoException(condition, elseVal);
+        }
+
+        @Override
+        public String toString() {
+            return "CaseEvaluator[condition=" + condition + ", elseVal=" + elseVal + ']';
+        }
+    }
 }