Add support for compression level in ZstdDictTrainer

luben · Jan 27, 2025 · adb1ea2 · adb1ea2
1 parent 9b39b59
commit adb1ea2
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java b/src/main/java/com/github/luben/zstd/ZstdDictTrainer.java
@@ -10,12 +10,18 @@ public class ZstdDictTrainer {
     private final List<Integer> sampleSizes;
     private final int dictSize;
     private long filledSize;
+    private int level;
 
     public ZstdDictTrainer(int sampleSize, int dictSize) {
+        this(sampleSize, dictSize, Zstd.defaultCompressionLevel());
+    }
+
+    public ZstdDictTrainer(int sampleSize, int dictSize, int level) {
         trainingSamples = ByteBuffer.allocateDirect(sampleSize);
         sampleSizes =  new ArrayList<Integer>();
         this.allocatedSize = sampleSize;
         this.dictSize = dictSize;
+        this.level = level;
     }
 
     public synchronized boolean addSample(byte[] sample) {
@@ -34,7 +40,7 @@ public ByteBuffer trainSamplesDirect() throws ZstdException {
 
     public synchronized ByteBuffer trainSamplesDirect(boolean legacy) throws ZstdException {
         ByteBuffer dictBuffer = ByteBuffer.allocateDirect(dictSize);
-        long l = Zstd.trainFromBufferDirect(trainingSamples, copyToIntArray(sampleSizes), dictBuffer, legacy);
+        long l = Zstd.trainFromBufferDirect(trainingSamples, copyToIntArray(sampleSizes), dictBuffer, legacy, level);
         if (Zstd.isError(l)) {
             dictBuffer.limit(0);
             throw new ZstdException(l);