Skip iterating DISI when reading metric values

docs/changelog/133365.yaml

@@ -0,0 +1,5 @@
+pr: 133365
+summary: Skip iterating DISI when reading metric values
+area: Codec
+type: enhancement
+issues: []

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

@@ -43,6 +43,7 @@
 import org.apache.lucene.util.compress.LZ4;
 import org.apache.lucene.util.packed.DirectMonotonicReader;
 import org.apache.lucene.util.packed.PackedInts;
+import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.codec.tsdb.TSDBDocValuesEncoder;
 import org.elasticsearch.index.mapper.BlockDocValuesReader;
@@ -388,6 +389,7 @@ public BlockLoader.Block tryRead(
                 BlockLoader.BlockFactory factory,
                 BlockLoader.Docs docs,
                 int offset,
+                boolean nullsFiltered,
                 BlockDocValuesReader.ToDouble toDouble
             ) throws IOException {
                 assert toDouble == null;
@@ -468,6 +470,7 @@ public BlockLoader.Block tryRead(
             BlockLoader.BlockFactory factory,
             BlockLoader.Docs docs,
             int offset,
+            boolean nullsFiltered,
             BlockDocValuesReader.ToDouble toDouble
         ) throws IOException {
             return null;
@@ -520,6 +523,7 @@ public BlockLoader.Block tryRead(
             BlockLoader.BlockFactory factory,
             BlockLoader.Docs docs,
             int offset,
+            boolean nullsFiltered,
             BlockDocValuesReader.ToDouble toDouble
         ) throws IOException {
             return null;
@@ -532,7 +536,7 @@ BlockLoader.Block tryRead(BlockLoader.SingletonLongBuilder builder, BlockLoader.
         }
     }
 
-    abstract static class BaseSparseNumericValues extends NumericDocValues {
+    abstract static class BaseSparseNumericValues extends NumericDocValues implements BlockLoader.OptionalColumnAtATimeReader {
         protected final IndexedDISI disi;
 
         BaseSparseNumericValues(IndexedDISI disi) {
@@ -563,6 +567,17 @@ public final int docID() {
         public final long cost() {
             return disi.cost();
         }
+
+        @Override
+        public BlockLoader.Block tryRead(
+            BlockLoader.BlockFactory factory,
+            BlockLoader.Docs docs,
+            int offset,
+            boolean nullsFiltered,
+            BlockDocValuesReader.ToDouble toDouble
+        ) throws IOException {
+            return null;
+        }
     }
 
     abstract static class BaseSortedSetDocValues extends SortedSetDocValues {
@@ -1386,16 +1401,11 @@ public BlockLoader.Block tryRead(
                     BlockLoader.BlockFactory factory,
                     BlockLoader.Docs docs,
                     int offset,
+                    boolean nullsFiltered,
                     BlockDocValuesReader.ToDouble toDouble
                 ) throws IOException {
-                    if (toDouble != null) {
-                        try (BlockLoader.SingletonDoubleBuilder builder = factory.singletonDoubles(docs.count() - offset)) {
-                            SingletonLongToDoubleDelegate delegate = new SingletonLongToDoubleDelegate(builder, toDouble);
-                            return tryRead(delegate, docs, offset);
-                        }
-                    }
-                    try (BlockLoader.SingletonLongBuilder builder = factory.singletonLongs(docs.count() - offset)) {
-                        return tryRead(builder, docs, offset);
+                    try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, docs.count() - offset)) {
+                        return tryRead(singletonLongBuilder, docs, offset);
                     }
                 }
 
@@ -1484,6 +1494,7 @@ static boolean isDense(int firstDocId, int lastDocId, int length) {
             );
             return new BaseSparseNumericValues(disi) {
                 private final TSDBDocValuesEncoder decoder = new TSDBDocValuesEncoder(ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE);
+                private IndexedDISI lookAheadDISI;
                 private long currentBlockIndex = -1;
                 private final long[] currentBlock = new long[ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE];
 
@@ -1507,6 +1518,74 @@ public long longValue() throws IOException {
                     }
                     return currentBlock[blockInIndex];
                 }
+
+                @Override
+                public BlockLoader.Block tryRead(
+                    BlockLoader.BlockFactory factory,
+                    BlockLoader.Docs docs,
+                    int offset,
+                    boolean nullsFiltered,
+                    BlockDocValuesReader.ToDouble toDouble
+                ) throws IOException {
+                    if (nullsFiltered == false) {
+                        return null;
+                    }
+                    final int firstDoc = docs.get(offset);
+                    if (disi.advanceExact(firstDoc) == false) {
+                        assert false : "nullsFiltered is true, but doc [" + firstDoc + "] has no value";
+                        throw new IllegalStateException("nullsFiltered is true, but doc [" + firstDoc + "] has no value");
+                    }
+                    if (lookAheadDISI == null) {
+                        lookAheadDISI = new IndexedDISI(
+                            data,
+                            entry.docsWithFieldOffset,
+                            entry.docsWithFieldLength,
+                            entry.jumpTableEntryCount,
+                            entry.denseRankPower,
+                            entry.numValues
+                        );
+                    }
+                    final int lastDoc = docs.get(docs.count() - 1);
+                    if (lookAheadDISI.advanceExact(lastDoc) == false) {
+                        assert false : "nullsFiltered is true, but doc [" + lastDoc + "] has no value";
+                        throw new IllegalStateException("nullsFiltered is true, but doc [" + lastDoc + "] has no value");
+                    }
+                    // Assumes docIds are unique - if the number of value indices between the first
+                    // and last doc equals the doc count, all values can be read and converted in bulk
+                    // TODO: Pass docCount attr for enrich and lookup.
+                    final int firstIndex = disi.index();
+                    final int lastIndex = lookAheadDISI.index();
+                    final int valueCount = lastIndex - firstIndex + 1;
+                    if (valueCount != docs.count()) {
+                        return null;
+                    }
+                    if (Assertions.ENABLED) {
+                        for (int i = 0; i < docs.count(); i++) {
+                            final int doc = docs.get(i + offset);
+                            assert disi.advanceExact(doc) : "nullsFiltered is true, but doc [" + doc + "] has no value";
+                            assert disi.index() == firstIndex + i : "unexpected disi index " + (firstIndex + i) + "!=" + disi.index();
+                        }
+                    }
+                    try (var singletonLongBuilder = singletonLongBuilder(factory, toDouble, valueCount)) {
+                        for (int i = 0; i < valueCount;) {
+                            final int index = firstIndex + i;
+                            final int blockIndex = index >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT;
+                            final int blockStartIndex = index & ES819TSDBDocValuesFormat.NUMERIC_BLOCK_MASK;
+                            if (blockIndex != currentBlockIndex) {
+                                assert blockIndex > currentBlockIndex : blockIndex + "<=" + currentBlockIndex;
+                                if (currentBlockIndex + 1 != blockIndex) {
+                                    valuesData.seek(indexReader.get(blockIndex));
+                                }
+                                currentBlockIndex = blockIndex;
+                                decoder.decode(valuesData, currentBlock);
+                            }
+                            final int count = Math.min(ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE - blockStartIndex, valueCount - i);
+                            singletonLongBuilder.appendLongs(currentBlock, blockStartIndex, count);
+                            i += count;
+                        }
+                        return singletonLongBuilder.build();
+                    }
+                }
             };
         }
     }
@@ -1802,11 +1881,22 @@ public BlockLoader.Builder endPositionEntry() {
         public void close() {}
     }
 
+    static BlockLoader.SingletonLongBuilder singletonLongBuilder(
+        BlockLoader.BlockFactory factory,
+        BlockDocValuesReader.ToDouble toDouble,
+        int valueCount
+    ) {
+        if (toDouble != null) {
+            return new SingletonLongToDoubleDelegate(factory.singletonDoubles(valueCount), toDouble);
+        } else {
+            return factory.singletonLongs(valueCount);
+        }
+    }
+
     // Block builder that consumes long values and converts them to double using the provided converter function.
     static final class SingletonLongToDoubleDelegate implements BlockLoader.SingletonLongBuilder {
         private final BlockLoader.SingletonDoubleBuilder doubleBuilder;
         private final BlockDocValuesReader.ToDouble toDouble;
-        private final double[] buffer = new double[ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE];
 
         // The passed builder is used to store the converted double values and produce the final block containing them.
         SingletonLongToDoubleDelegate(BlockLoader.SingletonDoubleBuilder doubleBuilder, BlockDocValuesReader.ToDouble toDouble) {
@@ -1821,11 +1911,7 @@ public BlockLoader.SingletonLongBuilder appendLong(long value) {
 
         @Override
         public BlockLoader.SingletonLongBuilder appendLongs(long[] values, int from, int length) {
-            assert length <= buffer.length : "length " + length + " > " + buffer.length;
-            for (int i = 0; i < length; i++) {
-                buffer[i] = toDouble.convert(values[from + i]);
-            }
-            doubleBuilder.appendDoubles(buffer, 0, length);
+            doubleBuilder.appendLongs(toDouble, values, from, length);
             return this;
         }
 
@@ -1850,7 +1936,9 @@ public BlockLoader.Builder endPositionEntry() {
         }
 
         @Override
-        public void close() {}
+        public void close() {
+            doubleBuilder.close();
+        }
     }
 
 }

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

@@ -137,7 +137,7 @@ static class SingletonLongs extends BlockDocValuesReader implements NumericDocVa
         @Override
         public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
             if (numericDocValues instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
-                BlockLoader.Block result = direct.tryRead(factory, docs, offset, null);
+                BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, null);
                 if (result != null) {
                     return result;
                 }
@@ -409,7 +409,7 @@ static class SingletonDoubles extends BlockDocValuesReader implements NumericDoc
         @Override
         public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boolean nullsFiltered) throws IOException {
             if (docValues instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
-                BlockLoader.Block result = direct.tryRead(factory, docs, offset, toDouble);
+                BlockLoader.Block result = direct.tryRead(factory, docs, offset, nullsFiltered, toDouble);
                 if (result != null) {
                     return result;
                 }
@@ -736,7 +736,7 @@ public BlockLoader.Block read(BlockFactory factory, Docs docs, int offset, boole
                 return readSingleDoc(factory, docs.get(offset));
             }
             if (ordinals instanceof BlockLoader.OptionalColumnAtATimeReader direct) {
-                BlockLoader.Block block = direct.tryRead(factory, docs, offset, null);
+                BlockLoader.Block block = direct.tryRead(factory, docs, offset, nullsFiltered, null);
                 if (block != null) {
                     return block;
                 }

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

@@ -66,10 +66,18 @@ interface OptionalColumnAtATimeReader {
          * Attempts to read the values of all documents in {@code docs}
          * Returns {@code null} if unable to load the values.
          *
+         * @param nullsFiltered if {@code true}, then target docs are guaranteed to have a value for the field.
+         *                      see {@link ColumnAtATimeReader#read(BlockFactory, Docs, int, boolean)}
          * @param toDouble a function to convert long values to double, or null if no conversion is needed/supported
          */
         @Nullable
-        BlockLoader.Block tryRead(BlockFactory factory, Docs docs, int offset, BlockDocValuesReader.ToDouble toDouble) throws IOException;
+        BlockLoader.Block tryRead(
+            BlockFactory factory,
+            Docs docs,
+            int offset,
+            boolean nullsFiltered,
+            BlockDocValuesReader.ToDouble toDouble
+        ) throws IOException;
     }
 
     interface RowStrideReader extends Reader {
@@ -562,6 +570,8 @@ interface SingletonDoubleBuilder extends Builder {
         SingletonDoubleBuilder appendDouble(double value);
 
         SingletonDoubleBuilder appendDoubles(double[] values, int from, int length);
+
+        SingletonDoubleBuilder appendLongs(BlockDocValuesReader.ToDouble toDouble, long[] values, int from, int length);
     }
 
     interface LongBuilder extends Builder {