ESQL: Optimize LENGTH for high cardinality keyword (#137586)

Optimize the fused block loader for the `LENGTH` function when applied
to high cardinality keywords. Previously it was loading values doc by
doc, thus jumping around the ordinal dictionary. Now it performs the
lookup on the ords dictionary in ascending order. This drops a fairly
simple per test from ~`360ms` to ~`95ms`. Ascending is indeed much
faster.

Author: nik9000

server/src/main/java/org/elasticsearch/index/mapper/blockloader/docvalues/Utf8CodePointsFromOrdsBlockLoader.java

@@ -78,15 +78,13 @@ public String toString() {
     /**
      * Loads low cardinality singleton ordinals in using a cache of code point counts.
      * <p>
-     *     If we haven't cached the counts for all ordinals then the process looks like:
+     *     It's very important to look up ordinals in ascending order. So, if we haven't
+     *     cached the counts for all ordinals, then the process looks like:
      * </p>
      * <ol>
-     *     <li>Build an int[] containing the ordinals</li>
-     *     <li>
-     *         Sort a copy of the int[] and load the cache for each ordinal. The sorting
-     *         is important here because ordinals are faster to resolved in ascending order.
-     *     </li>
-     *     <li>Walk the int[] in order, reading from the cache to build the page</li>
+     *     <li>Build an {@code int[]} containing the ordinals</li>
+     *     <li>Sort a copy of the {@code int[]} and load the count into the cache for each ordinal.</li>
+     *     <li>Walk the unsorted {@code int[]} reading from the cache to build the page</li>
      * </ol>
      * <p>
      *     If we <strong>have</strong> cached the counts for all ordinals we load the
@@ -178,6 +176,9 @@ private int[] readOrds(BlockFactory factory, Docs docs, int offset) throws IOExc
             }
         }
 
+        /**
+         * Fill the cache for all ords. We skip values {@code -1} which represent "no data".
+         */
         private void fillCache(BlockFactory factory, int[] ords) throws IOException {
             factory.adjustBreaker(RamUsageEstimator.sizeOf(ords));
             try {
@@ -196,6 +197,10 @@ private void fillCache(BlockFactory factory, int[] ords) throws IOException {
             }
         }
 
+        /**
+         * Build the results for a list of documents directly from the cache. We use this
+         * if we're sure that all ordinals we're going to load are already cached.
+         */
         private Block buildFromFilledCache(BlockFactory factory, Docs docs, int offset) throws IOException {
             int count = docs.count() - offset;
             try (IntBuilder builder = factory.ints(count)) {
@@ -211,6 +216,10 @@ private Block buildFromFilledCache(BlockFactory factory, Docs docs, int offset)
             }
         }
 
+        /**
+         * Get the count of code points at the ord, reading from the cache if possible.
+         * The {@code ord} must be {@code >= 0} or this will fail.
+         */
         private int codePointsAtOrd(int ord) throws IOException {
             if (cache[ord] >= 0) {
                 return cache[ord];
@@ -253,7 +262,7 @@ public Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFilt
                 return buildFromFilledCache(factory, docs, offset);
             }
 
-            int[] ords = readOrds(factory, docs, offset);
+            int[] ords = readOrds(ordinals, warnings, factory, docs, offset);
             try {
                 fillCache(factory, ords);
                 return buildFromCache(factory, cache, ords);
@@ -297,36 +306,9 @@ private Block blockForSingleDoc(BlockFactory factory, int docId) throws IOExcept
             return factory.constantNulls(1);
         }
 
-        private int[] readOrds(BlockFactory factory, Docs docs, int offset) throws IOException {
-            int count = docs.count() - offset;
-            long size = sizeOfArray(count);
-            factory.adjustBreaker(size);
-            int[] ords = null;
-            try {
-                ords = new int[docs.count() - offset];
-                for (int i = offset; i < docs.count(); i++) {
-                    int doc = docs.get(i);
-                    if (ordinals.advanceExact(doc) == false) {
-                        ords[i] = -1;
-                        continue;
-                    }
-                    if (ordinals.docValueCount() != 1) {
-                        registerSingleValueWarning(warnings);
-                        ords[i] = -1;
-                        continue;
-                    }
-                    ords[i] = Math.toIntExact(ordinals.nextOrd());
-                }
-                int[] result = ords;
-                ords = null;
-                return result;
-            } finally {
-                if (ords != null) {
-                    factory.adjustBreaker(-size);
-                }
-            }
-        }
-
+        /**
+         * Fill the cache for all ords. We skip values {@code -1} which represent "no data".
+         */
         private void fillCache(BlockFactory factory, int[] ords) throws IOException {
             factory.adjustBreaker(RamUsageEstimator.sizeOf(ords));
             try {
@@ -345,6 +327,10 @@ private void fillCache(BlockFactory factory, int[] ords) throws IOException {
             }
         }
 
+        /**
+         * Build the results for a list of documents directly from the cache. We use this
+         * if we're sure that all ordinals we're going to load are already cached.
+         */
         private Block buildFromFilledCache(BlockFactory factory, Docs docs, int offset) throws IOException {
             int count = docs.count() - offset;
             try (IntBuilder builder = factory.ints(count)) {
@@ -365,6 +351,10 @@ private Block buildFromFilledCache(BlockFactory factory, Docs docs, int offset)
             }
         }
 
+        /**
+         * Get the count of code points at the ord, reading from the cache if possible.
+         * The {@code ord} must be {@code >= 0} or this will fail.
+         */
         private int codePointsAtOrd(int ord) throws IOException {
             if (cache[ord] >= 0) {
                 return cache[ord];
@@ -378,8 +368,21 @@ private int codePointsAtOrd(int ord) throws IOException {
     }
 
     /**
-     * Loads a count of utf-8 code points for each ordinal doc by doc, without a cache. We use this when there
+     * Loads a count of utf-8 code points for each ordinal without a cache. We use this when there
      * are many unique doc values and the cache hit rate is unlikely to be high.
+     * <p>
+     *     It's very important to read values in sorted order so we:
+     * </p>
+     * <ul>
+     *     <li>Load the ordinals into an {@code int[]}, using -1 for "empty" values</li>
+     *     <li>Create a sorted copy of the {@code int[]}</li>
+     *     <li>Compact the sorted {@code int[]}s into a sorted list of unique, non "empty" ordinals</li>
+     *     <li>Count the code points for each of the sorted, compacted ordinals</li>
+     *     <li>
+     *         Walk the original ordinals {@code int[]} which are in doc order, building a {@link Block}
+     *         of counts.
+     *     </li>
+     * </ul>
      */
     private static class ImmediateOrdinals extends BlockDocValuesReader {
         private final Warnings warnings;
@@ -395,11 +398,32 @@ public Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFilt
             if (docs.count() - offset == 1) {
                 return blockForSingleDoc(factory, docs.get(offset));
             }
-            try (IntBuilder builder = factory.ints(docs.count() - offset)) {
-                for (int i = offset; i < docs.count(); i++) {
-                    read(docs.get(i), builder);
+
+            int[] ords = readOrds(ordinals, warnings, factory, docs, offset);
+            int[] sortedOrds = null;
+            int[] counts = null;
+            try {
+                sortedOrds = sortedOrds(factory, ords);
+                int compactedLength = compactSorted(sortedOrds);
+                counts = counts(factory, sortedOrds, compactedLength);
+                try (IntBuilder builder = factory.ints(ords.length)) {
+                    for (int ord : ords) {
+                        if (ord >= 0) {
+                            builder.appendInt(counts[Arrays.binarySearch(sortedOrds, 0, compactedLength, ord)]);
+                        } else {
+                            builder.appendNull();
+                        }
+                    }
+                    return builder.build();
+                }
+            } finally {
+                factory.adjustBreaker(-RamUsageEstimator.shallowSizeOf(ords));
+                if (sortedOrds != null) {
+                    factory.adjustBreaker(-RamUsageEstimator.shallowSizeOf(sortedOrds));
+                }
+                if (counts != null) {
+                    factory.adjustBreaker(-RamUsageEstimator.shallowSizeOf(counts));
                 }
-                return builder.build();
             }
         }
 
@@ -442,11 +466,93 @@ private Block blockForSingleDoc(BlockFactory factory, int docId) throws IOExcept
             return factory.constantNulls(1);
         }
 
+        /**
+         * Builds a sorted copy of the loaded ordinals.
+         */
+        private int[] sortedOrds(BlockFactory factory, int[] ords) {
+            factory.adjustBreaker(RamUsageEstimator.sizeOf(ords));
+            int[] sortedOrds = ords.clone();
+            Arrays.sort(sortedOrds);
+            return sortedOrds;
+        }
+
+        /**
+         * Compacts the array of sorted ordinals into an array of populated ({@code >= 0}), unique ordinals.
+         * @return the length of the unique array
+         */
+        private int compactSorted(int[] sortedOrds) {
+            int c = 0;
+            int i = 0;
+            while (i < sortedOrds.length && sortedOrds[i] < 0) {
+                i++;
+            }
+            while (i < sortedOrds.length) {
+                if (false == (i > 0 && sortedOrds[i - 1] == sortedOrds[i])) {
+                    sortedOrds[c++] = sortedOrds[i];
+                }
+                i++;
+            }
+            return c;
+        }
+
+        private int[] counts(BlockFactory factory, int[] compactedSortedOrds, int compactedLength) throws IOException {
+            long size = sizeOfArray(compactedLength);
+            factory.adjustBreaker(size);
+            int[] counts = new int[compactedLength];
+            for (int i = 0; i < counts.length; i++) {
+                counts[i] = codePointsAtOrd(compactedSortedOrds[i]);
+            }
+            return counts;
+        }
+
+        /**
+         * Get the count of code points at the ord.
+         * The {@code ord} must be {@code >= 0} or this will fail.
+         */
         private int codePointsAtOrd(long ord) throws IOException {
             return UnicodeUtil.codePointCount(ordinals.lookupOrd(ord));
         }
     }
 
+    /**
+     * Load an ordinal for each position. Three cases:
+     * <ul>
+     *     <li>There is a single ordinal at this position - load the ordinals value in to the array</li>
+     *     <li>There are no values at this position - load a -1 - we'll skip loading that later</li>
+     *     <li>There are <strong>many</strong> values at this position - load a -1 which we'll skip like above - and emit a warning</li>
+     * </ul>
+     */
+    private static int[] readOrds(SortedSetDocValues ordinals, Warnings warnings, BlockFactory factory, Docs docs, int offset)
+        throws IOException {
+        int count = docs.count() - offset;
+        long size = sizeOfArray(count);
+        factory.adjustBreaker(size);
+        int[] ords = null;
+        try {
+            ords = new int[docs.count() - offset];
+            for (int i = offset; i < docs.count(); i++) {
+                int doc = docs.get(i);
+                if (ordinals.advanceExact(doc) == false) {
+                    ords[i] = -1;
+                    continue;
+                }
+                if (ordinals.docValueCount() != 1) {
+                    registerSingleValueWarning(warnings);
+                    ords[i] = -1;
+                    continue;
+                }
+                ords[i] = Math.toIntExact(ordinals.nextOrd());
+            }
+            int[] result = ords;
+            ords = null;
+            return result;
+        } finally {
+            if (ords != null) {
+                factory.adjustBreaker(-size);
+            }
+        }
+    }
+
     private static Block buildFromCache(BlockFactory factory, int[] cache, int[] ords) {
         try (IntBuilder builder = factory.ints(ords.length)) {
             for (int ord : ords) {