Push down loading of singleton dense double based field types to the … (#133397)

* Push down loading of singleton dense double based field types to the es819 doc value codec

* Update docs/changelog/133397.yaml

* use `SingletonDoubleBuilder`

* add tests

* [CI] Auto commit changes from spotless

* Update test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java

Co-authored-by: Martijn van Groningen <[email protected]>

* add assert

---------

Co-authored-by: elasticsearchmachine <[email protected]>
Co-authored-by: Martijn van Groningen <[email protected]>

Author: 3 people

docs/changelog/133397.yaml

@@ -0,0 +1,5 @@
+pr: 133397
+summary: Push down loading of singleton dense double based field types to the …
+area: "Codec"
+type: enhancement
+issues: []

modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/ScaledFloatFieldMapperTests.java

@@ -78,6 +78,11 @@ protected void registerParameters(ParameterChecker checker) throws IOException {
         checker.registerUpdateCheck(b -> b.field("coerce", false), m -> assertFalse(((ScaledFloatFieldMapper) m).coerce()));
     }
 
+    @Override
+    protected boolean supportsBulkDoubleBlockReading() {
+        return true;
+    }
+
     public void testExistsQueryDocValuesDisabled() throws IOException {
         MapperService mapperService = createMapperService(fieldMapping(b -> {
             minimalMapping(b);

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

@@ -45,6 +45,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.BlockDocValuesReader;
 import org.elasticsearch.index.mapper.BlockLoader;
 
 import java.io.IOException;
@@ -383,7 +384,13 @@ public long cost() {
             }
 
             @Override
-            public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+            public BlockLoader.Block tryRead(
+                BlockLoader.BlockFactory factory,
+                BlockLoader.Docs docs,
+                int offset,
+                BlockDocValuesReader.ToDouble toDouble
+            ) throws IOException {
+                assert toDouble == null;
                 if (ords instanceof BaseDenseNumericValues denseOrds) {
                     var block = tryReadAHead(factory, docs, offset);
                     if (block != null) {
@@ -457,7 +464,12 @@ public TermsEnum termsEnum() throws IOException {
         }
 
         @Override
-        public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+        public BlockLoader.Block tryRead(
+            BlockLoader.BlockFactory factory,
+            BlockLoader.Docs docs,
+            int offset,
+            BlockDocValuesReader.ToDouble toDouble
+        ) throws IOException {
             return null;
         }
 
@@ -504,7 +516,12 @@ public final long cost() {
         }
 
         @Override
-        public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+        public BlockLoader.Block tryRead(
+            BlockLoader.BlockFactory factory,
+            BlockLoader.Docs docs,
+            int offset,
+            BlockDocValuesReader.ToDouble toDouble
+        ) throws IOException {
             return null;
         }
 
@@ -1365,7 +1382,18 @@ public long longValue() throws IOException {
                 }
 
                 @Override
-                public BlockLoader.Block tryRead(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset) throws IOException {
+                public BlockLoader.Block tryRead(
+                    BlockLoader.BlockFactory factory,
+                    BlockLoader.Docs docs,
+                    int offset,
+                    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);
                     }
@@ -1774,4 +1802,55 @@ public BlockLoader.Builder endPositionEntry() {
         public void close() {}
     }
 
+    // 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) {
+            this.doubleBuilder = doubleBuilder;
+            this.toDouble = toDouble;
+        }
+
+        @Override
+        public BlockLoader.SingletonLongBuilder appendLong(long value) {
+            throw new UnsupportedOperationException();
+        }
+
+        @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);
+            return this;
+        }
+
+        @Override
+        public BlockLoader.Block build() {
+            return doubleBuilder.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

@@ -122,7 +122,12 @@ public AllReader reader(LeafReaderContext context) throws IOException {
         }
     }
 
-    static class SingletonLongs extends BlockDocValuesReader {
+    // Used for testing.
+    interface NumericDocValuesAccessor {
+        NumericDocValues numericDocValues();
+    }
+
+    static class SingletonLongs extends BlockDocValuesReader implements NumericDocValuesAccessor {
         final NumericDocValues numericDocValues;
 
         SingletonLongs(NumericDocValues numericDocValues) {
@@ -132,7 +137,7 @@ static class SingletonLongs extends BlockDocValuesReader {
         @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);
+                BlockLoader.Block result = direct.tryRead(factory, docs, offset, null);
                 if (result != null) {
                     return result;
                 }
@@ -169,6 +174,11 @@ public int docId() {
         public String toString() {
             return "BlockDocValuesReader.SingletonLongs";
         }
+
+        @Override
+        public NumericDocValues numericDocValues() {
+            return numericDocValues;
+        }
     }
 
     static class Longs extends BlockDocValuesReader {
@@ -387,7 +397,7 @@ public AllReader reader(LeafReaderContext context) throws IOException {
         }
     }
 
-    private static class SingletonDoubles extends BlockDocValuesReader {
+    static class SingletonDoubles extends BlockDocValuesReader implements NumericDocValuesAccessor {
         private final NumericDocValues docValues;
         private final ToDouble toDouble;
 
@@ -398,6 +408,12 @@ private static class SingletonDoubles extends BlockDocValuesReader {
 
         @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);
+                if (result != null) {
+                    return result;
+                }
+            }
             try (BlockLoader.DoubleBuilder builder = factory.doublesFromDocValues(docs.count() - offset)) {
                 for (int i = offset; i < docs.count(); i++) {
                     int doc = docs.get(i);
@@ -430,9 +446,14 @@ public int docId() {
         public String toString() {
             return "BlockDocValuesReader.SingletonDoubles";
         }
+
+        @Override
+        public NumericDocValues numericDocValues() {
+            return docValues;
+        }
     }
 
-    private static class Doubles extends BlockDocValuesReader {
+    static class Doubles extends BlockDocValuesReader {
         private final SortedNumericDocValues docValues;
         private final ToDouble toDouble;
 
@@ -715,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);
+                BlockLoader.Block block = direct.tryRead(factory, docs, offset, null);
                 if (block != null) {
                     return block;
                 }

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

@@ -65,9 +65,11 @@ interface OptionalColumnAtATimeReader {
         /**
          * Attempts to read the values of all documents in {@code docs}
          * Returns {@code null} if unable to load the values.
+         *
+         * @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) throws IOException;
+        BlockLoader.Block tryRead(BlockFactory factory, Docs docs, int offset, BlockDocValuesReader.ToDouble toDouble) throws IOException;
     }
 
     interface RowStrideReader extends Reader {
@@ -435,6 +437,17 @@ interface BlockFactory {
          */
         SingletonLongBuilder singletonLongs(int expectedCount);
 
+        /**
+         * Build a specialized builder for singleton dense double 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.
+         */
+        SingletonDoubleBuilder singletonDoubles(int expectedCount);
+
         /**
          * Build a builder to load only {@code null}s.
          */
@@ -537,12 +550,20 @@ interface IntBuilder extends Builder {
      * 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);
     }
 
+    /**
+     * Specialized builder for collecting dense arrays of double values.
+     */
+    interface SingletonDoubleBuilder extends Builder {
+        SingletonDoubleBuilder appendDouble(double value);
+
+        SingletonDoubleBuilder appendDoubles(double[] values, int from, int length);
+    }
+
     interface LongBuilder extends Builder {
         /**
          * Appends a long to the current entry.