decouple column serializer compression closers from SegmentWriteoutMe…

…dium to optionally allow serializers to release direct memory allocated for compression earlier than when segment is completed (apache#16076)
monetate · Mar 11, 2024 · 313da98 · 313da98
1 parent 8084f22
commit 313da98
Show file tree

Hide file tree

Showing 39 changed files with 224 additions and 117 deletions.
diff --git a/...arks/src/test/java/org/apache/druid/benchmark/compression/BaseColumnarLongsBenchmark.java b/...arks/src/test/java/org/apache/druid/benchmark/compression/BaseColumnarLongsBenchmark.java
@@ -240,7 +240,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
             "lz4-longs",
             ByteOrder.LITTLE_ENDIAN,
             CompressionFactory.LongEncodingStrategy.LONGS,
-            CompressionStrategy.LZ4
+            CompressionStrategy.LZ4,
+            writeOutMedium.getCloser()
         );
         break;
       case "lz4-auto":
@@ -250,7 +251,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
             "lz4-auto",
             ByteOrder.LITTLE_ENDIAN,
             CompressionFactory.LongEncodingStrategy.AUTO,
-            CompressionStrategy.LZ4
+            CompressionStrategy.LZ4,
+            writeOutMedium.getCloser()
         );
         break;
       case "none-longs":
@@ -260,7 +262,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
             "none-longs",
             ByteOrder.LITTLE_ENDIAN,
             CompressionFactory.LongEncodingStrategy.LONGS,
-            CompressionStrategy.NONE
+            CompressionStrategy.NONE,
+            writeOutMedium.getCloser()
         );
         break;
       case "none-auto":
@@ -270,7 +273,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
             "none-auto",
             ByteOrder.LITTLE_ENDIAN,
             CompressionFactory.LongEncodingStrategy.AUTO,
-            CompressionStrategy.NONE
+            CompressionStrategy.NONE,
+            writeOutMedium.getCloser()
         );
         break;
       case "zstd-longs":
@@ -280,7 +284,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
                 "zstd-longs",
                 ByteOrder.LITTLE_ENDIAN,
                 CompressionFactory.LongEncodingStrategy.LONGS,
-                CompressionStrategy.ZSTD
+                CompressionStrategy.ZSTD,
+                writeOutMedium.getCloser()
         );
         break;
       case "zstd-auto":
@@ -290,7 +295,8 @@ static int encodeToFile(long[] vals, String encoding, FileChannel output)throws
                 "zstd-auto",
                 ByteOrder.LITTLE_ENDIAN,
                 CompressionFactory.LongEncodingStrategy.AUTO,
-                CompressionStrategy.ZSTD
+                CompressionStrategy.ZSTD,
+                writeOutMedium.getCloser()
         );
         break;
       default:

diff --git a/...t/java/org/apache/druid/benchmark/compression/FloatCompressionBenchmarkFileGenerator.java b/...t/java/org/apache/druid/benchmark/compression/FloatCompressionBenchmarkFileGenerator.java
@@ -29,6 +29,7 @@
 import org.apache.druid.segment.generator.ColumnValueGenerator;
 import org.apache.druid.segment.generator.GeneratorColumnSchema;
 import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMedium;
+import org.apache.druid.segment.writeout.SegmentWriteOutMedium;
 
 import java.io.BufferedReader;
 import java.io.File;
@@ -155,12 +156,14 @@ public static void main(String[] args) throws IOException
         compFile.delete();
         File dataFile = new File(dir, entry.getKey());
 
+        SegmentWriteOutMedium segmentWriteOutMedium = new OffHeapMemorySegmentWriteOutMedium();
         ColumnarFloatsSerializer writer = CompressionFactory.getFloatSerializer(
             "float-benchmark",
-            new OffHeapMemorySegmentWriteOutMedium(),
+            segmentWriteOutMedium,
             "float",
             ByteOrder.nativeOrder(),
-            compression
+            compression,
+            segmentWriteOutMedium.getCloser()
         );
         try (
             BufferedReader br = Files.newBufferedReader(dataFile.toPath(), StandardCharsets.UTF_8);

diff --git a/...st/java/org/apache/druid/benchmark/compression/LongCompressionBenchmarkFileGenerator.java b/...st/java/org/apache/druid/benchmark/compression/LongCompressionBenchmarkFileGenerator.java
@@ -29,6 +29,7 @@
 import org.apache.druid.segment.generator.ColumnValueGenerator;
 import org.apache.druid.segment.generator.GeneratorColumnSchema;
 import org.apache.druid.segment.writeout.OffHeapMemorySegmentWriteOutMedium;
+import org.apache.druid.segment.writeout.SegmentWriteOutMedium;
 
 import java.io.BufferedReader;
 import java.io.File;
@@ -148,13 +149,15 @@ public static void main(String[] args) throws IOException
           compFile.delete();
           File dataFile = new File(dir, entry.getKey());
 
+          SegmentWriteOutMedium segmentWriteOutMedium = new OffHeapMemorySegmentWriteOutMedium();
           ColumnarLongsSerializer writer = CompressionFactory.getLongSerializer(
               "long-benchmark",
-              new OffHeapMemorySegmentWriteOutMedium(),
+              segmentWriteOutMedium,
               "long",
               ByteOrder.nativeOrder(),
               encoding,
-              compression
+              compression,
+              segmentWriteOutMedium.getCloser()
           );
           try (
               BufferedReader br = Files.newBufferedReader(dataFile.toPath(), StandardCharsets.UTF_8);

diff --git a/.../java/org/apache/druid/compressedbigdecimal/CompressedBigDecimalLongColumnSerializer.java b/.../java/org/apache/druid/compressedbigdecimal/CompressedBigDecimalLongColumnSerializer.java
@@ -49,21 +49,26 @@ public class CompressedBigDecimalLongColumnSerializer implements GenericColumnSe
    */
   public static CompressedBigDecimalLongColumnSerializer create(
       SegmentWriteOutMedium segmentWriteOutMedium,
-      String filenameBase)
+      String filenameBase
+  )
   {
     return new CompressedBigDecimalLongColumnSerializer(
         CompressedVSizeColumnarIntsSerializer.create(
             "dummy",
             segmentWriteOutMedium,
             String.format(Locale.ROOT, "%s.scale", filenameBase),
             16,
-            CompressionStrategy.LZ4),
+            CompressionStrategy.LZ4,
+            segmentWriteOutMedium.getCloser()
+        ),
         V3CompressedVSizeColumnarMultiIntsSerializer.create(
             "dummy",
             segmentWriteOutMedium,
             String.format(Locale.ROOT, "%s.magnitude", filenameBase),
             Integer.MAX_VALUE,
-            CompressionStrategy.LZ4));
+            CompressionStrategy.LZ4
+        )
+    );
   }
 
   private final CompressedVSizeColumnarIntsSerializer scaleWriter;

diff --git a/processing/src/main/java/org/apache/druid/segment/DictionaryEncodedColumnMerger.java b/processing/src/main/java/org/apache/druid/segment/DictionaryEncodedColumnMerger.java
@@ -423,7 +423,8 @@ protected void setupEncodedValueWriter() throws IOException
             segmentWriteOutMedium,
             filenameBase,
             cardinality,
-            compressionStrategy
+            compressionStrategy,
+            segmentWriteOutMedium.getCloser()
         );
       } else {
         encodedValueSerializer = new VSizeColumnarIntsSerializer(segmentWriteOutMedium, cardinality);

diff --git a/processing/src/main/java/org/apache/druid/segment/DoubleColumnSerializer.java b/processing/src/main/java/org/apache/druid/segment/DoubleColumnSerializer.java
@@ -72,7 +72,8 @@ public void open() throws IOException
         segmentWriteOutMedium,
         StringUtils.format("%s.double_column", filenameBase),
         byteOrder,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
   }

diff --git a/processing/src/main/java/org/apache/druid/segment/DoubleColumnSerializerV2.java b/processing/src/main/java/org/apache/druid/segment/DoubleColumnSerializerV2.java
@@ -99,7 +99,8 @@ public void open() throws IOException
         segmentWriteOutMedium,
         StringUtils.format("%s.double_column", filenameBase),
         byteOrder,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
     nullValueBitmapWriter = new ByteBufferWriter<>(

diff --git a/processing/src/main/java/org/apache/druid/segment/FloatColumnSerializer.java b/processing/src/main/java/org/apache/druid/segment/FloatColumnSerializer.java
@@ -72,7 +72,8 @@ public void open() throws IOException
         segmentWriteOutMedium,
         StringUtils.format("%s.float_column", filenameBase),
         byteOrder,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
   }

diff --git a/processing/src/main/java/org/apache/druid/segment/FloatColumnSerializerV2.java b/processing/src/main/java/org/apache/druid/segment/FloatColumnSerializerV2.java
@@ -99,7 +99,8 @@ public void open() throws IOException
         segmentWriteOutMedium,
         StringUtils.format("%s.float_column", filenameBase),
         byteOrder,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
     nullValueBitmapWriter = new ByteBufferWriter<>(

diff --git a/processing/src/main/java/org/apache/druid/segment/IndexMergerV9.java b/processing/src/main/java/org/apache/druid/segment/IndexMergerV9.java
@@ -202,7 +202,6 @@ private File makeIndexFiles(
         mergers.add(
             handler.makeMerger(
                 indexSpec,
-
                 segmentWriteOutMedium,
                 dimFormats.get(i).toColumnCapabilities(),
                 progress,

diff --git a/processing/src/main/java/org/apache/druid/segment/LongColumnSerializer.java b/processing/src/main/java/org/apache/druid/segment/LongColumnSerializer.java
@@ -80,7 +80,8 @@ public void open() throws IOException
         StringUtils.format("%s.long_column", filenameBase),
         byteOrder,
         encoding,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
   }

diff --git a/processing/src/main/java/org/apache/druid/segment/LongColumnSerializerV2.java b/processing/src/main/java/org/apache/druid/segment/LongColumnSerializerV2.java
@@ -105,7 +105,8 @@ public void open() throws IOException
         StringUtils.format("%s.long_column", filenameBase),
         byteOrder,
         encoding,
-        compression
+        compression,
+        segmentWriteOutMedium.getCloser()
     );
     writer.open();
     nullValueBitmapWriter = new ByteBufferWriter<>(

diff --git a/...ing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarDoublesSerializer.java b/...ing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarDoublesSerializer.java
@@ -56,19 +56,20 @@ public class BlockLayoutColumnarDoublesSerializer implements ColumnarDoublesSeri
       SegmentWriteOutMedium segmentWriteOutMedium,
       String filenameBase,
       ByteOrder byteOrder,
-      CompressionStrategy compression
+      CompressionStrategy compression,
+      Closer closer
   )
   {
     this.columnName = columnName;
     this.flattener = GenericIndexedWriter.ofCompressedByteBuffers(
         segmentWriteOutMedium,
         filenameBase,
         compression,
-        CompressedPools.BUFFER_SIZE
+        CompressedPools.BUFFER_SIZE,
+        closer
     );
     this.compression = compression;
     CompressionStrategy.Compressor compressor = compression.getCompressor();
-    Closer closer = segmentWriteOutMedium.getCloser();
     this.endBuffer = compressor.allocateInBuffer(CompressedPools.BUFFER_SIZE, closer).order(byteOrder);
   }
 

diff --git a/...sing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarFloatsSerializer.java b/...sing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarFloatsSerializer.java
@@ -56,19 +56,20 @@ public class BlockLayoutColumnarFloatsSerializer implements ColumnarFloatsSerial
       SegmentWriteOutMedium segmentWriteOutMedium,
       String filenameBase,
       ByteOrder byteOrder,
-      CompressionStrategy compression
+      CompressionStrategy compression,
+      Closer closer
   )
   {
     this.columnName = columnName;
     this.flattener = GenericIndexedWriter.ofCompressedByteBuffers(
         segmentWriteOutMedium,
         filenameBase,
         compression,
-        CompressedPools.BUFFER_SIZE
+        CompressedPools.BUFFER_SIZE,
+        closer
     );
     this.compression = compression;
     CompressionStrategy.Compressor compressor = compression.getCompressor();
-    Closer closer = segmentWriteOutMedium.getCloser();
     this.endBuffer = compressor.allocateInBuffer(CompressedPools.BUFFER_SIZE, closer).order(byteOrder);
   }
 

diff --git a/...ssing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarLongsSerializer.java b/...ssing/src/main/java/org/apache/druid/segment/data/BlockLayoutColumnarLongsSerializer.java
@@ -19,13 +19,13 @@
 
 package org.apache.druid.segment.data;
 
+import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.io.smoosh.FileSmoosher;
 import org.apache.druid.segment.CompressedPools;
 import org.apache.druid.segment.serde.MetaSerdeHelper;
 import org.apache.druid.segment.writeout.SegmentWriteOutMedium;
 
 import javax.annotation.Nullable;
-
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -59,17 +59,24 @@ public class BlockLayoutColumnarLongsSerializer implements ColumnarLongsSerializ
       String filenameBase,
       ByteOrder byteOrder,
       CompressionFactory.LongEncodingWriter writer,
-      CompressionStrategy compression
+      CompressionStrategy compression,
+      Closer closer
   )
   {
     this.columnName = columnName;
     this.sizePer = writer.getBlockSize(CompressedPools.BUFFER_SIZE);
     int bufferSize = writer.getNumBytes(sizePer);
-    this.flattener = GenericIndexedWriter.ofCompressedByteBuffers(segmentWriteOutMedium, filenameBase, compression, bufferSize);
+    this.flattener = GenericIndexedWriter.ofCompressedByteBuffers(
+        segmentWriteOutMedium,
+        filenameBase,
+        compression,
+        bufferSize,
+        closer
+    );
     this.writer = writer;
     this.compression = compression;
     CompressionStrategy.Compressor compressor = compression.getCompressor();
-    endBuffer = compressor.allocateInBuffer(writer.getNumBytes(sizePer), segmentWriteOutMedium.getCloser()).order(byteOrder);
+    endBuffer = compressor.allocateInBuffer(writer.getNumBytes(sizePer), closer).order(byteOrder);
     writer.setBuffer(endBuffer);
     numInsertedForNextFlush = sizePer;
   }

diff --git a/processing/src/main/java/org/apache/druid/segment/data/CompressedBlockSerializer.java b/processing/src/main/java/org/apache/druid/segment/data/CompressedBlockSerializer.java
@@ -20,6 +20,7 @@
 package org.apache.druid.segment.data;
 
 import org.apache.druid.io.Channels;
+import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.io.smoosh.FileSmoosher;
 import org.apache.druid.segment.CompressedPools;
 import org.apache.druid.segment.serde.MetaSerdeHelper;
@@ -60,18 +61,16 @@ public class CompressedBlockSerializer implements Serializer
 
   public CompressedBlockSerializer(
       SegmentWriteOutMedium segmentWriteOutMedium,
-
       CompressionStrategy compression,
-      int blockSize
+      int blockSize,
+      Closer closer
   )
   {
     this.segmentWriteOutMedium = segmentWriteOutMedium;
     this.compression = compression;
     this.compressor = compression.getCompressor();
-    this.uncompressedDataBuffer = compressor.allocateInBuffer(blockSize, segmentWriteOutMedium.getCloser())
-                                            .order(ByteOrder.nativeOrder());
-    this.compressedDataBuffer = compressor.allocateOutBuffer(blockSize, segmentWriteOutMedium.getCloser())
-                                          .order(ByteOrder.nativeOrder());
+    this.uncompressedDataBuffer = compressor.allocateInBuffer(blockSize, closer).order(ByteOrder.nativeOrder());
+    this.compressedDataBuffer = compressor.allocateOutBuffer(blockSize, closer).order(ByteOrder.nativeOrder());
   }
 
   public void open() throws IOException

diff --git a/processing/src/main/java/org/apache/druid/segment/data/CompressedColumnarIntsSerializer.java b/processing/src/main/java/org/apache/druid/segment/data/CompressedColumnarIntsSerializer.java
@@ -58,39 +58,40 @@ public class CompressedColumnarIntsSerializer extends SingleValueColumnarIntsSer
       final String filenameBase,
       final int chunkFactor,
       final ByteOrder byteOrder,
-      final CompressionStrategy compression
+      final CompressionStrategy compression,
+      final Closer closer
   )
   {
     this(
         columnName,
-        segmentWriteOutMedium,
         chunkFactor,
         byteOrder,
         compression,
         GenericIndexedWriter.ofCompressedByteBuffers(
             segmentWriteOutMedium,
             filenameBase,
             compression,
-            chunkFactor * Integer.BYTES
-        )
+            chunkFactor * Integer.BYTES,
+            closer
+        ),
+        closer
     );
   }
 
   CompressedColumnarIntsSerializer(
       final String columnName,
-      final SegmentWriteOutMedium segmentWriteOutMedium,
       final int chunkFactor,
       final ByteOrder byteOrder,
       final CompressionStrategy compression,
-      final GenericIndexedWriter<ByteBuffer> flattener
+      final GenericIndexedWriter<ByteBuffer> flattener,
+      final Closer closer
   )
   {
     this.columnName = columnName;
     this.chunkFactor = chunkFactor;
     this.compression = compression;
     this.flattener = flattener;
     CompressionStrategy.Compressor compressor = compression.getCompressor();
-    Closer closer = segmentWriteOutMedium.getCloser();
     this.endBuffer = compressor.allocateInBuffer(chunkFactor * Integer.BYTES, closer).order(byteOrder);
     this.numInserted = 0;
   }