Push compute engine value loading down to tsdb codec

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/BlockAwareNumericDocValues.java

@@ -0,0 +1,18 @@
+/*
+ * 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;
+
+public abstract class BlockAwareNumericDocValues extends NumericDocValues {
+
+    public abstract SingletonDocValuesBlockLoader getSingletonBlockLoader();
+
+}

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

@@ -0,0 +1,18 @@
+/*
+ * 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.SortedDocValues;
+
+public abstract class BlockAwareSortedDocValues extends SortedDocValues {
+
+    public abstract SingletonDocValuesBlockLoader getSingletonBlockLoader();
+
+}

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

@@ -0,0 +1,24 @@
+/*
+ * 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.elasticsearch.index.mapper.BlockLoader;
+
+import java.io.IOException;
+
+public interface SingletonDocValuesBlockLoader {
+
+    void loadBlock(BlockLoader.SingletonLongBuilder builder, BlockLoader.Docs docs, int offset) throws IOException;
+
+    void loadBlock(BlockLoader.TSSingletonOrdinalsBuilder builder, BlockLoader.Docs docs, int offset) throws IOException;
+
+    int docID();
+
+}

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

@@ -116,7 +116,7 @@ public AllReader reader(LeafReaderContext context) throws IOException {
         }
     }
 
-    private static class SingletonLongs extends BlockDocValuesReader {
+    static class SingletonLongs extends BlockDocValuesReader {
         private final NumericDocValues numericDocValues;
 
         SingletonLongs(NumericDocValues numericDocValues) {
@@ -632,7 +632,7 @@ public String toString() {
         }
     }
 
-    private static class SingletonOrdinals extends BlockDocValuesReader {
+    public static class SingletonOrdinals extends BlockDocValuesReader {
         private final SortedDocValues ordinals;
 
         SingletonOrdinals(SortedDocValues ordinals) {

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

@@ -400,6 +400,8 @@ interface BlockFactory {
          */
         LongBuilder longs(int expectedCount);
 
+        BlockLoader.SingletonLongBuilder singletonLongs(int expectedCount);
+
         /**
          * Build a builder to load only {@code null}s.
          */
@@ -421,6 +423,8 @@ interface BlockFactory {
          */
         SingletonOrdinalsBuilder singletonOrdinalsBuilder(SortedDocValues ordinals, int count);
 
+        TSSingletonOrdinalsBuilder tsSingletonOrdinalsBuilder(boolean isPrimaryIndexSortField, SortedDocValues ordinals, int count);
+
         /**
          * Build a reader for reading {@link SortedSetDocValues}
          */
@@ -505,13 +509,28 @@ interface LongBuilder extends Builder {
         LongBuilder appendLong(long value);
     }
 
+    interface SingletonLongBuilder extends Builder {
+
+        SingletonLongBuilder appendLong(long value);
+
+        SingletonLongBuilder appendLongs(long[] values, int from, int length);
+    }
+
     interface SingletonOrdinalsBuilder extends Builder {
         /**
          * Appends an ordinal to the builder.
          */
         SingletonOrdinalsBuilder appendOrd(int value);
     }
 
+    interface TSSingletonOrdinalsBuilder extends Builder {
+
+        TSSingletonOrdinalsBuilder appendOrd(long value);
+
+        TSSingletonOrdinalsBuilder appendOrds(long[] values, int from, int length);
+
+    }
+
     interface SortedSetOrdinalsBuilder extends Builder {
         /**
          * Appends an ordinal to the builder.

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

@@ -1012,8 +1012,13 @@ public Function<byte[], Number> pointReaderIfPossible() {
 
         @Override
         public BlockLoader blockLoader(BlockLoaderContext blContext) {
+            var indexMode = blContext.indexSettings().getMode();
             if (hasDocValues()) {
-                return new BlockDocValuesReader.LongsBlockLoader(name());
+                if (name().equals("@timestamp") && (indexMode == IndexMode.TIME_SERIES || indexMode == IndexMode.LOGSDB)) {
+                    return new TimestampBlockLoader();
+                } else {
+                    return new BlockDocValuesReader.LongsBlockLoader(name());
+                }
             }
 
             // Multi fields don't have fallback synthetic source.

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

@@ -794,7 +794,11 @@ NamedAnalyzer normalizer() {
         @Override
         public BlockLoader blockLoader(BlockLoaderContext blContext) {
             if (hasDocValues() && (blContext.fieldExtractPreference() != FieldExtractPreference.STORED || isSyntheticSource)) {
-                return new BlockDocValuesReader.BytesRefsFromOrdsBlockLoader(name());
+                if (isDimension) {
+                    return new TSDimensionBlockLoader(name());
+                } else {
+                    return new BlockDocValuesReader.BytesRefsFromOrdsBlockLoader(name());
+                }
             }
             if (isStored()) {
                 return new BlockStoredFieldsReader.BytesFromBytesRefsBlockLoader(name());

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

@@ -0,0 +1,103 @@
+/*
+ * 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.mapper;
+
+import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.index.SortedSetDocValues;
+import org.elasticsearch.index.codec.tsdb.es819.BlockAwareSortedDocValues;
+import org.elasticsearch.index.codec.tsdb.es819.SingletonDocValuesBlockLoader;
+import org.elasticsearch.search.fetch.StoredFieldsSpec;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class TSDimensionBlockLoader implements BlockLoader {
+
+    private final String fieldName;
+    private final BlockLoader fallback;
+
+    public TSDimensionBlockLoader(String fieldName) {
+        this.fieldName = fieldName;
+        this.fallback = new BlockDocValuesReader.BytesRefsFromOrdsBlockLoader(fieldName);
+    }
+
+    @Override
+    public Builder builder(BlockFactory factory, int expectedCount) {
+        return factory.bytesRefs(expectedCount);
+    }
+
+    @Override
+    public ColumnAtATimeReader columnAtATimeReader(LeafReaderContext context) throws IOException {
+        var singleton = DocValues.unwrapSingleton(context.reader().getSortedSetDocValues(fieldName));
+        if (singleton instanceof BlockAwareSortedDocValues b) {
+            var loader = b.getSingletonBlockLoader();
+            if (loader != null) {
+                return new TSDimensions(b, loader);
+            }
+        }
+        return fallback.columnAtATimeReader(context);
+    }
+
+    @Override
+    public RowStrideReader rowStrideReader(LeafReaderContext context) throws IOException {
+        return fallback.rowStrideReader(context);
+    }
+
+    @Override
+    public StoredFieldsSpec rowStrideStoredFieldSpec() {
+        return StoredFieldsSpec.NO_REQUIREMENTS;
+    }
+
+    @Override
+    public boolean supportsOrdinals() {
+        return true;
+    }
+
+    @Override
+    public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException {
+        return DocValues.getSortedSet(context.reader(), fieldName);
+    }
+
+    public String toString() {
+        return "TSIDBlockLoader[" + fieldName + "]";
+    }
+
+    public static final class TSDimensions implements ColumnAtATimeReader {
+        private final Thread creationThread;
+        private final SortedDocValues sorted;
+        private final SingletonDocValuesBlockLoader blockLoader;
+
+        TSDimensions(BlockAwareSortedDocValues sorted, SingletonDocValuesBlockLoader blockLoader) {
+            this.creationThread = Thread.currentThread();
+            this.sorted = Objects.requireNonNull(sorted);
+            this.blockLoader = Objects.requireNonNull(blockLoader);
+        }
+
+        @Override
+        public Block read(BlockFactory factory, Docs docs, int offset) throws IOException {
+            try (var builder = factory.tsSingletonOrdinalsBuilder(false, sorted, docs.count() - offset)) {
+                blockLoader.loadBlock(builder, docs, offset);
+                return builder.build();
+            }
+        }
+
+        @Override
+        public boolean canReuse(int startingDocID) {
+            return creationThread == Thread.currentThread() && blockLoader.docID() <= startingDocID;
+        }
+
+        @Override
+        public String toString() {
+            return "TSDimensionBlockLoader.TSDimensions";
+        }
+    }
+}

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

@@ -0,0 +1,87 @@
+/*
+ * 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.mapper;
+
+import org.apache.lucene.index.DocValues;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedDocValues;
+import org.apache.lucene.index.SortedSetDocValues;
+import org.elasticsearch.index.codec.tsdb.es819.BlockAwareSortedDocValues;
+import org.elasticsearch.index.codec.tsdb.es819.SingletonDocValuesBlockLoader;
+import org.elasticsearch.search.fetch.StoredFieldsSpec;
+
+import java.io.IOException;
+
+public final class TSIDBlockLoader implements BlockLoader {
+
+    private static final String FIELD_NAME = TimeSeriesIdFieldMapper.NAME;
+
+    @Override
+    public Builder builder(BlockFactory factory, int expectedCount) {
+        return factory.bytesRefs(expectedCount);
+    }
+
+    @Override
+    public ColumnAtATimeReader columnAtATimeReader(LeafReaderContext context) throws IOException {
+        var singleton = context.reader().getSortedDocValues(FIELD_NAME);
+        return new TSIDs((BlockAwareSortedDocValues) singleton);
+    }
+
+    @Override
+    public RowStrideReader rowStrideReader(LeafReaderContext context) throws IOException {
+        var singleton = context.reader().getSortedDocValues(FIELD_NAME);
+        return new BlockDocValuesReader.SingletonOrdinals(singleton);
+    }
+
+    @Override
+    public StoredFieldsSpec rowStrideStoredFieldSpec() {
+        return StoredFieldsSpec.NO_REQUIREMENTS;
+    }
+
+    @Override
+    public boolean supportsOrdinals() {
+        return true;
+    }
+
+    @Override
+    public SortedSetDocValues ordinals(LeafReaderContext context) throws IOException {
+        return DocValues.getSortedSet(context.reader(), FIELD_NAME);
+    }
+
+    public static final class TSIDs implements ColumnAtATimeReader {
+        private final Thread creationThread;
+        private final SortedDocValues sorted;
+        private final SingletonDocValuesBlockLoader blockLoader;
+
+        TSIDs(BlockAwareSortedDocValues sorted) {
+            this.creationThread = Thread.currentThread();
+            this.sorted = sorted;
+            this.blockLoader = sorted.getSingletonBlockLoader();
+        }
+
+        @Override
+        public Block read(BlockFactory factory, Docs docs, int offset) throws IOException {
+            try (TSSingletonOrdinalsBuilder builder = factory.tsSingletonOrdinalsBuilder(true, sorted, docs.count() - offset)) {
+                blockLoader.loadBlock(builder, docs, offset);
+                return builder.build();
+            }
+        }
+
+        @Override
+        public boolean canReuse(int startingDocID) {
+            return creationThread == Thread.currentThread() && blockLoader.docID() <= startingDocID;
+        }
+
+        @Override
+        public String toString() {
+            return "TSIDBlockLoader.TSIDs";
+        }
+    }
+}

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

@@ -135,7 +135,7 @@ public Query termQuery(Object value, SearchExecutionContext context) {
 
         @Override
         public BlockLoader blockLoader(BlockLoaderContext blContext) {
-            return new BlockDocValuesReader.BytesRefsFromOrdsBlockLoader(name());
+            return new TSIDBlockLoader();
         }
     }