Support pre-defined compression dictionaries (#14253)

Summary:
... in addition to those derived from samples. This could be useful when trade-offs favor an offline trained dictionary that's good for the whole work load, which can involve heavy-weight training, vs. on-the-fly training on samples for each file, which has limitations.

This involves some breaking changes to some deeper parts of the new compression API. I'm not concerned about performance because this doesn't touch the per-block parts of the API, just the per-file parts.

Bonus: change to
CompressionManagerWrapper::FindCompatibleCompressionManager to implement what is likely the preferred behavior.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/14253

Test Plan: unit test included

Reviewed By: hx235

Differential Revision: D91082208

Pulled By: pdillinger

fbshipit-source-id: 1442db65e15c9435437204c19787c96f7a40a207

include/rocksdb/advanced_compression.h

@@ -11,6 +11,8 @@
 
 #pragma once
 
+#include <variant>
+
 #include "rocksdb/cache.h"
 #include "rocksdb/compression_type.h"
 #include "rocksdb/data_structure.h"
@@ -56,7 +58,64 @@ class Decompressor;
 // because RocksDB is not exception-safe. This could cause undefined behavior
 // including data loss, unreported corruption, deadlocks, and more.
 class Compressor {
- public:
+ public:  // Auxiliary types
+  // No dictionary should be used (for a given block type).
+  struct DictDisabled {};
+
+  // A recommendation for dictionary compression by collecting samples from
+  // blocks. The caller should collect up to `max_sample_bytes` of sample data
+  // and pass it to MaybeCloneSpecialized() to create a specialized compressor.
+  struct DictSampling {
+    // Maximum total bytes of sample data to collect from blocks.
+    // This controls how much data is buffered before dictionary training.
+    size_t max_sample_bytes = 0;
+  };
+
+  // A pre-defined dictionary that is recommended or specified for direct use
+  // with MaybeCloneSpecialized(), without any sampling.
+  struct DictPreDefined {
+    // The owned raw/serialized dictionary bytes. Recommend std::move to
+    // MaybeCloneSpecialized()
+    std::string dict_data;
+  };
+
+  // The result type for GetDictGuidance() - indicates how dictionary
+  // compression should be configured for a given block type.
+  using DictConfig = std::variant<DictDisabled, DictSampling, DictPreDefined>;
+
+  // Sample data collected from blocks for dictionary training.
+  struct DictSamples {
+    // All the sample input blocks stored contiguously
+    std::string sample_data;
+    // The lengths of each of the sample blocks in `sample_data`
+    std::vector<size_t> sample_lens;
+
+    bool empty() const { return sample_data.empty(); }
+    bool Verify() const {
+      size_t total_len = 0;
+      for (auto len : sample_lens) {
+        total_len += len;
+      }
+      return total_len == sample_data.size();
+    }
+  };
+
+  // Arguments for MaybeCloneSpecialized() - provides either samples, a
+  // pre-defined dictionary, or indicates no dictionary should be used.
+  // NOTE: DictPreDefined here is the same type as above, allowing the
+  // pre-defined dictionary from GetDictGuidance() to be passed through.
+  using DictConfigArgs =
+      std::variant<DictDisabled, DictSamples, DictPreDefined>;
+
+  // A WorkingArea is an optional structure (both for callers and
+  // implementations) that can enable optimizing repeated compressions by
+  // reusing working space or thread-local tracking of statistics or trends.
+  // This enables use of ZSTD context, for example.
+  //
+  // EXTENSIBLE or reinterpret_cast-able by custom Compressor implementations
+  struct WorkingArea {};
+
+ public:  // Functions
   Compressor() = default;
   virtual ~Compressor() = default;
 
@@ -69,15 +128,17 @@ class Compressor {
     return id;
   }
 
-  // Returns the max total bytes of for all sampled blocks for creating the data
-  // dictionary, or zero indicating dictionary compression should not be
-  // used/configured. This will typically be called after
-  // CompressionManager::GetCompressor() to see if samples should be accumulated
-  // and passed to MaybeCloneSpecialized().
-  virtual size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const {
+  // Returns the recommended dictionary configuration for the given block type.
+  // See the comments on DictConfig and variants for details.
+  //
+  // NOTE: This may be called on the "base" Compressor returned by
+  // CompressionManager, which is not yet configured with a dictionary,
+  // or it can be skipped by callers not intending to handle dictionary
+  // compression.
+  virtual DictConfig GetDictGuidance(CacheEntryRole block_type) const {
     // Default implementation: no dictionary
     (void)block_type;
-    return 0;
+    return DictDisabled{};
   }
 
   // Returns the serialized form of the data dictionary associated with this
@@ -94,52 +155,32 @@ class Compressor {
   // needed to implement MaybeCloneSpecialized() in wrapper compressors.
   virtual std::unique_ptr<Compressor> Clone() const = 0;
 
-  // Utility struct for providing sample data for the compression dictionary.
-  // Potentially extensible by callers of Compressor (but not recommended)
-  struct DictSampleArgs {
-    // All the sample input blocks stored contiguously
-    std::string sample_data;
-    // The lengths of each of the sample blocks in `sample_data`
-    std::vector<size_t> sample_lens;
-
-    bool empty() { return sample_data.empty(); }
-    bool Verify() {
-      size_t total_len = 0;
-      for (auto len : sample_lens) {
-        total_len += len;
-      }
-      return total_len == sample_data.size();
-    }
-  };
-
   // Create potential variants of the same Compressor that might be
   // (a) optimized for a particular block type (does not affect correct
   //     decompression), and/or
-  // (b) configured to use a compression dictionary, based on the given
-  //     samples (decompression must provide the dictionary from
-  //     GetSerializedDict())
+  // (b) configured to use a compression dictionary based on the provided
+  //     configuration (samples or pre-defined dictionary). See the comments on
+  //     DictConfigArgs and its variants for detail.
+  //
   // Return of nullptr indicates no specialization exists or was attempted
-  // and the caller is best to use the current Compressor for the desired
-  // scenario. Using CacheEntryRole:kMisc for block_type generally means
-  // "unspecified", and both parameters are merely suggestions. The exact
-  // dictionary associated with a returned compressor must be read from
-  // GetSerializedDict().
+  // and the caller should use the current Compressor for the desired scenario.
+  // Using CacheEntryRole::kMisc for block_type generally means "unspecified".
+  //
+  // The exact dictionary associated with a returned compressor must be read
+  // from GetSerializedDict().
   virtual std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const {
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const {
     // Default implementation: no specialization
     (void)block_type;
-    (void)dict_samples;
-    // Caller should have checked GetMaxSampleSizeIfWantDict before attempting
-    // to provide dictionary samples
-    assert(dict_samples.empty());
+    (void)dict_config;
     return nullptr;
   }
 
   // A convenience function when a clone is needed and may or may not be
   // specialized.
   std::unique_ptr<Compressor> CloneMaybeSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const {
-    auto clone = MaybeCloneSpecialized(block_type, std::move(dict_samples));
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const {
+    auto clone = MaybeCloneSpecialized(block_type, std::move(dict_config));
     if (clone == nullptr) {
       clone = Clone();
       assert(clone != nullptr);
@@ -147,14 +188,6 @@ class Compressor {
     return clone;
   }
 
-  // A WorkingArea is an optional structure (both for callers and
-  // implementations) that can enable optimizing repeated compressions by
-  // reusing working space or thread-local tracking of statistics or trends.
-  // This enables use of ZSTD context, for example.
-  //
-  // EXTENSIBLE or reinterpret_cast-able by custom Compressor implementations
-  struct WorkingArea {};
-
   // To allow for flexible re-use / reclaimation, we have explicit Get and
   // Release functions, and usually wrap in a special RAII smart pointer.
   // For example, a WorkingArea could be saved/recycled in thread-local or
@@ -423,6 +456,12 @@ class CompressionManager
   // which is valid at the discretion of the CompressionManager. Returning
   // nullptr should normally be the result if preferred == kNoCompression.
   //
+  // Compressors returned here are configured WITHOUT a dictionary, so that
+  // it's always possible to get correct compression->decompression results
+  // if not opting-in to dictionary handling. The compressors may recommend
+  // dictionary usage via GetDictGuidance() and creating a modified Compressor
+  // for that. See Compressor::GetDictGuidance() etc. for details.
+  //
   // These functions must be thread-safe.
 
   // Get a compressor for an SST file.
@@ -477,8 +516,8 @@ class CompressorWrapper : public Compressor {
   CompressorWrapper(const CompressorWrapper&) = delete;
   CompressorWrapper& operator=(const CompressorWrapper&) = delete;
 
-  size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override {
-    return wrapped_->GetMaxSampleSizeIfWantDict(block_type);
+  DictConfig GetDictGuidance(CacheEntryRole block_type) const override {
+    return wrapped_->GetDictGuidance(block_type);
   }
 
   Slice GetSerializedDict() const override {
@@ -496,9 +535,9 @@ class CompressorWrapper : public Compressor {
   // when the wrapped Compressor uses the default implementation of
   // MaybeCloneSpecialized(). This needs to be overridden if not.
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const override {
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const override {
     auto clone =
-        wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+        wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_config));
     // Assert default no-op MaybeCloneSpecialized()
     assert(clone == nullptr);
     return clone;
@@ -592,8 +631,15 @@ class CompressionManagerWrapper : public CompressionManager {
 
   std::shared_ptr<CompressionManager> FindCompatibleCompressionManager(
       Slice compatibility_name) override {
+    // NOTE: We expect that the wrapped CompressionManager will generally
+    // be preferred if compatible, so the default implementation here does
+    // not purely defer to the wrapped instance
+    if (compatibility_name == CompatibilityName()) {
+      return shared_from_this();
+    } else {
       return wrapped_->FindCompatibleCompressionManager(compatibility_name);
     }
+  }
 
   bool SupportsCompressionType(CompressionType type) const override {
     return wrapped_->SupportsCompressionType(type);

table/block_based/block_based_table_builder.cc

@@ -113,9 +113,9 @@ FilterBlockBuilder* CreateFilterBlockBuilder(
 // A convenience function for populating the Compressor* fields; see ~Rep()
 Compressor* MaybeCloneSpecialized(
     Compressor* compressor, CacheEntryRole block_type,
-    Compressor::DictSampleArgs&& dict_samples = {}) {
+    Compressor::DictConfigArgs&& dict_config = Compressor::DictDisabled{}) {
   auto specialized =
-      compressor->MaybeCloneSpecialized(block_type, std::move(dict_samples));
+      compressor->MaybeCloneSpecialized(block_type, std::move(dict_config));
   if (specialized) {
     // Caller is responsible for freeing when distinct
     return specialized.release();
@@ -833,7 +833,8 @@ struct BlockBasedTableBuilder::Rep {
   RelaxedAtomic<uint64_t> sampled_output_fast_data_bytes{0};
   uint32_t compression_parallel_threads;
   int max_compressed_bytes_per_kb;
-  size_t max_dict_sample_bytes = 0;
+  // Dictionary guidance for data blocks (from GetDictGuidance())
+  Compressor::DictConfig data_block_dict_guidance;
 
   // *** Compressors & decompressors - Yes, it seems like a lot here but ***
   // *** these are distinct fields to minimize extra conditionals and    ***
@@ -1122,9 +1123,12 @@ struct BlockBasedTableBuilder::Rep {
         index_block_working_area.compress =
             index_block_compressor->ObtainWorkingArea();
       }
-      max_dict_sample_bytes = basic_compressor->GetMaxSampleSizeIfWantDict(
-          CacheEntryRole::kDataBlock);
-      if (max_dict_sample_bytes > 0) {
+      data_block_dict_guidance =
+          basic_compressor->GetDictGuidance(CacheEntryRole::kDataBlock);
+      if (auto* sampling =
+              std::get_if<Compressor::DictSampling>(&data_block_dict_guidance);
+          sampling != nullptr && sampling->max_sample_bytes > 0) {
+        // Sampling mode: collect samples up to max_sample_bytes
         state = State::kBuffered;
         if (tbo.target_file_size == 0) {
           buffer_limit = tbo.compression_opts.max_dict_buffer_bytes;
@@ -1134,7 +1138,22 @@ struct BlockBasedTableBuilder::Rep {
           buffer_limit = std::min(tbo.target_file_size,
                                   tbo.compression_opts.max_dict_buffer_bytes);
         }
+      } else if (auto* predef = std::get_if<Compressor::DictPreDefined>(
+                     &data_block_dict_guidance);
+                 predef != nullptr && !predef->dict_data.empty()) {
+        // Pre-defined dictionary mode: use it immediately, no buffering
+        data_block_compressor = MaybeCloneSpecialized(
+            basic_compressor.get(), CacheEntryRole::kDataBlock,
+            Compressor::DictPreDefined{std::string{predef->dict_data}});
+        data_block_working_area.compress =
+            data_block_compressor->ObtainWorkingArea();
       } else {
+        assert(std::holds_alternative<Compressor::DictSampling>(
+                   data_block_dict_guidance) ||
+               std::holds_alternative<Compressor::DictPreDefined>(
+                   data_block_dict_guidance) ||
+               std::holds_alternative<Compressor::DictDisabled>(
+                   data_block_dict_guidance));
         // No distinct data block compressor using dictionary, but
         // implementation might still want to specialize for data blocks
         data_block_compressor = MaybeCloneSpecialized(
@@ -2632,13 +2651,17 @@ void BlockBasedTableBuilder::MaybeEnterUnbuffered(
       kPrimeGenerator % static_cast<uint64_t>(kNumBlocksBuffered));
   const size_t kInitSampleIdx = kNumBlocksBuffered / 2;
 
-  Compressor::DictSampleArgs samples;
+  Compressor::DictSamples samples;
   size_t buffer_idx = kInitSampleIdx;
-  for (size_t i = 0; i < kNumBlocksBuffered &&
-                     samples.sample_data.size() < r->max_dict_sample_bytes;
+  // Get max_sample_bytes from the DictSampling guidance
+  auto* sampling =
+      std::get_if<Compressor::DictSampling>(&r->data_block_dict_guidance);
+  assert(sampling != nullptr);
+  size_t max_sample_bytes = sampling->max_sample_bytes;
+  for (size_t i = 0;
+       i < kNumBlocksBuffered && samples.sample_data.size() < max_sample_bytes;
        ++i) {
-    size_t copy_len =
-        std::min(r->max_dict_sample_bytes - samples.sample_data.size(),
+    size_t copy_len = std::min(max_sample_bytes - samples.sample_data.size(),
                                r->data_block_buffers[buffer_idx].size());
     samples.sample_data.append(r->data_block_buffers[buffer_idx], 0, copy_len);
     samples.sample_lens.emplace_back(copy_len);

test_util/testutil.h

@@ -796,9 +796,9 @@ struct CompressorCustomAlg : public CompressorWrapper {
   }
 
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const override {
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const override {
     auto clone =
-        wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples));
+        wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_config));
     return std::make_unique<CompressorCustomAlg>(std::move(clone));
   }
 

util/auto_tune_compressor.cc

@@ -64,9 +64,9 @@ std::unique_ptr<Compressor> AutoSkipCompressorWrapper::Clone() const {
 }
 
 std::unique_ptr<Compressor> AutoSkipCompressorWrapper::MaybeCloneSpecialized(
-    CacheEntryRole block_type, DictSampleArgs&& dict_samples) const {
+    CacheEntryRole block_type, DictConfigArgs&& dict_config) const {
   auto clone =
-      wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples));
+      wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_config));
   return std::make_unique<AutoSkipCompressorWrapper>(std::move(clone), opts_);
 }
 
@@ -189,11 +189,10 @@ const char* CostAwareCompressor::Name() const { return "CostAwareCompressor"; }
 std::unique_ptr<Compressor> CostAwareCompressor::Clone() const {
   return std::make_unique<CostAwareCompressor>(opts_);
 }
-size_t CostAwareCompressor::GetMaxSampleSizeIfWantDict(
+Compressor::DictConfig CostAwareCompressor::GetDictGuidance(
     CacheEntryRole block_type) const {
   auto idx = allcompressors_index_.back();
-  return allcompressors_[idx.first][idx.second]->GetMaxSampleSizeIfWantDict(
-      block_type);
+  return allcompressors_[idx.first][idx.second]->GetDictGuidance(block_type);
 }
 
 Slice CostAwareCompressor::GetSerializedDict() const {
@@ -205,12 +204,12 @@ CompressionType CostAwareCompressor::GetPreferredCompressionType() const {
   return kZSTD;
 }
 std::unique_ptr<Compressor> CostAwareCompressor::MaybeCloneSpecialized(
-    CacheEntryRole block_type, DictSampleArgs&& dict_samples) const {
+    CacheEntryRole block_type, DictConfigArgs&& dict_config) const {
   // TODO: full dictionary compression support. Currently this just falls
   // back on a non-multi compressor when asked to use a dictionary.
   auto idx = allcompressors_index_.back();
   return allcompressors_[idx.first][idx.second]->MaybeCloneSpecialized(
-      block_type, std::move(dict_samples));
+      block_type, std::move(dict_config));
 }
 Status CostAwareCompressor::CompressBlock(Slice uncompressed_data,
                                           char* compressed_output,

util/auto_tune_compressor.h

@@ -66,7 +66,7 @@ class AutoSkipCompressorWrapper : public CompressorWrapper {
 
   std::unique_ptr<Compressor> Clone() const override;
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const override;
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const override;
   Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                        size_t* compressed_output_size,
                        CompressionType* out_compression_type,
@@ -153,12 +153,12 @@ class CostAwareCompressor : public Compressor {
   explicit CostAwareCompressor(const CompressionOptions& opts);
   const char* Name() const override;
   std::unique_ptr<Compressor> Clone() const override;
-  size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override;
+  DictConfig GetDictGuidance(CacheEntryRole block_type) const override;
   Slice GetSerializedDict() const override;
   CompressionType GetPreferredCompressionType() const override;
   ManagedWorkingArea ObtainWorkingArea() override;
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const override;
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const override;
 
   Status CompressBlock(Slice uncompressed_data, char* compressed_output,
                        size_t* compressed_output_size,

util/compression.cc

@@ -222,9 +222,11 @@ class CompressorWithSimpleDictBase : public CompressorBase {
                                         std::string&& dict_data = {})
       : CompressorBase(opts), dict_data_(std::move(dict_data)) {}
 
-  size_t GetMaxSampleSizeIfWantDict(
-      CacheEntryRole /*block_type*/) const override {
-    return opts_.max_dict_bytes;
+  DictConfig GetDictGuidance(CacheEntryRole /*block_type*/) const override {
+    if (opts_.max_dict_bytes == 0) {
+      return DictDisabled{};
+    }
+    return DictSampling{opts_.max_dict_bytes};
   }
 
   // NOTE: empty dict is equivalent to no dict
@@ -236,13 +238,21 @@ class CompressorWithSimpleDictBase : public CompressorBase {
 
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
       CacheEntryRole /*block_type*/,
-      DictSampleArgs&& dict_samples) const final override {
-    assert(dict_samples.Verify());
-    if (dict_samples.empty()) {
-      // Nothing to specialize on
+      DictConfigArgs&& dict_config) const final override {
+    if (auto* samples = std::get_if<DictSamples>(&dict_config)) {
+      assert(samples->Verify());
+      if (samples->empty()) {
         return nullptr;
+      }
+      return CloneForDict(std::move(samples->sample_data));
+    } else if (auto* predef = std::get_if<DictPreDefined>(&dict_config)) {
+      if (predef->dict_data.empty()) {
+        return nullptr;
+      }
+      return CloneForDict(std::move(predef->dict_data));
     } else {
-      return CloneForDict(std::move(dict_samples.sample_data));
+      assert(std::holds_alternative<DictDisabled>(dict_config));
+      return nullptr;
     }
   }
 
@@ -858,14 +868,15 @@ class BuiltinZSTDCompressorV2 final : public CompressorBase {
                                                      std::move(dict_copy));
   }
 
-  size_t GetMaxSampleSizeIfWantDict(
-      CacheEntryRole /*block_type*/) const override {
+  DictConfig GetDictGuidance(CacheEntryRole /*block_type*/) const override {
     if (opts_.max_dict_bytes == 0) {
       // Dictionary compression disabled
-      return 0;
+      return DictDisabled{};
     } else {
-      return opts_.zstd_max_train_bytes > 0 ? opts_.zstd_max_train_bytes
+      size_t max_sample_bytes = opts_.zstd_max_train_bytes > 0
+                                    ? opts_.zstd_max_train_bytes
                                     : opts_.max_dict_bytes;
+      return DictSampling{max_sample_bytes};
     }
   }
 
@@ -974,31 +985,49 @@ class BuiltinZSTDCompressorV2 final : public CompressorBase {
 
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
       CacheEntryRole /*block_type*/,
-      DictSampleArgs&& dict_samples) const override {
-    assert(dict_samples.Verify());
-    if (dict_samples.empty()) {
-      // Nothing to specialize on
+      DictConfigArgs&& dict_config) const override {
+    // Handle DictDisabled
+    // TODO: use holds_alternative
+    if (auto* disabled = std::get_if<DictDisabled>(&dict_config)) {
+      (void)disabled;
       return nullptr;
     }
+
     std::string dict_data;
+
+    // Handle DictPreDefined - use the pre-defined dictionary directly
+    if (auto* predef = std::get_if<DictPreDefined>(&dict_config)) {
+      if (predef->dict_data.empty()) {
+        return nullptr;
+      }
+      dict_data = std::move(predef->dict_data);
+    }
+
+    // Handle DictSamples - train dictionary from samples
+    if (auto* samples = std::get_if<DictSamples>(&dict_config)) {
+      assert(samples->Verify());
+      if (samples->empty()) {
+        return nullptr;
+      }
       // Migrated from BlockBasedTableBuilder::EnterUnbuffered()
       if (opts_.zstd_max_train_bytes > 0) {
-      assert(dict_samples.sample_data.size() <= opts_.zstd_max_train_bytes);
+        assert(samples->sample_data.size() <= opts_.zstd_max_train_bytes);
         if (opts_.use_zstd_dict_trainer) {
-        dict_data = ZSTD_TrainDictionary(dict_samples.sample_data,
-                                         dict_samples.sample_lens,
-                                         opts_.max_dict_bytes);
+          dict_data = ZSTD_TrainDictionary(
+              samples->sample_data, samples->sample_lens, opts_.max_dict_bytes);
         } else {
-        dict_data = ZSTD_FinalizeDictionary(dict_samples.sample_data,
-                                            dict_samples.sample_lens,
-                                            opts_.max_dict_bytes, opts_.level);
+          dict_data = ZSTD_FinalizeDictionary(
+              samples->sample_data, samples->sample_lens, opts_.max_dict_bytes,
+              opts_.level);
         }
       } else {
-      assert(dict_samples.sample_data.size() <= opts_.max_dict_bytes);
+        assert(samples->sample_data.size() <= opts_.max_dict_bytes);
         // ZSTD "raw content dictionary" - "Any buffer is a valid raw content
         // dictionary." Or similar for other compressions.
-      dict_data = std::move(dict_samples.sample_data);
+        dict_data = std::move(samples->sample_data);
       }
+    }
+
     CompressionDict dict{std::move(dict_data), kZSTD, opts_.level};
     return std::make_unique<BuiltinZSTDCompressorV2>(opts_, std::move(dict));
   }

util/compression_test.cc

@@ -1362,9 +1362,9 @@ TEST_P(DBCompressionTestMaybeParallel, CompressionManagerWrapper) {
 
     std::unique_ptr<Compressor> MaybeCloneSpecialized(
         CacheEntryRole block_type,
-        DictSampleArgs&& dict_samples) const override {
+        DictConfigArgs&& dict_config) const override {
       std::unique_ptr<Compressor> result = std::make_unique<MyCompressor>(
-          wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_samples)));
+          wrapped_->CloneMaybeSpecialized(block_type, std::move(dict_config)));
       if (block_type == CacheEntryRole::kDataBlock) {
         result = std::make_unique<CheckDataBlockCompressorWrapper>(
             std::move(result));
@@ -2138,6 +2138,232 @@ TEST_F(DBCompressionCostPredictor, CostAwareCompressorManager) {
   ASSERT_OK(Flush());
 }
 
+// Test pre-defined dictionary compression with a custom CompressionManager
+TEST_F(DBCompressionTest, PreDefinedDictionaryCompression) {
+  if (!ZSTD_Supported()) {
+    ROCKSDB_GTEST_BYPASS("ZSTD compression not supported");
+    return;
+  }
+
+  // A custom compressor that returns a pre-defined dictionary
+  class PreDefinedDictCompressor : public CompressorWrapper {
+   public:
+    explicit PreDefinedDictCompressor(std::unique_ptr<Compressor> wrapped,
+                                      std::string dict_data)
+        : CompressorWrapper(std::move(wrapped)),
+          predefined_dict_(std::move(dict_data)) {}
+
+    const char* Name() const override { return "PreDefinedDictCompressor"; }
+
+    DictConfig GetDictGuidance(CacheEntryRole block_type) const override {
+      if (block_type == CacheEntryRole::kDataBlock &&
+          !predefined_dict_.empty()) {
+        return DictPreDefined{/*copy*/ predefined_dict_};
+      }
+      return DictDisabled{};
+    }
+
+    std::unique_ptr<Compressor> Clone() const override {
+      return std::make_unique<PreDefinedDictCompressor>(wrapped_->Clone(),
+                                                        predefined_dict_);
+    }
+
+    std::unique_ptr<Compressor> MaybeCloneSpecialized(
+        CacheEntryRole block_type,
+        DictConfigArgs&& dict_config) const override {
+      // Delegate to wrapped compressor for dictionary handling
+      auto specialized =
+          wrapped_->MaybeCloneSpecialized(block_type, std::move(dict_config));
+      if (specialized) {
+        return specialized;
+      }
+      return nullptr;
+    }
+
+   private:
+    std::string predefined_dict_;
+  };
+
+  // Custom CompatibilityName so the builtin compression manager won't be used
+  static const char* kTestCompatibilityName = "PreDefinedDictTest";
+
+  class PreDefinedDictManager : public CompressionManagerWrapper {
+   public:
+    explicit PreDefinedDictManager(std::shared_ptr<CompressionManager> wrapped,
+                                   std::string dict_data)
+        : CompressionManagerWrapper(std::move(wrapped)),
+          predefined_dict_(std::move(dict_data)) {}
+
+    const char* Name() const override { return "PreDefinedDictManager"; }
+
+    const char* CompatibilityName() const override {
+      return kTestCompatibilityName;
+    }
+
+    std::unique_ptr<Compressor> GetCompressorForSST(
+        const FilterBuildingContext& context, const CompressionOptions& opts,
+        CompressionType preferred) override {
+      auto base = wrapped_->GetCompressorForSST(context, opts, preferred);
+      if (base) {
+        return std::make_unique<PreDefinedDictCompressor>(std::move(base),
+                                                          predefined_dict_);
+      }
+      return nullptr;
+    }
+
+   private:
+    std::string predefined_dict_;
+  };
+
+  // A broken manager that ignores the dictionary when decompressing.
+  // This simulates a buggy decompressor that doesn't properly apply the
+  // dictionary, causing ZSTD to produce wrong output when decompressing
+  // dictionary-compressed data.
+  class BrokenDictManager : public CompressionManagerWrapper {
+   public:
+    explicit BrokenDictManager(std::shared_ptr<CompressionManager> wrapped)
+        : CompressionManagerWrapper(std::move(wrapped)) {}
+
+    const char* Name() const override { return "BrokenDictManager"; }
+
+    const char* CompatibilityName() const override {
+      return kTestCompatibilityName;
+    }
+
+    std::shared_ptr<Decompressor> GetDecompressor() override {
+      return std::make_shared<IgnoreDictDecompressor>(
+          wrapped_->GetDecompressor());
+    }
+
+    std::shared_ptr<Decompressor> GetDecompressorOptimizeFor(
+        CompressionType optimize_for_type) override {
+      return std::make_shared<IgnoreDictDecompressor>(
+          wrapped_->GetDecompressorOptimizeFor(optimize_for_type));
+    }
+
+    std::shared_ptr<Decompressor> GetDecompressorForTypes(
+        const CompressionType* types_begin,
+        const CompressionType* types_end) override {
+      return std::make_shared<IgnoreDictDecompressor>(
+          wrapped_->GetDecompressorForTypes(types_begin, types_end));
+    }
+
+   private:
+    // A decompressor that stores the dictionary (for GetSerializedDict) but
+    // ignores it during decompression, causing ZSTD to produce garbage
+    class IgnoreDictDecompressor : public DecompressorWrapper {
+     public:
+      explicit IgnoreDictDecompressor(std::shared_ptr<Decompressor> wrapped)
+          : DecompressorWrapper(std::move(wrapped)) {}
+
+      IgnoreDictDecompressor(std::shared_ptr<Decompressor> wrapped,
+                             std::string dict)
+          : DecompressorWrapper(std::move(wrapped)),
+            dict_(std::move(dict)),
+            dict_slice_(dict_) {}
+
+      const char* Name() const override { return "IgnoreDictDecompressor"; }
+
+      const Slice& GetSerializedDict() const override { return dict_slice_; }
+
+      Status MaybeCloneForDict(const Slice& serialized_dict,
+                               std::unique_ptr<Decompressor>* out) override {
+        // Store the dict but don't actually use it for decompression
+        *out = std::make_unique<IgnoreDictDecompressor>(
+            wrapped_,
+            std::string(serialized_dict.data(), serialized_dict.size()));
+        return Status::OK();
+      }
+
+     private:
+      std::string dict_;
+      Slice dict_slice_;
+    };
+  };
+
+  // Create a dictionary that will be heavily referenced. The key insight is
+  // that ZSTD dictionary compression works by finding matches between the input
+  // data and the dictionary content. To force ZSTD to create dictionary
+  // references, we need to use data that contains exact copies of dictionary
+  // content.
+  Random rnd(42);
+
+  // Create a dictionary with recognizable patterns
+  std::string predefined_dict;
+  std::vector<std::string> dict_patterns;
+  for (int i = 0; i < 50; i++) {
+    std::string pattern = rnd.RandomString(200);
+    dict_patterns.push_back(pattern);
+    predefined_dict += pattern;
+  }
+  // Total dict size: 50 * 200 = 10000 bytes
+  size_t kDictSize = predefined_dict.size();
+
+  auto mgr = std::make_shared<PreDefinedDictManager>(
+      GetBuiltinV2CompressionManager(), predefined_dict);
+
+  Options options = CurrentOptions();
+  options.compression = kZSTD;
+  options.compression_opts.max_dict_bytes = static_cast<int>(kDictSize);
+  options.compression_manager = mgr;
+  options.statistics = ROCKSDB_NAMESPACE::CreateDBStatistics();
+  BlockBasedTableOptions bbto;
+  bbto.enable_index_compression = true;
+  // Need format_version >= 7 for custom CompatibilityName
+  bbto.format_version = 7;
+  // Need dictionary block load statistics
+  bbto.block_cache = NewLRUCache(1 << 20);
+  bbto.cache_index_and_filter_blocks = true;
+  options.table_factory.reset(NewBlockBasedTableFactory(bbto));
+  DestroyAndReopen(options);
+
+  // Write data that uses the same patterns from the dictionary.
+  // This forces ZSTD to create back-references to the dictionary.
+  std::vector<std::string> expected_values;
+  for (int i = 0; i < 100; i++) {
+    std::string value;
+    // Compose value from random dictionary patterns - same content as dict
+    for (int j = 0; j < 5; j++) {
+      value +=
+          dict_patterns[rnd.Uniform(static_cast<int>(dict_patterns.size()))];
+    }
+    expected_values.push_back(value);
+    ASSERT_OK(Put(Key(i), value));
+  }
+  ASSERT_OK(Flush());
+
+  // Verify dictionary was used by checking that dict bytes were inserted
+  ASSERT_GE(
+      TestGetTickerCount(options, BLOCK_CACHE_COMPRESSION_DICT_BYTES_INSERT),
+      predefined_dict.size());
+
+  // Read back data and verify correctness
+  for (int i = 0; i < 100; i++) {
+    std::string value;
+    ASSERT_OK(db_->Get(ReadOptions(), Key(i), &value));
+    ASSERT_EQ(value, expected_values[i]);
+  }
+
+  // Now re-open with a broken decompressor that ignores dictionary.
+  // This should result in corruption on read because ZSTD will fail to
+  // decompress data that references the missing dictionary content.
+  Close();
+  auto broken_mgr =
+      std::make_shared<BrokenDictManager>(GetBuiltinV2CompressionManager());
+  options.compression_manager = broken_mgr;
+  // New block cache to ensure dictionary is re-loaded, because the
+  // dictionary block in cache is actually associated with a decompressor
+  bbto.block_cache = NewLRUCache(1 << 20);
+  options.table_factory.reset(NewBlockBasedTableFactory(bbto));
+  ASSERT_OK(TryReopen(options));
+
+  // Read should fail with corruption because the decompressor ignores
+  // the dictionary, causing ZSTD to produce garbage output
+  std::string value;
+  ASSERT_EQ(db_->Get(ReadOptions(), Key(0), &value).code(),
+            Status::kCorruption);
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 int main(int argc, char** argv) {
   ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();

util/simple_mixed_compressor.cc

@@ -28,9 +28,9 @@ MultiCompressorWrapper::MultiCompressorWrapper(const CompressionOptions& opts)
   }
 }
 
-size_t MultiCompressorWrapper::GetMaxSampleSizeIfWantDict(
+Compressor::DictConfig MultiCompressorWrapper::GetDictGuidance(
     CacheEntryRole block_type) const {
-  return compressors_.back()->GetMaxSampleSizeIfWantDict(block_type);
+  return compressors_.back()->GetDictGuidance(block_type);
 }
 
 Slice MultiCompressorWrapper::GetSerializedDict() const {
@@ -46,11 +46,11 @@ Compressor::ManagedWorkingArea MultiCompressorWrapper::ObtainWorkingArea() {
 }
 
 std::unique_ptr<Compressor> MultiCompressorWrapper::MaybeCloneSpecialized(
-    CacheEntryRole block_type, DictSampleArgs&& dict_samples) const {
+    CacheEntryRole block_type, DictConfigArgs&& dict_config) const {
   // TODO: full dictionary compression support. Currently this just falls
   // back on a non-multi compressor when asked to use a dictionary.
   return compressors_.back()->MaybeCloneSpecialized(block_type,
-                                                    std::move(dict_samples));
+                                                    std::move(dict_config));
 }
 
 // RandomMixedCompressor implementation

util/simple_mixed_compressor.h

@@ -19,12 +19,12 @@ class MultiCompressorWrapper : public Compressor {
  public:
   explicit MultiCompressorWrapper(const CompressionOptions& opts);
 
-  size_t GetMaxSampleSizeIfWantDict(CacheEntryRole block_type) const override;
+  DictConfig GetDictGuidance(CacheEntryRole block_type) const override;
   Slice GetSerializedDict() const override;
   CompressionType GetPreferredCompressionType() const override;
   ManagedWorkingArea ObtainWorkingArea() override;
   std::unique_ptr<Compressor> MaybeCloneSpecialized(
-      CacheEntryRole block_type, DictSampleArgs&& dict_samples) const override;
+      CacheEntryRole block_type, DictConfigArgs&& dict_config) const override;
 
  protected:
   const CompressionOptions opts_;