[#24020] DocDB: Vector LSM

Summary:
Initial implementation of Vector LSM.

Introduces Vector LSM interface and some basic stubs.
Implements insert of vector batch in multiple thread.
Search implementation is incomplete.
Jira: DB-12907

Test Plan: Jenkins

Reviewers: mbautin, xCluster, hsunder

Reviewed By: mbautin

Subscribers: ybase

Tags: #jenkins-ready

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

src/yb/docdb/CMakeLists.txt

@@ -81,6 +81,7 @@ set(DOCDB_SRCS
         vector_index.cc
         vector_index_read.cc
         vector_index_update.cc
+        vector_lsm.cc
         wait_queue.cc
         )
 

src/yb/docdb/vector_lsm.cc

@@ -0,0 +1,349 @@
+// Copyright (c) YugabyteDB, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied.  See the License for the specific language governing permissions and limitations
+// under the License.
+//
+
+#include "yb/docdb/vector_lsm.h"
+
+#include <boost/intrusive/list.hpp>
+
+#include "yb/rocksdb/env.h"
+
+#include "yb/rpc/thread_pool.h"
+
+#include "yb/util/flags.h"
+#include "yb/util/file_util.h"
+#include "yb/util/shared_lock.h"
+#include "yb/util/unique_lock.h"
+
+DEFINE_RUNTIME_uint64(vector_index_task_size, 16,
+                      "Number of vectors in a single insert subtask for vector index. "
+                      "When the number of elements in the batch is more than this times the task "
+                      "size multiple tasks will be used to insert batch.");
+
+DEFINE_RUNTIME_uint64(vector_index_max_insert_tasks, 1000,
+                      "Limit for number of insert subtask for vector index. "
+                      "When this limit is reached, new inserts are blocked until necessary amount "
+                      "of tasks is released.");
+
+DEFINE_RUNTIME_uint64(vector_index_task_pool_size, 1000,
+                      "Pool size of insert subtasks for vector index. "
+                      "Pool is just used to avoid memory allocations and does not limit the total "
+                      "number of tasks.");
+
+namespace yb::docdb {
+
+using vectorindex::IndexableVectorType;
+using vectorindex::ValidDistanceResultType;
+using vectorindex::VertexId;
+using vectorindex::VertexWithDistance;
+namespace bi = boost::intrusive;
+
+namespace {
+
+std::string GetIndexChunkPath(std::string storage_dir, size_t chunk_index) {
+  return JoinPathSegments(storage_dir, Format("vectorindex_$0", chunk_index));
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+class VectorLSMInsertTask :
+    public rpc::ThreadPoolTask,
+    public bi::list_base_hook<bi::link_mode<bi::normal_link>> {
+ public:
+  using Types = VectorLSMTypes<Vector, DistanceResult>;
+  using InsertRegistry = typename Types::InsertRegistry;
+  using Chunk = typename Types::Chunk;
+
+  explicit VectorLSMInsertTask(InsertRegistry& registry) : registry_(registry) {
+  }
+
+  void Bind(Chunk& chunk) {
+    chunk_ = &chunk;
+  }
+
+  void Add(VertexId vertex_id, Vector vector) {
+    vectors_.emplace_back(vertex_id, std::move(vector));
+  }
+
+  void Run() override {
+    for (const auto& [vertex_id, vector] : vectors_) {
+      // TODO(lsm) handle failure
+      CHECK_OK(chunk_->Insert(vertex_id, vector));
+    }
+  }
+
+  void Done(const Status& status) override;
+ private:
+  InsertRegistry& registry_;
+  Chunk* chunk_;
+  std::vector<std::pair<VertexId, Vector>> vectors_;
+};
+
+} // namespace
+
+// Registry for all active Vector LSM insert subtasks.
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+class VectorLSMInsertRegistry {
+ public:
+  using Types = VectorLSMTypes<Vector, DistanceResult>;
+  using Chunk = typename Types::Chunk;
+  using InsertTask = VectorLSMInsertTask<Vector, DistanceResult>;
+  using InsertTaskList = boost::intrusive::list<InsertTask>;
+  using InsertTaskPtr = std::unique_ptr<InsertTask>;
+
+  explicit VectorLSMInsertRegistry(rpc::ThreadPool& thread_pool)
+      : thread_pool_(thread_pool) {}
+
+  InsertTaskList AllocateTasks(Chunk& chunk, size_t num_tasks) EXCLUDES(mutex_) {
+    InsertTaskList result;
+    {
+      UniqueLock lock(mutex_);
+      while (allocated_tasks_ + num_tasks >= FLAGS_vector_index_max_insert_tasks) {
+        // TODO(lsm) Pass timeout here.
+        allocated_tasks_cond_.wait(GetLockForCondition(lock));
+      }
+      allocated_tasks_ += num_tasks;
+      for (size_t left = num_tasks; left-- > 0;) {
+        InsertTaskPtr task;
+        if (task_pool_.empty()) {
+          task = std::make_unique<InsertTask>(*this);
+        } else {
+          task = std::move(task_pool_.back());
+          task_pool_.pop_back();
+        }
+        task->Bind(chunk);
+        result.push_back(*task.release());
+      }
+    }
+    return result;
+  }
+
+  void ExecuteTasks(InsertTaskList& list) EXCLUDES(mutex_) {
+    for (auto& task : list) {
+      thread_pool_.Enqueue(&task);
+    }
+    std::lock_guard lock(mutex_);
+    active_tasks_.splice(active_tasks_.end(), list);
+  }
+
+  void TaskDone(InsertTask* raw_task) EXCLUDES(mutex_) {
+    InsertTaskPtr task(raw_task);
+    std::lock_guard lock(mutex_);
+    --allocated_tasks_;
+    allocated_tasks_cond_.notify_all();
+    active_tasks_.erase(active_tasks_.iterator_to(*raw_task));
+    if (task_pool_.size() < FLAGS_vector_index_task_pool_size) {
+      task_pool_.push_back(std::move(task));
+    }
+  }
+ private:
+  rpc::ThreadPool& thread_pool_;
+  std::mutex mutex_;
+  std::condition_variable allocated_tasks_cond_;
+  size_t allocated_tasks_ GUARDED_BY(mutex_);
+  InsertTaskList active_tasks_ GUARDED_BY(mutex_);
+  std::vector<InsertTaskPtr> task_pool_ GUARDED_BY(mutex_);
+};
+
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+void VectorLSMInsertTask<Vector, DistanceResult>::Done(const Status& status) {
+  vectors_.clear();
+  registry_.TaskDone(this);
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+VectorLSM<Vector, DistanceResult>::VectorLSM(Options options)
+    : options_(std::move(options)), env_(rocksdb::Env::Default()),
+      insert_registry_(new InsertRegistry(*options_.insert_thread_pool)) {
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+Status VectorLSM<Vector, DistanceResult>::Open() {
+  return env_->CreateDirIfMissing(options_.storage_dir);
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+Status VectorLSM<Vector, DistanceResult>::Insert(std::vector<InsertEntry> entries) {
+  ChunkPtr chunk;
+  {
+    std::lock_guard lock(mutex_);
+    if (!mutable_chunk_) {
+      RETURN_NOT_OK(CreateNewMutableChunk());
+    } else if (entries_in_mutable_chunks_ &&
+               entries_in_mutable_chunks_ + entries.size() > options_.points_per_chunk) {
+      // TODO(lsm) Handle size of entries greater than points_per_chunk.
+      RETURN_NOT_OK(RollChunk());
+    }
+    chunk = mutable_chunk_;
+  }
+
+  size_t num_tasks = ceil_div<size_t>(entries.size(), FLAGS_vector_index_task_size);
+  size_t entries_per_task = ceil_div(entries.size(), num_tasks);
+
+  auto tasks = insert_registry_->AllocateTasks(*chunk, num_tasks);
+  auto tasks_it = tasks.begin();
+  size_t index_in_task = std::numeric_limits<size_t>::max();
+  for (auto& [vertex_id, base_table_key, v] : entries) {
+    if (index_in_task++ >= entries_per_task) {
+      ++tasks_it;
+      index_in_task = 0;
+    }
+    tasks_it->Add(vertex_id, std::move(v));
+    options_.key_value_storage->StoreBaseTableKey(vertex_id, base_table_key);
+  }
+  insert_registry_->ExecuteTasks(tasks);
+
+  return Status::OK();
+}
+
+// Merges results from multiple chunks.
+// results_with_chunk - already merged results from chunks on previous steps.
+// chunk_results - results from a new chunk - chunk_index.
+// max_num_results - limit for number of results.
+// Expects that results_with_chunk and chunk_results already ordered by distance.
+template<ValidDistanceResultType DistanceResult>
+void MergeChunkResults(
+    std::vector<std::pair<VertexWithDistance<DistanceResult>, int>>& results_with_chunk,
+    const std::vector<VertexWithDistance<DistanceResult>>& chunk_results,
+    int chunk_index,
+    size_t max_num_results) {
+  // Store the current size of the existing results.
+  auto old_size = results_with_chunk.size();
+
+  // Calculate the combined size of the existing and new chunk results.
+  auto sum_size = old_size + chunk_results.size();
+
+  // Resize the results vector to hold the merged results, but cap it at max_num_results.
+  results_with_chunk.resize(std::min(sum_size, max_num_results));
+
+  // Initialize iterators for traversing both results:
+  // lhi: left-hand iterator for the existing results, starting from the end.
+  // rhi: right-hand iterator for the new chunk results, starting from the end.
+  auto lhi = make_signed(old_size) - 1;
+  auto rhi = make_signed(chunk_results.size()) - 1;
+
+  // out_index: keeps track of where to insert the next result in the merged vector.
+  auto out_index = sum_size;
+
+  // Merge the results from both vectors while both have elements left.
+  while (lhi >= 0 && rhi >= 0) {
+    --out_index;
+
+    // Compare elements from both sides, insert the larger one at the correct position.
+    if (results_with_chunk[lhi].first <= chunk_results[rhi]) {
+      if (out_index < results_with_chunk.size()) {
+        // Insert the new chunk result and associate it with the current chunk index.
+        results_with_chunk[out_index] = {chunk_results[rhi], chunk_index};
+      }
+      --rhi;  // Move to the next element in the new chunk results.
+    } else {
+      if (out_index < results_with_chunk.size()) {
+        // Otherwise, move the element from the existing results.
+        results_with_chunk[out_index] = results_with_chunk[lhi];
+      }
+      --lhi;  // Move to the next element in the existing results.
+    }
+  }
+
+  // If any elements remain in the new chunk results, insert them.
+  while (rhi >= 0) {
+    --out_index;
+    results_with_chunk[out_index] = {chunk_results[rhi], chunk_index};
+    --rhi;
+  }
+
+  // At this point, the existing results don't need further processing
+  // because they were already sorted in the correct position or ignored.
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+auto VectorLSM<Vector, DistanceResult>::Search(
+    const Vector& query_vector, const SearchOptions& options) const ->
+    Result<typename VectorLSM<Vector, DistanceResult>::SearchResults> {
+  if (options.max_num_results == 0) {
+    return SearchResults();
+  }
+  const int kMutableChunkIndex = -1;
+  using IntermediateResult = std::pair<vectorindex::VertexWithDistance<DistanceResult>, int>;
+  std::vector<IntermediateResult> results_with_chunk;
+
+  ChunkPtr mutable_chunk;
+  // TODO(lsm) Optimize memory allocation.
+  decltype(immutable_chunks_) immutable_chunks;
+  {
+    SharedLock lock(mutex_);
+    mutable_chunk = mutable_chunk_;
+    immutable_chunks = immutable_chunks_;
+  }
+
+  if (mutable_chunk) {
+    auto mutable_results = mutable_chunk->Search(query_vector, options.max_num_results);
+    for (auto& result : mutable_results) {
+      results_with_chunk.push_back({
+          result,
+          kMutableChunkIndex
+      });
+    }
+  }
+
+  {
+    int chunk_index = 0;
+    for (const auto& chunk : immutable_chunks) {
+      auto chunk_results = chunk->Search(query_vector, options.max_num_results);
+      MergeChunkResults(results_with_chunk, chunk_results, chunk_index, options.max_num_results);
+      ++chunk_index;
+    }
+  }
+
+  SearchResults final_results;
+  final_results.reserve(results_with_chunk.size());
+  for (const auto& item : results_with_chunk) {
+    auto [vertex_id, distance] = item.first;
+    auto base_table_key = VERIFY_RESULT(options_.key_value_storage->ReadBaseTableKey(vertex_id));
+    final_results.push_back({
+      .distance = distance,
+      .base_table_key = std::move(base_table_key)
+    });
+  }
+
+  return final_results;
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+Status VectorLSM<Vector, DistanceResult>::RollChunk() {
+  auto chunk_index = current_chunk_serial_no_++;
+  const auto chunk_path = GetIndexChunkPath(options_.storage_dir, chunk_index);
+  // TODO(lsm) Implement storing and loading data to file.
+  // RETURN_NOT_OK(mutable_chunk_->AttachToFile(chunk_path));
+
+  immutable_chunks_.push_back(std::move(mutable_chunk_));
+
+  return CreateNewMutableChunk();
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+Status VectorLSM<Vector, DistanceResult>::CreateNewMutableChunk() {
+  mutable_chunk_ = options_.chunk_factory();
+  RETURN_NOT_OK(mutable_chunk_->Reserve(options_.points_per_chunk));
+  entries_in_mutable_chunks_ = 0;
+  return Status::OK();
+}
+
+template<IndexableVectorType Vector, ValidDistanceResultType DistanceResult>
+size_t VectorLSM<Vector, DistanceResult>::TEST_num_immutable_chunks() const {
+  SharedLock lock(mutex_);
+  return immutable_chunks_.size();
+}
+
+YB_INSTANTIATE_TEMPLATE_FOR_ALL_VECTOR_AND_DISTANCE_RESULT_TYPES(VectorLSM);
+
+}  // namespace yb::docdb