Push compute engine value loading for longs down to tsdb codec. (#132622)

This is the first of many changes that pushes loading of field values to the es819 doc values codec in case of logsdb/tsdb and when the field supports it.

This change first targets reading field values in bulk mode at codec level when doc values type is numeric doc values or sorted doc values, there is only one value per document, and the field is dense (all documents have a value). Multivalued and sparse fields are more complex to support bulk reading for, but it is possible.

With this change, the following field types will support bulk read mode at codec level under the described conditions: long, date, geo_point, point and unsigned_long.

Other number types like integer, short, double, float, scaled_float will be supported in a followup, but would be similar to long based fields, but required an additional conversion step to either an int or float vector.

This change originates from #132460 (which adds bulk reading to `@timestamp`, `_tsid` and dimension fields) and is basically the timestamp support part of it. In another followup, support for single valued, dense sorted (set) doc values will be added for field like _tsid.

Relates to #128445

Author: martijnvg

docs/changelog/132622.yaml

@@ -0,0 +1,5 @@
+pr: 132622
+summary: Push down compute engine value loading of long based singleton numeric doc value to the es819 tsdb doc values codec.
+area: "Codec"
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/index/codec/tsdb/DocValuesForUtil.java

@@ -16,7 +16,7 @@
 
 import java.io.IOException;
 
-public class DocValuesForUtil {
+public final class DocValuesForUtil {
     private static final int BITS_IN_FOUR_BYTES = 4 * Byte.SIZE;
     private static final int BITS_IN_FIVE_BYTES = 5 * Byte.SIZE;
     private static final int BITS_IN_SIX_BYTES = 6 * Byte.SIZE;

server/src/main/java/org/elasticsearch/index/codec/tsdb/TSDBDocValuesEncoder.java

@@ -54,7 +54,7 @@
  *
  * Of course, decoding follows the opposite order with respect to encoding.
  */
-public class TSDBDocValuesEncoder {
+public final class TSDBDocValuesEncoder {
     private final DocValuesForUtil forUtil;
     private final int numericBlockSize;
 

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.tsdb.es819;
+
+import org.apache.lucene.index.NumericDocValues;
+import org.elasticsearch.index.mapper.BlockLoader;
+
+import java.io.IOException;
+
+/**
+ * An es819 doc values specialization that allows bulk loading of values that is optimized in the context of compute engine.
+ */
+public abstract class BulkNumericDocValues extends NumericDocValues {
+
+    /**
+     * Reads the values of all documents in {@code docs}.
+     */
+    public abstract BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException;
+
+}

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

@@ -43,6 +43,7 @@
 import org.apache.lucene.util.packed.PackedInts;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.codec.tsdb.TSDBDocValuesEncoder;
+import org.elasticsearch.index.mapper.BlockLoader;
 
 import java.io.IOException;
 
@@ -1140,7 +1141,7 @@ public long longValue() {
         final int bitsPerOrd = maxOrd >= 0 ? PackedInts.bitsRequired(maxOrd - 1) : -1;
         if (entry.docsWithFieldOffset == -1) {
             // dense
-            return new NumericDocValues() {
+            return new BulkNumericDocValues() {
 
                 private final int maxDoc = ES819TSDBDocValuesProducer.this.maxDoc;
                 private int doc = -1;
@@ -1197,6 +1198,53 @@ public long longValue() throws IOException {
                     }
                     return currentBlock[blockInIndex];
                 }
+
+                @Override
+                public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+                    assert maxOrd == -1 : "unexpected maxOrd[" + maxOrd + "]";
+                    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;
+                                decoder.decode(valuesData, currentBlock);
+                            }
+
+                            // 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();
+                    }
+                }
+
+                static boolean isDense(int firstDocId, int lastDocId, int length) {
+                    // This does not detect duplicate docids (e.g [1, 1, 2, 4] would be detected as dense),
+                    // this can happen with enrich or lookup. However this codec isn't used for enrich / lookup.
+                    // This codec is only used in the context of logsdb and tsdb, so this is fine here.
+                    return lastDocId - firstDocId == length - 1;
+                }
+
             };
         } else {
             final IndexedDISI disi = new IndexedDISI(

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

@@ -21,6 +21,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.ByteArrayStreamInput;
 import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.codec.tsdb.es819.BulkNumericDocValues;
 import org.elasticsearch.index.mapper.BlockLoader.BlockFactory;
 import org.elasticsearch.index.mapper.BlockLoader.BooleanBuilder;
 import org.elasticsearch.index.mapper.BlockLoader.Builder;
@@ -84,6 +85,7 @@ public boolean supportsOrdinals() {
         public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException {
             throw new UnsupportedOperationException();
         }
+
     }
 
     public static class LongsBlockLoader extends DocValuesBlockLoader {
@@ -116,15 +118,18 @@ public AllReader reader(LeafReaderContext context) throws IOException {
         }
     }
 
-    private static class SingletonLongs extends BlockDocValuesReader {
-        private final NumericDocValues numericDocValues;
+    static class SingletonLongs extends BlockDocValuesReader {
+        final NumericDocValues numericDocValues;
 
         SingletonLongs(NumericDocValues numericDocValues) {
             this.numericDocValues = numericDocValues;
         }
 
         @Override
         public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset) throws IOException {
+            if (numericDocValues instanceof BulkNumericDocValues bulkDv) {
+                return bulkDv.read(factory, docs, offset);
+            }
             try (BlockLoader.LongBuilder builder = factory.longsFromDocValues(docs.count() - offset)) {
                 int lastDoc = -1;
                 for (int i = offset; i < docs.count(); i++) {
@@ -164,7 +169,7 @@ public String toString() {
         }
     }
 
-    private static class Longs extends BlockDocValuesReader {
+    static class Longs extends BlockDocValuesReader {
         private final SortedNumericDocValues numericDocValues;
         private int docID = -1;
 

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

@@ -400,6 +400,17 @@ interface BlockFactory {
          */
         LongBuilder longs(int expectedCount);
 
+        /**
+         * Build a specialized builder for singleton dense long based fields with the following constraints:
+         * <ul>
+         *     <li>Only one value per document can be collected</li>
+         *     <li>No more than expectedCount values can be collected</li>
+         * </ul>
+         *
+         * @param expectedCount The maximum number of values to be collected.
+         */
+        SingletonLongBuilder singletonLongs(int expectedCount);
+
         /**
          * Build a builder to load only {@code null}s.
          */
@@ -498,6 +509,16 @@ interface IntBuilder extends Builder {
         IntBuilder appendInt(int value);
     }
 
+    /**
+     * Specialized builder for collecting dense arrays of long values.
+     */
+    interface SingletonLongBuilder extends Builder {
+
+        SingletonLongBuilder appendLong(long value);
+
+        SingletonLongBuilder appendLongs(long[] values, int from, int length);
+    }
+
     interface LongBuilder extends Builder {
         /**
          * Appends a long to the current entry.