@nxtedition/rocksdb 8.2.7 → 9.0.0

package/deps/rocksdb/rocksdb/table/block_based/block_based_table_reader_test.cc

@@ -38,8 +38,8 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
   // complexity: human readable strings are easier to compress than random
   // strings.
   static std::map<std::string, std::string> GenerateKVMap(
-      int num_block = 100,
-      bool mixed_with_human_readable_string_value = false) {
+      int num_block = 100, bool mixed_with_human_readable_string_value = false,
+      size_t ts_sz = 0) {
     std::map<std::string, std::string> kv;
 
     Random rnd(101);
@@ -58,7 +58,13 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
         } else {
           v = rnd.RandomString(256);
         }
-        kv[std::string(k)] = v;
+        if (ts_sz > 0) {
+          std::string user_key;
+          AppendKeyWithMinTimestamp(&user_key, std::string(k), ts_sz);
+          kv[user_key] = v;
+        } else {
+          kv[std::string(k)] = v;
+        }
         key++;
       }
     }
@@ -80,21 +86,29 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
 
   // Creates a table with the specificied key value pairs (kv).
   void CreateTable(const std::string& table_name,
+                   const ImmutableOptions& ioptions,
                    const CompressionType& compression_type,
-                   const std::map<std::string, std::string>& kv) {
+                   const std::map<std::string, std::string>& kv,
+                   uint32_t compression_parallel_threads = 1,
+                   uint32_t compression_dict_bytes = 0) {
     std::unique_ptr<WritableFileWriter> writer;
     NewFileWriter(table_name, &writer);
 
-    // Create table builder.
-    ImmutableOptions ioptions(options_);
-    InternalKeyComparator comparator(options_.comparator);
+    InternalKeyComparator comparator(ioptions.user_comparator);
     ColumnFamilyOptions cf_options;
+    cf_options.prefix_extractor = options_.prefix_extractor;
     MutableCFOptions moptions(cf_options);
+    CompressionOptions compression_opts;
+    compression_opts.parallel_threads = compression_parallel_threads;
+    // Enable compression dictionary and set a buffering limit that is the same
+    // as each block's size.
+    compression_opts.max_dict_bytes = compression_dict_bytes;
+    compression_opts.max_dict_buffer_bytes = compression_dict_bytes;
     IntTblPropCollectorFactories factories;
     std::unique_ptr<TableBuilder> table_builder(
         options_.table_factory->NewTableBuilder(
             TableBuilderOptions(ioptions, moptions, comparator, &factories,
-                                compression_type, CompressionOptions(),
+                                compression_type, compression_opts,
                                 0 /* column_family_id */,
                                 kDefaultColumnFamilyName, -1 /* level */),
             writer.get()));
@@ -114,11 +128,17 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
                                 const std::string& table_name,
                                 std::unique_ptr<BlockBasedTable>* table,
                                 bool prefetch_index_and_filter_in_cache = true,
-                                Status* status = nullptr) {
+                                Status* status = nullptr,
+                                bool user_defined_timestamps_persisted = true) {
     const MutableCFOptions moptions(options_);
-    TableReaderOptions table_reader_options =
-        TableReaderOptions(ioptions, moptions.prefix_extractor, EnvOptions(),
-                           comparator, 0 /* block_protection_bytes_per_key */);
+    TableReaderOptions table_reader_options = TableReaderOptions(
+        ioptions, moptions.prefix_extractor, EnvOptions(), comparator,
+        0 /* block_protection_bytes_per_key */, false /* _skip_filters */,
+        false /* _immortal */, false /* _force_direct_prefetch */,
+        -1 /* _level */, nullptr /* _block_cache_tracer */,
+        0 /* _max_file_size_for_l0_meta_pin */, "" /* _cur_db_session_id */,
+        0 /* _cur_file_num */, {} /* _unique_id */, 0 /* _largest_seqno */,
+        0 /* _tail_size */, user_defined_timestamps_persisted);
 
     std::unique_ptr<RandomAccessFileReader> file;
     NewFileReader(table_name, foptions, &file);
@@ -126,9 +146,11 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
     uint64_t file_size = 0;
     ASSERT_OK(env_->GetFileSize(Path(table_name), &file_size));
 
+    ReadOptions read_opts;
+    read_opts.verify_checksums = true;
     std::unique_ptr<TableReader> general_table;
     Status s = options_.table_factory->NewTableReader(
-        ReadOptions(), table_reader_options, std::move(file), file_size,
+        read_opts, table_reader_options, std::move(file), file_size,
         &general_table, prefetch_index_and_filter_in_cache);
 
     if (s.ok()) {
@@ -147,6 +169,11 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
   std::shared_ptr<FileSystem> fs_;
   Options options_;
 
+  std::string ToInternalKey(const std::string& key) {
+    InternalKey internal_key(key, 0, ValueType::kTypeValue);
+    return internal_key.Encode().ToString();
+  }
+
  private:
   void WriteToFile(const std::string& content, const std::string& filename) {
     std::unique_ptr<FSWritableFile> f;
@@ -173,21 +200,31 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
     reader->reset(new RandomAccessFileReader(std::move(f), path,
                                              env_->GetSystemClock().get()));
   }
-
-  std::string ToInternalKey(const std::string& key) {
-    InternalKey internal_key(key, 0, ValueType::kTypeValue);
-    return internal_key.Encode().ToString();
-  }
 };
 
+// Param 1: compression type
+// Param 2: whether to use direct reads
+// Param 3: Block Based Table Index type
+// Param 4: BBTO no_block_cache option
+// Param 5: test mode for the user-defined timestamp feature
+// Param 6: number of parallel compression threads
+// Param 7: CompressionOptions.max_dict_bytes and
+//          CompressionOptions.max_dict_buffer_bytes to enable/disable
+//          compression dictionary.
 class BlockBasedTableReaderTest
     : public BlockBasedTableReaderBaseTest,
       public testing::WithParamInterface<std::tuple<
-          CompressionType, bool, BlockBasedTableOptions::IndexType, bool>> {
+          CompressionType, bool, BlockBasedTableOptions::IndexType, bool,
+          test::UserDefinedTimestampTestMode, uint32_t, uint32_t>> {
  protected:
   void SetUp() override {
     compression_type_ = std::get<0>(GetParam());
     use_direct_reads_ = std::get<1>(GetParam());
+    test::UserDefinedTimestampTestMode udt_test_mode = std::get<4>(GetParam());
+    udt_enabled_ = test::IsUDTEnabled(udt_test_mode);
+    persist_udt_ = test::ShouldPersistUDT(udt_test_mode);
+    compression_parallel_threads_ = std::get<5>(GetParam());
+    compression_dict_bytes_ = std::get<6>(GetParam());
     BlockBasedTableReaderBaseTest::SetUp();
   }
 
@@ -195,24 +232,45 @@ class BlockBasedTableReaderTest
     BlockBasedTableOptions opts;
     opts.index_type = std::get<2>(GetParam());
     opts.no_block_cache = std::get<3>(GetParam());
+    opts.filter_policy.reset(NewBloomFilterPolicy(10, false));
+    opts.partition_filters =
+        opts.index_type ==
+        BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch;
     options_.table_factory.reset(
         static_cast<BlockBasedTableFactory*>(NewBlockBasedTableFactory(opts)));
+    options_.prefix_extractor =
+        std::shared_ptr<const SliceTransform>(NewFixedPrefixTransform(3));
   }
 
   CompressionType compression_type_;
   bool use_direct_reads_;
+  bool udt_enabled_;
+  bool persist_udt_;
+  uint32_t compression_parallel_threads_;
+  uint32_t compression_dict_bytes_;
 };
 
 // Tests MultiGet in both direct IO and non-direct IO mode.
 // The keys should be in cache after MultiGet.
 TEST_P(BlockBasedTableReaderTest, MultiGet) {
+  Options options;
+  ReadOptions read_opts;
+  std::string dummy_ts(sizeof(uint64_t), '\0');
+  Slice read_timestamp = dummy_ts;
+  if (udt_enabled_) {
+    options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+    read_opts.timestamp = &read_timestamp;
+  }
+  options.persist_user_defined_timestamps = persist_udt_;
+  size_t ts_sz = options.comparator->timestamp_size();
   std::map<std::string, std::string> kv =
       BlockBasedTableReaderBaseTest::GenerateKVMap(
           100 /* num_block */,
-          true /* mixed_with_human_readable_string_value */);
+          true /* mixed_with_human_readable_string_value */, ts_sz);
 
   // Prepare keys, values, and statuses for MultiGet.
   autovector<Slice, MultiGetContext::MAX_BATCH_SIZE> keys;
+  autovector<Slice, MultiGetContext::MAX_BATCH_SIZE> keys_without_timestamps;
   autovector<PinnableSlice, MultiGetContext::MAX_BATCH_SIZE> values;
   autovector<Status, MultiGetContext::MAX_BATCH_SIZE> statuses;
   {
@@ -221,27 +279,40 @@ TEST_P(BlockBasedTableReaderTest, MultiGet) {
     auto it = kv.begin();
     for (int i = 0; i < MultiGetContext::MAX_BATCH_SIZE; i++) {
       keys.emplace_back(it->first);
+      if (ts_sz > 0) {
+        Slice ukey_without_ts = StripTimestampFromUserKey(it->first, ts_sz);
+        keys_without_timestamps.push_back(ukey_without_ts);
+      } else {
+        keys_without_timestamps.emplace_back(it->first);
+      }
       values.emplace_back();
       statuses.emplace_back();
       std::advance(it, step);
     }
   }
 
-  std::string table_name =
-      "BlockBasedTableReaderTest" + CompressionTypeToString(compression_type_);
-  CreateTable(table_name, compression_type_, kv);
+  std::string table_name = "BlockBasedTableReaderTest_MultiGet" +
+                           CompressionTypeToString(compression_type_);
 
-  std::unique_ptr<BlockBasedTable> table;
-  Options options;
   ImmutableOptions ioptions(options);
+  CreateTable(table_name, ioptions, compression_type_, kv,
+              compression_parallel_threads_, compression_dict_bytes_);
+
+  std::unique_ptr<BlockBasedTable> table;
   FileOptions foptions;
   foptions.use_direct_reads = use_direct_reads_;
   InternalKeyComparator comparator(options.comparator);
-  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
+
+  ASSERT_OK(
+      table->VerifyChecksum(read_opts, TableReaderCaller::kUserVerifyChecksum));
 
   // Ensure that keys are not in cache before MultiGet.
   for (auto& key : keys) {
-    ASSERT_FALSE(table->TEST_KeyInCache(ReadOptions(), key));
+    std::string ikey = ToInternalKey(key.ToString());
+    ASSERT_FALSE(table->TEST_KeyInCache(read_opts, ikey));
   }
 
   // Prepare MultiGetContext.
@@ -249,26 +320,26 @@ TEST_P(BlockBasedTableReaderTest, MultiGet) {
   autovector<KeyContext, MultiGetContext::MAX_BATCH_SIZE> key_context;
   autovector<KeyContext*, MultiGetContext::MAX_BATCH_SIZE> sorted_keys;
   for (size_t i = 0; i < keys.size(); ++i) {
-    get_context.emplace_back(BytewiseComparator(), nullptr, nullptr, nullptr,
+    get_context.emplace_back(options.comparator, nullptr, nullptr, nullptr,
                              GetContext::kNotFound, keys[i], &values[i],
                              nullptr, nullptr, nullptr, nullptr,
                              true /* do_merge */, nullptr, nullptr, nullptr,
                              nullptr, nullptr, nullptr);
-    key_context.emplace_back(nullptr, keys[i], &values[i], nullptr, nullptr,
-                             &statuses.back());
+    key_context.emplace_back(nullptr, keys_without_timestamps[i], &values[i],
+                             nullptr, nullptr, &statuses.back());
     key_context.back().get_context = &get_context.back();
   }
   for (auto& key_ctx : key_context) {
     sorted_keys.emplace_back(&key_ctx);
   }
-  MultiGetContext ctx(&sorted_keys, 0, sorted_keys.size(), 0, ReadOptions(),
+  MultiGetContext ctx(&sorted_keys, 0, sorted_keys.size(), 0, read_opts,
                       fs_.get(), nullptr);
 
   // Execute MultiGet.
   MultiGetContext::Range range = ctx.GetMultiGetRange();
   PerfContext* perf_ctx = get_perf_context();
   perf_ctx->Reset();
-  table->MultiGet(ReadOptions(), &range, nullptr);
+  table->MultiGet(read_opts, &range, nullptr);
 
   ASSERT_GE(perf_ctx->block_read_count - perf_ctx->index_block_read_count -
                 perf_ctx->filter_block_read_count -
@@ -281,11 +352,78 @@ TEST_P(BlockBasedTableReaderTest, MultiGet) {
   }
   // Check that keys are in cache after MultiGet.
   for (size_t i = 0; i < keys.size(); i++) {
-    ASSERT_TRUE(table->TEST_KeyInCache(ReadOptions(), keys[i]));
+    std::string ikey = ToInternalKey(keys[i].ToString());
+    ASSERT_TRUE(table->TEST_KeyInCache(read_opts, ikey));
     ASSERT_EQ(values[i].ToString(), kv[keys[i].ToString()]);
   }
 }
 
+TEST_P(BlockBasedTableReaderTest, NewIterator) {
+  Options options;
+  ReadOptions read_opts;
+  std::string dummy_ts(sizeof(uint64_t), '\0');
+  Slice read_timestamp = dummy_ts;
+  if (udt_enabled_) {
+    options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+    read_opts.timestamp = &read_timestamp;
+  }
+  options.persist_user_defined_timestamps = persist_udt_;
+  size_t ts_sz = options.comparator->timestamp_size();
+  std::map<std::string, std::string> kv =
+      BlockBasedTableReaderBaseTest::GenerateKVMap(
+          100 /* num_block */,
+          true /* mixed_with_human_readable_string_value */, ts_sz);
+
+  std::string table_name = "BlockBasedTableReaderTest_NewIterator" +
+                           CompressionTypeToString(compression_type_);
+
+  ImmutableOptions ioptions(options);
+  CreateTable(table_name, ioptions, compression_type_, kv,
+              compression_parallel_threads_, compression_dict_bytes_);
+
+  std::unique_ptr<BlockBasedTable> table;
+  FileOptions foptions;
+  foptions.use_direct_reads = use_direct_reads_;
+  InternalKeyComparator comparator(options.comparator);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
+  ASSERT_OK(
+      table->VerifyChecksum(read_opts, TableReaderCaller::kUserVerifyChecksum));
+
+  std::unique_ptr<InternalIterator> iter;
+  iter.reset(table->NewIterator(
+      read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
+      /*skip_filters=*/false, TableReaderCaller::kUncategorized));
+
+  // Test forward scan.
+  ASSERT_TRUE(!iter->Valid());
+  iter->SeekToFirst();
+  ASSERT_OK(iter->status());
+  for (auto kv_iter = kv.begin(); kv_iter != kv.end(); kv_iter++) {
+    std::string ikey = ToInternalKey(kv_iter->first);
+    ASSERT_EQ(iter->key().ToString(), ikey);
+    ASSERT_EQ(iter->value().ToString(), kv_iter->second);
+    iter->Next();
+    ASSERT_OK(iter->status());
+  }
+  ASSERT_TRUE(!iter->Valid());
+  ASSERT_OK(iter->status());
+
+  // Test backward scan.
+  iter->SeekToLast();
+  ASSERT_OK(iter->status());
+  for (auto kv_iter = kv.rbegin(); kv_iter != kv.rend(); kv_iter++) {
+    std::string ikey = ToInternalKey(kv_iter->first);
+    ASSERT_EQ(iter->key().ToString(), ikey);
+    ASSERT_EQ(iter->value().ToString(), kv_iter->second);
+    iter->Prev();
+    ASSERT_OK(iter->status());
+  }
+  ASSERT_TRUE(!iter->Valid());
+  ASSERT_OK(iter->status());
+}
+
 class ChargeTableReaderTest
     : public BlockBasedTableReaderBaseTest,
       public testing::WithParamInterface<
@@ -374,7 +512,8 @@ class ChargeTableReaderTest
     std::size_t approx_table_reader_mem = 0;
 
     std::string table_name = "table_for_approx_table_reader_mem";
-    CreateTable(table_name, compression_type_, kv_);
+    ImmutableOptions ioptions(options_);
+    CreateTable(table_name, ioptions, compression_type_, kv_);
 
     std::unique_ptr<BlockBasedTable> table;
     Status s;
@@ -424,13 +563,14 @@ TEST_P(ChargeTableReaderTest, Basic) {
   std::size_t opened_table_reader_num = 0;
   std::string table_name;
   std::vector<std::unique_ptr<BlockBasedTable>> tables;
+  ImmutableOptions ioptions(options_);
   // Keep creating BlockBasedTableReader till hiting the memory limit based on
   // cache capacity and creation fails (when charge_table_reader_ ==
   // kEnabled) or reaching a specfied big number of table readers (when
   // charge_table_reader_ == kDisabled)
   while (s.ok() && opened_table_reader_num < max_table_reader_num_uncapped) {
     table_name = "table_" + std::to_string(opened_table_reader_num);
-    CreateTable(table_name, compression_type_, kv_);
+    CreateTable(table_name, ioptions, compression_type_, kv_);
     tables.push_back(std::unique_ptr<BlockBasedTable>());
     NewBlockBasedTableReader(
         FileOptions(), ImmutableOptions(options_),
@@ -465,7 +605,7 @@ TEST_P(ChargeTableReaderTest, Basic) {
       --opened_table_reader_num;
     }
     table_name = "table_for_successful_table_reader_open";
-    CreateTable(table_name, compression_type_, kv_);
+    CreateTable(table_name, ioptions, compression_type_, kv_);
     tables.push_back(std::unique_ptr<BlockBasedTable>());
     NewBlockBasedTableReader(
         FileOptions(), ImmutableOptions(options_),
@@ -491,28 +631,42 @@ class BlockBasedTableReaderTestVerifyChecksum
 };
 
 TEST_P(BlockBasedTableReaderTestVerifyChecksum, ChecksumMismatch) {
+  Options options;
+  ReadOptions read_opts;
+  std::string dummy_ts(sizeof(uint64_t), '\0');
+  Slice read_timestamp = dummy_ts;
+  if (udt_enabled_) {
+    options.comparator = test::BytewiseComparatorWithU64TsWrapper();
+    read_opts.timestamp = &read_timestamp;
+  }
+  options.persist_user_defined_timestamps = persist_udt_;
+  size_t ts_sz = options.comparator->timestamp_size();
   std::map<std::string, std::string> kv =
-      BlockBasedTableReaderBaseTest::GenerateKVMap(800 /* num_block */);
+      BlockBasedTableReaderBaseTest::GenerateKVMap(
+          800 /* num_block */,
+          false /* mixed_with_human_readable_string_value=*/, ts_sz);
 
+  options.statistics = CreateDBStatistics();
+  ImmutableOptions ioptions(options);
   std::string table_name =
       "BlockBasedTableReaderTest" + CompressionTypeToString(compression_type_);
-  CreateTable(table_name, compression_type_, kv);
+  CreateTable(table_name, ioptions, compression_type_, kv,
+              compression_parallel_threads_, compression_dict_bytes_);
 
   std::unique_ptr<BlockBasedTable> table;
-  Options options;
-  options.statistics = CreateDBStatistics();
-  ImmutableOptions ioptions(options);
   FileOptions foptions;
   foptions.use_direct_reads = use_direct_reads_;
   InternalKeyComparator comparator(options.comparator);
-  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
 
   // Use the top level iterator to find the offset/size of the first
   // 2nd level index block and corrupt the block
   IndexBlockIter iiter_on_stack;
   BlockCacheLookupContext context{TableReaderCaller::kUserVerifyChecksum};
   InternalIteratorBase<IndexValue>* iiter = table->NewIndexIterator(
-      ReadOptions(), /*disable_prefix_seek=*/false, &iiter_on_stack,
+      read_opts, /*need_upper_bound_check=*/false, &iiter_on_stack,
       /*get_context=*/nullptr, &context);
   std::unique_ptr<InternalIteratorBase<IndexValue>> iiter_unique_ptr;
   if (iiter != &iiter_on_stack) {
@@ -529,11 +683,13 @@ TEST_P(BlockBasedTableReaderTestVerifyChecksum, ChecksumMismatch) {
   ASSERT_OK(test::CorruptFile(options.env, Path(table_name),
                               static_cast<int>(handle.offset()), 128));
 
-  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table);
+  NewBlockBasedTableReader(foptions, ioptions, comparator, table_name, &table,
+                           true /* bool prefetch_index_and_filter_in_cache */,
+                           nullptr /* status */, persist_udt_);
   ASSERT_EQ(0,
             options.statistics->getTickerCount(BLOCK_CHECKSUM_MISMATCH_COUNT));
-  Status s = table->VerifyChecksum(ReadOptions(),
-                                   TableReaderCaller::kUserVerifyChecksum);
+  Status s =
+      table->VerifyChecksum(read_opts, TableReaderCaller::kUserVerifyChecksum);
   ASSERT_EQ(1,
             options.statistics->getTickerCount(BLOCK_CHECKSUM_MISMATCH_COUNT));
   ASSERT_EQ(s.code(), Status::kCorruption);
@@ -541,14 +697,25 @@ TEST_P(BlockBasedTableReaderTestVerifyChecksum, ChecksumMismatch) {
 
 // Param 1: compression type
 // Param 2: whether to use direct reads
-// Param 3: Block Based Table Index type
+// Param 3: Block Based Table Index type, partitioned filters are also enabled
+//          when index type is kTwoLevelIndexSearch
 // Param 4: BBTO no_block_cache option
+// Param 5: test mode for the user-defined timestamp feature
+// Param 6: number of parallel compression threads
+// Param 7: CompressionOptions.max_dict_bytes and
+//          CompressionOptions.max_dict_buffer_bytes. This enable/disables
+//          compression dictionary.
 INSTANTIATE_TEST_CASE_P(
-    MultiGet, BlockBasedTableReaderTest,
+    BlockBasedTableReaderTest, BlockBasedTableReaderTest,
     ::testing::Combine(
         ::testing::ValuesIn(GetSupportedCompressions()), ::testing::Bool(),
-        ::testing::Values(BlockBasedTableOptions::IndexType::kBinarySearch),
-        ::testing::Values(false)));
+        ::testing::Values(
+            BlockBasedTableOptions::IndexType::kBinarySearch,
+            BlockBasedTableOptions::IndexType::kHashSearch,
+            BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch,
+            BlockBasedTableOptions::IndexType::kBinarySearchWithFirstKey),
+        ::testing::Values(false), ::testing::ValuesIn(test::GetUDTTestModes()),
+        ::testing::Values(1, 2), ::testing::Values(0, 4096)));
 INSTANTIATE_TEST_CASE_P(
     VerifyChecksum, BlockBasedTableReaderTestVerifyChecksum,
     ::testing::Combine(
@@ -556,7 +723,8 @@ INSTANTIATE_TEST_CASE_P(
         ::testing::Values(false),
         ::testing::Values(
             BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch),
-        ::testing::Values(true)));
+        ::testing::Values(true), ::testing::ValuesIn(test::GetUDTTestModes()),
+        ::testing::Values(1, 2), ::testing::Values(0)));
 
 }  // namespace ROCKSDB_NAMESPACE
 

package/deps/rocksdb/rocksdb/table/block_based/block_builder.cc

@@ -48,10 +48,13 @@ BlockBuilder::BlockBuilder(
     int block_restart_interval, bool use_delta_encoding,
     bool use_value_delta_encoding,
     BlockBasedTableOptions::DataBlockIndexType index_type,
-    double data_block_hash_table_util_ratio)
+    double data_block_hash_table_util_ratio, size_t ts_sz,
+    bool persist_user_defined_timestamps, bool is_user_key)
     : block_restart_interval_(block_restart_interval),
       use_delta_encoding_(use_delta_encoding),
       use_value_delta_encoding_(use_value_delta_encoding),
+      strip_ts_sz_(persist_user_defined_timestamps ? 0 : ts_sz),
+      is_user_key_(is_user_key),
       restarts_(1, 0),  // First restart point is at offset 0
       counter_(0),
       finished_(false) {
@@ -96,6 +99,9 @@ size_t BlockBuilder::EstimateSizeAfterKV(const Slice& key,
   // Note: this is an imprecise estimate as it accounts for the whole key size
   // instead of non-shared key size.
   estimate += key.size();
+  if (strip_ts_sz_ > 0) {
+    estimate -= strip_ts_sz_;
+  }
   // In value delta encoding we estimate the value delta size as half the full
   // value size since only the size field of block handle is encoded.
   estimate +=
@@ -168,13 +174,13 @@ void BlockBuilder::AddWithLastKey(const Slice& key, const Slice& value,
   // or Reset. This is more convenient for the caller and we can be more
   // clever inside BlockBuilder. On this hot code path, we want to avoid
   // conditional jumps like `buffer_.empty() ? ... : ...` so we can use a
-  // fast min operation instead, with an assertion to be sure our logic is
-  // sound.
+  // fast arithmetic operation instead, with an assertion to be sure our logic
+  // is sound.
   size_t buffer_size = buffer_.size();
   size_t last_key_size = last_key_param.size();
-  assert(buffer_size == 0 || buffer_size >= last_key_size);
+  assert(buffer_size == 0 || buffer_size >= last_key_size - strip_ts_sz_);
 
-  Slice last_key(last_key_param.data(), std::min(buffer_size, last_key_size));
+  Slice last_key(last_key_param.data(), last_key_size * (buffer_size > 0));
 
   AddWithLastKeyImpl(key, value, last_key, delta_value, buffer_size);
 }
@@ -187,6 +193,15 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key,
   assert(!finished_);
   assert(counter_ <= block_restart_interval_);
   assert(!use_value_delta_encoding_ || delta_value);
+  std::string key_buf;
+  std::string last_key_buf;
+  const Slice key_to_persist = MaybeStripTimestampFromKey(&key_buf, key);
+  // For delta key encoding, the first key in each restart interval doesn't have
+  // a last key to share bytes with.
+  const Slice last_key_persisted =
+      last_key.size() == 0
+          ? last_key
+          : MaybeStripTimestampFromKey(&last_key_buf, last_key);
   size_t shared = 0;  // number of bytes shared with prev key
   if (counter_ >= block_restart_interval_) {
     // Restart compression
@@ -195,10 +210,10 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key,
     counter_ = 0;
   } else if (use_delta_encoding_) {
     // See how much sharing to do with previous string
-    shared = key.difference_offset(last_key);
+    shared = key_to_persist.difference_offset(last_key_persisted);
   }
 
-  const size_t non_shared = key.size() - shared;
+  const size_t non_shared = key_to_persist.size() - shared;
 
   if (use_value_delta_encoding_) {
     // Add "<shared><non_shared>" to buffer_
@@ -212,7 +227,7 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key,
   }
 
   // Add string delta to buffer_ followed by value
-  buffer_.append(key.data() + shared, non_shared);
+  buffer_.append(key_to_persist.data() + shared, non_shared);
   // Use value delta encoding only when the key has shared bytes. This would
   // simplify the decoding, where it can figure which decoding to use simply by
   // looking at the shared bytes size.
@@ -222,7 +237,12 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key,
     buffer_.append(value.data(), value.size());
   }
 
+  // TODO(yuzhangyu): make user defined timestamp work with block hash index.
   if (data_block_hash_index_builder_.Valid()) {
+    // Only data blocks should be using `kDataBlockBinaryAndHash` index type.
+    // And data blocks should always be built with internal keys instead of
+    // user keys.
+    assert(!is_user_key_);
     data_block_hash_index_builder_.Add(ExtractUserKey(key),
                                        restarts_.size() - 1);
   }
@@ -231,4 +251,17 @@ inline void BlockBuilder::AddWithLastKeyImpl(const Slice& key,
   estimate_ += buffer_.size() - buffer_size;
 }
 
+const Slice BlockBuilder::MaybeStripTimestampFromKey(std::string* key_buf,
+                                                     const Slice& key) {
+  Slice stripped_key = key;
+  if (strip_ts_sz_ > 0) {
+    if (is_user_key_) {
+      stripped_key.remove_suffix(strip_ts_sz_);
+    } else {
+      StripTimestampFromInternalKey(key_buf, key, strip_ts_sz_);
+      stripped_key = *key_buf;
+    }
+  }
+  return stripped_key;
+}
 }  // namespace ROCKSDB_NAMESPACE

package/deps/rocksdb/rocksdb/table/block_based/block_builder.h

@@ -28,7 +28,10 @@ class BlockBuilder {
                         bool use_value_delta_encoding = false,
                         BlockBasedTableOptions::DataBlockIndexType index_type =
                             BlockBasedTableOptions::kDataBlockBinarySearch,
-                        double data_block_hash_table_util_ratio = 0.75);
+                        double data_block_hash_table_util_ratio = 0.75,
+                        size_t ts_sz = 0,
+                        bool persist_user_defined_timestamps = true,
+                        bool is_user_key = false);
 
   // Reset the contents as if the BlockBuilder was just constructed.
   void Reset();
@@ -83,11 +86,32 @@ class BlockBuilder {
                                  const Slice* const delta_value,
                                  size_t buffer_size);
 
+  inline const Slice MaybeStripTimestampFromKey(std::string* key_buf,
+                                                const Slice& key);
+
   const int block_restart_interval_;
   // TODO(myabandeh): put it into a separate IndexBlockBuilder
   const bool use_delta_encoding_;
   // Refer to BlockIter::DecodeCurrentValue for format of delta encoded values
   const bool use_value_delta_encoding_;
+  // Size in bytes for the user-defined timestamp to strip in a user key.
+  // This is non-zero if there is user-defined timestamp in the user key and it
+  // should not be persisted.
+  const size_t strip_ts_sz_;
+  // Whether the keys provided to build this block are user keys. If not,
+  // the keys are internal keys. This will affect how timestamp stripping is
+  // done for the key if `persisted_user_defined_timestamps_` is false and
+  // `ts_sz_` is non-zero.
+  // The timestamp stripping only applies to the keys added to the block. If the
+  // value contains user defined timestamp that needed to be stripped too, such
+  // as the `first_internal_key` in an `IndexValue` for an index block, the
+  // value part for a range deletion entry, their timestamp should be stripped
+  // before calling `BlockBuilder::Add`.
+  // Timestamp stripping only applies to data block and index blocks including
+  // index block for data blocks, index block for partitioned filter blocks,
+  // index block for partitioned index blocks. In summary, this only applies to
+  // block whose key are real user keys or internal keys created from user keys.
+  const bool is_user_key_;
 
   std::string buffer_;              // Destination buffer
   std::vector<uint32_t> restarts_;  // Restart points