Add binary doc value compression with variable doc count blocks (elastic#137139)

Add compression for binary doc values using Zstd and blocks with a variable number of values.

Block-wise LZ4 compression for binary doc values was previously added to Lucene in LUCENE-9211. This was subsequently removed in LUCENE-9378 due to query performance issues. We investigated adding to adding the original Lucene implementation to ES in elastic#112416 and elastic#105301. This previous approach used a constant number of values per block (specifically 32 values). This is nice because it makes it very easy to map a given value index (eg docId for dense values) to the block containing it with blockId = docId / 32. Unfortunately, if values are very large we cannot reduce the number of values per block and (de)compressing a block could cause an OOM. Also, since this is a concern, we have to keep the number of values lower than ideal.

This PR instead stores a variable number of documents per block. It stores a minimum of 1 document per block and stops adding values when the size of a block exceeds a threshold. Like the previous version it stores an array of address for the start of each block. Additionally, it stores a parallel array with the value index at the start of each block. When looking up a given value index, if it is not in the current block, we binary search the array of value index starts to find the blockId containing the value. Then look up the address of the block.

Author: parkertimmins

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,13 @@ 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_ZSTD_LEVEL_1,
+            true
+        );
         config.setCodec(new Elasticsearch92Lucene103Codec() {
             @Override
             public DocValuesFormat getDocValuesFormatForField(String field) {

server/src/main/java/org/elasticsearch/index/codec/PerFieldFormatSupplier.java

@@ -41,6 +41,7 @@
 public class PerFieldFormatSupplier {
 
     private static final Set<String> INCLUDE_META_FIELDS;
+    private static final Set<String> EXCLUDE_MAPPER_TYPES;
 
     static {
         // TODO: should we just allow all fields to use tsdb doc values codec?
@@ -53,6 +54,7 @@ public class PerFieldFormatSupplier {
         // Don't the include _recovery_source_size and _recovery_source fields, since their values can be trimmed away in
         // RecoverySourcePruneMergePolicy, which leads to inconsistencies between merge stats and actual values.
         INCLUDE_META_FIELDS = Collections.unmodifiableSet(includeMetaField);
+        EXCLUDE_MAPPER_TYPES = Set.of("geo_shape");
     }
 
     private static final DocValuesFormat docValuesFormat = new Lucene90DocValuesFormat();
@@ -145,6 +147,10 @@ boolean useTSDBDocValuesFormat(final String field) {
             return false;
         }
 
+        if (excludeMapperTypes(field)) {
+            return false;
+        }
+
         return mapperService != null
             && mapperService.getIndexSettings().useTimeSeriesDocValuesFormat()
             && mapperService.getIndexSettings().isES87TSDBCodecEnabled();
@@ -154,4 +160,29 @@ private boolean excludeFields(String fieldName) {
         return fieldName.startsWith("_") && INCLUDE_META_FIELDS.contains(fieldName) == false;
     }
 
+    private boolean excludeMapperTypes(String fieldName) {
+        var typeName = getMapperType(fieldName);
+        if (typeName == null) {
+            return false;
+        }
+        return EXCLUDE_MAPPER_TYPES.contains(getMapperType(fieldName));
+    }
+
+    private boolean isTimeSeriesModeIndex() {
+        return mapperService != null && IndexMode.TIME_SERIES == mapperService.getIndexSettings().getMode();
+    }
+
+    private boolean isLogsModeIndex() {
+        return mapperService != null && IndexMode.LOGSDB == mapperService.getIndexSettings().getMode();
+    }
+
+    String getMapperType(final String field) {
+        if (mapperService != null) {
+            Mapper mapper = mapperService.mappingLookup().getMapper(field);
+            if (mapper != null) {
+                return mapper.typeName();
+            }
+        }
+        return null;
+    }
 }

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

@@ -0,0 +1,65 @@
+/*
+ * 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_ZSTD_LEVEL_1((byte) 1, new ZstdCompressionMode(1));
+
+    public final byte code;
+    private final CompressionMode compressionMode;
+
+    private static final BinaryDVCompressionMode[] values = new BinaryDVCompressionMode[values().length];
+    static {
+        for (BinaryDVCompressionMode mode : values()) {
+            values[mode.code] = mode;
+        }
+    }
+
+    BinaryDVCompressionMode(byte code, CompressionMode compressionMode) {
+        this.code = code;
+        this.compressionMode = compressionMode;
+    }
+
+    public static BinaryDVCompressionMode fromMode(byte code) {
+        if (code < 0 || code >= values.length) {
+            throw new IllegalStateException("unknown compression mode [" + code + "]");
+        }
+        return values[code];
+    }
+
+    public CompressionMode compressionMode() {
+        if (compressionMode == null) {
+            throw new UnsupportedOperationException("BinaryDVCompressionMode [" + code + "] does not support compression");
+        }
+        return compressionMode;
+    }
+
+    public record BlockHeader(boolean isCompressed) {
+        static final byte IS_COMPRESSED = 0x1;
+
+        public static BlockHeader fromByte(byte header) {
+            boolean isCompressed = (header & IS_COMPRESSED) != 0;
+            return new BlockHeader(isCompressed);
+        }
+
+        public byte toByte() {
+            byte header = 0;
+            if (isCompressed) {
+                header |= IS_COMPRESSED;
+            }
+            return header;
+        }
+    }
+}

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

@@ -0,0 +1,142 @@
+/*
+ * 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;
+
+final class BlockMetadataAccumulator implements Closeable {
+
+    private final DelayedOffsetAccumulator blockAddressAcc;
+    private final DelayedOffsetAccumulator blockDocRangeAcc;
+
+    BlockMetadataAccumulator(Directory dir, IOContext context, IndexOutput data, long addressesStart) throws IOException {
+        boolean success = false;
+        try {
+            blockDocRangeAcc = new DelayedOffsetAccumulator(dir, context, data, "block-doc-ranges", 0);
+            blockAddressAcc = new DelayedOffsetAccumulator(dir, context, data, "block-addresses", addressesStart);
+            success = true;
+        } finally {
+            if (success == false) {
+                IOUtils.closeWhileHandlingException(this); // self-close because constructor caller can't
+            }
+        }
+    }
+
+    public void addDoc(long numDocsInBlock, long blockLenInBytes) throws IOException {
+        blockDocRangeAcc.addDoc(numDocsInBlock);
+        blockAddressAcc.addDoc(blockLenInBytes);
+    }
+
+    public void build(IndexOutput meta, IndexOutput data) throws IOException {
+        long dataAddressesStart = data.getFilePointer();
+        blockAddressAcc.build(meta, data);
+        long dataDocRangeStart = data.getFilePointer();
+        long addressesLength = dataDocRangeStart - dataAddressesStart;
+        meta.writeLong(addressesLength);
+
+        meta.writeLong(dataDocRangeStart);
+        blockDocRangeAcc.build(meta, data);
+        long docRangesLen = data.getFilePointer() - dataDocRangeStart;
+        meta.writeLong(docRangesLen);
+    }
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.closeWhileHandlingException(blockAddressAcc, blockDocRangeAcc);
+    }
+
+    /**
+     * 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.
+     */
+    static 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
+                }
+            }
+        }
+
+        void addDoc(long delta) throws IOException {
+            tempOutput.writeVLong(delta);
+            numValues++;
+        }
+
+        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()));
+            }
+        }
+    }
+}