Add parameterized forward index compression codecs

Author: xiangfu0

pinot-common/src/test/java/org/apache/pinot/common/utils/config/TableConfigSerDeUtilsTest.java

@@ -28,6 +28,7 @@
 import java.util.Map;
 import org.apache.pinot.common.tier.TierFactory;
 import org.apache.pinot.spi.config.table.CompletionConfig;
+import org.apache.pinot.spi.config.table.CompressionCodecSpec;
 import org.apache.pinot.spi.config.table.DedupConfig;
 import org.apache.pinot.spi.config.table.FieldConfig;
 import org.apache.pinot.spi.config.table.HashFunction;
@@ -240,7 +241,9 @@ public void testSerDe()
               Lists.newArrayList(FieldConfig.IndexType.INVERTED, FieldConfig.IndexType.RANGE), null, properties),
           new FieldConfig("column2", null, Collections.emptyList(), null, null),
           new FieldConfig("column3", FieldConfig.EncodingType.RAW, Collections.emptyList(),
-              FieldConfig.CompressionCodec.SNAPPY, null));
+              FieldConfig.CompressionCodec.SNAPPY, null),
+          new FieldConfig.Builder("column4").withEncodingType(FieldConfig.EncodingType.RAW)
+              .withCompressionCodecSpec(CompressionCodecSpec.fromString("zstd(3)")).build());
       TableConfig tableConfig = tableConfigBuilder.setFieldConfigList(fieldConfigList).build();
 
       checkFieldConfig(tableConfig);
@@ -546,7 +549,7 @@ private void checkInstanceAssignmentConfig(TableConfig tableConfig) {
   private void checkFieldConfig(TableConfig tableConfig) {
     List<FieldConfig> fieldConfigList = tableConfig.getFieldConfigList();
     assertNotNull(fieldConfigList);
-    assertEquals(fieldConfigList.size(), 3);
+    assertEquals(fieldConfigList.size(), 4);
 
     FieldConfig firstFieldConfig = fieldConfigList.get(0);
     assertEquals(firstFieldConfig.getName(), "column1");
@@ -571,6 +574,14 @@ private void checkFieldConfig(TableConfig tableConfig) {
     assertNull(thirdFieldConfig.getIndexType());
     assertEquals(thirdFieldConfig.getCompressionCodec(), FieldConfig.CompressionCodec.SNAPPY);
     assertNull(thirdFieldConfig.getProperties());
+
+    FieldConfig fourthFieldConfig = fieldConfigList.get(3);
+    assertEquals(fourthFieldConfig.getName(), "column4");
+    assertEquals(fourthFieldConfig.getEncodingType(), FieldConfig.EncodingType.RAW);
+    assertEquals(fourthFieldConfig.getCompressionCodec(), FieldConfig.CompressionCodec.ZSTANDARD);
+    assertEquals(fourthFieldConfig.getCompressionCodecSpec(), CompressionCodecSpec.of(
+        FieldConfig.CompressionCodec.ZSTANDARD, 3));
+    assertNull(fourthFieldConfig.getProperties());
   }
 
   private void checkTableConfigWithUpsertConfig(TableConfig tableConfig) {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/compression/ChunkCompressorFactory.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.segment.local.io.compression;
 
+import javax.annotation.Nullable;
 import org.apache.pinot.segment.spi.compression.ChunkCompressionType;
 import org.apache.pinot.segment.spi.compression.ChunkCompressor;
 import org.apache.pinot.segment.spi.compression.ChunkDecompressor;
@@ -39,7 +40,7 @@ private ChunkCompressorFactory() {
    * @return Compressor for the specified type.
    */
   public static ChunkCompressor getCompressor(ChunkCompressionType compressionType) {
-    return getCompressor(compressionType, false);
+    return getCompressor(compressionType, null, false);
   }
 
   /**
@@ -51,30 +52,49 @@ public static ChunkCompressor getCompressor(ChunkCompressionType compressionType
    * @return Compressor for the specified type.
    */
   public static ChunkCompressor getCompressor(ChunkCompressionType compressionType, boolean upgradeToLengthPrefixed) {
+    return getCompressor(compressionType, null, upgradeToLengthPrefixed);
+  }
+
+  /**
+   * Returns the chunk compressor for the specified name and optional compression level.
+   */
+  public static ChunkCompressor getCompressor(ChunkCompressionType compressionType, @Nullable Integer compressionLevel,
+      boolean upgradeToLengthPrefixed) {
     switch (compressionType) {
 
       case PASS_THROUGH:
+        validateUnsupportedLevel(compressionType, compressionLevel);
         return PassThroughCompressor.INSTANCE;
 
       case SNAPPY:
+        validateUnsupportedLevel(compressionType, compressionLevel);
         return SnappyCompressor.INSTANCE;
 
       case ZSTANDARD:
-        return ZstandardCompressor.INSTANCE;
+        return compressionLevel != null ? new ZstandardCompressor(compressionLevel) : ZstandardCompressor.INSTANCE;
 
       case LZ4:
+        if (compressionLevel != null) {
+          return upgradeToLengthPrefixed ? new LZ4WithLengthCompressor(compressionLevel)
+              : new LZ4Compressor(compressionLevel);
+        }
         return upgradeToLengthPrefixed ? LZ4WithLengthCompressor.INSTANCE : LZ4Compressor.INSTANCE;
 
       case LZ4_LENGTH_PREFIXED:
+        if (compressionLevel != null) {
+          return new LZ4WithLengthCompressor(compressionLevel);
+        }
         return LZ4WithLengthCompressor.INSTANCE;
 
       case GZIP:
-        return new GzipCompressor();
+        return compressionLevel != null ? new GzipCompressor(compressionLevel) : new GzipCompressor();
 
       case DELTA:
+        validateUnsupportedLevel(compressionType, compressionLevel);
         return DeltaCompressor.INSTANCE;
 
       case DELTADELTA:
+        validateUnsupportedLevel(compressionType, compressionLevel);
         return DeltaDeltaCompressor.INSTANCE;
 
       default:
@@ -118,4 +138,12 @@ public static ChunkDecompressor getDecompressor(ChunkCompressionType compression
         throw new IllegalArgumentException("Illegal decompressor name " + compressionType);
     }
   }
+
+  private static void validateUnsupportedLevel(ChunkCompressionType compressionType,
+      @Nullable Integer compressionLevel) {
+    if (compressionLevel != null) {
+      throw new IllegalArgumentException(
+          String.format("Compression type %s does not support an explicit compression level", compressionType));
+    }
+  }
 }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/compression/GzipCompressor.java

@@ -36,6 +36,10 @@ public GzipCompressor() {
     _compressor = new Deflater();
   }
 
+  public GzipCompressor(int compressionLevel) {
+    _compressor = new Deflater(compressionLevel);
+  }
+
   @Override
   public int compress(ByteBuffer inUncompressed, ByteBuffer outCompressed)
       throws IOException {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/compression/LZ4Compressor.java

@@ -35,13 +35,20 @@ class LZ4Compressor implements ChunkCompressor {
 
   static final LZ4Compressor INSTANCE = new LZ4Compressor();
 
+  private final net.jpountz.lz4.LZ4Compressor _compressor;
+
   private LZ4Compressor() {
+    _compressor = LZ4_FACTORY.fastCompressor();
+  }
+
+  LZ4Compressor(int compressionLevel) {
+    _compressor = LZ4_FACTORY.highCompressor(compressionLevel);
   }
 
   @Override
   public int compress(ByteBuffer inUncompressed, ByteBuffer outCompressed)
       throws IOException {
-    LZ4_FACTORY.fastCompressor().compress(inUncompressed, outCompressed);
+    _compressor.compress(inUncompressed, outCompressed);
     // When the compress method returns successfully,
     // dstBuf's position() will be set to its current position() plus the compressed size of the data.
     // and srcBuf's position() will be set to its limit()
@@ -52,7 +59,7 @@ public int compress(ByteBuffer inUncompressed, ByteBuffer outCompressed)
 
   @Override
   public int maxCompressedSize(int uncompressedSize) {
-    return LZ4_FACTORY.fastCompressor().maxCompressedLength(uncompressedSize);
+    return _compressor.maxCompressedLength(uncompressedSize);
   }
 
   @Override

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/compression/LZ4WithLengthCompressor.java

@@ -39,6 +39,10 @@ private LZ4WithLengthCompressor() {
     _compressor = new LZ4CompressorWithLength(LZ4Compressor.LZ4_FACTORY.fastCompressor());
   }
 
+  LZ4WithLengthCompressor(int compressionLevel) {
+    _compressor = new LZ4CompressorWithLength(LZ4Compressor.LZ4_FACTORY.highCompressor(compressionLevel));
+  }
+
   @Override
   public int compress(ByteBuffer inUncompressed, ByteBuffer outCompressed)
       throws IOException {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/compression/ZstandardCompressor.java

@@ -33,13 +33,20 @@ class ZstandardCompressor implements ChunkCompressor {
 
   static final ZstandardCompressor INSTANCE = new ZstandardCompressor();
 
+  private final int _compressionLevel;
+
   private ZstandardCompressor() {
+    _compressionLevel = Zstd.defaultCompressionLevel();
+  }
+
+  ZstandardCompressor(int compressionLevel) {
+    _compressionLevel = compressionLevel;
   }
 
   @Override
   public int compress(ByteBuffer inUncompressed, ByteBuffer outCompressed)
       throws IOException {
-    int compressedSize = Zstd.compress(outCompressed, inUncompressed);
+    int compressedSize = Zstd.compress(outCompressed, inUncompressed, _compressionLevel);
     // When the compress method returns successfully,
     // dstBuf's position() will be set to its current position() plus the compressed size of the data.
     // and srcBuf's position() will be set to its limit()

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/BaseChunkForwardIndexWriter.java

@@ -25,6 +25,7 @@
 import java.io.RandomAccessFile;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
+import javax.annotation.Nullable;
 import org.apache.pinot.segment.local.io.compression.ChunkCompressorFactory;
 import org.apache.pinot.segment.spi.compression.ChunkCompressionType;
 import org.apache.pinot.segment.spi.compression.ChunkCompressor;
@@ -88,11 +89,18 @@ public abstract class BaseChunkForwardIndexWriter implements Closeable {
   protected BaseChunkForwardIndexWriter(File file, ChunkCompressionType compressionType, int totalDocs,
       int numDocsPerChunk, long chunkSize, int sizeOfEntry, int version, boolean fixed)
       throws IOException {
+    this(file, compressionType, null, totalDocs, numDocsPerChunk, chunkSize, sizeOfEntry, version, fixed);
+  }
+
+  protected BaseChunkForwardIndexWriter(File file, ChunkCompressionType compressionType,
+      @Nullable Integer compressionLevel, int totalDocs, int numDocsPerChunk, long chunkSize, int sizeOfEntry,
+      int version, boolean fixed)
+      throws IOException {
     Preconditions.checkArgument(version == 2 || version == 3 || (fixed && (version == 4 || version == 5)),
         "Illegal version: %s for %s bytes values", version, fixed ? "fixed" : "variable");
     Preconditions.checkArgument(chunkSize <= Integer.MAX_VALUE, "Chunk size limited to 2GB");
     _chunkSize = (int) chunkSize;
-    _chunkCompressor = ChunkCompressorFactory.getCompressor(compressionType);
+    _chunkCompressor = ChunkCompressorFactory.getCompressor(compressionType, compressionLevel, false);
     _headerEntryChunkOffsetSize = version == 2 ? Integer.BYTES : Long.BYTES;
     _dataOffset = writeHeader(compressionType, totalDocs, numDocsPerChunk, sizeOfEntry, version);
     _chunkBuffer = ByteBuffer.allocateDirect(_chunkSize);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/FixedByteChunkForwardIndexWriter.java

@@ -21,6 +21,7 @@
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.pinot.segment.spi.compression.ChunkCompressionType;
 
@@ -48,7 +49,13 @@ public class FixedByteChunkForwardIndexWriter extends BaseChunkForwardIndexWrite
   public FixedByteChunkForwardIndexWriter(File file, ChunkCompressionType compressionType, int totalDocs,
       int numDocsPerChunk, int sizeOfEntry, int writerVersion)
       throws IOException {
-    super(file, compressionType, totalDocs, normalizeDocsPerChunk(writerVersion, numDocsPerChunk),
+    this(file, compressionType, null, totalDocs, numDocsPerChunk, sizeOfEntry, writerVersion);
+  }
+
+  public FixedByteChunkForwardIndexWriter(File file, ChunkCompressionType compressionType,
+      @Nullable Integer compressionLevel, int totalDocs, int numDocsPerChunk, int sizeOfEntry, int writerVersion)
+      throws IOException {
+    super(file, compressionType, compressionLevel, totalDocs, normalizeDocsPerChunk(writerVersion, numDocsPerChunk),
         (long) sizeOfEntry * normalizeDocsPerChunk(writerVersion, numDocsPerChunk), sizeOfEntry, writerVersion, true);
     _chunkDataOffset = 0;
   }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkForwardIndexWriter.java

@@ -22,6 +22,7 @@
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.math.BigDecimal;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.pinot.segment.local.utils.ArraySerDeUtils;
 import org.apache.pinot.segment.spi.compression.ChunkCompressionType;
@@ -66,7 +67,14 @@ public class VarByteChunkForwardIndexWriter extends BaseChunkForwardIndexWriter
   public VarByteChunkForwardIndexWriter(File file, ChunkCompressionType compressionType, int totalDocs,
       int numDocsPerChunk, int lengthOfLongestEntry, int writerVersion)
       throws IOException {
-    super(file, compressionType, totalDocs, numDocsPerChunk,
+    this(file, compressionType, null, totalDocs, numDocsPerChunk, lengthOfLongestEntry, writerVersion);
+  }
+
+  public VarByteChunkForwardIndexWriter(File file, ChunkCompressionType compressionType,
+      @Nullable Integer compressionLevel, int totalDocs, int numDocsPerChunk, int lengthOfLongestEntry,
+      int writerVersion)
+      throws IOException {
+    super(file, compressionType, compressionLevel, totalDocs, numDocsPerChunk,
         numDocsPerChunk * (CHUNK_HEADER_ENTRY_ROW_OFFSET_SIZE + (long) lengthOfLongestEntry),
         // chunkSize
         lengthOfLongestEntry, writerVersion, false);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/VarByteChunkForwardIndexWriterV4.java

@@ -27,6 +27,7 @@
 import java.nio.ByteOrder;
 import java.nio.channels.FileChannel;
 import java.nio.charset.StandardCharsets;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.NotThreadSafe;
 import org.apache.commons.io.FileUtils;
 import org.apache.pinot.segment.local.io.compression.ChunkCompressorFactory;
@@ -95,10 +96,16 @@ public class VarByteChunkForwardIndexWriterV4 implements VarByteChunkWriter {
 
   public VarByteChunkForwardIndexWriterV4(File file, ChunkCompressionType compressionType, int chunkSize)
       throws IOException {
+    this(file, compressionType, null, chunkSize);
+  }
+
+  public VarByteChunkForwardIndexWriterV4(File file, ChunkCompressionType compressionType,
+      @Nullable Integer compressionLevel, int chunkSize)
+      throws IOException {
     _dataBuffer = new File(file.getParentFile(), file.getName() + DATA_BUFFER_SUFFIX);
     _output = new RandomAccessFile(file, "rw");
     _dataChannel = new RandomAccessFile(_dataBuffer, "rw").getChannel();
-    _chunkCompressor = ChunkCompressorFactory.getCompressor(compressionType, true);
+    _chunkCompressor = ChunkCompressorFactory.getCompressor(compressionType, compressionLevel, true);
     _chunkBuffer = ByteBuffer.allocateDirect(chunkSize).order(ByteOrder.LITTLE_ENDIAN);
     _compressionBuffer =
         ByteBuffer.allocateDirect(_chunkCompressor.maxCompressedSize(chunkSize)).order(ByteOrder.LITTLE_ENDIAN);