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: #14253

Test Plan: unit test included

Reviewed By: hx235

Differential Revision: D91082208

Pulled By: pdillinger

fbshipit-source-id: 1442db65e15c9435437204c19787c96f7a40a207

Author: pdillinger

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,67 +155,39 @@ 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);
     }
     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,7 +631,14 @@ class CompressionManagerWrapper : public CompressionManager {
 
   std::shared_ptr<CompressionManager> FindCompatibleCompressionManager(
       Slice compatibility_name) override {
-    return wrapped_->FindCompatibleCompressionManager(compatibility_name);
+    // 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 {

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,14 +2651,18 @@ 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(),
-                 r->data_block_buffers[buffer_idx].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,