add uniform block tracking

Author: joshkang97

include/rocksdb/table_properties.h

@@ -50,6 +50,7 @@ struct TablePropertiesNames {
   static const std::string kRawKeySize;
   static const std::string kRawValueSize;
   static const std::string kNumDataBlocks;
+  static const std::string kNumUniformBlocks;
   static const std::string kNumEntries;
   static const std::string kNumFilterEntries;
   static const std::string kDeletedKeys;
@@ -245,6 +246,8 @@ struct TableProperties {
   uint64_t raw_value_size = 0;
   // the number of blocks in this table
   uint64_t num_data_blocks = 0;
+  // the number of uniform blocks in this table
+  uint64_t num_uniform_blocks = 0;
   // the number of entries in this table
   uint64_t num_entries = 0;
   // the number of unique entries (keys or prefixes) added to filters

table/block_based/block_based_table_builder.cc

@@ -2432,6 +2432,10 @@ void BlockBasedTableBuilder::WriteIndexBlock(
       // The last index_block_handle will be for the partition index block
     }
   }
+  if (LIKELY(ok())) {
+    rep_->props.num_uniform_blocks =
+        rep_->index_builder->NumUniformIndexBlocks();
+  }
   // If success and need to record in metaindex rather than footer...
   if (LIKELY(ok()) && !FormatVersionUsesIndexHandleInFooter(
                           rep_->table_options.format_version)) {

table/block_based/block_builder.cc

@@ -105,6 +105,7 @@ BlockBuilder::BlockBuilder(
       restarts_(1, 0),  // First restart point is at offset 0
       counter_(0),
       finished_(false),
+      is_uniform_(false),
       uniform_cv_threshold_(uniform_cv_threshold),
       statistics_(statistics),
       use_separated_kv_storage_(use_separated_kv_storage) {
@@ -131,6 +132,7 @@ void BlockBuilder::Reset() {
               (use_separated_kv_storage_ ? sizeof(uint32_t) : 0);
   counter_ = 0;
   finished_ = false;
+  is_uniform_ = false;
   last_key_.clear();
   if (data_block_hash_index_builder_.Valid()) {
     data_block_hash_index_builder_.Reset();
@@ -185,7 +187,7 @@ size_t BlockBuilder::EstimateSizeAfterKV(const Slice& key,
 }
 
 Slice BlockBuilder::Finish() {
-  bool is_uniform = ScanForUniformity();
+  is_uniform_ = ScanForUniformity();
 
   // Append restart array
   size_t values_buffer_offset = buffer_.size();
@@ -201,7 +203,7 @@ Slice BlockBuilder::Finish() {
   DataBlockFooter footer;
   footer.num_restarts = static_cast<uint32_t>(restarts_.size());
   footer.index_type = BlockBasedTableOptions::kDataBlockBinarySearch;
-  footer.is_uniform = is_uniform;
+  footer.is_uniform = is_uniform_;
   if (data_block_hash_index_builder_.Valid() &&
       CurrentSizeEstimate() <= kMaxBlockSizeSupportedByHashIndex) {
     data_block_hash_index_builder_.Finish(buffer_);

table/block_based/block_builder.h

@@ -86,6 +86,10 @@ class BlockBuilder {
 
   std::string& MutableBuffer() { return buffer_; }
 
+  // Returns true if the most recently Finish()'d block was marked uniform.
+  // REQUIRES: Finish() has been called.
+  bool IsUniform() const { return is_uniform_; }
+
  private:
   inline void AddWithLastKeyImpl(const Slice& key, const Slice& value,
                                  const Slice& last_key,
@@ -128,6 +132,7 @@ class BlockBuilder {
   size_t estimate_;
   int counter_;    // Number of entries emitted since restart
   bool finished_;  // Has Finish() been called?
+  bool is_uniform_;  // Was the last Finish()'d block uniform?
   std::string last_key_;
   DataBlockHashIndexBuilder data_block_hash_index_builder_;
   const double uniform_cv_threshold_;

table/block_based/block_test.cc

@@ -803,6 +803,7 @@ TEST_P(IndexBlockTest, IndexValueEncodingTest) {
     bool expect_uniform = reader.NumRestarts() >= 3 &&
                           keyDistribution() == KeyDistribution::kUniform;
     EXPECT_EQ(reader.IsUniform(), expect_uniform);
+    EXPECT_EQ(builder.IsUniform(), expect_uniform);
   }
 
   // read block contents randomly

table/block_based/index_builder.cc

@@ -346,9 +346,12 @@ Status PartitionedIndexBuilder::Finish(
   if (UNLIKELY(entries_.empty())) {
     if (must_use_separator_with_seq_.LoadRelaxed()) {
       index_blocks->index_block_contents = index_block_builder_.Finish();
+      num_uniform_index_blocks_ += index_block_builder_.IsUniform() ? 1 : 0;
     } else {
       index_blocks->index_block_contents =
           index_block_builder_without_seq_.Finish();
+      num_uniform_index_blocks_ +=
+          index_block_builder_without_seq_.IsUniform() ? 1 : 0;
     }
     top_level_index_size_ = index_blocks->index_block_contents.size();
     index_size_ += top_level_index_size_;
@@ -361,6 +364,7 @@ Status PartitionedIndexBuilder::Finish(
     entry.value->must_use_separator_with_seq_.StoreRelaxed(
         must_use_separator_with_seq_.LoadRelaxed());
     auto s = entry.value->Finish(index_blocks);
+    num_uniform_index_blocks_ += entry.value->NumUniformIndexBlocks();
     index_size_ += index_blocks->index_block_contents.size();
     finishing_indexes_ = true;
     return s.ok() ? Status::Incomplete() : s;

table/block_based/index_builder.h

@@ -160,6 +160,9 @@ class IndexBuilder {
   // Get the size for index block. Must be called after ::Finish.
   virtual size_t IndexSize() const = 0;
 
+  // Get the number of uniform index blocks. Must be called after ::Finish.
+  virtual uint64_t NumUniformIndexBlocks() const { return 0; }
+
   // Returns an estimate of the current index size based on the builder's state.
   // Implementations should cache the estimate and update it via
   // UpdateIndexSizeEstimate() to avoid recalculating on every key add,
@@ -434,16 +437,22 @@ class ShortenedIndexBuilder : public IndexBuilder {
                 const BlockHandle& /*last_partition_block_handle*/) override {
     if (must_use_separator_with_seq_.LoadRelaxed()) {
       index_blocks->index_block_contents = index_block_builder_.Finish();
+      is_uniform_ = index_block_builder_.IsUniform();
     } else {
       index_blocks->index_block_contents =
           index_block_builder_without_seq_.Finish();
+      is_uniform_ = index_block_builder_without_seq_.IsUniform();
     }
     index_size_ = index_blocks->index_block_contents.size();
     return Status::OK();
   }
 
   size_t IndexSize() const override { return index_size_; }
 
+  uint64_t NumUniformIndexBlocks() const override {
+    return is_uniform_ ? 1 : 0;
+  }
+
   uint64_t CurrentIndexSizeEstimate() const override {
     return estimated_index_size_.LoadRelaxed();
   }
@@ -480,6 +489,7 @@ class ShortenedIndexBuilder : public IndexBuilder {
   const bool include_first_key_;
   BlockBasedTableOptions::IndexShorteningMode shortening_mode_;
   BlockHandle last_encoded_handle_ = BlockHandle::NullBlockHandle();
+  bool is_uniform_ = false;
   std::string current_block_first_internal_key_;
   uint64_t num_index_entries_ = 0;
   // Cache for index size estimate to avoid recalculating in hot path
@@ -607,6 +617,10 @@ class HashIndexBuilder : public IndexBuilder {
            prefix_meta_block_.size();
   }
 
+  uint64_t NumUniformIndexBlocks() const override {
+    return primary_index_builder_.NumUniformIndexBlocks();
+  }
+
   uint64_t CurrentIndexSizeEstimate() const override { return 0; }
 
   bool separator_is_key_plus_seq() override {
@@ -685,6 +699,10 @@ class PartitionedIndexBuilder : public IndexBuilder {
   size_t TopLevelIndexSize(uint64_t) const { return top_level_index_size_; }
   size_t NumPartitions() const;
 
+  uint64_t NumUniformIndexBlocks() const override {
+    return num_uniform_index_blocks_;
+  }
+
   // Returns a cached estimate of the current index size. This
   // estimate is updated when data blocks are added.
   uint64_t CurrentIndexSizeEstimate() const override {
@@ -724,6 +742,8 @@ class PartitionedIndexBuilder : public IndexBuilder {
   size_t top_level_index_size_ = 0;
   // Set after ::Finish is called
   size_t partition_cnt_ = 0;
+  // Accumulated across all Finish() calls
+  uint64_t num_uniform_index_blocks_ = 0;
 
   void MakeNewSubIndexBuilder();
   void UpdateIndexSizeEstimate() override;

table/meta_blocks.cc

@@ -97,6 +97,7 @@ void PropertyBlockBuilder::AddTableProperty(const TableProperties& props) {
   Add(TablePropertiesNames::kMergeOperands, props.num_merge_operands);
   Add(TablePropertiesNames::kNumRangeDeletions, props.num_range_deletions);
   Add(TablePropertiesNames::kNumDataBlocks, props.num_data_blocks);
+  Add(TablePropertiesNames::kNumUniformBlocks, props.num_uniform_blocks);
   Add(TablePropertiesNames::kFilterSize, props.filter_size);
   Add(TablePropertiesNames::kFormatVersion, props.format_version);
   Add(TablePropertiesNames::kFixedKeyLen, props.fixed_key_len);
@@ -291,6 +292,8 @@ Status ParsePropertiesBlock(
        &new_table_properties->raw_value_size},
       {TablePropertiesNames::kNumDataBlocks,
        &new_table_properties->num_data_blocks},
+      {TablePropertiesNames::kNumUniformBlocks,
+       &new_table_properties->num_uniform_blocks},
       {TablePropertiesNames::kNumEntries, &new_table_properties->num_entries},
       {TablePropertiesNames::kNumFilterEntries,
        &new_table_properties->num_filter_entries},

table/table_properties.cc

@@ -47,6 +47,8 @@ std::string TableProperties::ToString(const std::string& prop_delim,
   // Basic Info
   AppendProperty(result, "# data blocks", num_data_blocks, prop_delim,
                  kv_delim);
+  AppendProperty(result, "# uniform blocks", num_uniform_blocks, prop_delim,
+                 kv_delim);
   AppendProperty(result, "# entries", num_entries, prop_delim, kv_delim);
   AppendProperty(result, "# deletions", num_deletions, prop_delim, kv_delim);
   AppendProperty(result, "# merge operands", num_merge_operands, prop_delim,
@@ -192,6 +194,7 @@ void TableProperties::Add(const TableProperties& tp) {
   raw_key_size += tp.raw_key_size;
   raw_value_size += tp.raw_value_size;
   num_data_blocks += tp.num_data_blocks;
+  num_uniform_blocks += tp.num_uniform_blocks;
   num_entries += tp.num_entries;
   num_filter_entries += tp.num_filter_entries;
   num_deletions += tp.num_deletions;
@@ -215,6 +218,7 @@ TableProperties::GetAggregatablePropertiesAsMap() const {
   rv["raw_key_size"] = raw_key_size;
   rv["raw_value_size"] = raw_value_size;
   rv["num_data_blocks"] = num_data_blocks;
+  rv["num_uniform_blocks"] = num_uniform_blocks;
   rv["num_entries"] = num_entries;
   rv["num_filter_entries"] = num_filter_entries;
   rv["num_deletions"] = num_deletions;
@@ -280,6 +284,8 @@ const std::string TablePropertiesNames::kRawValueSize =
     "rocksdb.raw.value.size";
 const std::string TablePropertiesNames::kNumDataBlocks =
     "rocksdb.num.data.blocks";
+const std::string TablePropertiesNames::kNumUniformBlocks =
+    "rocksdb.num.uniform.blocks";
 const std::string TablePropertiesNames::kNumEntries = "rocksdb.num.entries";
 const std::string TablePropertiesNames::kNumFilterEntries =
     "rocksdb.num.filter_entries";
@@ -381,6 +387,10 @@ static std::unordered_map<std::string, OptionTypeInfo>
          {offsetof(struct TableProperties, num_data_blocks),
           OptionType::kUInt64T, OptionVerificationType::kNormal,
           OptionTypeFlags::kNone}},
+        {"num_uniform_blocks",
+         {offsetof(struct TableProperties, num_uniform_blocks),
+          OptionType::kUInt64T, OptionVerificationType::kNormal,
+          OptionTypeFlags::kNone}},
         {"num_entries",
          {offsetof(struct TableProperties, num_entries), OptionType::kUInt64T,
           OptionVerificationType::kNormal, OptionTypeFlags::kNone}},