Speed up loading dense singleton keyword fields (#132994)

With this change both sorted set and number doc values use the same bulk loading for values/ordinals.

This PR supersedes (#132715) that also sped up loading dense singleton keyword fields, but duplicated the bulk encoding logic.

Author: martijnvg

server/src/main/java/org/elasticsearch/index/codec/tsdb/es819/ES819TSDBDocValuesProducer.java

@@ -345,6 +345,10 @@ public SortedDocValues getSorted(FieldInfo field) throws IOException {
     }
 
     private SortedDocValues getSorted(SortedEntry entry, boolean valuesSorted) throws IOException {
+        if (entry.ordsEntry.docsWithFieldOffset == -2) {
+            return DocValues.emptySorted();
+        }
+
         final NumericDocValues ords = getNumeric(entry.ordsEntry, entry.termsDictEntry.termsDictSize);
         return new BaseSortedDocValues(entry) {
 
@@ -380,7 +384,25 @@ public long cost() {
 
             @Override
             public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
-                if (valuesSorted && ords instanceof BaseDenseNumericValues denseOrds) {
+                if (ords instanceof BaseDenseNumericValues denseOrds) {
+                    var block = tryReadAHead(factory, docs, offset);
+                    if (block != null) {
+                        return block;
+                    }
+                    // Falling back to tryRead(...) is safe here, given that current block index wasn't altered by looking ahead.
+                    try (var builder = factory.singletonOrdinalsBuilder(this, docs.count() - offset, true)) {
+                        BlockLoader.SingletonLongBuilder delegate = new SingletonLongToSingletonOrdinalDelegate(builder);
+                        var result = denseOrds.tryRead(delegate, docs, offset);
+                        if (result != null) {
+                            return result;
+                        }
+                    }
+                }
+                return null;
+            }
+
+            BlockLoader.Block tryReadAHead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+                if (ords instanceof BaseDenseNumericValues denseOrds && (valuesSorted || entry.termsDictEntry.termsDictSize == 1)) {
                     int firstDoc = docs.get(offset);
                     denseOrds.advanceExact(firstDoc);
                     long startValue = denseOrds.longValue();
@@ -438,10 +460,18 @@ public TermsEnum termsEnum() throws IOException {
         public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
             return null;
         }
+
+        BlockLoader.Block tryReadAHead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+            return null;
+        }
     }
 
     abstract static class BaseDenseNumericValues extends NumericDocValues implements BlockLoader.OptionalColumnAtATimeReader {
         abstract long lookAheadValueAt(int targetDoc) throws IOException;
+
+        BlockLoader.Block tryRead(BlockLoader.SingletonLongBuilder builder, BlockLoader.Docs docs, int offset) throws IOException {
+            return null;
+        }
     }
 
     abstract static class BaseSortedSetDocValues extends SortedSetDocValues {
@@ -1256,41 +1286,49 @@ public long longValue() throws IOException {
 
                 @Override
                 public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
-                    assert maxOrd == -1 : "unexpected maxOrd[" + maxOrd + "]";
+                    try (BlockLoader.SingletonLongBuilder builder = factory.singletonLongs(docs.count() - offset)) {
+                        return tryRead(builder, docs, offset);
+                    }
+                }
+
+                @Override
+                BlockLoader.Block tryRead(BlockLoader.SingletonLongBuilder builder, BlockLoader.Docs docs, int offset) throws IOException {
                     final int docsCount = docs.count();
                     doc = docs.get(docsCount - 1);
-                    try (BlockLoader.SingletonLongBuilder builder = factory.singletonLongs(docs.count() - offset)) {
-                        for (int i = offset; i < docsCount;) {
-                            int index = docs.get(i);
-                            final int blockIndex = index >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT;
-                            final int blockInIndex = index & ES819TSDBDocValuesFormat.NUMERIC_BLOCK_MASK;
-                            if (blockIndex != currentBlockIndex) {
-                                assert blockIndex > currentBlockIndex : blockIndex + " < " + currentBlockIndex;
-                                // no need to seek if the loading block is the next block
-                                if (currentBlockIndex + 1 != blockIndex) {
-                                    valuesData.seek(indexReader.get(blockIndex));
-                                }
-                                currentBlockIndex = blockIndex;
+                    for (int i = offset; i < docsCount;) {
+                        int index = docs.get(i);
+                        final int blockIndex = index >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT;
+                        final int blockInIndex = index & ES819TSDBDocValuesFormat.NUMERIC_BLOCK_MASK;
+                        if (blockIndex != currentBlockIndex) {
+                            assert blockIndex > currentBlockIndex : blockIndex + " < " + currentBlockIndex;
+                            // no need to seek if the loading block is the next block
+                            if (currentBlockIndex + 1 != blockIndex) {
+                                valuesData.seek(indexReader.get(blockIndex));
+                            }
+                            currentBlockIndex = blockIndex;
+                            if (bitsPerOrd == -1) {
                                 decoder.decode(valuesData, currentBlock);
+                            } else {
+                                decoder.decodeOrdinals(valuesData, currentBlock, bitsPerOrd);
                             }
+                        }
 
-                            // Try to append more than just one value:
-                            // Instead of iterating over docs and find the max length, take an optimistic approach to avoid as
-                            // many comparisons as there are remaining docs and instead do at most 7 comparisons:
-                            int length = 1;
-                            int remainingBlockLength = Math.min(ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE - blockInIndex, docsCount - i);
-                            for (int newLength = remainingBlockLength; newLength > 1; newLength = newLength >> 1) {
-                                int lastIndex = i + newLength - 1;
-                                if (isDense(index, docs.get(lastIndex), newLength)) {
-                                    length = newLength;
-                                    break;
-                                }
+                        // Try to append more than just one value:
+                        // Instead of iterating over docs and find the max length, take an optimistic approach to avoid as
+                        // many comparisons as there are remaining docs and instead do at most 7 comparisons:
+                        int length = 1;
+                        int remainingBlockLength = Math.min(ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE - blockInIndex, docsCount - i);
+                        for (int newLength = remainingBlockLength; newLength > 1; newLength = newLength >> 1) {
+                            int lastIndex = i + newLength - 1;
+                            if (isDense(index, docs.get(lastIndex), newLength)) {
+                                length = newLength;
+                                break;
                             }
-                            builder.appendLongs(currentBlock, blockInIndex, length);
-                            i += length;
                         }
-                        return builder.build();
+                        builder.appendLongs(currentBlock, blockInIndex, length);
+                        i += length;
                     }
+                    return builder.build();
                 }
 
                 @Override
@@ -1624,4 +1662,59 @@ private static class TermsDictEntry {
         int maxBlockLength;
     }
 
+    static final class SingletonLongToSingletonOrdinalDelegate implements BlockLoader.SingletonLongBuilder {
+        private final BlockLoader.SingletonOrdinalsBuilder builder;
+
+        SingletonLongToSingletonOrdinalDelegate(BlockLoader.SingletonOrdinalsBuilder builder) {
+            this.builder = builder;
+        }
+
+        @Override
+        public BlockLoader.SingletonLongBuilder appendLong(long value) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public BlockLoader.SingletonLongBuilder appendLongs(long[] values, int from, int length) {
+            // Unfortunately, no array copy here...
+            // Since we need to loop here, let's also keep track of min/max.
+            int minOrd = Integer.MAX_VALUE;
+            int maxOrd = Integer.MIN_VALUE;
+            int counter = 0;
+            int[] convertedOrds = new int[length];
+            int end = from + length;
+            for (int j = from; j < end; j++) {
+                int ord = Math.toIntExact(values[j]);
+                convertedOrds[counter++] = ord;
+                minOrd = Math.min(minOrd, ord);
+                maxOrd = Math.max(maxOrd, ord);
+            }
+            builder.appendOrds(convertedOrds, 0, length, minOrd, maxOrd);
+            return this;
+        }
+
+        @Override
+        public BlockLoader.Block build() {
+            return builder.build();
+        }
+
+        @Override
+        public BlockLoader.Builder appendNull() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public BlockLoader.Builder beginPositionEntry() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public BlockLoader.Builder endPositionEntry() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void close() {}
+    }
+
 }

server/src/main/java/org/elasticsearch/index/mapper/BlockDocValuesReader.java

@@ -756,7 +756,7 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset) throw
                     return block;
                 }
             }
-            try (var builder = factory.singletonOrdinalsBuilder(ordinals, docs.count() - offset)) {
+            try (var builder = factory.singletonOrdinalsBuilder(ordinals, docs.count() - offset, false)) {
                 for (int i = offset; i < docs.count(); i++) {
                     int doc = docs.get(i);
                     if (doc < ordinals.docID()) {

server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java

@@ -447,7 +447,7 @@ interface BlockFactory {
         /**
          * Build a reader for reading {@link SortedDocValues}
          */
-        SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count);
+        SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count, boolean isDense);
 
         /**
          * Build a reader for reading {@link SortedSetDocValues}
@@ -548,6 +548,8 @@ interface SingletonOrdinalsBuilder extends Builder {
          * Appends an ordinal to the builder.
          */
         SingletonOrdinalsBuilder appendOrd(int value);
+
+        SingletonOrdinalsBuilder appendOrds(int[] values, int from, int length, int minOrd, int maxOrd);
     }
 
     interface SortedSetOrdinalsBuilder extends Builder {