Force caller to pass comparator in MultiScanArgs (#13970)

Xingbo Wang · facebook-github-bot · commit 5a1ff2cb1469 · 2025-09-18T15:18:18.000-07:00
Summary: Force caller of MultiScanArgs to pass comparator. Pass comparator from CF handle to MultiScanArgs in NewMultiScan. Expand MultiScanArgs unit test with different comparator. Pull Request resolved: #13970 Test Plan: unit test Reviewed By: cbi42 Differential Revision: D82739270 Pulled By: xingbowang fbshipit-source-id: e709f4a333ad547c0ba6d24d8fb2b22e50e8a12f
diff --git a/db/version_set.cc b/db/version_set.cc
@@ -1101,6 +1101,27 @@ class LevelIterator final : public InternalIterator {
     read_seq_ = read_seq;
   }
 
+  inline bool FileHasMultiScanArg(size_t file_index) {
+    if (file_to_scan_opts_.get()) {
+      auto it = file_to_scan_opts_->find(file_index);
+      if (it != file_to_scan_opts_->end()) {
+        return !it->second.empty();
+      }
+    }
+    return false;
+  }
+
+  MultiScanArgs& GetMultiScanArgForFile(size_t file_index) {
+    auto multi_scan_args_it = file_to_scan_opts_->find(file_index);
+    if (multi_scan_args_it == file_to_scan_opts_->end()) {
+      auto ret = file_to_scan_opts_->emplace(
+          file_index, MultiScanArgs(user_comparator_.user_comparator()));
+      multi_scan_args_it = ret.first;
+      assert(ret.second);
+    }
+    return multi_scan_args_it->second;
+  }
+
   void Prepare(const MultiScanArgs* so) override {
     // We assume here that scan_opts is sorted such that
     // scan_opts[0].range.start < scan_opts[1].range.start, and non overlapping
@@ -1109,6 +1130,9 @@ class LevelIterator final : public InternalIterator {
     }
     scan_opts_ = so;
 
+    // Verify comparator is consistent
+    assert(so->GetComparator() == user_comparator_.user_comparator());
+
     file_to_scan_opts_ = std::make_unique<ScanOptionsMap>();
     for (size_t k = 0; k < scan_opts_->size(); k++) {
       const ScanOptions& opt = scan_opts_->GetScanRanges().at(k);
@@ -1157,8 +1181,8 @@ class LevelIterator final : public InternalIterator {
       // 3. [  S  ] ...... [  E  ]
       for (auto i = fstart; i <= fend; i++) {
         if (i < flevel_->num_files) {
-          (*file_to_scan_opts_)[i].insert(start.value(), end.value(),
-                                          opt.property_bag);
+          auto args = GetMultiScanArgForFile(i);
+          args.insert(start.value(), end.value(), opt.property_bag);
         }
       }
     }
@@ -1562,9 +1586,10 @@ bool LevelIterator::SkipEmptyFileForward() {
     if (file_iter_.iter() != nullptr) {
       // If we are doing prepared scan opts then we should seek to the values
       // specified by the scan opts
-      if (scan_opts_ && (*file_to_scan_opts_)[file_index_].size()) {
+
+      if (scan_opts_ && FileHasMultiScanArg(file_index_)) {
         const ScanOptions& opts =
-            file_to_scan_opts_->at(file_index_).GetScanRanges().front();
+            GetMultiScanArgForFile(file_index_).GetScanRanges().front();
         if (opts.range.start.has_value()) {
           InternalKey target(*opts.range.start.AsPtr(), kMaxSequenceNumber,
                              kValueTypeForSeek);
@@ -1621,9 +1646,8 @@ void LevelIterator::SetFileIterator(InternalIterator* iter) {
 
   InternalIterator* old_iter = file_iter_.Set(iter);
   if (iter && scan_opts_) {
-    if (file_to_scan_opts_.get() &&
-        file_to_scan_opts_->find(file_index_) != file_to_scan_opts_->end()) {
-      const MultiScanArgs& new_opts = file_to_scan_opts_->at(file_index_);
+    if (FileHasMultiScanArg(file_index_)) {
+      const MultiScanArgs& new_opts = GetMultiScanArgForFile(file_index_);
       file_iter_.Prepare(&new_opts);
     } else {
       file_iter_.Prepare(scan_opts_);
diff --git a/db_stress_tool/db_stress_test_base.cc b/db_stress_tool/db_stress_test_base.cc
@@ -1692,7 +1692,8 @@ Status StressTest::TestMultiScan(ThreadState* thread,
 
   std::vector<std::string> start_key_strs;
   std::vector<std::string> end_key_strs;
-  MultiScanArgs scan_opts;
+  // TODO support reverse BytewiseComparator in the stress test
+  MultiScanArgs scan_opts(BytewiseComparator());
   scan_opts.use_async_io = FLAGS_multiscan_use_async_io;
   start_key_strs.reserve(num_scans);
   end_key_strs.reserve(num_scans);
diff --git a/include/rocksdb/db.h b/include/rocksdb/db.h
@@ -1118,11 +1118,11 @@ class DB {
   //    // Check ex.what()
   //  }
   virtual std::unique_ptr<MultiScan> NewMultiScan(
-      const ReadOptions& /*options*/, ColumnFamilyHandle* /*column_family*/,
+      const ReadOptions& /*options*/, ColumnFamilyHandle* column_family,
       const MultiScanArgs& /*scan_opts*/) {
     std::unique_ptr<Iterator> iter(NewErrorIterator(Status::NotSupported()));
-    std::unique_ptr<MultiScan> ms_iter =
-        std::make_unique<MultiScan>(std::move(iter));
+    std::unique_ptr<MultiScan> ms_iter = std::make_unique<MultiScan>(
+        column_family->GetComparator(), std::move(iter));
     return ms_iter;
   }
 
diff --git a/include/rocksdb/multi_scan.h b/include/rocksdb/multi_scan.h
@@ -155,8 +155,9 @@ class MultiScan {
   MultiScan(const ReadOptions& read_options, const MultiScanArgs& scan_opts,
             DB* db, ColumnFamilyHandle* cfh);
 
-  explicit MultiScan(std::unique_ptr<Iterator>&& db_iter)
-      : db_iter_(std::move(db_iter)) {}
+  explicit MultiScan(const Comparator* comp,
+                     std::unique_ptr<Iterator>&& db_iter)
+      : scan_opts_(comp), db_iter_(std::move(db_iter)) {}
 
   class MultiScanIterator {
    public:
diff --git a/include/rocksdb/options.h b/include/rocksdb/options.h
@@ -1782,8 +1782,7 @@ struct ScanOptions {
 class MultiScanArgs {
  public:
   // Constructor that takes a comparator
-  explicit MultiScanArgs(const Comparator* comparator = BytewiseComparator())
-      : comp_(comparator) {}
+  explicit MultiScanArgs(const Comparator* comparator) : comp_(comparator) {}
 
   // Copy Constructor
   MultiScanArgs(const MultiScanArgs& other) {
@@ -1855,6 +1854,8 @@ class MultiScanArgs {
     return original_ranges_;
   }
 
+  const Comparator* GetComparator() const { return comp_; }
+
   uint64_t io_coalesce_threshold = 16 << 10;  // 16KB by default
 
   // Maximum size (in bytes) for the data blocks loaded by a MultiScan.
diff --git a/table/block_based/block_based_table_reader_test.cc b/table/block_based/block_based_table_reader_test.cc
@@ -52,7 +52,8 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
   // user defined timestamps and different sequence number to differentiate them
   static std::vector<std::pair<std::string, std::string>> GenerateKVMap(
       int num_block = 2, bool mixed_with_human_readable_string_value = false,
-      size_t ts_sz = 0, bool same_key_diff_ts = false) {
+      size_t ts_sz = 0, bool same_key_diff_ts = false,
+      const Comparator* comparator = BytewiseComparator()) {
     std::vector<std::pair<std::string, std::string>> kv;
 
     SequenceNumber seq_no = 0;
@@ -100,6 +101,10 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
         }
       }
     }
+    auto comparator_name = std::string(comparator->Name());
+    if (comparator_name.find("Reverse") != std::string::npos) {
+      std::reverse(kv.begin(), kv.end());
+    }
     return kv;
   }
 
@@ -128,6 +133,7 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
 
     InternalKeyComparator comparator(ioptions.user_comparator);
     ColumnFamilyOptions cf_options;
+    cf_options.comparator = ioptions.user_comparator;
     cf_options.prefix_extractor = options_.prefix_extractor;
     MutableCFOptions moptions(cf_options);
     CompressionOptions compression_opts;
@@ -255,11 +261,13 @@ class BlockBasedTableReaderBaseTest : public testing::Test {
 //          generate keys with different user provided key, same user-defined
 //          timestamps (if udt enabled), same sequence number. This test mode is
 //          used for testing `Get`, `MultiGet`, and `NewIterator`.
+// Param 9: test both the default comparator and a reverse comparator.
 class BlockBasedTableReaderTest
     : public BlockBasedTableReaderBaseTest,
-      public testing::WithParamInterface<std::tuple<
-          CompressionType, bool, BlockBasedTableOptions::IndexType, bool,
-          test::UserDefinedTimestampTestMode, uint32_t, uint32_t, bool>> {
+      public testing::WithParamInterface<
+          std::tuple<CompressionType, bool, BlockBasedTableOptions::IndexType,
+                     bool, test::UserDefinedTimestampTestMode, uint32_t,
+                     uint32_t, bool, const Comparator*>> {
  protected:
   void SetUp() override {
     compression_type_ = std::get<0>(GetParam());
@@ -270,6 +278,7 @@ class BlockBasedTableReaderTest
     compression_parallel_threads_ = std::get<5>(GetParam());
     compression_dict_bytes_ = std::get<6>(GetParam());
     same_key_diff_ts_ = std::get<7>(GetParam());
+    comparator_ = std::get<8>(GetParam());
     BlockBasedTableReaderBaseTest::SetUp();
   }
 
@@ -295,6 +304,7 @@ class BlockBasedTableReaderTest
   uint32_t compression_parallel_threads_;
   uint32_t compression_dict_bytes_;
   bool same_key_diff_ts_;
+  const Comparator* comparator_{};
 };
 
 class BlockBasedTableReaderGetTest : public BlockBasedTableReaderTest {};
@@ -1022,14 +1032,16 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       SCOPED_TRACE(std::string("use_async_io=") + std::to_string(use_async_io));
       Options options;
       options.statistics = CreateDBStatistics();
+      options.comparator = comparator_;
       std::shared_ptr<FileSystem> fs = options.env->GetFileSystem();
       ReadOptions read_opts;
       read_opts.fill_cache = fill_cache;
       size_t ts_sz = options.comparator->timestamp_size();
       std::vector<std::pair<std::string, std::string>> kv =
           BlockBasedTableReaderBaseTest::GenerateKVMap(
               100 /* num_block */,
-              true /* mixed_with_human_readable_string_value */, ts_sz);
+              true /* mixed_with_human_readable_string_value */, ts_sz,
+              same_key_diff_ts_, comparator_);
       std::string table_name = "BlockBasedTableReaderTest_NewIterator" +
                                CompressionTypeToString(compression_type_) +
                                "_async" + std::to_string(use_async_io);
@@ -1052,7 +1064,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
           read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
           /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-      MultiScanArgs scan_options(BytewiseComparator());
+      MultiScanArgs scan_options(comparator_);
       scan_options.use_async_io = use_async_io;
       scan_options.insert(ExtractUserKey(kv[0].first),
                           ExtractUserKey(kv[kEntriesPerBlock].first));
@@ -1087,7 +1099,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       iter.reset(table->NewIterator(
           read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
           /*skip_filters=*/false, TableReaderCaller::kUncategorized));
-      scan_options = MultiScanArgs(BytewiseComparator());
+      scan_options = MultiScanArgs(comparator_);
       scan_options.insert(ExtractUserKey(kv[70 * kEntriesPerBlock].first),
                           ExtractUserKey(kv[75 * kEntriesPerBlock].first));
       scan_options.insert(ExtractUserKey(kv[90 * kEntriesPerBlock].first),
@@ -1125,7 +1137,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       // From reads above, blocks 70-75 and 90-95 already in cache
       // So we should read 50-70 76-89 96-99 in three I/Os.
       // If fill_cache is false, then we'll do one giant I/O.
-      scan_options = MultiScanArgs(BytewiseComparator());
+      scan_options = MultiScanArgs(comparator_);
       scan_options.use_async_io = use_async_io;
       scan_options.insert(ExtractUserKey(kv[50 * kEntriesPerBlock].first));
       read_count_before =
@@ -1165,7 +1177,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       iter.reset(table->NewIterator(
           read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
           /*skip_filters=*/false, TableReaderCaller::kUncategorized));
-      scan_options = MultiScanArgs(BytewiseComparator());
+      scan_options = MultiScanArgs(comparator_);
       scan_options.use_async_io = use_async_io;
       scan_options.insert(ExtractUserKey(kv[10 * kEntriesPerBlock].first),
                           ExtractUserKey(kv[20 * kEntriesPerBlock].first));
@@ -1195,7 +1207,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       iter.reset(table->NewIterator(
           read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
           /*skip_filters=*/false, TableReaderCaller::kUncategorized));
-      scan_options = MultiScanArgs(BytewiseComparator());
+      scan_options = MultiScanArgs(comparator_);
       scan_options.use_async_io = use_async_io;
       scan_options.insert(ExtractUserKey(kv[10 * kEntriesPerBlock].first));
       scan_options.insert(ExtractUserKey(kv[11 * kEntriesPerBlock].first));
@@ -1226,14 +1238,15 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
     return;
   }
   Options options;
+  options.comparator = comparator_;
   ReadOptions read_opts;
   size_t ts_sz = options.comparator->timestamp_size();
 
   // Generate data that spans multiple blocks
   std::vector<std::pair<std::string, std::string>> kv =
       BlockBasedTableReaderBaseTest::GenerateKVMap(
           20 /* num_block */, true /* mixed_with_human_readable_string_value */,
-          ts_sz);
+          ts_sz, same_key_diff_ts_, comparator_);
 
   std::string table_name = "BlockBasedTableReaderTest_PrefetchSizeLimit" +
                            CompressionTypeToString(compression_type_);
@@ -1259,7 +1272,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
         read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
         /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-    MultiScanArgs scan_options(BytewiseComparator());
+    MultiScanArgs scan_options(comparator_);
     scan_options.max_prefetch_size = 1024;  // less than block size
     scan_options.insert(ExtractUserKey(kv[0].first),
                         ExtractUserKey(kv[5].first));
@@ -1279,7 +1292,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
         read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
         /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-    MultiScanArgs scan_options(BytewiseComparator());
+    MultiScanArgs scan_options(comparator_);
     scan_options.max_prefetch_size = 9 * 1024;  // 9KB - 2 blocks with buffer
     scan_options.insert(ExtractUserKey(kv[1 * kEntriesPerBlock].first),
                         ExtractUserKey(kv[8 * kEntriesPerBlock].first));
@@ -1310,7 +1323,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
         read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
         /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-    MultiScanArgs scan_options(BytewiseComparator());
+    MultiScanArgs scan_options(comparator_);
     scan_options.max_prefetch_size = 3 * 4 * 1024 + 1024;  // 3 blocks + 1KB
     scan_options.insert(ExtractUserKey(kv[0].first),
                         ExtractUserKey(kv[5 * kEntriesPerBlock].first));
@@ -1336,7 +1349,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
         read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
         /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-    MultiScanArgs scan_options(BytewiseComparator());
+    MultiScanArgs scan_options(comparator_);
     scan_options.max_prefetch_size = 5 * 4 * 1024 + 1024;  // 5 blocks + 1KB
     // Will read 5 entries from first scan range, and 4 blocks from the second
     // scan range
@@ -1373,7 +1386,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrefetchSizeLimit) {
         read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
         /*skip_filters=*/false, TableReaderCaller::kUncategorized));
 
-    MultiScanArgs scan_options(BytewiseComparator());
+    MultiScanArgs scan_options(comparator_);
     scan_options.max_prefetch_size = 10 * 1024 * 1024;  // 10MB
     scan_options.insert(ExtractUserKey(kv[0].first),
                         ExtractUserKey(kv[5].first));
@@ -1440,7 +1453,8 @@ INSTANTIATE_TEST_CASE_P(
             BlockBasedTableOptions::IndexType::kBinarySearchWithFirstKey),
         ::testing::Values(false), ::testing::ValuesIn(test::GetUDTTestModes()),
         ::testing::Values(1, 2), ::testing::Values(0, 4096),
-        ::testing::Values(false)));
+        ::testing::Values(false),
+        ::testing::Values(BytewiseComparator(), ReverseBytewiseComparator())));
 INSTANTIATE_TEST_CASE_P(
     BlockBasedTableReaderGetTest, BlockBasedTableReaderGetTest,
     ::testing::Combine(
@@ -1452,7 +1466,8 @@ INSTANTIATE_TEST_CASE_P(
             BlockBasedTableOptions::IndexType::kBinarySearchWithFirstKey),
         ::testing::Values(false), ::testing::ValuesIn(test::GetUDTTestModes()),
         ::testing::Values(1, 2), ::testing::Values(0, 4096),
-        ::testing::Values(false, true)));
+        ::testing::Values(false, true),
+        ::testing::Values(BytewiseComparator(), ReverseBytewiseComparator())));
 INSTANTIATE_TEST_CASE_P(
     StrictCapacityLimitReaderTest, StrictCapacityLimitReaderTest,
     ::testing::Combine(
@@ -1461,7 +1476,8 @@ INSTANTIATE_TEST_CASE_P(
             BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch),
         ::testing::Values(false), ::testing::ValuesIn(test::GetUDTTestModes()),
         ::testing::Values(1, 2), ::testing::Values(0),
-        ::testing::Values(false, true)));
+        ::testing::Values(false, true),
+        ::testing::Values(BytewiseComparator(), ReverseBytewiseComparator())));
 INSTANTIATE_TEST_CASE_P(
     VerifyChecksum, BlockBasedTableReaderTestVerifyChecksum,
     ::testing::Combine(
@@ -1470,8 +1486,8 @@ INSTANTIATE_TEST_CASE_P(
         ::testing::Values(
             BlockBasedTableOptions::IndexType::kTwoLevelIndexSearch),
         ::testing::Values(true), ::testing::ValuesIn(test::GetUDTTestModes()),
-        ::testing::Values(1, 2), ::testing::Values(0),
-        ::testing::Values(false)));
+        ::testing::Values(1, 2), ::testing::Values(0), ::testing::Values(false),
+        ::testing::Values(BytewiseComparator(), ReverseBytewiseComparator())));
 
 }  // namespace ROCKSDB_NAMESPACE
 
diff --git a/table/table_test.cc b/table/table_test.cc
@@ -8056,7 +8056,7 @@ TEST_F(UserDefinedIndexTest, IngestTest) {
   ro.iterate_upper_bound = nullptr;
   iter.reset(db->NewIterator(ro, cfh));
   ASSERT_NE(iter, nullptr);
-  MultiScanArgs scan_opts;
+  MultiScanArgs scan_opts(options.comparator);
   std::unordered_map<std::string, std::string> property_bag;
   property_bag["count"] = std::to_string(25);
   scan_opts.insert(Slice("key20"), std::optional(property_bag));
@@ -8146,7 +8146,7 @@ TEST_F(UserDefinedIndexTest, EmptyRangeTest) {
 
   ro.table_index_factory = user_defined_index_factory.get();
   std::vector<int> key_counts;
-  MultiScanArgs scan_opts;
+  MultiScanArgs scan_opts(options.comparator);
   std::unordered_map<std::string, std::string> property_bag;
   property_bag["count"] = std::to_string(5);
   // Empty scans
@@ -8405,7 +8405,7 @@ TEST_F(UserDefinedIndexTest, ConfigTest) {
   ro.table_index_factory = user_defined_index_factory.get();
   std::unique_ptr<Iterator> iter(db->NewIterator(ro, cfh));
   ASSERT_NE(iter, nullptr);
-  MultiScanArgs scan_opts;
+  MultiScanArgs scan_opts(options.comparator);
   std::unordered_map<std::string, std::string> property_bag;
   property_bag["count"] = std::to_string(25);
   scan_opts.insert(Slice("key20"), std::optional(property_bag));
diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc
diff --git a/unreleased_history/public_api_changes/MultiScanArgs_contructor_parameter_change.md b/unreleased_history/public_api_changes/MultiScanArgs_contructor_parameter_change.md
