User defined index reader and iterator (facebook#13727)

Summary:
Pull Request resolved: facebook#13727

Add UserDefinedIndexReader and UserDefinedIndexIterator. The BlockBasedTable reads the user defined index meta block during open, verifies the checksum, pins in cache or heap depending on configuration, and allocates a UserDefinedIndexReader object with the contents. Similar to the builder, an IndexReader wrapper is allocated. The wrapper forwards the calls to the native reader and/or user defined index reader as appropriate.

A new option, table_index_name, in ReadOptions specifies the index to use when creating a new Iterator.

Reviewed By: pdillinger

Differential Revision: D76165694

fbshipit-source-id: c30bde4c5ce91ea3dc9ad302e73fe4963c1ed457

Author: anand1976

include/rocksdb/options.h

@@ -57,6 +57,7 @@ class Statistics;
 class InternalKeyComparator;
 class WalFilter;
 class FileSystem;
+class UserDefinedIndexFactory;
 
 struct Options;
 struct DbPath;
@@ -2069,6 +2070,17 @@ struct ReadOptions {
   // Default: false
   bool auto_refresh_iterator_with_snapshot = false;
 
+  // EXPERIMENTAL
+  //
+  // Specify an alternate index to use in the SST files instead of the native
+  // block based table index. The table_factory used for the column family
+  // must support building/reading this index.
+  //
+  // Currently, only forward scans are supported. For forward scans, only Seek()
+  // is supported. SeekToFirst() is not supported. If the caller wishes to scan
+  // from start to end, the native index must be used.
+  const UserDefinedIndexFactory* table_index_factory = nullptr;
+
   // *** END options only relevant to iterators or scans ***
 
   // *** BEGIN options for RocksDB internal use only ***

include/rocksdb/user_defined_index.h

@@ -11,7 +11,9 @@
 
 #include <string>
 
+#include "rocksdb/advanced_iterator.h"
 #include "rocksdb/customizable.h"
+#include "rocksdb/options.h"
 #include "rocksdb/slice.h"
 #include "rocksdb/status.h"
 
@@ -23,7 +25,8 @@ inline const std::string kUserDefinedIndexPrefix =
 
 // This is a public API for user-defined index builders.
 // It allows users to define their own index format and build custom
-// indexes during table building.
+// indexes during table building. Currently, only a monolithic index
+// block is supported (no partitioned index).
 
 // The interface for building user-defined index.
 class UserDefinedIndexBuilder {
@@ -52,6 +55,9 @@ class UserDefinedIndexBuilder {
   // @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
   //                           the last one in the table
+  // @block_handle: offset/size of the data block referenced by this index
+  //                entry. This should be stored along with the index entry
+  //                key
   // @separator_scratch: a scratch buffer to back a computed separator between
   //                     those, as needed. May be modified on each call.
   // @return: the key or separator stored in the index, which could be
@@ -76,13 +82,57 @@ class UserDefinedIndexBuilder {
   virtual Status Finish(Slice* index_contents) = 0;
 };
 
+// The interface for iterating the user defined index. This will be
+// instantiated and used by a scan to iterate through the index entries
+// covered by the scan.
+class UserDefinedIndexIterator {
+ public:
+  virtual ~UserDefinedIndexIterator() = default;
+
+  // Given the target key, position the index iterator at the index entry
+  // with the smallest key >= target. The result must be updated with the
+  // index key, and the bound_check_result. The bound_check_result should
+  // be set to kOutOfBound if no block satisfies the target key and
+  // termination criteria, kInbound if the data block is definitely fully
+  // within bounds, or kUnknown if the data block could be partially
+  // within bounds.
+  virtual Status SeekAndGetResult(const Slice& target,
+                                  IterateResult* result) = 0;
+
+  // Advance to the next index entry. The result must be populated similar
+  // to SeekAndGetResult.
+  virtual Status NextAndGetResult(IterateResult* result) = 0;
+
+  // Return the BlockHandle in the current index entry
+  virtual UserDefinedIndexBuilder::BlockHandle value() = 0;
+};
+
+// A reader interface for the user defined index
+class UserDefinedIndexReader {
+ public:
+  virtual ~UserDefinedIndexReader() = default;
+
+  // Allocate an iterator that will be used by RocksDB to perform scans
+  virtual std::unique_ptr<UserDefinedIndexIterator> NewIterator(
+      const ReadOptions& read_options) = 0;
+
+  // The memory usage of the index, including the size of the raw contents and
+  // any other heap data structures allocated by the reader
+  virtual size_t ApproximateMemoryUsage() const = 0;
+};
+
 // Factory for creating user-defined index builders.
 class UserDefinedIndexFactory : public Customizable {
  public:
   virtual ~UserDefinedIndexFactory() = default;
 
   // Create a new builder for user-defined index.
   virtual UserDefinedIndexBuilder* NewBuilder() const = 0;
+
+  // Create a new user defined index reader given the contents of the index
+  // block
+  virtual std::unique_ptr<UserDefinedIndexReader> NewReader(
+      Slice& index_block) const = 0;
 };
 
 }  // namespace ROCKSDB_NAMESPACE

table/block_based/block_based_table_reader.cc

@@ -59,6 +59,7 @@
 #include "table/block_based/hash_index_reader.h"
 #include "table/block_based/partitioned_filter_block.h"
 #include "table/block_based/partitioned_index_reader.h"
+#include "table/block_based/user_defined_index_wrapper.h"
 #include "table/block_fetcher.h"
 #include "table/format.h"
 #include "table/get_context.h"
@@ -114,6 +115,7 @@ INSTANTIATE_BLOCKLIKE_TEMPLATES(Block_kIndex);
 INSTANTIATE_BLOCKLIKE_TEMPLATES(Block_kFilterPartitionIndex);
 INSTANTIATE_BLOCKLIKE_TEMPLATES(Block_kRangeDeletion);
 INSTANTIATE_BLOCKLIKE_TEMPLATES(Block_kMetaIndex);
+INSTANTIATE_BLOCKLIKE_TEMPLATES(Block_kUserDefinedIndex);
 
 }  // namespace ROCKSDB_NAMESPACE
 
@@ -1319,6 +1321,34 @@ Status BlockBasedTable::PrefetchIndexAndFilterBlocks(
   if (!s.ok()) {
     return s;
   }
+  if (table_options.user_defined_index_factory != nullptr) {
+    std::string udi_name(table_options.user_defined_index_factory->Name());
+    BlockHandle udi_block_handle;
+
+    // Should we use FindOptionalMetaBlock here?
+    s = FindMetaBlock(meta_iter, kUserDefinedIndexPrefix + udi_name,
+                      &udi_block_handle);
+    if (!s.ok()) {
+      return s;
+    }
+    // Read the block, and allocate on heap or pin in cache. The UDI block is
+    // not compressed. RetrieveBlock will verify the checksum.
+    s = RetrieveBlock(prefetch_buffer, ro, udi_block_handle,
+                      rep_->decompressor.get(), &rep_->udi_block,
+                      /*get_context=*/nullptr, lookup_context,
+                      /*for_compaction=*/false, use_cache, /*async_read=*/false,
+                      /*use_block_cache_for_lookup=*/false);
+    if (!s.ok()) {
+      return s;
+    }
+    assert(!rep_->udi_block.IsEmpty());
+
+    std::unique_ptr<UserDefinedIndexReader> udi_reader =
+        table_options.user_defined_index_factory->NewReader(
+            rep_->udi_block.GetValue()->data);
+    index_reader = std::make_unique<UserDefinedIndexReaderWrapper>(
+        udi_name, std::move(index_reader), std::move(udi_reader));
+  }
 
   rep_->index_reader = std::move(index_reader);
 
@@ -1766,8 +1796,7 @@ BlockBasedTable::MaybeReadBlockAndLoadToCache(
         ro.fill_cache) {
       Statistics* statistics = rep_->ioptions.stats;
       const bool maybe_compressed =
-          TBlocklike::kBlockType != BlockType::kFilter &&
-          TBlocklike::kBlockType != BlockType::kCompressionDictionary &&
+          BlockTypeMaybeCompressed(TBlocklike::kBlockType) &&
           rep_->decompressor;
       // This flag, if true, tells BlockFetcher to return the uncompressed
       // block when ReadBlockContents() is called.
@@ -1911,6 +1940,7 @@ BlockBasedTable::SaveLookupContextOrTraceRecord(
       trace_block_type = TraceType::kBlockTraceRangeDeletionBlock;
       break;
     case BlockType::kIndex:
+    case BlockType::kUserDefinedIndex:
       trace_block_type = TraceType::kBlockTraceIndexBlock;
       break;
     default:
@@ -2003,9 +2033,7 @@ WithBlocklikeCheck<Status, TBlocklike> BlockBasedTable::RetrieveBlock(
   }
 
   const bool maybe_compressed =
-      TBlocklike::kBlockType != BlockType::kFilter &&
-      TBlocklike::kBlockType != BlockType::kCompressionDictionary &&
-      rep_->decompressor;
+      BlockTypeMaybeCompressed(TBlocklike::kBlockType) && rep_->decompressor;
   std::unique_ptr<TBlocklike> block;
 
   {

table/block_based/block_based_table_reader.h

@@ -228,11 +228,15 @@ class BlockBasedTable : public TableReader {
 
     // Create an iterator for index access. If iter is null, then a new object
     // is created on the heap, and the callee will have the ownership.
-    // If a non-null iter is passed in, it will be used, and the returned value
-    // is either the same as iter or a new on-heap object that
-    // wraps the passed iter. In the latter case the return value points
-    // to a different object then iter, and the callee has the ownership of the
-    // returned object.
+    // If a non-null iter is passed in, it may be used, and the returned value
+    // is either the same as iter or a new on-heap object.
+    // In the latter case the return value points to a different object then
+    // iter, and the callee has the ownership of the returned object.
+    //
+    // Under all circumstances, the caller MUST use the returned iterator
+    // for further operations. If the returned iterator != iter, then the
+    // caller MUST ensure that iter stays in scope until the returned
+    // iterator is destroyed.
     virtual InternalIteratorBase<IndexValue>* NewIterator(
         const ReadOptions& read_options, bool disable_prefix_seek,
         IndexBlockIter* iter, GetContext* get_context,
@@ -544,6 +548,12 @@ class BlockBasedTable : public TableReader {
 
   bool TimestampMayMatch(const ReadOptions& read_options) const;
 
+  bool BlockTypeMaybeCompressed(BlockType type) const {
+    return type != BlockType::kFilter &&
+           type != BlockType::kCompressionDictionary &&
+           type != BlockType::kUserDefinedIndex;
+  }
+
   // A cumulative data block file read in MultiGet lower than this size will
   // use a stack buffer
   static constexpr size_t kMultiGetReadStackBufSize = 8192;
@@ -689,6 +699,8 @@ struct BlockBasedTable::Rep {
   std::unique_ptr<CacheReservationManager::CacheReservationHandle>
       table_reader_cache_res_handle = nullptr;
 
+  CachableEntry<Block_kUserDefinedIndex> udi_block;
+
   SequenceNumber get_global_seqno(BlockType block_type) const {
     return (block_type == BlockType::kFilterPartitionIndex ||
             block_type == BlockType::kCompressionDictionary)

table/block_based/block_cache.cc

@@ -46,6 +46,12 @@ void BlockCreateContext::Create(std::unique_ptr<Block_kMetaIndex>* parsed_out,
       protection_bytes_per_key);
 }
 
+void BlockCreateContext::Create(
+    std::unique_ptr<Block_kUserDefinedIndex>* parsed_out,
+    BlockContents&& block) {
+  parsed_out->reset(new Block_kUserDefinedIndex(std::move(block)));
+}
+
 void BlockCreateContext::Create(
     std::unique_ptr<ParsedFullFilterBlock>* parsed_out, BlockContents&& block) {
   parsed_out->reset(new ParsedFullFilterBlock(

table/block_based/block_cache.h

@@ -67,6 +67,16 @@ class Block_kMetaIndex : public Block {
   static constexpr BlockType kBlockType = BlockType::kMetaIndex;
 };
 
+class Block_kUserDefinedIndex : public BlockContents {
+ public:
+  static constexpr CacheEntryRole kCacheEntryRole = CacheEntryRole::kIndexBlock;
+  static constexpr BlockType kBlockType = BlockType::kUserDefinedIndex;
+
+  explicit Block_kUserDefinedIndex(BlockContents&& other)
+      : BlockContents(std::move(other)) {}
+  const Slice& ContentSlice() const { return data; }
+};
+
 struct BlockCreateContext : public Cache::CreateContext {
   BlockCreateContext() {}
   BlockCreateContext(const BlockBasedTableOptions* _table_options,
@@ -126,6 +136,8 @@ struct BlockCreateContext : public Cache::CreateContext {
               BlockContents&& block);
   void Create(std::unique_ptr<Block_kMetaIndex>* parsed_out,
               BlockContents&& block);
+  void Create(std::unique_ptr<Block_kUserDefinedIndex>* parsed_out,
+              BlockContents&& block);
   void Create(std::unique_ptr<ParsedFullFilterBlock>* parsed_out,
               BlockContents&& block);
   void Create(std::unique_ptr<DecompressorDict>* parsed_out,