Use the comparator from the sst file table properties in sst_dump_tool (#9491)

Summary:
We introduced a new Comparator for timestamp in user keys. In the sst_dump_tool by default we use BytewiseComparator to read sst files. This change allows us to read comparator_name from table properties in meta data block and use it to read.

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

Test Plan:
added unittests for new functionality.
make check
![image](https://user-images.githubusercontent.com/4923556/152915444-28b88a1f-7b4e-47d0-815f-7011552bd9a2.png)
![image](https://user-images.githubusercontent.com/4923556/152916196-bea3d2a1-a3d5-4362-b911-036131b83e8d.png)

Reviewed By: riversand963

Differential Revision: D33993614

Pulled By: satyajanga

fbshipit-source-id: 4b5cf938e6d2cb3931d763bef5baccc900b8c536

HISTORY.md

@@ -37,6 +37,7 @@
 
 ## New Features
 * Introduced an option `BlockBasedTableBuilder::detect_filter_construct_corruption` for detecting corruption during Bloom Filter (format_version >= 5) and Ribbon Filter construction.
+* Improved the SstDumpTool to read the comparator from table properties and use it to read the SST File.
 
 ## 6.29.0 (01/21/2022)
 Note: The next release will be major release 7.0. See https://github.com/facebook/rocksdb/issues/9390 for more info.

db/compaction/compaction_iterator_test.cc

@@ -1079,7 +1079,7 @@ INSTANTIATE_TEST_CASE_P(CompactionIteratorWithAllowIngestBehindTestInstance,
 class CompactionIteratorTsGcTest : public CompactionIteratorTest {
  public:
   CompactionIteratorTsGcTest()
-      : CompactionIteratorTest(test::ComparatorWithU64Ts()) {}
+      : CompactionIteratorTest(test::BytewiseComparatorWithU64TsWrapper()) {}
 };
 
 TEST_P(CompactionIteratorTsGcTest, NoKeyEligibleForGC) {

db/compaction/compaction_job_test.cc

@@ -1292,7 +1292,8 @@ class CompactionJobTimestampTest : public CompactionJobTestBase {
  public:
   CompactionJobTimestampTest()
       : CompactionJobTestBase(test::PerThreadDBPath("compaction_job_ts_test"),
-                              test::ComparatorWithU64Ts(), test::EncodeInt) {}
+                              test::BytewiseComparatorWithU64TsWrapper(),
+                              test::EncodeInt) {}
 };
 
 TEST_F(CompactionJobTimestampTest, GCDisabled) {

db/db_test2.cc

@@ -6837,7 +6837,7 @@ TEST_F(DBTest2, GetLatestSeqAndTsForKey) {
   options.max_write_buffer_size_to_maintain = 64 << 10;
   options.create_if_missing = true;
   options.disable_auto_compactions = true;
-  options.comparator = test::ComparatorWithU64Ts();
+  options.comparator = test::BytewiseComparatorWithU64TsWrapper();
   options.statistics = CreateDBStatistics();
 
   Reopen(options);

db/db_with_timestamp_basic_test.cc

@@ -206,7 +206,7 @@ TEST_F(DBBasicTestWithTimestamp, SanityChecks) {
 
   Options options1 = CurrentOptions();
   options1.env = env_;
-  options1.comparator = test::ComparatorWithU64Ts();
+  options1.comparator = test::BytewiseComparatorWithU64TsWrapper();
   options1.merge_operator = MergeOperators::CreateStringAppendTESTOperator();
   assert(options1.comparator &&
          options1.comparator->timestamp_size() == sizeof(uint64_t));

db/db_with_timestamp_compaction_test.cc

@@ -54,7 +54,7 @@ TEST_F(TimestampCompatibleCompactionTest, UserKeyCrossFileBoundary) {
   Options options = CurrentOptions();
   options.env = env_;
   options.compaction_style = kCompactionStyleLevel;
-  options.comparator = test::ComparatorWithU64Ts();
+  options.comparator = test::BytewiseComparatorWithU64TsWrapper();
   options.level0_file_num_compaction_trigger = 3;
   constexpr size_t kNumKeysPerFile = 101;
   options.memtable_factory.reset(

db/flush_job_test.cc

@@ -530,7 +530,7 @@ class FlushJobTimestampTest : public FlushJobTestBase {
  public:
   FlushJobTimestampTest()
       : FlushJobTestBase(test::PerThreadDBPath("flush_job_ts_gc_test"),
-                         test::ComparatorWithU64Ts()) {}
+                         test::BytewiseComparatorWithU64TsWrapper()) {}
 
   void AddKeyValueToMemtable(MemTable* memtable, std::string key, uint64_t ts,
                              SequenceNumber seq, ValueType value_type,

db/version_set_test.cc

@@ -700,7 +700,8 @@ TEST_F(VersionStorageInfoTest, ForcedBlobGC) {
 class VersionStorageInfoTimestampTest : public VersionStorageInfoTestBase {
  public:
   VersionStorageInfoTimestampTest()
-      : VersionStorageInfoTestBase(test::ComparatorWithU64Ts()) {}
+      : VersionStorageInfoTestBase(test::BytewiseComparatorWithU64TsWrapper()) {
+  }
   ~VersionStorageInfoTimestampTest() override {}
   std::string Timestamp(uint64_t ts) const {
     std::string ret;
@@ -1967,7 +1968,7 @@ class VersionSetWithTimestampTest : public VersionSetTest {
   void SetUp() override {
     NewDB();
     Options options;
-    options.comparator = test::ComparatorWithU64Ts();
+    options.comparator = test::BytewiseComparatorWithU64TsWrapper();
     cfd_ = CreateColumnFamily(kNewCfName, options);
     EXPECT_NE(nullptr, cfd_);
     EXPECT_NE(nullptr, cfd_->GetLatestMutableCFOptions());

db/write_batch_test.cc

@@ -950,7 +950,8 @@ Status CheckTimestampsInWriteBatch(
 }  // namespace
 
 TEST_F(WriteBatchTest, SanityChecks) {
-  ColumnFamilyHandleImplDummy cf0(0, test::ComparatorWithU64Ts());
+  ColumnFamilyHandleImplDummy cf0(0,
+                                  test::BytewiseComparatorWithU64TsWrapper());
   ColumnFamilyHandleImplDummy cf4(4);
 
   WriteBatch wb(0, 0, 0, /*default_cf_ts_sz=*/sizeof(uint64_t));
@@ -998,8 +999,10 @@ TEST_F(WriteBatchTest, UpdateTimestamps) {
   std::vector<std::string> key_strs(num_of_keys, std::string(key_size, '\0'));
 
   ColumnFamilyHandleImplDummy cf0(0);
-  ColumnFamilyHandleImplDummy cf4(4, test::ComparatorWithU64Ts());
-  ColumnFamilyHandleImplDummy cf5(5, test::ComparatorWithU64Ts());
+  ColumnFamilyHandleImplDummy cf4(4,
+                                  test::BytewiseComparatorWithU64TsWrapper());
+  ColumnFamilyHandleImplDummy cf5(5,
+                                  test::BytewiseComparatorWithU64TsWrapper());
 
   const std::unordered_map<uint32_t, const Comparator*> cf_to_ucmps = {
       {0, cf0.GetComparator()},

db_stress_tool/db_stress_test_base.cc

@@ -19,6 +19,7 @@
 #include "rocksdb/sst_file_manager.h"
 #include "rocksdb/types.h"
 #include "rocksdb/utilities/object_registry.h"
+#include "test_util/testutil.h"
 #include "util/cast_util.h"
 #include "utilities/backupable/backupable_db_impl.h"
 #include "utilities/fault_injection_fs.h"
@@ -2854,7 +2855,7 @@ void StressTest::Reopen(ThreadState* thread) {
 
 void StressTest::CheckAndSetOptionsForUserTimestamp() {
   assert(FLAGS_user_timestamp_size > 0);
-  const Comparator* const cmp = test::ComparatorWithU64Ts();
+  const Comparator* const cmp = test::BytewiseComparatorWithU64TsWrapper();
   assert(cmp);
   if (FLAGS_user_timestamp_size != cmp->timestamp_size()) {
     fprintf(stderr,

table/sst_file_dumper.cc

@@ -125,7 +125,7 @@ Status SstFileDumper::GetTableReader(const std::string& file_path) {
                               nullptr);
       file_.reset(new RandomAccessFileReader(std::move(file), file_path));
     }
-    options_.comparator = &internal_comparator_;
+
     // For old sst format, ReadTableProperties might fail but file can be read
     if (ReadTableProperties(magic_number, file_.get(), file_size,
                             (magic_number == kBlockBasedTableMagicNumber)
@@ -133,9 +133,24 @@ Status SstFileDumper::GetTableReader(const std::string& file_path) {
                                 : nullptr)
             .ok()) {
       s = SetTableOptionsByMagicNumber(magic_number);
+      if (s.ok()) {
+        if (table_properties_ && !table_properties_->comparator_name.empty()) {
+          ConfigOptions config_options;
+          const Comparator* user_comparator = nullptr;
+          s = Comparator::CreateFromString(config_options,
+                                           table_properties_->comparator_name,
+                                           &user_comparator);
+          if (s.ok()) {
+            assert(user_comparator);
+            internal_comparator_ =
+                InternalKeyComparator(user_comparator, /*named=*/true);
+          }
+        }
+      }
     } else {
       s = SetOldTableOptions();
     }
+    options_.comparator = internal_comparator_.user_comparator();
   }
 
   if (s.ok()) {

table/sst_file_reader_test.cc

@@ -215,7 +215,7 @@ class SstFileReaderTimestampTest : public testing::Test {
 
     options_.env = env;
 
-    options_.comparator = test::ComparatorWithU64Ts();
+    options_.comparator = test::BytewiseComparatorWithU64TsWrapper();
 
     sst_name_ = test::PerThreadDBPath("sst_file_ts");
   }

test_util/testutil.cc

@@ -118,59 +118,6 @@ class Uint64ComparatorImpl : public Comparator {
 
   void FindShortSuccessor(std::string* /*key*/) const override { return; }
 };
-
-// A test implementation of comparator with 64-bit integer timestamp.
-class ComparatorWithU64TsImpl : public Comparator {
- public:
-  ComparatorWithU64TsImpl()
-      : Comparator(/*ts_sz=*/sizeof(uint64_t)),
-        cmp_without_ts_(BytewiseComparator()) {
-    assert(cmp_without_ts_);
-    assert(cmp_without_ts_->timestamp_size() == 0);
-  }
-  const char* Name() const override { return "ComparatorWithU64Ts"; }
-  void FindShortSuccessor(std::string*) const override {}
-  void FindShortestSeparator(std::string*, const Slice&) const override {}
-  int Compare(const Slice& a, const Slice& b) const override {
-    int ret = CompareWithoutTimestamp(a, b);
-    size_t ts_sz = timestamp_size();
-    if (ret != 0) {
-      return ret;
-    }
-    // Compare timestamp.
-    // For the same user key with different timestamps, larger (newer) timestamp
-    // comes first.
-    return -CompareTimestamp(ExtractTimestampFromUserKey(a, ts_sz),
-                             ExtractTimestampFromUserKey(b, ts_sz));
-  }
-  using Comparator::CompareWithoutTimestamp;
-  int CompareWithoutTimestamp(const Slice& a, bool a_has_ts, const Slice& b,
-                              bool b_has_ts) const override {
-    const size_t ts_sz = timestamp_size();
-    assert(!a_has_ts || a.size() >= ts_sz);
-    assert(!b_has_ts || b.size() >= ts_sz);
-    Slice lhs = a_has_ts ? StripTimestampFromUserKey(a, ts_sz) : a;
-    Slice rhs = b_has_ts ? StripTimestampFromUserKey(b, ts_sz) : b;
-    return cmp_without_ts_->Compare(lhs, rhs);
-  }
-  int CompareTimestamp(const Slice& ts1, const Slice& ts2) const override {
-    assert(ts1.size() == sizeof(uint64_t));
-    assert(ts2.size() == sizeof(uint64_t));
-    uint64_t lhs = DecodeFixed64(ts1.data());
-    uint64_t rhs = DecodeFixed64(ts2.data());
-    if (lhs < rhs) {
-      return -1;
-    } else if (lhs > rhs) {
-      return 1;
-    } else {
-      return 0;
-    }
-  }
-
- private:
-  const Comparator* cmp_without_ts_{nullptr};
-};
-
 }  // namespace
 
 const Comparator* Uint64Comparator() {
@@ -178,9 +125,13 @@ const Comparator* Uint64Comparator() {
   return &uint64comp;
 }
 
-const Comparator* ComparatorWithU64Ts() {
-  static ComparatorWithU64TsImpl comp_with_u64_ts;
-  return &comp_with_u64_ts;
+const Comparator* BytewiseComparatorWithU64TsWrapper() {
+  ConfigOptions config_options;
+  const Comparator* user_comparator = nullptr;
+  Status s = Comparator::CreateFromString(
+      config_options, "leveldb.BytewiseComparator.u64ts", &user_comparator);
+  s.PermitUncheckedError();
+  return user_comparator;
 }
 
 void CorruptKeyType(InternalKey* ikey) {

test_util/testutil.h

@@ -113,6 +113,9 @@ class SimpleSuffixReverseComparator : public Comparator {
 // endian machines.
 extern const Comparator* Uint64Comparator();
 
+// A wrapper api for getting the ComparatorWithU64Ts<BytewiseComparator>
+extern const Comparator* BytewiseComparatorWithU64TsWrapper();
+
 class StringSink : public FSWritableFile {
  public:
   std::string contents_;
@@ -794,8 +797,6 @@ class ChanglingCompactionFilterFactory : public CompactionFilterFactory {
 // number of entries exceeds a threshold.
 extern MemTableRepFactory* NewSpecialSkipListFactory(int num_entries_per_flush);
 
-extern const Comparator* ComparatorWithU64Ts();
-
 CompressionType RandomCompressionType(Random* rnd);
 
 void RandomCompressionTypeVector(const size_t count,

tools/db_bench_tool.cc

@@ -4155,7 +4155,7 @@ class Benchmark {
         fprintf(stderr, "Only 64 bits timestamps are supported.\n");
         exit(1);
       }
-      options.comparator = ROCKSDB_NAMESPACE::test::ComparatorWithU64Ts();
+      options.comparator = test::BytewiseComparatorWithU64TsWrapper();
     }
 
     // Integrated BlobDB

tools/sst_dump_test.cc

@@ -6,7 +6,6 @@
 // Copyright (c) 2012 The LevelDB Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file. See the AUTHORS file for names of contributors.
-
 #ifndef ROCKSDB_LITE
 
 #include <stdint.h>
@@ -33,6 +32,12 @@ static std::string MakeKey(int i) {
   return key.Encode().ToString();
 }
 
+static std::string MakeKeyWithTimeStamp(int i, uint64_t ts) {
+  char buf[100];
+  snprintf(buf, sizeof(buf), "k_%04d", i);
+  return test::KeyStr(ts, std::string(buf), /*seq=*/0, kTypeValue);
+}
+
 static std::string MakeValue(int i) {
   char buf[100];
   snprintf(buf, sizeof(buf), "v_%04d", i);
@@ -116,8 +121,21 @@ class SSTDumpToolTest : public testing::Test {
 
     // Populate slightly more than 1K keys
     uint32_t num_keys = kNumKey;
-    for (uint32_t i = 0; i < num_keys; i++) {
-      tb->Add(MakeKey(i), MakeValue(i));
+    const char* comparator_name = ikc.user_comparator()->Name();
+    if (strcmp(comparator_name, ReverseBytewiseComparator()->Name()) == 0) {
+      for (int32_t i = num_keys; i >= 0; i--) {
+        tb->Add(MakeKey(i), MakeValue(i));
+      }
+    } else if (strcmp(comparator_name,
+                      test::BytewiseComparatorWithU64TsWrapper()->Name()) ==
+               0) {
+      for (uint32_t i = 0; i < num_keys; i++) {
+        tb->Add(MakeKeyWithTimeStamp(i, 100 + i), MakeValue(i));
+      }
+    } else {
+      for (uint32_t i = 0; i < num_keys; i++) {
+        tb->Add(MakeKey(i), MakeValue(i));
+      }
     }
     ASSERT_OK(tb->Finish());
     ASSERT_OK(file_writer->Close());
@@ -161,6 +179,54 @@ TEST_F(SSTDumpToolTest, EmptyFilter) {
   }
 }
 
+TEST_F(SSTDumpToolTest, SstDumpReverseBytewiseComparator) {
+  Options opts;
+  opts.env = env();
+  opts.comparator = ReverseBytewiseComparator();
+  BlockBasedTableOptions table_opts;
+  table_opts.filter_policy.reset(
+      ROCKSDB_NAMESPACE::NewBloomFilterPolicy(10, false));
+  opts.table_factory.reset(new BlockBasedTableFactory(table_opts));
+  std::string file_path =
+      MakeFilePath("rocksdb_sst_reverse_bytewise_comparator.sst");
+  createSST(opts, file_path);
+
+  char* usage[3];
+  PopulateCommandArgs(file_path, "--command=raw", usage);
+
+  ROCKSDB_NAMESPACE::SSTDumpTool tool;
+  ASSERT_TRUE(!tool.Run(3, usage, opts));
+
+  cleanup(opts, file_path);
+  for (int i = 0; i < 3; i++) {
+    delete[] usage[i];
+  }
+}
+
+TEST_F(SSTDumpToolTest, SstDumpComparatorWithU64Ts) {
+  Options opts;
+  opts.env = env();
+  opts.comparator = test::BytewiseComparatorWithU64TsWrapper();
+  BlockBasedTableOptions table_opts;
+  table_opts.filter_policy.reset(
+      ROCKSDB_NAMESPACE::NewBloomFilterPolicy(10, false));
+  opts.table_factory.reset(new BlockBasedTableFactory(table_opts));
+  std::string file_path =
+      MakeFilePath("rocksdb_sst_comparator_with_u64_ts.sst");
+  createSST(opts, file_path);
+
+  char* usage[3];
+  PopulateCommandArgs(file_path, "--command=raw", usage);
+
+  ROCKSDB_NAMESPACE::SSTDumpTool tool;
+  ASSERT_TRUE(!tool.Run(3, usage, opts));
+
+  cleanup(opts, file_path);
+  for (int i = 0; i < 3; i++) {
+    delete[] usage[i];
+  }
+}
+
 TEST_F(SSTDumpToolTest, FilterBlock) {
   Options opts;
   opts.env = env();

util/comparator.cc

@@ -14,7 +14,9 @@
 #include <algorithm>
 #include <memory>
 #include <mutex>
+#include <sstream>
 
+#include "db/dbformat.h"
 #include "port/port.h"
 #include "rocksdb/convenience.h"
 #include "rocksdb/slice.h"
@@ -216,6 +218,75 @@ class ReverseBytewiseComparatorImpl : public BytewiseComparatorImpl {
     return -a.compare(b);
   }
 };
+
+// EXPERIMENTAL
+// Comparator with 64-bit integer timestamp.
+// We did not performance test this yet.
+template <typename TComparator>
+class ComparatorWithU64TsImpl : public Comparator {
+  static_assert(std::is_base_of<Comparator, TComparator>::value,
+                "template type must be a inherited type of comparator");
+
+ public:
+  explicit ComparatorWithU64TsImpl() : Comparator(/*ts_sz=*/sizeof(uint64_t)) {
+    assert(cmp_without_ts_.timestamp_size() == 0);
+  }
+
+  static const char* kClassName() {
+    static std::string class_name = kClassNameInternal();
+    return class_name.c_str();
+  }
+
+  const char* Name() const override { return kClassName(); }
+
+  void FindShortSuccessor(std::string*) const override {}
+  void FindShortestSeparator(std::string*, const Slice&) const override {}
+  int Compare(const Slice& a, const Slice& b) const override {
+    int ret = CompareWithoutTimestamp(a, b);
+    size_t ts_sz = timestamp_size();
+    if (ret != 0) {
+      return ret;
+    }
+    // Compare timestamp.
+    // For the same user key with different timestamps, larger (newer) timestamp
+    // comes first.
+    return -CompareTimestamp(ExtractTimestampFromUserKey(a, ts_sz),
+                             ExtractTimestampFromUserKey(b, ts_sz));
+  }
+  using Comparator::CompareWithoutTimestamp;
+  int CompareWithoutTimestamp(const Slice& a, bool a_has_ts, const Slice& b,
+                              bool b_has_ts) const override {
+    const size_t ts_sz = timestamp_size();
+    assert(!a_has_ts || a.size() >= ts_sz);
+    assert(!b_has_ts || b.size() >= ts_sz);
+    Slice lhs = a_has_ts ? StripTimestampFromUserKey(a, ts_sz) : a;
+    Slice rhs = b_has_ts ? StripTimestampFromUserKey(b, ts_sz) : b;
+    return cmp_without_ts_.Compare(lhs, rhs);
+  }
+  int CompareTimestamp(const Slice& ts1, const Slice& ts2) const override {
+    assert(ts1.size() == sizeof(uint64_t));
+    assert(ts2.size() == sizeof(uint64_t));
+    uint64_t lhs = DecodeFixed64(ts1.data());
+    uint64_t rhs = DecodeFixed64(ts2.data());
+    if (lhs < rhs) {
+      return -1;
+    } else if (lhs > rhs) {
+      return 1;
+    } else {
+      return 0;
+    }
+  }
+
+ private:
+  static std::string kClassNameInternal() {
+    std::stringstream ss;
+    ss << TComparator::kClassName() << ".u64ts";
+    return ss.str();
+  }
+
+  TComparator cmp_without_ts_;
+};
+
 }// namespace
 
 const Comparator* BytewiseComparator() {
@@ -228,6 +299,11 @@ const Comparator* ReverseBytewiseComparator() {
   return &rbytewise;
 }
 
+const Comparator* BytewiseComparatorWithU64Ts() {
+  static ComparatorWithU64TsImpl<BytewiseComparatorImpl> comp_with_u64_ts;
+  return &comp_with_u64_ts;
+}
+
 #ifndef ROCKSDB_LITE
 static int RegisterBuiltinComparators(ObjectLibrary& library,
                                       const std::string& /*arg*/) {
@@ -241,7 +317,12 @@ static int RegisterBuiltinComparators(ObjectLibrary& library,
       [](const std::string& /*uri*/,
          std::unique_ptr<const Comparator>* /*guard */,
          std::string* /* errmsg */) { return ReverseBytewiseComparator(); });
-  return 2;
+  library.AddFactory<const Comparator>(
+      ComparatorWithU64TsImpl<BytewiseComparatorImpl>::kClassName(),
+      [](const std::string& /*uri*/,
+         std::unique_ptr<const Comparator>* /*guard */,
+         std::string* /* errmsg */) { return BytewiseComparatorWithU64Ts(); });
+  return 3;
 }
 #endif  // ROCKSDB_LITE
 
@@ -265,6 +346,9 @@ Status Comparator::CreateFromString(const ConfigOptions& config_options,
     *result = BytewiseComparator();
   } else if (id == ReverseBytewiseComparatorImpl::kClassName()) {
     *result = ReverseBytewiseComparator();
+  } else if (id ==
+             ComparatorWithU64TsImpl<BytewiseComparatorImpl>::kClassName()) {
+    *result = BytewiseComparatorWithU64Ts();
   } else if (value.empty()) {
     // No Id and no options.  Clear the object
     *result = nullptr;