Allow passing comparator in UDI (facebook#14001)

Summary:
Pass the comparator to UDI interface for both reader and builder.

Pull Request resolved: facebook#14001

Test Plan: Unit test

Reviewed By: anand1976

Differential Revision: D83339943

Pulled By: xingbowang

fbshipit-source-id: 7f6541776b0995260e28224329f0cca37f13b3d4

Author: Xingbo Wang

include/rocksdb/user_defined_index.h

@@ -30,8 +30,7 @@ inline const std::string kUserDefinedIndexPrefix =
 //
 // 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.
+// SstFileWriter are supported.
 
 // The interface for building user-defined index.
 class UserDefinedIndexBuilder {
@@ -145,6 +144,11 @@ class UserDefinedIndexReader {
   virtual size_t ApproximateMemoryUsage() const = 0;
 };
 
+// Options for user defined index
+struct UserDefinedIndexOption {
+  const Comparator* comparator = BytewiseComparator();
+};
+
 // Factory for creating user-defined index builders.
 class UserDefinedIndexFactory : public Customizable {
  public:
@@ -163,6 +167,21 @@ class UserDefinedIndexFactory : public Customizable {
   // block
   virtual std::unique_ptr<UserDefinedIndexReader> NewReader(
       Slice& index_block) const = 0;
+
+  // New API for allowing customized comparator
+  virtual Status NewBuilder(
+      const UserDefinedIndexOption& /*option*/,
+      std::unique_ptr<UserDefinedIndexBuilder>& builder) const {
+    builder.reset(NewBuilder());
+    return Status::OK();
+  };
+
+  virtual Status NewReader(
+      const UserDefinedIndexOption& /*option*/, Slice& index_block,
+      std::unique_ptr<UserDefinedIndexReader>& reader) const {
+    reader = NewReader(index_block);
+    return Status::OK();
+  };
 };
 
 }  // namespace ROCKSDB_NAMESPACE

table/block_based/block_based_table_builder.cc

@@ -1191,13 +1191,20 @@ struct BlockBasedTableBuilder::Rep {
             Status::InvalidArgument("user_defined_index_factory not supported "
                                     "with parallel compression"));
       } else {
-        std::unique_ptr<UserDefinedIndexBuilder> user_defined_index_builder(
-            table_options.user_defined_index_factory->NewBuilder());
-        if (user_defined_index_builder != nullptr) {
-          index_builder = std::make_unique<UserDefinedIndexBuilderWrapper>(
-              std::string(table_options.user_defined_index_factory->Name()),
-              std::move(index_builder), std::move(user_defined_index_builder),
-              &internal_comparator, ts_sz, persist_user_defined_timestamps);
+        std::unique_ptr<UserDefinedIndexBuilder> user_defined_index_builder;
+        UserDefinedIndexOption udi_options;
+        udi_options.comparator = internal_comparator.user_comparator();
+        auto s = table_options.user_defined_index_factory->NewBuilder(
+            udi_options, user_defined_index_builder);
+        if (!s.ok()) {
+          SetStatus(s);
+        } else {
+          if (user_defined_index_builder != nullptr) {
+            index_builder = std::make_unique<UserDefinedIndexBuilderWrapper>(
+                std::string(table_options.user_defined_index_factory->Name()),
+                std::move(index_builder), std::move(user_defined_index_builder),
+                &internal_comparator, ts_sz, persist_user_defined_timestamps);
+          }
         }
       }
     }

table/block_based/block_based_table_reader.cc

@@ -1369,15 +1369,20 @@ Status BlockBasedTable::PrefetchIndexAndFilterBlocks(
       if (s.ok()) {
         assert(!rep_->udi_block.IsEmpty());
 
-        std::unique_ptr<UserDefinedIndexReader> udi_reader =
-            table_options.user_defined_index_factory->NewReader(
-                rep_->udi_block.GetValue()->data);
-        if (udi_reader) {
-          index_reader = std::make_unique<UserDefinedIndexReaderWrapper>(
-              udi_name, std::move(index_reader), std::move(udi_reader));
-        } else {
-          s = Status::Corruption("Failed to create UDI reader for " + udi_name +
-                                 " in file " + rep_->file->file_name());
+        std::unique_ptr<UserDefinedIndexReader> udi_reader;
+        UserDefinedIndexOption udi_option;
+        udi_option.comparator = rep_->internal_comparator.user_comparator();
+        s = table_options.user_defined_index_factory->NewReader(
+            udi_option, rep_->udi_block.GetValue()->data, udi_reader);
+        if (s.ok()) {
+          if (udi_reader) {
+            index_reader = std::make_unique<UserDefinedIndexReaderWrapper>(
+                udi_name, std::move(index_reader), std::move(udi_reader));
+          } else {
+            s = Status::Corruption("Failed to create UDI reader for " +
+                                   udi_name + " in file " +
+                                   rep_->file->file_name());
+          }
         }
       }
     }