Fix omega build and optimize regressions

Author: driPyf

src/core/algorithm/hnsw_rabitq/CMakeLists.txt

@@ -15,7 +15,7 @@ cc_library(
     NAME core_knn_hnsw_rabitq
     STATIC SHARED STRICT ALWAYS_LINK
     SRCS *.cc
-    LIBS core_framework core_utility rabitqlib sparsehash
+    LIBS core_framework core_utility core_knn_cluster rabitqlib sparsehash
     INCS . ${PROJECT_ROOT_DIR}/src ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
     VERSION "${PROXIMA_ZVEC_VERSION}"
   )

src/core/mixed_reducer/mixed_streamer_reducer.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "mixed_streamer_reducer.h"
+#include <algorithm>
 #include <ailego/pattern/defer.h>
 #include <utility/sparse_utility.h>
 #include <zvec/ailego/utility/file_helper.h>
@@ -141,15 +142,16 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) {
   ailego::ElapsedTime timer;
 
 
-  std::vector<int> add_results(num_of_add_threads_, -1);
+  const size_t add_thread_count = enable_pk_rewrite_ ? 1 : num_of_add_threads_;
+  std::vector<int> add_results(add_thread_count, -1);
   auto add_group = thread_pool_->make_group();
 
   std::vector<int> read_results(streamers_.size(), -1);
   // TODO: use id instead of key
   uint32_t id_offset = 0, next_id = 0;
 
   if (is_sparse_) {
-    for (size_t i = 0; i < num_of_add_threads_; i++) {
+    for (size_t i = 0; i < add_thread_count; i++) {
       add_group->submit(ailego::Closure::New(
           this, &MixedStreamerReducer::add_sparse_vec, &add_results[i]));
     }
@@ -162,7 +164,7 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) {
 
     sparse_mt_list_.done();
   } else {
-    for (size_t i = 0; i < num_of_add_threads_; i++) {
+    for (size_t i = 0; i < add_thread_count; i++) {
       add_group->submit(ailego::Closure::New(
           this, &MixedStreamerReducer::add_vec, &add_results[i]));
       // add_vec(&add_results[i]);
@@ -304,6 +306,7 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index,
 
   IndexProvider::Pointer provider = streamer->create_provider();
   IndexProvider::Iterator::Pointer iterator = provider->create_iterator();
+  std::vector<std::pair<uint32_t, std::vector<uint8_t>>> pending_items;
 
   while (iterator->is_valid()) {
     if (stop_flag_ != nullptr && stop_flag_->load(std::memory_order_relaxed)) {
@@ -332,13 +335,19 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index,
       memcpy(bytes.data(), iterator->data(), bytes.size());
     }
 
-    // TODO: use id instead of key
-    if (!mt_list_.produce(VectorItem((*next_id)++, std::move(bytes)))) {
-      LOG_ERROR("Produce vector to queue failed. key[%lu]",
-                (size_t)iterator->key());
+    pending_items.emplace_back(iterator->key() + id_offset, std::move(bytes));
+    iterator->next();
+  }
+
+  std::sort(pending_items.begin(), pending_items.end(),
+            [](const auto &lhs, const auto &rhs) {
+              return lhs.first < rhs.first;
+            });
+  for (auto &item : pending_items) {
+    if (!mt_list_.produce(VectorItem((*next_id)++, std::move(item.second)))) {
+      LOG_ERROR("Produce vector to queue failed. key[%u]", item.first);
       return IndexError_Runtime;
     }
-    iterator->next();
   }
   return 0;
 }
@@ -508,6 +517,7 @@ int MixedStreamerReducer::read_sparse_vec(size_t source_streamer_index,
       streamer->create_sparse_provider();
   IndexStreamer::SparseProvider::Iterator::Pointer iterator =
       provider->create_iterator();
+  std::vector<SparseVectorItem> pending_items;
 
   while (iterator->is_valid()) {
     if (stop_flag_ != nullptr && stop_flag_->load(std::memory_order_relaxed)) {
@@ -547,15 +557,24 @@ int MixedStreamerReducer::read_sparse_vec(size_t source_streamer_index,
     memcpy(sparse_indices.data(), iterator->sparse_indices(),
            sparse_indices.size() * sizeof(uint32_t));
 
-    // TODO: use id instead of key
-    if (!sparse_mt_list_.produce(SparseVectorItem((*next_id)++,
-                                                  std::move(sparse_indices),
-                                                  std::move(sparse_values)))) {
+    pending_items.emplace_back(iterator->key() + id_offset,
+                               std::move(sparse_indices),
+                               std::move(sparse_values));
+    iterator->next();
+  }
+
+  std::sort(pending_items.begin(), pending_items.end(),
+            [](const SparseVectorItem &lhs, const SparseVectorItem &rhs) {
+              return lhs.pkey_ < rhs.pkey_;
+            });
+  for (auto &item : pending_items) {
+    if (!sparse_mt_list_.produce(SparseVectorItem(
+            (*next_id)++, std::move(item.sparse_indices_),
+            std::move(item.sparse_values_)))) {
       LOG_ERROR("Produce vector to queue failed. key[%lu]",
-                (size_t)iterator->key());
+                static_cast<size_t>(item.pkey_));
       return IndexError_Runtime;
     }
-    iterator->next();
   }
   return 0;
 }

src/db/collection.cc

@@ -834,60 +834,10 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) {
     return Status::OK();
   }
 
-  // Step 1: Build vector indexes if not ready
-  // This ensures indexes are built even for single segments that won't be
-  // compacted
-  std::vector<SegmentTask::Ptr> index_build_tasks;
-  for (auto &segment : persist_segments) {
-    if (!segment->all_vector_index_ready()) {
-      // Build all vector indexes for this segment
-      index_build_tasks.push_back(SegmentTask::CreateCreateVectorIndexTask(
-          CreateVectorIndexTask{segment, "", nullptr, options.concurrency_}));
-    }
-  }
-
-  if (!index_build_tasks.empty()) {
-    LOG_INFO("Building vector indexes for %zu segments",
-             index_build_tasks.size());
-    auto s = execute_tasks(index_build_tasks);
-    CHECK_RETURN_STATUS(s);
-
-    // Update segment metadata
-    std::lock_guard write_lock(write_mtx_);
-    Version new_version = version_manager_->get_current_version();
-
-    for (auto &task : index_build_tasks) {
-      auto task_info = task->task_info();
-      if (std::holds_alternative<CreateVectorIndexTask>(task_info)) {
-        auto create_index_task = std::get<CreateVectorIndexTask>(task_info);
-        s = new_version.update_persisted_segment_meta(
-            create_index_task.output_segment_meta_);
-        CHECK_RETURN_STATUS(s);
-      }
-    }
-
-    s = version_manager_->apply(new_version);
-    CHECK_RETURN_STATUS(s);
-    s = version_manager_->flush();
-    CHECK_RETURN_STATUS(s);
-
-    // Reload indexes in segments
-    for (auto &task : index_build_tasks) {
-      auto task_info = task->task_info();
-      if (std::holds_alternative<CreateVectorIndexTask>(task_info)) {
-        auto create_index_task = std::get<CreateVectorIndexTask>(task_info);
-        s = create_index_task.input_segment_->reload_vector_index(
-            *schema_, create_index_task.output_segment_meta_,
-            create_index_task.output_vector_indexers_,
-            create_index_task.output_quant_vector_indexers_);
-        CHECK_RETURN_STATUS(s);
-      }
-    }
-
-    LOG_INFO("Completed building vector indexes");
-  }
-
-  // Step 2: build segment compact task
+  // Build optimize tasks once so compacted segments are merged directly from
+  // their current per-segment sources. Pre-building filtered vector indexes for
+  // every persisted segment would shift source row ids before compaction and
+  // break alignment with scalar row-id filters.
   auto delete_store_clone = delete_store_->clone();
   auto tasks =
       build_compact_task(schema_, persist_segments, options.concurrency_,

src/db/index/CMakeLists.txt

@@ -3,11 +3,19 @@ include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
 
 cc_library(
     NAME zvec_index STATIC STRICT
-    SRCS *.cc segment/*.cc column/vector_column/*.cc column/inverted_column/*.cc storage/*.cc storage/wal/*.cc common/*.cc
+    SRCS *.cc
+         segment/*.cc
+         column/vector_column/*.cc
+         column/inverted_column/*.cc
+         storage/*.cc
+         storage/wal/*.cc
+         common/*.cc
+         ../training/*.cc
     LIBS zvec_common
          zvec_proto
          rocksdb
          core_interface
+         omega
          Arrow::arrow_static
          Arrow::arrow_compute
          Arrow::arrow_dataset

tests/core/interface/omega_training_session_test.cc

@@ -34,15 +34,31 @@ TEST(OmegaTrainingSessionTest, ConsumeArtifactsAggregatesRecordsAndGtCmps) {
   first.training_query_id = 0;
   first.total_cmps = 13;
   first.gt_cmps_per_rank = {3, 7, 11};
-  first.records.push_back(
-      TrainingRecord{0, 1, 3, 0.1f, 0.2f, std::vector<float>(7, 1.0f), 1});
+  TrainingRecord first_record;
+  first_record.query_id = 0;
+  first_record.hops_visited = 1;
+  first_record.cmps_visited = 3;
+  first_record.dist_1st = 0.1f;
+  first_record.dist_start = 0.2f;
+  first_record.traversal_window_stats = {1.0f, 1.0f, 1.0f, 1.0f,
+                                         1.0f, 1.0f, 1.0f};
+  first_record.label = 1;
+  first.records.push_back(first_record);
 
   QueryTrainingArtifacts second;
   second.training_query_id = 2;
   second.total_cmps = 21;
   second.gt_cmps_per_rank = {5, 9, 15};
-  second.records.push_back(
-      TrainingRecord{2, 4, 8, 0.3f, 0.4f, std::vector<float>(7, 2.0f), 0});
+  TrainingRecord second_record;
+  second_record.query_id = 2;
+  second_record.hops_visited = 4;
+  second_record.cmps_visited = 8;
+  second_record.dist_1st = 0.3f;
+  second_record.dist_start = 0.4f;
+  second_record.traversal_window_stats = {2.0f, 2.0f, 2.0f, 2.0f,
+                                          2.0f, 2.0f, 2.0f};
+  second_record.label = 0;
+  second.records.push_back(second_record);
 
   session.CollectQueryArtifacts(std::move(first));
   session.CollectQueryArtifacts(std::move(second));

tests/db/sqlengine/mock_segment.h

@@ -504,6 +504,10 @@ class MockSegment : public Segment {
     return Status::OK();
   }
 
+  Status retrain_omega_model() override {
+    return Status::OK();
+  }
+
   Status destroy() override {
     return Status::OK();
   }