Add binary doc value compression with variable doc count blocks

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.cluster.metadata.DataStream;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.index.codec.Elasticsearch92Lucene103Codec;
+import org.elasticsearch.index.codec.tsdb.BinaryDVCompressionMode;
 import org.elasticsearch.index.codec.tsdb.es819.ES819TSDBDocValuesFormat;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -257,7 +258,12 @@ private static IndexWriterConfig createIndexWriterConfig(boolean optimizedMergeE
         );
         config.setLeafSorter(DataStream.TIMESERIES_LEAF_READERS_SORTER);
         config.setMergePolicy(new LogByteSizeMergePolicy());
-        var docValuesFormat = new ES819TSDBDocValuesFormat(4096, 512, optimizedMergeEnabled);
+        var docValuesFormat = new ES819TSDBDocValuesFormat(
+            4096,
+            512,
+            optimizedMergeEnabled,
+            BinaryDVCompressionMode.COMPRESSED_WITH_ZSTD_LEVEL_1
+        );
         config.setCodec(new Elasticsearch92Lucene103Codec() {
             @Override
             public DocValuesFormat getDocValuesFormatForField(String field) {

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

@@ -0,0 +1,35 @@
+/*
+ * 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;
+
+import org.apache.lucene.codecs.compressing.CompressionMode;
+import org.elasticsearch.index.codec.zstd.ZstdCompressionMode;
+
+public enum BinaryDVCompressionMode {
+
+    NO_COMPRESS((byte) 0, null),
+    COMPRESSED_WITH_ZSTD_LEVEL_1((byte) 1, new ZstdCompressionMode(1));
+
+    public final byte code;
+    public final CompressionMode compressionMode;
+
+    BinaryDVCompressionMode(byte code, CompressionMode compressionMode) {
+        this.code = code;
+        this.compressionMode = compressionMode;
+    }
+
+    public static BinaryDVCompressionMode fromMode(byte mode) {
+        return switch (mode) {
+            case 0 -> NO_COMPRESS;
+            case 1 -> COMPRESSED_WITH_ZSTD_LEVEL_1;
+            default -> throw new IllegalStateException("unknown compression mode [" + mode + "]");
+        };
+    }
+}

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

@@ -0,0 +1,100 @@
+/*
+ * 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.codecs.CodecUtil;
+import org.apache.lucene.store.ChecksumIndexInput;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexOutput;
+import org.apache.lucene.util.packed.DirectMonotonicWriter;
+import org.elasticsearch.core.IOUtils;
+
+import java.io.Closeable;
+import java.io.IOException;
+
+/**
+ *  Like OffsetsAccumulator builds offsets and stores in a DirectMonotonicWriter. But write to temp file
+ *  rather than directly to a DirectMonotonicWriter because the number of values is unknown. If number of
+ *  values is known prefer OffsetsWriter.
+ */
+final class DelayedOffsetAccumulator implements Closeable {
+    private final Directory dir;
+    private final long startOffset;
+
+    private int numValues = 0;
+    private final IndexOutput tempOutput;
+    private final String suffix;
+
+    DelayedOffsetAccumulator(Directory dir, IOContext context, IndexOutput data, String suffix, long startOffset) throws IOException {
+        this.dir = dir;
+        this.startOffset = startOffset;
+        this.suffix = suffix;
+
+        boolean success = false;
+        try {
+            tempOutput = dir.createTempOutput(data.getName(), suffix, context);
+            CodecUtil.writeHeader(tempOutput, ES819TSDBDocValuesFormat.META_CODEC + suffix, ES819TSDBDocValuesFormat.VERSION_CURRENT);
+            success = true;
+        } finally {
+            if (success == false) {
+                IOUtils.closeWhileHandlingException(this); // self-close because constructor caller can't
+            }
+        }
+    }
+
+    public void addDoc(long delta) throws IOException {
+        tempOutput.writeVLong(delta);
+        numValues++;
+    }
+
+    public void build(IndexOutput meta, IndexOutput data) throws IOException {
+        CodecUtil.writeFooter(tempOutput);
+        IOUtils.close(tempOutput);
+
+        // write the offsets info to the meta file by reading from temp file
+        try (ChecksumIndexInput tempInput = dir.openChecksumInput(tempOutput.getName());) {
+            CodecUtil.checkHeader(
+                tempInput,
+                ES819TSDBDocValuesFormat.META_CODEC + suffix,
+                ES819TSDBDocValuesFormat.VERSION_CURRENT,
+                ES819TSDBDocValuesFormat.VERSION_CURRENT
+            );
+            Throwable priorE = null;
+            try {
+                final DirectMonotonicWriter writer = DirectMonotonicWriter.getInstance(
+                    meta,
+                    data,
+                    numValues + 1,
+                    ES819TSDBDocValuesFormat.DIRECT_MONOTONIC_BLOCK_SHIFT
+                );
+
+                long offset = startOffset;
+                writer.add(offset);
+                for (int i = 0; i < numValues; ++i) {
+                    offset += tempInput.readVLong();
+                    writer.add(offset);
+                }
+                writer.finish();
+            } catch (Throwable e) {
+                priorE = e;
+            } finally {
+                CodecUtil.checkFooter(tempInput, priorE);
+            }
+        }
+    }
+
+    @Override
+    public void close() throws IOException {
+        if (tempOutput != null) {
+            IOUtils.close(tempOutput, () -> dir.deleteFile(tempOutput.getName()));
+        }
+    }
+}