[yugabyte#24947] DocDB: Fetch intents DB when querying vector index

Summary:
The vector index is updated only when transaction is applied.
So we could get into situation when transaction was committed, but not yet applied.
In this case we should scan intents DB for updated vectors and take them into account when populating read result.
Jira: DB-14085

Test Plan: PgVectorIndexTest.NotApplied/*

Reviewers: arybochkin

Reviewed By: arybochkin

Subscribers: ybase, yql

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D40042

Author: spolitov

src/yb/docdb/doc_reader.cc

@@ -220,7 +220,7 @@ Result<DocHybridTime> GetTableTombstoneTime(
   auto iter = CreateIntentAwareIterator(
       doc_db, BloomFilterMode::USE_BLOOM_FILTER, table_id, rocksdb::kDefaultQueryId, txn_op_context,
       read_operation_data.WithStatistics(nullptr));
-  iter->Seek(table_id);
+  iter->Seek(table_id, SeekFilter::kAll);
   const auto& entry_data = VERIFY_RESULT_REF(iter->Fetch());
   if (!entry_data || !entry_data.value.FirstByteIs(dockv::ValueEntryTypeAsChar::kTombstone) ||
       entry_data.key != table_id) {
@@ -253,7 +253,7 @@ Result<std::optional<SubDocument>> TEST_GetSubDocument(
   DOCDB_DEBUG_LOG("GetSubDocument for key $0 @ $1", sub_doc_key.ToDebugHexString(),
                   iter->read_time().ToString());
 
-  iter->Seek(sub_doc_key);
+  iter->Seek(sub_doc_key, SeekFilter::kAll);
   const auto& fetched = VERIFY_RESULT_REF(iter->Fetch());
   if (!fetched || !fetched.key.starts_with(sub_doc_key)) {
     return std::nullopt;
@@ -1845,7 +1845,7 @@ Result<DocReaderResult> DocDBTableReader::Get(
   // It means that other columns have NULL values, so if such column present, then
   // we should return row consisting of NULLs.
   // Here we check if there are columns values not listed in projection.
-  data_.iter->Seek(root_doc_key->AsSlice());
+  data_.iter->Seek(root_doc_key->AsSlice(), SeekFilter::kAll);
   const auto& new_fetched_entry = VERIFY_RESULT_REF(data_.iter->Fetch());
   if (!new_fetched_entry) {
     return DocReaderResult::kNotFound;

src/yb/docdb/doc_reader_redis.cc

@@ -194,7 +194,7 @@ Status BuildSubDocument(
                 ? doc_value.timestamp()
                 : write_time.hybrid_time().GetPhysicalValueMicros());
         if (!data.high_index->CanInclude(current_values_observed)) {
-          iter->SeekOutOfSubDoc(&key_copy);
+          iter->SeekOutOfSubDoc(SeekFilter::kAll, &key_copy);
           DCHECK(iter->Fetch().ok()); // Enforce call to Fetch in debug mode
           return Status::OK();
         }
@@ -203,7 +203,7 @@ Status BuildSubDocument(
         }
         (*num_values_observed)++;
         VLOG(3) << "SeekOutOfSubDoc: " << SubDocKey::DebugSliceToString(key);
-        iter->SeekOutOfSubDoc(&key_copy);
+        iter->SeekOutOfSubDoc(SeekFilter::kAll, &key_copy);
         DCHECK(iter->Fetch().ok()); // Enforce call to Fetch in debug mode
         return Status::OK();
       } else {
@@ -221,7 +221,7 @@ Status BuildSubDocument(
           iter, data.Adjusted(key, &descendant), low_ts,
           num_values_observed));
     }
-    iter->Revalidate();
+    iter->Revalidate(SeekFilter::kAll);
     if (descendant.value_type() == ValueEntryType::kInvalid) {
       // The document was not found in this level (maybe a tombstone was encountered).
       continue;
@@ -426,7 +426,7 @@ Status GetRedisSubDocument(
   if (seek_fwd_suffices) {
     db_iter->SeekForward(key_slice);
   } else {
-    db_iter->Seek(key_slice);
+    db_iter->Seek(key_slice, SeekFilter::kAll);
   }
   {
     auto temp_key = data.subdocument_key;
@@ -470,7 +470,7 @@ Status GetRedisSubDocument(
     upperbound_buffer.Append(key_slice.WithoutPrefix(upperbound_buffer.size()));
     upperbound_buffer.PushBack(dockv::KeyEntryTypeAsChar::kHighest);
     IntentAwareIteratorUpperboundScope upperbound_scope2(upperbound_buffer.AsSlice(), db_iter);
-    db_iter->Revalidate();
+    db_iter->Revalidate(SeekFilter::kAll);
     RETURN_NOT_OK(BuildSubDocument(db_iter, data, max_overwrite_ht,
                                    &num_values_observed));
     *data.doc_found = data.result->value_type() != ValueEntryType::kInvalid;
@@ -518,7 +518,7 @@ Status GetRedisSubDocument(
   }
   // Make sure the iterator is placed outside the whole document in the end.
   key_bytes.Truncate(dockey_size);
-  db_iter->SeekOutOfSubDoc(&key_bytes);
+  db_iter->SeekOutOfSubDoc(SeekFilter::kAll, &key_bytes);
   return Status::OK();
 }
 

src/yb/docdb/doc_rowwise_iterator.cc

@@ -153,9 +153,16 @@ void DocRowwiseIterator::InitResult() {
   }
 }
 
+void DocRowwiseIterator::Refresh(SeekFilter seek_filter) {
+  done_ = false;
+  seek_filter_ = seek_filter;
+}
+
 inline void DocRowwiseIterator::Seek(Slice key) {
   VLOG_WITH_FUNC(3) << " Seeking to " << key << "/" << dockv::DocKey::DebugSliceToString(key);
 
+  DCHECK(!done_);
+
   prev_doc_found_ = DocReaderResult::kNotFound;
 
   // We do not have values before dockv::KeyEntryTypeAsChar::kNullLow, but there is
@@ -164,18 +171,18 @@ inline void DocRowwiseIterator::Seek(Slice key) {
   // Another option would be changing kLowest value to kNullLow. But there are much more scenarios
   // that could be affected and should be tested.
   if (!key.empty() && key[0] >= dockv::KeyEntryTypeAsChar::kNullLow) {
-    db_iter_->Seek(key, Full::kTrue);
+    db_iter_->Seek(key, seek_filter_, Full::kTrue);
     return;
   }
 
   auto shared_prefix = shared_key_prefix();
   if (!shared_prefix.empty()) {
-    db_iter_->Seek(shared_prefix, Full::kFalse);
+    db_iter_->Seek(shared_prefix, seek_filter_, Full::kFalse);
     return;
   }
 
   const auto null_low = dockv::KeyEntryTypeAsChar::kNullLow;
-  db_iter_->Seek(Slice(&null_low, 1), Full::kFalse);
+  db_iter_->Seek(Slice(&null_low, 1), seek_filter_, Full::kFalse);
 }
 
 inline void DocRowwiseIterator::SeekPrevDocKey(Slice key) {
@@ -190,18 +197,20 @@ inline void DocRowwiseIterator::SeekPrevDocKey(Slice key) {
 
 Status DocRowwiseIterator::AdvanceIteratorToNextDesiredRow(bool row_finished,
                                                            bool current_fetched_row_skipped) {
-  if (!IsFetchedRowStatic() &&
+  if (seek_filter_ == SeekFilter::kAll && !IsFetchedRowStatic() &&
       VERIFY_RESULT(scan_choices_->AdvanceToNextRow(&row_key_, db_iter_.get(),
                                                     current_fetched_row_skipped))) {
     return Status::OK();
   }
   if (!is_forward_scan_) {
     VLOG(4) << __PRETTY_FUNCTION__ << " setting as PrevDocKey";
+    RSTATUS_DCHECK_EQ(seek_filter_, SeekFilter::kAll, IllegalState,
+                      "Backward scan is not supported with this filter");
     db_iter_->PrevDocKey(row_key_);
   } else if (row_finished) {
-    db_iter_->Revalidate();
+    db_iter_->Revalidate(seek_filter_);
   } else {
-    db_iter_->SeekOutOfSubDoc(&row_key_);
+    db_iter_->SeekOutOfSubDoc(seek_filter_, &row_key_);
   }
 
   return Status::OK();

src/yb/docdb/doc_rowwise_iterator.h

@@ -80,6 +80,10 @@ class DocRowwiseIterator final : public DocRowwiseIteratorBase {
 
   void Seek(Slice key) override;
 
+  // Refreshes the iterator if it was in finished state.
+  // filter - filter mode that should be used with refreshed iterator.
+  void Refresh(SeekFilter filter);
+
   HybridTime TEST_MaxSeenHt() override;
 
   // key slice should point to block of memory, that contains kHighest after the end.
@@ -161,6 +165,8 @@ class DocRowwiseIterator final : public DocRowwiseIteratorBase {
   bool use_fast_backward_scan_ = false;
 
   DeadlineInfo deadline_info_;
+
+  SeekFilter seek_filter_ = SeekFilter::kAll;
 };
 
 }  // namespace docdb

src/yb/docdb/doc_write_batch.cc

@@ -110,7 +110,7 @@ Status DocWriteBatch::SeekToKeyPrefix(IntentAwareIterator* doc_iter, HasAncestor
   const auto prev_key_prefix_exact = current_entry_.found_exact_key_prefix;
 
   // Seek the value.
-  doc_iter->Seek(key_prefix_.AsSlice());
+  doc_iter->Seek(key_prefix_.AsSlice(), SeekFilter::kAll);
   VLOG_WITH_FUNC(4) << SubDocKey::DebugSliceToString(key_prefix_.AsSlice())
                     << ", prev_subdoc_ht: " << prev_subdoc_ht
                     << ", prev_key_prefix_exact: " << prev_key_prefix_exact

src/yb/docdb/docdb_fwd.h

@@ -27,6 +27,7 @@
 
 namespace yb::docdb {
 
+class BoundedRocksDbIterator;
 class ConsensusFrontier;
 class DeadlineInfo;
 class DocDBCompactionFilterFactory;
@@ -69,6 +70,7 @@ struct ObjectLockPrefix;
 struct PgsqlReadOperationData;
 struct ReadOperationData;
 struct VectorIndexInsertEntry;
+struct VectorIndexSearchResultEntry;
 
 using DocKeyHash = uint16_t;
 using DocReadContextPtr = std::shared_ptr<DocReadContext>;
@@ -88,6 +90,7 @@ using VectorIndexPtr = std::shared_ptr<VectorIndex>;
 using VectorIndexes = std::vector<VectorIndexPtr>;
 using VectorIndexesPtr = std::shared_ptr<VectorIndexes>;
 using VectorIndexInsertEntries = std::vector<VectorIndexInsertEntry>;
+using VectorIndexSearchResult = std::vector<VectorIndexSearchResultEntry>;
 
 YB_STRONGLY_TYPED_BOOL(SkipFlush);
 YB_STRONGLY_TYPED_BOOL(SkipSeek);

src/yb/docdb/docrowwiseiterator-test.cc

@@ -2012,7 +2012,7 @@ SubDocKey(DocKey([], ["row2", 22222]), [SystemColumnId(0); HT{ physical: 1000 }]
           if (use_seek_forward) {
             iter.SeekForward(kEncodedDocKey2);
           } else {
-            iter.Seek(kEncodedDocKey2);
+            iter.Seek(kEncodedDocKey2, SeekFilter::kAll);
           }
           if (VLOG_IS_ON(1)) {
             iter.DebugDump();

src/yb/docdb/intent_aware_iterator.cc

@@ -232,11 +232,12 @@ IntentAwareIterator::IntentAwareIterator(
 }
 
 void IntentAwareIterator::Seek(const dockv::DocKey &doc_key) {
-  Seek(doc_key.Encode(), Full::kFalse);
+  Seek(doc_key.Encode(), SeekFilter::kAll, Full::kFalse);
 }
 
-void IntentAwareIterator::Seek(Slice key, Full full) {
-  VLOG_WITH_FUNC(4) << "key: " << DebugDumpKeyToStr(key) << ", full: " << full;
+void IntentAwareIterator::Seek(Slice key, SeekFilter filter, Full full) {
+  VLOG_WITH_FUNC(4)
+      << "key: " << DebugDumpKeyToStr(key) << ", full: " << full << ", filter: " << filter;
   DOCDB_DEBUG_SCOPE_LOG(
       key.ToDebugString(),
       std::bind(&IntentAwareIterator::DebugDump, this));
@@ -246,7 +247,17 @@ void IntentAwareIterator::Seek(Slice key, Full full) {
 
   SeekTriggered();
 
-  SkipFutureRecords<Direction::kForward>(ROCKSDB_SEEK(&iter_, key));
+  [&] {
+    switch (filter) {
+      case SeekFilter::kAll:
+        SkipFutureRecords<Direction::kForward>(ROCKSDB_SEEK(&iter_, key));
+        return;
+      case SeekFilter::kIntentsOnly:
+        regular_entry_.Reset();
+        return;
+    }
+    FATAL_INVALID_ENUM_VALUE(SeekFilter, filter);
+  }();
   if (intent_iter_.Initialized()) {
     if (!SetIntentUpperbound()) {
       return;
@@ -348,7 +359,7 @@ void IntentAwareIterator::SeekPastSubKey(Slice key) {
   FillEntry();
 }
 
-void IntentAwareIterator::SeekOutOfSubDoc(KeyBytes* key_bytes) {
+void IntentAwareIterator::SeekOutOfSubDoc(SeekFilter filter, KeyBytes* key_bytes) {
   VLOG_WITH_FUNC(4) << DebugDumpKeyToStr(*key_bytes);
   if (!status_.ok()) {
     return;
@@ -358,7 +369,17 @@ void IntentAwareIterator::SeekOutOfSubDoc(KeyBytes* key_bytes) {
 
   auto prefix_len = intent_iter_.Initialized()
       ? IntentPrepareSeek(*key_bytes, KeyEntryTypeAsChar::kMaxByte) : 0;
-  SkipFutureRecords<Direction::kForward>(docdb::SeekOutOfSubKey(key_bytes, &iter_));
+  [&] {
+    switch (filter) {
+      case SeekFilter::kAll:
+        SkipFutureRecords<Direction::kForward>(docdb::SeekOutOfSubKey(key_bytes, &iter_));
+        return;
+      case SeekFilter::kIntentsOnly:
+        regular_entry_.Reset();
+        return;
+    }
+    FATAL_INVALID_ENUM_VALUE(SeekFilter, filter);
+  }();
   IntentSeekForward(prefix_len);
   FillEntry();
 }
@@ -645,7 +666,7 @@ void IntentAwareIterator::SeekToLatestSubDocKeyInternal() {
     return;
   }
   subdockey_slice.remove_suffix(1);
-  Seek(subdockey_slice);
+  Seek(subdockey_slice, SeekFilter::kAll);
 }
 
 void IntentAwareIterator::SeekToLatestDocKeyInternal() {
@@ -656,13 +677,23 @@ void IntentAwareIterator::SeekToLatestDocKeyInternal() {
   if (!HandleStatus(dockey_size)) {
     return;
   }
-  Seek(Slice(subdockey_slice.data(), *dockey_size));
+  Seek(Slice(subdockey_slice.data(), *dockey_size), SeekFilter::kAll);
 }
 
-void IntentAwareIterator::Revalidate() {
+void IntentAwareIterator::Revalidate(SeekFilter seek_filter) {
   VLOG_WITH_FUNC(4);
 
-  SkipFutureRecords<Direction::kForward>(iter_.Entry());
+  [&] {
+    switch (seek_filter) {
+      case SeekFilter::kAll:
+        SkipFutureRecords<Direction::kForward>(iter_.Entry());
+        return;
+      case SeekFilter::kIntentsOnly:
+        regular_entry_.Reset();
+        return;
+    }
+    FATAL_INVALID_ENUM_VALUE(SeekFilter, seek_filter);
+  }();
   if (intent_iter_.Initialized()) {
     if (!SetIntentUpperbound()) {
       return;

src/yb/docdb/intent_aware_iterator.h

@@ -62,12 +62,12 @@ class IntentAwareIterator final : public IntentAwareIteratorIf {
   IntentAwareIterator(const IntentAwareIterator& other) = delete;
   void operator=(const IntentAwareIterator& other) = delete;
 
-  void Revalidate();
+  void Revalidate(SeekFilter seek_filter);
 
   // Seek to the smallest key which is greater or equal than doc_key.
   void Seek(const dockv::DocKey& doc_key);
 
-  void Seek(Slice key, Full full = Full::kTrue) override;
+  void Seek(Slice key, SeekFilter seek_filter, Full full = Full::kTrue) override;
 
   // Seek forward to specified encoded key (it is responsibility of caller to make sure it
   // doesn't have hybrid time).
@@ -82,7 +82,7 @@ class IntentAwareIterator final : public IntentAwareIteratorIf {
   // For efficiency, this overload takes a non-const KeyBytes pointer avoids memory allocation by
   // using the KeyBytes buffer to prepare the key to seek to by appending an extra byte. The
   // appended byte is removed when the method returns.
-  void SeekOutOfSubDoc(dockv::KeyBytes* key_bytes) override;
+  void SeekOutOfSubDoc(SeekFilter seek_filter, dockv::KeyBytes* key_bytes) override;
 
   // Seek to last doc key.
   void SeekToLastDocKey();
@@ -139,10 +139,6 @@ class IntentAwareIterator final : public IntentAwareIteratorIf {
   // Returns HybridTime::kInvalid if no such record was found.
   Result<HybridTime> FindOldestRecord(Slice key_without_ht, HybridTime min_hybrid_time);
 
-  size_t NumberOfBytesAppendedDuringSeekForward() const {
-    return 1 + encoded_read_time_.global_limit.size();
-  }
-
   void DebugDump();
 
   std::string DebugPosToString() override;

src/yb/docdb/intent_aware_iterator_interface.h

@@ -21,6 +21,7 @@ namespace yb {
 namespace docdb {
 
 YB_DEFINE_ENUM(Direction, (kForward)(kBackward));
+YB_DEFINE_ENUM(SeekFilter, (kAll)(kIntentsOnly));
 YB_STRONGLY_TYPED_BOOL(Full);
 
 struct FetchedEntry {
@@ -75,7 +76,8 @@ class IntentAwareIteratorIf {
   // hybrid time).
   // full means that key was fully specified, and we could add intent type at the end of the key,
   // to skip read only intents.
-  virtual void Seek(Slice key, Full full = Full::kTrue) = 0;
+  // filter - allows to ignore regular records during seek.
+  virtual void Seek(Slice key, SeekFilter filter, Full full = Full::kTrue) = 0;
 
   // Seek forward to specified encoded key (it is responsibility of caller to make sure it
   // doesn't have hybrid time). For efficiency, the method that takes a non-const KeyBytes pointer
@@ -88,7 +90,8 @@ class IntentAwareIteratorIf {
   // time). For efficiency, the method takes a non-const KeyBytes pointer avoids memory allocation
   // by using the KeyBytes buffer to prepare the key to seek to by appending an extra byte. The
   // appended byte is removed when the method returns.
-  virtual void SeekOutOfSubDoc(dockv::KeyBytes* key_bytes) = 0;
+  // filter - allows to ignore regular records during seek.
+  virtual void SeekOutOfSubDoc(SeekFilter filter, dockv::KeyBytes* key_bytes) = 0;
 
   // Positions the iterator at the beginning of the DocKey found before the given encoded_doc_key.
   // If fast backward scan is enabled, the method positions the iterator at the end (at the last