ESQL: Keep ordinals in conversion functions (#125357) (#127933)

Make the conversion functions that process `BytesRef`s into `BytesRefs`
keep the `OrdinalBytesRefVector`s when processing. Let's use `TO_LOWER`
as an example. First, the performance numbers:
```
  (operation)  Mode   Score   Error ->  Score    Error Units
     to_lower  30.662 ± 6.163 -> 30.048 ±  0.479 ns/op
to_lower_ords  30.773 ± 0.370 ->  0.025 ±  0.001 ns/op
     to_upper  33.552 ± 0.529 -> 35.775 ±  1.799 ns/op
to_upper_ords  35.791 ± 0.658 ->  0.027 ±  0.001 ns/op
```
The test has a 8192 positions containing alternating `foo` and `bar`.
Running `TO_LOWER` via ordinals is super duper faster. No longer
`O(positions)` and now `O(unique_values)`.

Let's paint some pictures! `OrdinalBytesRefVector` is a lookup table.
Like this:
```
+-------+----------+
| bytes | ordinals |
| ----- | -------- |
|  FOO  | 0        |
|  BAR  | 1        |
|  BAZ  | 2        |
+-------+ 1        |
        | 1        |
        | 0        |
        +----------+
```

That lookup table is one block. When you read it you look up the
`ordinal` and match it to the `bytes`. Previously `TO_LOWER` would
process each value one at a time and make:
```
bytes
-----
 foo
 bar
 baz
 bar
 bar
 foo
```

So it'd run `TO_LOWER` once per `ordinal` and it'd make an ordinal
non-lookup table. With this change `TO_LOWER` will now make:
```
+-------+----------+
| bytes | ordinals |
| ----- | -------- |
|  foo  | 0        |
|  bar  | 1        |
|  baz  | 2        |
+-------+ 1        |
        | 1        |
        | 0        |
        +----------+
```
We don't even have to copy the `ordinals` - we can reuse those from the
input and just bump the reference count. That's why this goes from
`O(positions)` to `O(unique_values)`.

Author: nik9000

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

@@ -11,15 +11,19 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BooleanVector;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.compute.data.LongVector;
+import org.elasticsearch.compute.data.OrdinalBytesRefVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator;
@@ -40,9 +44,13 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
 import org.elasticsearch.xpack.esql.expression.function.scalar.string.RLike;
+import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToLower;
+import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToUpper;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
 import org.elasticsearch.xpack.esql.planner.Layout;
+import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
+import org.elasticsearch.xpack.esql.session.Configuration;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
 import org.openjdk.jmh.annotations.Fork;
@@ -56,8 +64,10 @@
 import org.openjdk.jmh.annotations.Warmup;
 
 import java.time.Duration;
+import java.time.ZoneOffset;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
@@ -106,7 +116,11 @@ public class EvalBenchmark {
             "long_equal_to_int",
             "mv_min",
             "mv_min_ascending",
-            "rlike" }
+            "rlike",
+            "to_lower",
+            "to_lower_ords",
+            "to_upper",
+            "to_upper_ords" }
     )
     public String operation;
 
@@ -214,6 +228,16 @@ private static EvalOperator.ExpressionEvaluator evaluator(String operation) {
                 RLike rlike = new RLike(Source.EMPTY, keywordField, new RLikePattern(".ar"));
                 yield EvalMapper.toEvaluator(FOLD_CONTEXT, rlike, layout(keywordField)).get(driverContext);
             }
+            case "to_lower", "to_lower_ords" -> {
+                FieldAttribute keywordField = keywordField();
+                ToLower toLower = new ToLower(Source.EMPTY, keywordField, configuration());
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, toLower, layout(keywordField)).get(driverContext);
+            }
+            case "to_upper", "to_upper_ords" -> {
+                FieldAttribute keywordField = keywordField();
+                ToUpper toUpper = new ToUpper(Source.EMPTY, keywordField, configuration());
+                yield EvalMapper.toEvaluator(FOLD_CONTEXT, toUpper, layout(keywordField)).get(driverContext);
+            }
             default -> throw new UnsupportedOperationException();
         };
     }
@@ -366,10 +390,35 @@ private static void checkExpected(String operation, Page actual) {
                     }
                 }
             }
+            case "to_lower" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
+            case "to_lower_ords" -> checkBytes(operation, actual, true, new BytesRef[] { new BytesRef("foo"), new BytesRef("bar") });
+            case "to_upper" -> checkBytes(operation, actual, false, new BytesRef[] { new BytesRef("FOO"), new BytesRef("BAR") });
+            case "to_upper_ords" -> checkBytes(operation, actual, true, new BytesRef[] { new BytesRef("FOO"), new BytesRef("BAR") });
             default -> throw new UnsupportedOperationException(operation);
         }
     }
 
+    private static void checkBytes(String operation, Page actual, boolean expectOrds, BytesRef[] expectedVals) {
+        BytesRef scratch = new BytesRef();
+        BytesRefVector v = actual.<BytesRefBlock>getBlock(1).asVector();
+        for (int i = 0; i < BLOCK_LENGTH; i++) {
+            BytesRef expected = expectedVals[i % 2];
+            BytesRef b = v.getBytesRef(i, scratch);
+            if (b.equals(expected) == false) {
+                throw new AssertionError("[" + operation + "] expected [" + expected + "] but was [" + b + "]");
+            }
+        }
+        if (expectOrds) {
+            if (v.asOrdinals() == null) {
+                throw new IllegalArgumentException("expected ords but got " + v);
+            }
+        } else {
+            if (v.asOrdinals() != null) {
+                throw new IllegalArgumentException("expected non-ords but got " + v);
+            }
+        }
+    }
+
     private static Page page(String operation) {
         return switch (operation) {
             case "abs", "add", "date_trunc", "equal_to_const" -> {
@@ -448,6 +497,16 @@ private static Page page(String operation) {
                 }
                 yield new Page(builder.build().asBlock());
             }
+            case "to_lower_ords", "to_upper_ords" -> {
+                var bytes = blockFactory.newBytesRefVectorBuilder(BLOCK_LENGTH);
+                bytes.appendBytesRef(new BytesRef("foo"));
+                bytes.appendBytesRef(new BytesRef("bar"));
+                var ordinals = blockFactory.newIntVectorFixedBuilder(BLOCK_LENGTH);
+                for (int i = 0; i < BLOCK_LENGTH; i++) {
+                    ordinals.appendInt(i % 2);
+                }
+                yield new Page(new OrdinalBytesRefVector(ordinals.build(), bytes.build()).asBlock());
+            }
             default -> throw new UnsupportedOperationException();
         };
     }
@@ -470,4 +529,21 @@ private static void run(String operation) {
             checkExpected(operation, output);
         }
     }
+
+    private static Configuration configuration() {
+        return new Configuration(
+            ZoneOffset.UTC,
+            Locale.ROOT,
+            null,
+            null,
+            null,
+            EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.get(Settings.EMPTY),
+            EsqlPlugin.QUERY_RESULT_TRUNCATION_DEFAULT_SIZE.get(Settings.EMPTY),
+            null,
+            false,
+            Map.of(),
+            0,
+            false
+        );
+    }
 }

docs/changelog/125357.yaml

@@ -0,0 +1,5 @@
+pr: 125357
+summary: Keep ordinals in conversion functions
+area: ES|QL
+type: enhancement
+issues: []

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

@@ -28,9 +28,12 @@
 import static org.elasticsearch.compute.gen.Types.ABSTRACT_CONVERT_FUNCTION_EVALUATOR;
 import static org.elasticsearch.compute.gen.Types.BLOCK;
 import static org.elasticsearch.compute.gen.Types.BYTES_REF;
+import static org.elasticsearch.compute.gen.Types.BYTES_REF_VECTOR_BUILDER;
 import static org.elasticsearch.compute.gen.Types.DRIVER_CONTEXT;
 import static org.elasticsearch.compute.gen.Types.EXPRESSION_EVALUATOR;
 import static org.elasticsearch.compute.gen.Types.EXPRESSION_EVALUATOR_FACTORY;
+import static org.elasticsearch.compute.gen.Types.INT_VECTOR;
+import static org.elasticsearch.compute.gen.Types.ORDINALS_BYTES_REF_VECTOR;
 import static org.elasticsearch.compute.gen.Types.SOURCE;
 import static org.elasticsearch.compute.gen.Types.VECTOR;
 import static org.elasticsearch.compute.gen.Types.blockType;
@@ -41,7 +44,7 @@ public class ConvertEvaluatorImplementer {
 
     private final TypeElement declarationType;
     private final EvaluatorImplementer.ProcessFunction processFunction;
-    private final String extraName;
+    private final boolean canProcessOrdinals;
     private final ClassName implementation;
     private final TypeName argumentType;
     private final List<TypeMirror> warnExceptions;
@@ -55,6 +58,10 @@ public ConvertEvaluatorImplementer(
     ) {
         this.declarationType = (TypeElement) processFunction.getEnclosingElement();
         this.processFunction = new EvaluatorImplementer.ProcessFunction(types, processFunction, warnExceptions);
+        this.canProcessOrdinals = warnExceptions.isEmpty()
+            && this.processFunction.returnType().equals(BYTES_REF)
+            && this.processFunction.args.get(0) instanceof EvaluatorImplementer.StandardProcessFunctionArg s
+            && s.type().equals(BYTES_REF);
 
         if (this.processFunction.args.get(0) instanceof EvaluatorImplementer.StandardProcessFunctionArg == false) {
             throw new IllegalArgumentException("first argument must be the field to process");
@@ -66,7 +73,6 @@ public ConvertEvaluatorImplementer(
             }
         }
 
-        this.extraName = extraName;
         this.argumentType = TypeName.get(processFunction.getParameters().get(0).asType());
         this.warnExceptions = warnExceptions;
 
@@ -102,6 +108,9 @@ private TypeSpec type() {
         builder.addMethod(evalValue(true));
         builder.addMethod(evalBlock());
         builder.addMethod(evalValue(false));
+        if (canProcessOrdinals) {
+            builder.addMethod(evalOrdinals());
+        }
         builder.addMethod(processFunction.toStringMethod(implementation));
         builder.addMethod(processFunction.close());
         builder.addType(factory());
@@ -132,6 +141,15 @@ private MethodSpec evalVector() {
 
         TypeName vectorType = vectorType(argumentType);
         builder.addStatement("$T vector = ($T) v", vectorType, vectorType);
+        if (canProcessOrdinals) {
+            builder.addStatement("$T ordinals = vector.asOrdinals()", ORDINALS_BYTES_REF_VECTOR);
+            builder.beginControlFlow("if (ordinals != null)");
+            {
+                builder.addStatement("return evalOrdinals(ordinals)");
+            }
+            builder.endControlFlow();
+        }
+
         builder.addStatement("int positionCount = v.getPositionCount()");
 
         String scratchPadName = argumentType.equals(BYTES_REF) ? "scratchPad" : null;
@@ -299,6 +317,31 @@ private MethodSpec evalValue(boolean forVector) {
         return builder.build();
     }
 
+    private MethodSpec evalOrdinals() {
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("evalOrdinals").addModifiers(Modifier.PRIVATE);
+        builder.addParameter(ORDINALS_BYTES_REF_VECTOR, "v").returns(BLOCK);
+
+        builder.addStatement("int positionCount = v.getDictionaryVector().getPositionCount()");
+        builder.addStatement("BytesRef scratchPad = new BytesRef()");
+        builder.beginControlFlow(
+            "try ($T builder = driverContext.blockFactory().newBytesRefVectorBuilder(positionCount))",
+            BYTES_REF_VECTOR_BUILDER
+        );
+        {
+            builder.beginControlFlow("for (int p = 0; p < positionCount; p++)");
+            {
+                builder.addStatement("builder.appendBytesRef($N)", evalValueCall("v.getDictionaryVector()", "p", "scratchPad"));
+            }
+            builder.endControlFlow();
+            builder.addStatement("$T ordinals = v.getOrdinalsVector()", INT_VECTOR);
+            builder.addStatement("ordinals.incRef()");
+            builder.addStatement("return new $T(ordinals, builder.build()).asBlock()", ORDINALS_BYTES_REF_VECTOR);
+        }
+        builder.endControlFlow();
+
+        return builder.build();
+    }
+
     private TypeSpec factory() {
         TypeSpec.Builder builder = TypeSpec.classBuilder("Factory");
         builder.addSuperinterface(EXPRESSION_EVALUATOR_FACTORY);

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

@@ -61,6 +61,7 @@ public class Types {
 
     static final ClassName BOOLEAN_VECTOR = ClassName.get(DATA_PACKAGE, "BooleanVector");
     static final ClassName BYTES_REF_VECTOR = ClassName.get(DATA_PACKAGE, "BytesRefVector");
+    static final ClassName ORDINALS_BYTES_REF_VECTOR = ClassName.get(DATA_PACKAGE, "OrdinalBytesRefVector");
     static final ClassName INT_VECTOR = ClassName.get(DATA_PACKAGE, "IntVector");
     static final ClassName LONG_VECTOR = ClassName.get(DATA_PACKAGE, "LongVector");
     static final ClassName DOUBLE_VECTOR = ClassName.get(DATA_PACKAGE, "DoubleVector");

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java

@@ -13,17 +13,23 @@
 import org.elasticsearch.compute.data.BlockUtils;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
 import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.compute.data.FloatBlock;
 import org.elasticsearch.compute.data.IntBlock;
 import org.elasticsearch.compute.data.LongBlock;
+import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
 import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasables;
 import org.hamcrest.Matcher;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
 import static org.elasticsearch.test.ESTestCase.between;
@@ -268,4 +274,67 @@ public static List<List<Object>> valuesAtPositions(Block block, int from, int to
         }
         return result;
     }
+
+    /**
+     * Convert all of the {@link Block}s in a page that contain {@link BytesRef}s into
+     * {@link OrdinalBytesRefBlock}s.
+     */
+    public static Page convertBytesRefsToOrdinals(Page page) {
+        Block[] blocks = new Block[page.getBlockCount()];
+        try {
+            for (int b = 0; b < page.getBlockCount(); b++) {
+                Block block = page.getBlock(b);
+                if (block.elementType() != ElementType.BYTES_REF) {
+                    blocks[b] = block;
+                    continue;
+                }
+                Map<BytesRef, Integer> dedupe = new HashMap<>();
+                BytesRefBlock bytesRefBlock = (BytesRefBlock) block;
+                try (
+                    IntBlock.Builder ordinals = block.blockFactory().newIntBlockBuilder(block.getPositionCount());
+                    BytesRefVector.Builder bytes = block.blockFactory().newBytesRefVectorBuilder(block.getPositionCount())
+                ) {
+                    BytesRef scratch = new BytesRef();
+                    for (int p = 0; p < block.getPositionCount(); p++) {
+                        int first = block.getFirstValueIndex(p);
+                        int count = block.getValueCount(p);
+                        if (count == 0) {
+                            ordinals.appendNull();
+                            continue;
+                        }
+                        if (count == 1) {
+                            BytesRef v = bytesRefBlock.getBytesRef(first, scratch);
+                            ordinals.appendInt(dedupe(dedupe, bytes, v));
+                            continue;
+                        }
+                        int end = first + count;
+                        ordinals.beginPositionEntry();
+                        for (int i = first; i < end; i++) {
+                            BytesRef v = bytesRefBlock.getBytesRef(i, scratch);
+                            ordinals.appendInt(dedupe(dedupe, bytes, v));
+                        }
+                        ordinals.endPositionEntry();
+                    }
+                    blocks[b] = new OrdinalBytesRefBlock(ordinals.build(), bytes.build());
+                    bytesRefBlock.decRef();
+                }
+            }
+            Page p = new Page(blocks);
+            Arrays.fill(blocks, null);
+            return p;
+        } finally {
+            Releasables.close(blocks);
+        }
+    }
+
+    private static int dedupe(Map<BytesRef, Integer> dedupe, BytesRefVector.Builder bytes, BytesRef v) {
+        Integer current = dedupe.get(v);
+        if (current != null) {
+            return current;
+        }
+        bytes.appendBytesRef(v);
+        int o = dedupe.size();
+        dedupe.put(v, o);
+        return o;
+    }
 }