Add ordinal range encode for tsid

benchmarks/src/main/java/org/elasticsearch/benchmark/index/codec/tsdb/TSDBDocValuesMergeBenchmark.java

@@ -258,7 +258,7 @@ private static IndexWriterConfig createIndexWriterConfig(boolean optimizedMergeE
         );
         config.setLeafSorter(DataStream.TIMESERIES_LEAF_READERS_SORTER);
         config.setMergePolicy(new LogByteSizeMergePolicy());
-        var docValuesFormat = new ES819TSDBDocValuesFormat(4096, optimizedMergeEnabled);
+        var docValuesFormat = new ES819TSDBDocValuesFormat(4096, 512, optimizedMergeEnabled);
         config.setCodec(new Elasticsearch900Lucene101Codec() {
 
             @Override

docs/changelog/133018.yaml

@@ -0,0 +1,5 @@
+pr: 133018
+summary: Add ordinal range encode for tsid
+area: TSDB
+type: enhancement
+issues: []

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

@@ -62,11 +62,14 @@ final class ES819TSDBDocValuesConsumer extends XDocValuesConsumer {
     final int maxDoc;
     private byte[] termsDictBuffer;
     private final int skipIndexIntervalSize;
+    private final int minDocsPerOrdinalForOrdinalRangeEncoding;
     final boolean enableOptimizedMerge;
+    private final int primarySortFieldNumber;
 
     ES819TSDBDocValuesConsumer(
         SegmentWriteState state,
         int skipIndexIntervalSize,
+        int minDocsPerOrdinalForOrdinalRangeEncoding,
         boolean enableOptimizedMerge,
         String dataCodec,
         String dataExtension,
@@ -75,6 +78,8 @@ final class ES819TSDBDocValuesConsumer extends XDocValuesConsumer {
     ) throws IOException {
         this.termsDictBuffer = new byte[1 << 14];
         this.dir = state.directory;
+        this.minDocsPerOrdinalForOrdinalRangeEncoding = minDocsPerOrdinalForOrdinalRangeEncoding;
+        this.primarySortFieldNumber = ES819TSDBDocValuesProducer.primarySortFieldNumber(state.segmentInfo, state.fieldInfos);
         this.context = state.context;
         boolean success = false;
         try {
@@ -124,6 +129,12 @@ public SortedNumericDocValues getSortedNumeric(FieldInfo field) throws IOExcepti
         writeField(field, producer, -1, null);
     }
 
+    private boolean shouldEncodeOrdinalRange(FieldInfo field, long maxOrd, int numDocsWithValue) {
+        return maxDoc > 1
+            && field.number == primarySortFieldNumber
+            && (numDocsWithValue / maxOrd) >= minDocsPerOrdinalForOrdinalRangeEncoding;
+    }
+
     private long[] writeField(FieldInfo field, TsdbDocValuesProducer valuesProducer, long maxOrd, OffsetsAccumulator offsetsAccumulator)
         throws IOException {
         int numDocsWithValue = 0;
@@ -149,19 +160,53 @@ private long[] writeField(FieldInfo field, TsdbDocValuesProducer valuesProducer,
         try {
             if (numValues > 0) {
                 assert numDocsWithValue > 0;
-                // Special case for maxOrd of 1, signal -1 that no blocks will be written
-                meta.writeInt(maxOrd != 1 ? ES819TSDBDocValuesFormat.DIRECT_MONOTONIC_BLOCK_SHIFT : -1);
                 final ByteBuffersDataOutput indexOut = new ByteBuffersDataOutput();
-                final DirectMonotonicWriter indexWriter = DirectMonotonicWriter.getInstance(
-                    meta,
-                    new ByteBuffersIndexOutput(indexOut, "temp-dv-index", "temp-dv-index"),
-                    1L + ((numValues - 1) >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT),
-                    ES819TSDBDocValuesFormat.DIRECT_MONOTONIC_BLOCK_SHIFT
-                );
+                DirectMonotonicWriter indexWriter = null;
 
                 final long valuesDataOffset = data.getFilePointer();
-                // Special case for maxOrd of 1, skip writing the blocks
-                if (maxOrd != 1) {
+                if (maxOrd == 1) {
+                    // Special case for maxOrd of 1, signal -1 that no blocks will be written
+                    meta.writeInt(-1);
+                } else if (shouldEncodeOrdinalRange(field, maxOrd, numDocsWithValue)) {
+                    // When a field is sorted, use ordinal range encode for long runs of the same ordinal.
+                    meta.writeInt(-2);
+                    meta.writeVInt(Math.toIntExact(maxOrd));
+                    meta.writeByte((byte) ES819TSDBDocValuesFormat.ORDINAL_RANGE_ENCODING_BLOCK_SHIFT);
+                    values = valuesProducer.getSortedNumeric(field);
+                    if (enableOptimizedMerge && numDocsWithValue < maxDoc) {
+                        disiAccumulator = new DISIAccumulator(dir, context, data, IndexedDISI.DEFAULT_DENSE_RANK_POWER);
+                    }
+                    DirectMonotonicWriter startDocs = DirectMonotonicWriter.getInstance(
+                        meta,
+                        data,
+                        maxOrd + 1,
+                        ES819TSDBDocValuesFormat.ORDINAL_RANGE_ENCODING_BLOCK_SHIFT
+                    );
+                    long lastOrd = 0;
+                    startDocs.add(0);
+                    for (int doc = values.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = values.nextDoc()) {
+                        if (disiAccumulator != null) {
+                            disiAccumulator.addDocId(doc);
+                        }
+                        if (offsetsAccumulator != null) {
+                            offsetsAccumulator.addDoc(1);
+                        }
+                        final long nextOrd = values.nextValue();
+                        if (nextOrd != lastOrd) {
+                            lastOrd = nextOrd;
+                            startDocs.add(doc);
+                        }
+                    }
+                    startDocs.add(maxDoc);
+                    startDocs.finish();
+                } else {
+                    indexWriter = DirectMonotonicWriter.getInstance(
+                        meta,
+                        new ByteBuffersIndexOutput(indexOut, "temp-dv-index", "temp-dv-index"),
+                        1L + ((numValues - 1) >>> ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SHIFT),
+                        ES819TSDBDocValuesFormat.DIRECT_MONOTONIC_BLOCK_SHIFT
+                    );
+                    meta.writeInt(DIRECT_MONOTONIC_BLOCK_SHIFT);
                     final long[] buffer = new long[ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE];
                     int bufferSize = 0;
                     final TSDBDocValuesEncoder encoder = new TSDBDocValuesEncoder(ES819TSDBDocValuesFormat.NUMERIC_BLOCK_SIZE);
@@ -204,8 +249,7 @@ private long[] writeField(FieldInfo field, TsdbDocValuesProducer valuesProducer,
                 }
 
                 final long valuesDataLength = data.getFilePointer() - valuesDataOffset;
-                if (maxOrd != 1) {
-                    // Special case for maxOrd of 1, indexWriter isn't really used, so no need to invoke finish() method.
+                if (indexWriter != null) {
                     indexWriter.finish();
                 }
                 final long indexDataOffset = data.getFilePointer();

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

@@ -104,21 +104,35 @@ private static boolean getOptimizedMergeEnabledDefault() {
         return Boolean.parseBoolean(System.getProperty(OPTIMIZED_MERGE_ENABLED_NAME, Boolean.TRUE.toString()));
     }
 
+    /**
+     * The default minimum number of documents per ordinal required to use ordinal range encoding.
+     * If the average number of documents per ordinal is below this threshold, it is more efficient to encode doc values in blocks.
+     * A much smaller value may be used in tests to exercise ordinal range encoding more frequently.
+     */
+    public static final int ORDINAL_RANGE_ENCODING_MIN_DOC_PER_ORDINAL = 512;
+
+    /**
+     * The block shift used in DirectMonotonicWriter when encoding the start docs of each ordinal with ordinal range encoding.
+     */
+    public static final int ORDINAL_RANGE_ENCODING_BLOCK_SHIFT = 12;
+
     final int skipIndexIntervalSize;
+    final int minDocsPerOrdinalForRangeEncoding;
     private final boolean enableOptimizedMerge;
 
     /** Default constructor. */
     public ES819TSDBDocValuesFormat() {
-        this(DEFAULT_SKIP_INDEX_INTERVAL_SIZE, OPTIMIZED_MERGE_ENABLE_DEFAULT);
+        this(DEFAULT_SKIP_INDEX_INTERVAL_SIZE, ORDINAL_RANGE_ENCODING_MIN_DOC_PER_ORDINAL, OPTIMIZED_MERGE_ENABLE_DEFAULT);
     }
 
     /** Doc values fields format with specified skipIndexIntervalSize. */
-    public ES819TSDBDocValuesFormat(int skipIndexIntervalSize, boolean enableOptimizedMerge) {
+    public ES819TSDBDocValuesFormat(int skipIndexIntervalSize, int minDocsPerOrdinalForRangeEncoding, boolean enableOptimizedMerge) {
         super(CODEC_NAME);
         if (skipIndexIntervalSize < 2) {
             throw new IllegalArgumentException("skipIndexIntervalSize must be > 1, got [" + skipIndexIntervalSize + "]");
         }
         this.skipIndexIntervalSize = skipIndexIntervalSize;
+        this.minDocsPerOrdinalForRangeEncoding = minDocsPerOrdinalForRangeEncoding;
         this.enableOptimizedMerge = enableOptimizedMerge;
     }
 
@@ -127,6 +141,7 @@ public DocValuesConsumer fieldsConsumer(SegmentWriteState state) throws IOExcept
         return new ES819TSDBDocValuesConsumer(
             state,
             skipIndexIntervalSize,
+            minDocsPerOrdinalForRangeEncoding,
             enableOptimizedMerge,
             DATA_CODEC,
             DATA_EXTENSION,