Make UDI interface consistently use the user key (#13865)

Summary:
The original intention of the User Defined Index interface was to use the user key. However, the implementation mixed user and internal key usage. This PR makes it consistent. It also clarifies the UDI contract.

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

Test Plan: Update tests in table_test.cc

Reviewed By: pdillinger

Differential Revision: D80050344

Pulled By: anand1976

fbshipit-source-id: ace47737d21684ec19709640a09e198cee2d98bd

include/rocksdb/user_defined_index.h

@@ -27,6 +27,11 @@ inline const std::string kUserDefinedIndexPrefix =
 // It allows users to define their own index format and build custom
 // indexes during table building. Currently, only a monolithic index
 // block is supported (no partitioned index).
+//
+// This is currently supported only for a restricted set of use cases. The
+// CF must be ingest only, and only files containing Puts generated by
+// SstFileWriter are supported. The user_comparator used for the CF must
+// be BytewiseComparator.
 
 // The interface for building user-defined index.
 class UserDefinedIndexBuilder {
@@ -51,6 +56,10 @@ class UserDefinedIndexBuilder {
   // The previous index entry key and the new index entry key cover
   // all the keys in the data block associated with the new index entry.
   //
+  // The last_key_in_current_block and first_key_in_next_block will be user
+  // keys, i.e the user key string, and optionally the user timestamp if one
+  // is configured, without a sequence number suffix.
+  //
   // Called before the OnKeyAdded() call for first_key_in_next_block.
   // @last_key_in_current_block: The last key in the current data block
   // @first_key_in_next_block: it will be nullptr if the entry being added is
@@ -72,6 +81,9 @@ class UserDefinedIndexBuilder {
   // override OnKeyAdded() if they need to collect additional information.
   // The type argument indicates whether the value is a full value or partial.
   // At the moment, only full values are supported.
+  //
+  // The key will be a user key. RocksDB guarantees that there will only be
+  // one entry for each key in the file/index.
   virtual void OnKeyAdded(const Slice& /*key*/, ValueType /*type*/,
                           const Slice& /*value*/) {}
 
@@ -100,6 +112,14 @@ class UserDefinedIndexIterator {
   // termination criteria, kInbound if the data block is definitely fully
   // within bounds, or kUnknown if the data block could be partially
   // within bounds.
+  // The UDI implementation needs to be careful about returning kOutOfBound.
+  // If a limit key is specified in ScanOptions, an implementation that
+  // does not store the first key in the block for the corresponding index
+  // entry cannot reliably determine if the block is out of bounds. It must
+  // compare against the previous index key to determine if the current block
+  // is out of bounds w.r.t the limit. Other termination criteria (specified
+  // in property_bag) may cause the scan to terminate earlier, in which case
+  // kOutOfBound can be returned earlier.
   virtual Status SeekAndGetResult(const Slice& target,
                                   IterateResult* result) = 0;
 

table/block_based/block_based_table_builder.cc

@@ -911,6 +911,12 @@ struct BlockBasedTableBuilder::Rep {
         SetStatus(
             Status::InvalidArgument("user_defined_index_factory not supported "
                                     "with parallel compression"));
+      } else if (ioptions.user_comparator != BytewiseComparator()) {
+        // TODO: Pass the user_comparator to the UDI and let it validate. Do
+        // it in a major release.
+        SetStatus(
+            Status::InvalidArgument("user_defined_index_factory only supported "
+                                    "with bytewise comparator"));
       } else {
         std::unique_ptr<UserDefinedIndexBuilder> user_defined_index_builder(
             table_options.user_defined_index_factory->NewBuilder());

table/block_based/user_defined_index_wrapper.h

@@ -46,9 +46,24 @@ class UserDefinedIndexBuilderWrapper : public IndexBuilder {
     handle.offset = block_handle.offset();
     handle.size = block_handle.size();
     // Forward the call to both index builders
-    user_defined_index_builder_->AddIndexEntry(last_key_in_current_block,
-                                               first_key_in_next_block, handle,
-                                               separator_scratch);
+    ParsedInternalKey pkey_last;
+    ParsedInternalKey pkey_first;
+    // There's no way to return an error here, so we remember the statsu and
+    // return it in Finish()
+    if (status_.ok()) {
+      status_ = ParseInternalKey(last_key_in_current_block, &pkey_last,
+                                 /*lof_err_key*/ false);
+    }
+    if (status_.ok() && first_key_in_next_block) {
+      status_ = ParseInternalKey(*first_key_in_next_block, &pkey_first,
+                                 /*lof_err_key*/ false);
+    }
+    if (status_.ok()) {
+      user_defined_index_builder_->AddIndexEntry(
+          pkey_last.user_key,
+          first_key_in_next_block ? &pkey_first.user_key : nullptr, handle,
+          separator_scratch);
+    }
     return internal_index_builder_->AddIndexEntry(
         last_key_in_current_block, first_key_in_next_block, block_handle,
         separator_scratch);
@@ -76,6 +91,12 @@ class UserDefinedIndexBuilderWrapper : public IndexBuilder {
 
     // Forward the call to both index builders
     internal_index_builder_->OnKeyAdded(key, value);
+
+    // Pass the user key to the UDI. We don't expect multiple entries with
+    // different sequence numbers for the same key in the file. RocksDB may
+    // enforce it in the future by allowing UDIs only for read only
+    // bulkloaded use cases, and only allow ingestion of files with
+    // sequence number 0.
     user_defined_index_builder_->OnKeyAdded(
         pkey.user_key, UserDefinedIndexBuilder::ValueType::kValue,
         value.value());
@@ -149,23 +170,41 @@ class UserDefinedIndexIteratorWrapper
     status_ = ParseInternalKey(target, &pkey, /*log_err_key=*/false);
     if (status_.ok()) {
       status_ = udi_iter_->SeekAndGetResult(pkey.user_key, &result_);
-      valid_ = status_.ok() &&
-               result_.bound_check_result == IterBoundCheck::kInbound;
+      if (status_.ok()) {
+        valid_ = result_.bound_check_result == IterBoundCheck::kInbound;
+        if (valid_) {
+          ikey_.Set(result_.key, 0, ValueType::kTypeValue);
+        }
+      }
+    } else {
+      valid_ = false;
     }
   }
 
   void Next() override {
     status_ = udi_iter_->NextAndGetResult(&result_);
-    valid_ =
-        status_.ok() && result_.bound_check_result == IterBoundCheck::kInbound;
+    if (status_.ok()) {
+      valid_ = result_.bound_check_result == IterBoundCheck::kInbound;
+      if (valid_) {
+        ikey_.Set(result_.key, 0, ValueType::kTypeValue);
+      }
+    } else {
+      valid_ = false;
+    }
   }
 
   bool NextAndGetResult(IterateResult* result) override {
     status_ = udi_iter_->NextAndGetResult(&result_);
-    valid_ =
-        status_.ok() && result_.bound_check_result == IterBoundCheck::kInbound;
     if (status_.ok()) {
-      *result = result_;
+      valid_ = result_.bound_check_result == IterBoundCheck::kInbound;
+      if (valid_) {
+        ikey_.Set(result_.key, 0, ValueType::kTypeValue);
+      }
+      if (status_.ok()) {
+        *result = result_;
+      }
+    } else {
+      valid_ = false;
     }
     return valid_;
   }
@@ -176,7 +215,7 @@ class UserDefinedIndexIteratorWrapper
 
   void Prev() override { status_ = Status::NotSupported("Prev not supported"); }
 
-  Slice key() const override { return result_.key; }
+  Slice key() const override { return Slice(*ikey_.const_rep()); }
 
   IndexValue value() const override {
     auto handle = udi_iter_->value();
@@ -196,6 +235,7 @@ class UserDefinedIndexIteratorWrapper
  private:
   std::unique_ptr<UserDefinedIndexIterator> udi_iter_;
   IterateResult result_;
+  InternalKey ikey_;
   Status status_;
   bool valid_;
 };

table/table_test.cc

@@ -7479,8 +7479,11 @@ class UserDefinedIndexTest : public BlockBasedTableTestBase {
                           const Slice* first_key_in_next_block,
                           const BlockHandle& block_handle,
                           std::string* separator_scratch) override {
+        EXPECT_EQ(last_key_in_current_block.size(), 5);
+        if (first_key_in_next_block) {
+          EXPECT_EQ(first_key_in_next_block->size(), 5);
+        }
         // Unused parameters
-        (void)first_key_in_next_block;
         (void)separator_scratch;
         entries_added_++;
         // Store the block handle for each key
@@ -7494,8 +7497,9 @@ class UserDefinedIndexTest : public BlockBasedTableTestBase {
         return last_key_in_current_block;
       }
 
-      void OnKeyAdded(const Slice& /*key*/, ValueType /*value*/,
+      void OnKeyAdded(const Slice& key, ValueType /*value*/,
                       const Slice& /*value*/) override {
+        EXPECT_EQ(key.size(), 5);
         // Track keys added to the index
         keys_added_++;
       }

unreleased_history/bug_fixes/udi_index_key_format.md

@@ -0,0 +1 @@
+Make the User Defined Index interface consistently use the user key format, fixing the previous mixed usage of internal and user key.