Restore compressionLevel in the dict builder

luben · Feb 20, 2025 · 6a6b1b1 · 6a6b1b1
1 parent cc5ca81
commit 6a6b1b1
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/src/main/native/dictBuilder/zdict.c b/src/main/native/dictBuilder/zdict.c
@@ -1105,15 +1105,18 @@ size_t ZDICT_trainFromBuffer_legacy(void* dictBuffer, size_t dictBufferCapacity,
 
 
 size_t ZDICT_trainFromBuffer(void* dictBuffer, size_t dictBufferCapacity,
-                             const void* samplesBuffer, const size_t* samplesSizes, unsigned nbSamples)
+                             const void* samplesBuffer, const size_t* samplesSizes, unsigned nbSamples, int compressionLevel)
 {
     ZDICT_fastCover_params_t params;
     DEBUGLOG(3, "ZDICT_trainFromBuffer");
     memset(&params, 0, sizeof(params));
     params.d = 8;
     params.steps = 4;
-    /* Use default level since no compression level information is available */
+    /* Use default level if no compression level information is available */
     params.zParams.compressionLevel = ZSTD_CLEVEL_DEFAULT;
+    if (compressionLevel != 0) {
+        params.zParams.compressionLevel = compressionLevel;
+    }
 #if defined(DEBUGLEVEL) && (DEBUGLEVEL>=1)
     params.zParams.notificationLevel = DEBUGLEVEL;
 #endif

diff --git a/src/main/native/jni_zdict.c b/src/main/native/jni_zdict.c
@@ -48,7 +48,7 @@ JNIEXPORT jlong Java_com_github_luben_zstd_Zstd_trainFromBuffer0
         params.zParams.compressionLevel = compressionLevel;
         size = ZDICT_trainFromBuffer_legacy(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples, params);
     } else {
-        size = ZDICT_trainFromBuffer(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples);
+        size = ZDICT_trainFromBuffer(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples, compressionLevel);
     }
     (*env)->ReleasePrimitiveArrayCritical(env, dictBuffer, dict_buff, 0);
     free(samples_buffer);
@@ -85,7 +85,7 @@ JNIEXPORT jlong Java_com_github_luben_zstd_Zstd_trainFromBufferDirect0
         params.zParams.compressionLevel = compressionLevel;
         size = ZDICT_trainFromBuffer_legacy(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples, params);
     } else {
-        size = ZDICT_trainFromBuffer(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples);
+        size = ZDICT_trainFromBuffer(dict_buff, dict_capacity, samples_buffer, samples_sizes, num_samples, compressionLevel);
     }
 E2: free(samples_sizes);
 E1: return size;

diff --git a/src/main/native/zdict.h b/src/main/native/zdict.h
@@ -209,7 +209,7 @@ extern "C" {
  */
 ZDICTLIB_API size_t ZDICT_trainFromBuffer(void* dictBuffer, size_t dictBufferCapacity,
                                     const void* samplesBuffer,
-                                    const size_t* samplesSizes, unsigned nbSamples);
+                                    const size_t* samplesSizes, unsigned nbSamples, int compressionLevel);
 
 typedef struct {
     int      compressionLevel;   /**< optimize for a specific zstd compression level; 0 means default */