Add support for data block hash index for column families with user-defined timestamps

db/db_with_timestamp_basic_test.cc

@@ -4795,6 +4795,74 @@ TEST_F(DBBasicTestWithTimestamp, TimestampFilterTableReadOnGet) {
   Close();
 }
 
+class DBBasicTestWithTimestampDataBlockHashIndex
+    : public DBBasicTestWithTimestampBase {
+ public:
+  DBBasicTestWithTimestampDataBlockHashIndex()
+      : DBBasicTestWithTimestampBase("data_block_hash_index_timestamp_test") {}
+ protected:
+  class TestComparator : public DBBasicTestWithTimestampBase::TestComparator {
+   public:
+    explicit TestComparator(size_t ts_sz)
+      : DBBasicTestWithTimestampBase::TestComparator(ts_sz) {}
+
+    // Override this method to allow data block hash index to be used
+    bool KeysAreBytewiseComparableOtherThanTimestamp() const override {
+      return true;
+    }
+  };
+};
+
+TEST_F(DBBasicTestWithTimestampDataBlockHashIndex, HashIndexWithTimestamp) {
+
+  Options options = CurrentOptions();
+
+  // Use a bytewise comparable, timestamp-aware comparator that overrides
+  // KeysAreBytewiseComparableOtherThanTimestamp() to return true,
+  // enabling use of data block hash index
+  const size_t kTimestampSize = Timestamp(0, 0).size();
+  DBBasicTestWithTimestampDataBlockHashIndex::TestComparator
+      comparator(kTimestampSize);
+  options.comparator = &comparator;
+
+  // Select data block hash index as this CF's data block index type
+  BlockBasedTableOptions table_options;
+  table_options.data_block_index_type =
+      BlockBasedTableOptions::kDataBlockBinaryAndHash;
+  options.table_factory.reset(NewBlockBasedTableFactory(table_options));
+
+
+  // Run sequence of $count Put()s and one Get() multiple times
+  // Original incompatibility between UDTs and Data Block Hash Index
+  // led to inconsistent behavior that would _usually_ cause this test to fail,
+  // so multiple iterations significantly reduces misleading passes.
+  // This test always passes when UDT and Data Block Hash Index are compatible.
+  for (uint32_t iteration = 0; iteration < 5; iteration++) {
+    DestroyAndReopen(options);
+
+    WriteOptions wopts;
+
+    const uint32_t count = 124;
+    for (uint32_t i = 0; i < count; i++) {
+      ASSERT_OK(db_->Put(wopts, "key", Timestamp(i, 0), std::to_string(i)));
+    }
+
+    ASSERT_OK(Flush());
+
+    ReadOptions ropts;
+    Slice ts = Timestamp(count, 0);
+    ropts.timestamp = &ts;
+
+    std::string value;
+    std::string expected_value = std::to_string(count - 1);
+
+    ASSERT_OK(db_->Get(ropts, "key", &value));
+
+    // When UDTs and hash index are incompatible, these values can mismatch
+    ASSERT_EQ(expected_value, value);
+  }
+}
+
 }  // namespace ROCKSDB_NAMESPACE
 
 int main(int argc, char** argv) {

include/rocksdb/comparator.h

@@ -114,6 +114,16 @@ class Comparator : public Customizable, public CompareInterface {
   // with the customized comparator.
   virtual bool CanKeysWithDifferentByteContentsBeEqual() const { return true; }
 
+  // return true if the user portion of the key is bytwise comparable other than
+  // any user defined timestamp.
+  // i.e. if the output of ExtractUserKeyAndStripTimestamp() on two keys results
+  // in distinct byte sequences, then this comparator considers them unequal.
+  // This is used to determine if DataBlockHashIndex is compatible
+  // with customized comparators that support user-defined timestamps.
+  virtual bool KeysAreBytewiseComparableOtherThanTimestamp() const {
+    return false;
+  }
+
   // if it is a wrapped comparator, may return the root one.
   // return itself it is not wrapped.
   virtual const Comparator* GetRootComparator() const { return this; }

table/block_based/block.cc

@@ -357,7 +357,7 @@ void MetaBlockIter::SeekImpl(const Slice& target) {
 //    with a smaller [ type | seqno ] (i.e. a larger seqno, or the same seqno
 //    but larger type).
 bool DataBlockIter::SeekForGetImpl(const Slice& target) {
-  Slice target_user_key = ExtractUserKey(target);
+  Slice target_user_key = ExtractUserKeyAndStripTimestamp(target, ts_sz_);
   uint32_t map_offset = restarts_ + num_restarts_ * sizeof(uint32_t);
   uint8_t entry =
       data_block_hash_index_->Lookup(data_, map_offset, target_user_key);

table/block_based/block_based_table_builder.cc

@@ -464,8 +464,10 @@ struct BlockBasedTableBuilder::Rep {
         data_block(table_options.block_restart_interval,
                    table_options.use_delta_encoding,
                    false /* use_value_delta_encoding */,
-                   tbo.internal_comparator.user_comparator()
-                           ->CanKeysWithDifferentByteContentsBeEqual()
+                   (tbo.internal_comparator.user_comparator()
+                    ->CanKeysWithDifferentByteContentsBeEqual() &&
+                   !tbo.internal_comparator.user_comparator()
+                   ->KeysAreBytewiseComparableOtherThanTimestamp())
                        ? BlockBasedTableOptions::kDataBlockBinarySearch
                        : table_options.data_block_index_type,
                    table_options.data_block_hash_table_util_ratio, ts_sz,

table/block_based/block_builder.cc

@@ -53,6 +53,7 @@ BlockBuilder::BlockBuilder(
       use_delta_encoding_(use_delta_encoding),
       use_value_delta_encoding_(use_value_delta_encoding),
       strip_ts_sz_(persist_user_defined_timestamps ? 0 : ts_sz),
+      ts_sz_(ts_sz),
       is_user_key_(is_user_key),
       restarts_(1, 0),  // First restart point is at offset 0
       counter_(0),
@@ -236,13 +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),
+    data_block_hash_index_builder_.Add(ExtractUserKeyAndStripTimestamp(key, ts_sz_),
                                        restarts_.size() - 1);
   }
 

table/block_based/block_builder.h

@@ -98,6 +98,7 @@ class BlockBuilder {
   // 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_;
+  const size_t 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

unreleased_history/public_api_changes/timestamp_compatible_data_block_hash_index.md

@@ -0,0 +1,4 @@
+Add `Comparator::KeysAreBytewiseComparableOtherThanTimestamp` to allow column
+families that use a timestamp-aware comparator with otherwise
+bytewise-comparable keys to enable
+`BlockBasedTableOptions::kDataBlockBinaryAndHash`