Add retained size to compressor

airlift · Sep 4, 2024 · 3cdea23 · 3cdea23
1 parent ea31a52
commit 3cdea23
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 2 deletions.
diff --git a/src/main/java/io/airlift/compress/v3/Compressor.java b/src/main/java/io/airlift/compress/v3/Compressor.java
@@ -28,4 +28,9 @@ public interface Compressor
      * @return number of bytes written to the output
      */
     int compress(MemorySegment input, MemorySegment output);
+
+    default int getRetainedSizeInBytes(int inputLength)
+    {
+        return 0;
+    }
 }
diff --git a/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java b/src/main/java/io/airlift/compress/v3/lz4/Lz4JavaCompressor.java
@@ -69,6 +69,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return Lz4RawCompressor.computeTableSize(inputLength);
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");

diff --git a/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java b/src/main/java/io/airlift/compress/v3/lz4/Lz4RawCompressor.java
@@ -300,7 +300,7 @@ private static long encodeRunLength(
         return output;
     }
 
-    private static int computeTableSize(int inputSize)
+    static int computeTableSize(int inputSize)
     {
         // smallest power of 2 larger than inputSize
         int target = Integer.highestOneBit(inputSize - 1) << 1;

diff --git a/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java b/src/main/java/io/airlift/compress/v3/lzo/LzoCompressor.java
@@ -71,6 +71,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return MAX_TABLE_SIZE;
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");

diff --git a/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java b/src/main/java/io/airlift/compress/v3/snappy/SnappyJavaCompressor.java
@@ -75,6 +75,12 @@ public int compress(MemorySegment input, MemorySegment output)
         }
     }
 
+    @Override
+    public int getRetainedSizeInBytes(int inputLength)
+    {
+        return SnappyRawCompressor.getHashTableSize(inputLength);
+    }
+
     private static void verifyRange(byte[] data, int offset, int length)
     {
         requireNonNull(data, "data is null");

diff --git a/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java b/src/main/java/io/airlift/compress/v3/snappy/SnappyRawCompressor.java
@@ -344,7 +344,7 @@ private static long emitCopy(Object outputBase, long output, long input, long ma
     }
 
     @SuppressWarnings("IllegalToken")
-    private static int getHashTableSize(int inputSize)
+    static int getHashTableSize(int inputSize)
     {
         // Use smaller hash table when input.size() is smaller, since we
         // fill the table, incurring O(hash table size) overhead for