diff --git a/cmake/svs.cmake b/cmake/svs.cmake index fa28df32c..160e761be 100644 --- a/cmake/svs.cmake +++ b/cmake/svs.cmake @@ -64,7 +64,7 @@ if(USE_SVS) if (CMAKE_CXX_COMPILER_ID MATCHES "Clang") if(GLIBC_VERSION VERSION_GREATER_EQUAL "2.31" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "21.0") if(GLIBCXX_VERSION VERSION_GREATER_EQUAL "11") - set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/v0.3.2/svs-shared-library-reduced-clang21-gcc11.tar.gz" CACHE STRING "SVS URL") + set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/nightly/svs-shared-library-0.5.0rc10-reduced-clang21-gcc11.tar.gz" CACHE STRING "SVS URL") else() message(STATUS "libstdc++ >= GCC 11 is required for Clang SVS binaries - disabling SVS_SHARED_LIB") set(SVS_SHARED_LIB OFF) @@ -76,14 +76,12 @@ if(USE_SVS) else() if(GLIBC_VERSION VERSION_GREATER_EQUAL "2.28") if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "14.0") - set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/v0.3.2/svs-shared-library-reduced-gcc14.tar.gz" CACHE STRING "SVS URL") + set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/nightly/svs-shared-library-0.5.0rc10-reduced-gcc14.tar.gz" CACHE STRING "SVS URL") else() - set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/v0.3.2/svs-shared-library-reduced.tar.gz" CACHE STRING "SVS URL") + set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/nightly/svs-shared-library-0.5.0rc10-reduced.tar.gz" CACHE STRING "SVS URL") endif() - elseif(GLIBC_VERSION VERSION_GREATER_EQUAL "2.26") - set(SVS_URL "https://github.com/intel/ScalableVectorSearch/releases/download/v0.3.0/svs-shared-library-0.3.0-reduced-glibc2_26.tar.gz" CACHE STRING "SVS URL") else() - message(STATUS "GLIBC >= 2.26 is required for GCC SVS binaries - disabling SVS_SHARED_LIB") + message(STATUS "GLIBC >= 2.28 is required for GCC SVS binaries - disabling SVS_SHARED_LIB") set(SVS_SHARED_LIB OFF) endif() endif() @@ -128,7 +126,7 @@ if(USE_SVS) # so whether it is available depends on which SVS this build ended up with rather than on the # platform. Detect it in the header instead of assuming: it is a method, not a file, so this # greps the header the LVQ check would have tested for existence. - set(SVS_DYNAMIC_INDEX_HEADER "svs/index/vamana/dynamic_index.h") + set(SVS_DYNAMIC_INDEX_HEADER "svs/concurrent/dynamic_index.h") set(SVS_HAS_REPLACE_EXTERNAL_ID 0) if(EXISTS "${svs_SOURCE_DIR}/include/${SVS_DYNAMIC_INDEX_HEADER}") file(READ "${svs_SOURCE_DIR}/include/${SVS_DYNAMIC_INDEX_HEADER}" SVS_DYNAMIC_INDEX_SRC) diff --git a/deps/ScalableVectorSearch b/deps/ScalableVectorSearch index a7e3494fa..1109f0cd2 160000 --- a/deps/ScalableVectorSearch +++ b/deps/ScalableVectorSearch @@ -1 +1 @@ -Subproject commit a7e3494faab9454f577b87ac934a96c78c45dd7b +Subproject commit 1109f0cd23f987753715cf4e635450f0d4462ec8 diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered.h b/src/VecSim/algorithms/hnsw/hnsw_tiered.h index 21ebaa342..c52385d71 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered.h @@ -27,9 +27,7 @@ struct HNSWSwapJob; /** * Definition of a job that inserts a new vector from flat into HNSW Index. */ -struct HNSWInsertJob : public AsyncJob { - labelType label; - idType id; +struct HNSWInsertJob : public TieredInsertJob { // The swap job of the id this vector replaces, when an async overwrite of a single-value label // earmarked that id for this vector instead of a fresh one. Set once, before the job can run - // which is only once that id's repairs are done (see `HNSWSwapJob::reuseInsertJob`). @@ -37,7 +35,7 @@ struct HNSWInsertJob : public AsyncJob { HNSWInsertJob(std::shared_ptr allocator, labelType label_, idType id_, JobCallback insertCb, VecSimIndex *index_) - : AsyncJob(allocator, HNSW_INSERT_VECTOR_JOB, insertCb, index_), label(label_), id(id_) {} + : TieredInsertJob(allocator, HNSW_INSERT_VECTOR_JOB, label_, id_, insertCb, index_) {} }; /** @@ -98,17 +96,10 @@ template class TieredHNSWIndex : public VecSimTieredIndex { private: /// Mappings from id/label to associated jobs, for invalidating and update ids if necessary. - // In MULTI, we can have more than one insert job pending per label. // **This map is protected with the flat buffer lock** - vecsim_stl::unordered_map> labelToInsertJobs; vecsim_stl::unordered_map> idToRepairJobs; vecsim_stl::unordered_map idToSwapJob; - // A mapping to hold invalid jobs, so we can dispose them upon index deletion. - vecsim_stl::unordered_map invalidJobs; - idType currInvalidJobId; // A unique arbitrary identifier for accessing invalid jobs - std::mutex invalidJobsLookupGuard; - // This threshold is tested upon deleting a label from HNSW, and once the number of deleted // vectors reached this limit, we apply swap jobs *only for vectors that has no more pending // repair jobs*, and are ready to be removed from the graph. @@ -218,13 +209,6 @@ class TieredHNSWIndex : public VecSimTieredIndex { inline HNSWIndex *getHNSWIndex() const; - // Helper function for deleting a vector from the flat buffer (after it has already been - // ingested into HNSW or deleted). This includes removing the corresponding insert job from the - // label-to-insert-jobs lookup. Also, since deletion a vector triggers swapping of the - // internal last id with the deleted vector id, here we update the pending insert job(s) for the - // last id (if needed). This should be called while *flat lock is held* (exclusive lock). - void updateInsertJobInternalId(idType prev_id, idType new_id, labelType label); - // Helper function for performing in place mark delete of vector(s) associated with a label // and creating the appropriate repair jobs for the effected connections. This should be called // while *HNSW shared lock is held* (shared locked). @@ -272,11 +256,6 @@ class TieredHNSWIndex : public VecSimTieredIndex { void insertVectorToHNSW(HNSWIndex *hnsw_index, labelType label, const void *blob); - // Set an insert/repair job as invalid, put the job pointer in the invalid jobs lookup under - // the current available id, increase it and return it (while holding invalidJobsLookupGuard). - // Returns the id that the job was stored under (to be set in the job id field). - idType setAndSaveInvalidJob(AsyncJob *job); - // Handle deletion of vector inplace considering that async deletion might occurred beforehand. int deleteLabelFromHNSWInplace(labelType label); @@ -304,36 +283,41 @@ class TieredHNSWIndex : public VecSimTieredIndex { #include "VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h" #endif +protected: + // Adapts HNSW's own data-guard mutex to Lockable + class HnswDataGuardLockable : public Lockable { + HNSWIndex *hnsw; + + public: + explicit HnswDataGuardLockable(HNSWIndex *hnsw) : hnsw(hnsw) {} + void lock() const override { hnsw->lockSharedIndexDataGuard(); } + void unlock() const override { hnsw->unlockSharedIndexDataGuard(); } + }; + HnswDataGuardLockable hnswDataGuardLockable{this->getHNSWIndex()}; + + ScopedLocks lockMainIndexForQuery() const override { + return ScopedLocks(this->mainIndexLockable); + } + + ScopedLocks lockIndexForSize() const override { + return ScopedLocks(this->flatIndexLockable, hnswDataGuardLockable); + } + + ScopedLocks lockIndexForCapacity() const override { + // No-op: No locking required for capacity() + return ScopedLocks(); + } + public: - class TieredHNSW_BatchIterator : public VecSimBatchIterator { + class TieredHNSW_BatchIterator : public TieredIndex_BatchIterator { private: const TieredHNSWIndex *index; std::shared_lock backend_index_lock; VecSimQueryParams *queryParams; - VecSimQueryResultContainer flat_results; - VecSimQueryResultContainer hnsw_results; - VecSimBatchIterator *flat_iterator; VecSimBatchIterator *hnsw_iterator; - // On single value indices, this set holds the IDs of the results that were returned from - // the flat buffer. - // On multi value indices, this set holds the IDs of all the results that were returned. - // The difference between the two cases is that on multi value indices, the same ID can - // appear in both indexes and results with different scores, and therefore we can't tell in - // advance when we expect a possibility of a duplicate. - // On single value indices, a duplicate may appear at the same batch (and we will handle it - // when merging the results) Or it may appear in a different batches, first from the flat - // buffer and then from the HNSW, in the cases where a better result if found later in HNSW - // because of the approximate nature of the algorithm. - vecsim_stl::unordered_set returned_results_set; - - private: - template - inline VecSimQueryReply *compute_current_batch(size_t n_res); - inline void filter_irrelevant_results(VecSimQueryResultContainer &); - public: TieredHNSW_BatchIterator(const void *query_vector, const TieredHNSWIndex *index, @@ -394,12 +378,9 @@ class TieredHNSWIndex : public VecSimTieredIndex { VecSimUpdateCode updateVectors(labelType label, const void *new_blobs, size_t n) override; size_t getNumMarkedDeleted() const override { return getHNSWIndex()->getNumMarkedDeleted(); } - size_t indexSize() const override; - size_t indexCapacity() const override; int updateVectorInPlace(labelType label, HNSWIndex *hnsw_index, const MemoryUtils::unique_blob &storage_blob); auto shouldDefrag(); - double getDistanceFrom_Unsafe(labelType label, const void *blob) const override; // Do nothing here, each tier (flat buffer and HNSW) should increase capacity for itself when // needed. VecSimIndexDebugInfo debugInfo() const override; @@ -468,16 +449,16 @@ class TieredHNSWIndex : public VecSimTieredIndex { /* Helper methods */ template void TieredHNSWIndex::executeInsertJobWrapper(AsyncJob *job) { - auto *insert_job = reinterpret_cast(job); - auto *job_index = reinterpret_cast *>(insert_job->index); + auto *insert_job = static_cast(job); + auto *job_index = static_cast *>(insert_job->index); job_index->executeInsertJob(insert_job); delete job; } template void TieredHNSWIndex::executeRepairJobWrapper(AsyncJob *job) { - auto *repair_job = reinterpret_cast(job); - auto *job_index = reinterpret_cast *>(repair_job->index); + auto *repair_job = static_cast(job); + auto *job_index = static_cast *>(repair_job->index); job_index->executeRepairJob(repair_job); delete job; } @@ -686,22 +667,6 @@ int TieredHNSWIndex::deleteLabelFromHNSW( return internal_ids.size(); } -template -void TieredHNSWIndex::updateInsertJobInternalId(idType prev_id, idType new_id, - labelType label) { - // Update the pending job id, due to a swap that was caused after the removal of new_id. - assert(new_id != INVALID_ID && prev_id != INVALID_ID); - auto it = this->labelToInsertJobs.find(label); - if (it != this->labelToInsertJobs.end()) { - // There is a pending job for the label of the swapped last id - update its id. - for (HNSWInsertJob *job_it : it->second) { - if (job_it->id == prev_id) { - job_it->id = new_id; - } - } - } -} - template template void TieredHNSWIndex::insertVectorToHNSW( @@ -862,16 +827,6 @@ void TieredHNSWIndex::insertVectorToHNSWReusingId( this->mainIndexGuard.unlock_shared(); } -template -idType TieredHNSWIndex::setAndSaveInvalidJob(AsyncJob *job) { - this->invalidJobsLookupGuard.lock(); - job->isValid = false; - idType curInvalidId = currInvalidJobId++; - this->invalidJobs.insert({curInvalidId, job}); - this->invalidJobsLookupGuard.unlock(); - return curInvalidId; -} - template int TieredHNSWIndex::deleteLabelFromHNSWInplace(labelType label) { auto *hnsw_index = this->getHNSWIndex(); @@ -941,40 +896,7 @@ void TieredHNSWIndex::executeInsertJob(HNSWInsertJob *job) { #endif // Remove the vector and the insert job from the flat buffer. - this->flatIndexGuard.lock(); - // The job might have been invalidated due to overwrite in the meantime. In this case, - // it was already deleted and the job has been evicted. Otherwise, we need to do it now. - if (job->isValid) { - // Remove the job pointer from the labelToInsertJobs mapping. - auto &jobs = labelToInsertJobs.at(job->label); - for (size_t i = 0; i < jobs.size(); i++) { - if (jobs[i]->id == job->id) { - jobs.erase(jobs.begin() + (long)i); - break; - } - } - if (labelToInsertJobs.at(job->label).empty()) { - labelToInsertJobs.erase(job->label); - } - // Remove the vector from the flat buffer. This may cause the last vector id to swap with - // the deleted id. Hold the label for the last id, so we can later on update its - // corresponding job id. Note that after calling deleteVectorById, the last id's label - // shouldn't be available, since it is removed from the lookup. - labelType last_vec_label = - this->frontendIndex->getVectorLabel(this->frontendIndex->indexSize() - 1); - int deleted = this->frontendIndex->deleteVectorById(job->label, job->id); - if (deleted && job->id != this->frontendIndex->indexSize()) { - // If the vector removal caused a swap with the last id, update the relevant insert job. - this->updateInsertJobInternalId(this->frontendIndex->indexSize(), job->id, - last_vec_label); - } - } else { - // Remove the current job from the invalid jobs' lookup, as we are about to delete it now. - this->invalidJobsLookupGuard.lock(); - this->invalidJobs.erase(job->id); - this->invalidJobsLookupGuard.unlock(); - } - this->flatIndexGuard.unlock(); + this->removeIngestedVectorFromFlat(job); // Reusing an id disposes of its element, which can complete another reserved id's repairs. this->submitDeferredReuseInsertJobs(); @@ -1061,10 +983,8 @@ TieredHNSWIndex::TieredHNSWIndex(HNSWIndex allocator) : VecSimTieredIndex(hnsw_index, bf_index, tiered_index_params, allocator), - labelToInsertJobs(this->allocator), idToRepairJobs(this->allocator), - idToSwapJob(this->allocator), invalidJobs(this->allocator), currInvalidJobId(0), - readySwapJobs(0), deferredReuseInsertJobs(this->allocator), - isQuantized(hnsw_index->usesQuantizedStorage()) { + idToRepairJobs(this->allocator), idToSwapJob(this->allocator), readySwapJobs(0), + deferredReuseInsertJobs(this->allocator), isQuantized(hnsw_index->usesQuantizedStorage()) { const size_t normalization_set_size = tiered_index_params.specificParams.tieredHnswParams.QuantNormalizationSetSize; if (isQuantized && normalization_set_size > 0) { @@ -1081,12 +1001,6 @@ TieredHNSWIndex::TieredHNSWIndex(HNSWIndex TieredHNSWIndex::~TieredHNSWIndex() { - // Delete all the pending insert jobs. - for (auto &jobs : this->labelToInsertJobs) { - for (auto *job : jobs.second) { - delete job; - } - } // Delete all the pending repair jobs. for (auto &jobs : this->idToRepairJobs) { for (auto *job : jobs.second) { @@ -1097,10 +1011,6 @@ TieredHNSWIndex::~TieredHNSWIndex() { for (auto &it : this->idToSwapJob) { delete it.second; } - // Delete all the pending invalid jobs. - for (auto &it : this->invalidJobs) { - delete it.second; - } } // Training writes stay in FLAT in both write modes, even beyond flatBufferLimit. @@ -1111,23 +1021,20 @@ int TieredHNSWIndex::addVectorDuringAccumulation(const void const auto [result, count, should_finalize] = [&] { std::lock_guard flat_lock(this->flatIndexGuard); idType id = this->frontendIndex->indexSize(); - HNSWInsertJob *job = nullptr; + TieredInsertJob *job = nullptr; if (!this->frontendIndex->isMultiValue() && this->frontendIndex->isLabelExists(label)) { id = static_cast *>(this->frontendIndex) ->getIdOfLabel(label); subtractFromSum( {this->frontendIndex->getDataByInternalId(id), this->frontendIndex->getDim()}); - job = labelToInsertJobs.at(label).front(); + job = this->labelToInsertJobs.at(label).front(); } const int result = this->frontendIndex->addVector(blob, label); addToSum({this->frontendIndex->getDataByInternalId(id), this->frontendIndex->getDim()}); if (!job) { job = new (this->allocator) HNSWInsertJob(this->allocator, label, id, executeInsertJobWrapper, this); - auto [it, inserted] = labelToInsertJobs.try_emplace(label, 1, job, this->allocator); - if (!inserted) { - it->second.push_back(job); - } + this->registerInsertJob(label, job); } const size_t count = this->frontendIndex->indexSize(); const bool should_finalize = count >= sqAccumulationState->normalizationSetSize; @@ -1143,8 +1050,8 @@ int TieredHNSWIndex::addVectorDuringAccumulation(const void template int TieredHNSWIndex::deleteVectorDuringAccumulation(labelType label) { std::lock_guard flat_lock(this->flatIndexGuard); - auto it = labelToInsertJobs.find(label); - if (it == labelToInsertJobs.end()) { + auto it = this->labelToInsertJobs.find(label); + if (it == this->labelToInsertJobs.end()) { return 0; } const int removed = it->second.size(); @@ -1153,10 +1060,10 @@ int TieredHNSWIndex::deleteVectorDuringAccumulation(labelTyp {this->frontendIndex->getDataByInternalId(job->id), this->frontendIndex->getDim()}); delete job; // Never submitted, so no worker can still reference this job. } - labelToInsertJobs.erase(it); + this->labelToInsertJobs.erase(it); auto updated_ids = this->frontendIndex->deleteVectorAndGetUpdatedIds(label); for (const auto &[new_id, previous] : updated_ids) { - updateInsertJobInternalId(previous.first, new_id, previous.second); + this->updateInsertJobInternalId(previous.first, new_id, previous.second); } return removed; } @@ -1183,7 +1090,7 @@ void TieredHNSWIndex::finalizeQuantizationAndSubmitJobs(size jobs.reserve(count); { std::shared_lock flat_lock(this->flatIndexGuard); - for (const auto &[label, pending] : labelToInsertJobs) { + for (const auto &[label, pending] : this->labelToInsertJobs) { jobs.insert(jobs.end(), pending.begin(), pending.end()); } } @@ -1212,20 +1119,6 @@ void TieredHNSWIndex::finalizeQuantizationAndSubmitJobs(size } } -template -size_t TieredHNSWIndex::indexSize() const { - std::shared_lock flat_index_lock(this->flatIndexGuard); - const size_t res = this->frontendIndex->indexSize(); - auto *hnsw_index = getHNSWIndex(); - auto index_data_lock = hnsw_index->acquireSharedIndexDataGuard(); - return res + hnsw_index->indexSize(); -} - -template -size_t TieredHNSWIndex::indexCapacity() const { - return this->backendIndex->indexCapacity() + this->frontendIndex->indexCapacity(); -} - template int TieredHNSWIndex::updateVectorInPlace( labelType label, HNSWIndex *hnsw_index, @@ -1322,16 +1215,10 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l this->flatIndexGuard.lock(); idType new_flat_id = this->frontendIndex->indexSize(); if (this->frontendIndex->isLabelExists(label) && !this->frontendIndex->isMultiValue()) { - // Overwrite the vector and invalidate its only pending job (since we are not in MULTI). - auto *old_job = this->labelToInsertJobs.at(label).at(0); - old_job->id = this->setAndSaveInvalidJob(old_job); - this->labelToInsertJobs.erase(label); ret = 0; // We are going to update the internal id that currently holds the vector associated with // the given label. - new_flat_id = - dynamic_cast *>(this->frontendIndex) - ->getIdOfLabel(label); + new_flat_id = this->invalidatePendingInsertJob(label); // If we are adding a new element (rather than updating an exiting one) we may need to // increase index capacity. } @@ -1340,13 +1227,7 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l auto *new_insert_job = new (this->allocator) HNSWInsertJob(this->allocator, label, new_flat_id, executeInsertJobWrapper, this); - // Construct the job vector only for a new label; multi-value labels append to the existing one. - auto [it, inserted] = - this->labelToInsertJobs.try_emplace(label, 1, new_insert_job, this->allocator); - if (!inserted) { - assert(this->backendIndex->isMultiValue()); - it->second.push_back(new_insert_job); - } + this->registerInsertJob(label, new_insert_job); this->flatIndexGuard.unlock(); // Here, a worker might ingest the previous vector that was stored under "label" @@ -1680,7 +1561,7 @@ VecSimRelabelCode TieredHNSWIndex::relabelVector(labelType o if (jobs_it != this->labelToInsertJobs.end()) { auto jobs = std::move(jobs_it->second); this->labelToInsertJobs.erase(jobs_it); - for (HNSWInsertJob *job : jobs) { + for (TieredInsertJob *job : jobs) { job->label = new_label; } this->labelToInsertJobs.emplace(new_label, std::move(jobs)); @@ -1711,50 +1592,6 @@ VecSimRelabelCode TieredHNSWIndex::relabelVector(labelType o return VecSimRelabel_OK; } -// `getDistanceFrom` returns the minimum distance between the given blob and the vector with the -// given label. If the label doesn't exist, the distance will be NaN. -// Therefore, it's better to just call `getDistanceFrom` on both indexes and return the minimum -// instead of checking if the label exists in each index. We first try to get the distance from the -// flat buffer, as vectors in the buffer might move to the Main while we're "between" the locks. -// Behavior for single (regular) index: -// 1. label doesn't exist in both indexes - return NaN -// 2. label exists in one of the indexes only - return the distance from that index (which is valid) -// 3. label exists in both indexes - return the value from the flat buffer (which is valid and equal -// to the value from the Main index), saving us from locking the Main index. -// Behavior for multi index: -// 1. label doesn't exist in both indexes - return NaN -// 2. label exists in one of the indexes only - return the distance from that index (which is valid) -// 3. label exists in both indexes - we may have some of the vectors with the same label in the flat -// buffer only and some in the Main index only (and maybe temporal duplications). -// So, we get the distance from both indexes and return the minimum. - -// IMPORTANT: this should be called when the *tiered index locks are locked for shared ownership*, -// along with HNSW index data guard lock. That is since the internal getDistanceFrom calls access -// the indexes' data, and it is not safe to run insert/delete operation in parallel. Also, we avoid -// acquiring the locks internally, since this is usually called for every vector individually, and -// the overhead of acquiring and releasing the locks is significant in that case. -template -double TieredHNSWIndex::getDistanceFrom_Unsafe(labelType label, - const void *blob) const { - // Try to get the distance from the flat buffer. - // If the label doesn't exist, the distance will be NaN. - auto flat_dist = this->frontendIndex->getDistanceFrom_Unsafe(label, blob); - - // Optimization. TODO: consider having different implementations for single and multi indexes, - // to avoid checking the index type on every query. - if (!this->backendIndex->isMultiValue() && !std::isnan(flat_dist)) { - // If the index is single value, and we got a valid distance from the flat buffer, - // we can return the distance without querying the Main index. - return flat_dist; - } - - // Try to get the distance from the Main index. - auto hnsw_dist = getHNSWIndex()->getDistanceFrom_Unsafe(label, blob); - - // Return the minimum distance that is not NaN. - return std::fmin(flat_dist, hnsw_dist); -} - //////////////////////////////////////////////////////////////////////////////////////////////////// // TieredHNSW_BatchIterator // //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -1775,11 +1612,10 @@ TieredHNSWIndex::TieredHNSW_BatchIterator::TieredHNSW_BatchI // copies: flat_iterator copy is created during TieredHNSW_BatchIterator construction When // TieredHNSW_BatchIterator::getNextResults() is called and hnsw_iterator is not initialized, it // retrieves the blob from flat_iterator - : VecSimBatchIterator(nullptr, queryParams ? queryParams->timeoutCtx : nullptr, - std::move(allocator)), + : TieredIndex_BatchIterator(nullptr, queryParams ? queryParams->timeoutCtx : nullptr, + std::move(allocator)), index(index), backend_index_lock(index->mainIndexGuard, std::defer_lock), - flat_results(this->allocator), hnsw_results(this->allocator), flat_iterator(UNINITIALIZED), - hnsw_iterator(UNINITIALIZED), returned_results_set(this->allocator) { + flat_iterator(UNINITIALIZED), hnsw_iterator(UNINITIALIZED) { { std::shared_lock flat_index_lock(this->index->flatIndexGuard); this->flat_iterator = @@ -1836,7 +1672,7 @@ VecSimQueryReply *TieredHNSWIndex::TieredHNSW_BatchIterator: this->flat_iterator->getQueryBlob(), queryParams); auto cur_hnsw_results = this->hnsw_iterator->getNextResults(n_res, BY_SCORE_THEN_ID); hnsw_code = cur_hnsw_results->code; - this->hnsw_results.swap(cur_hnsw_results->results); + this->backend_results.swap(cur_hnsw_results->results); VecSimQueryReply_Free(cur_hnsw_results); if (this->hnsw_iterator->isDepleted()) { delete this->hnsw_iterator; @@ -1864,19 +1700,19 @@ VecSimQueryReply *TieredHNSWIndex::TieredHNSW_BatchIterator: } } - while (this->hnsw_results.size() < n_res && this->hnsw_iterator != DEPLETED && + while (this->backend_results.size() < n_res && this->hnsw_iterator != DEPLETED && hnsw_code == VecSim_OK) { - auto tail = this->hnsw_iterator->getNextResults(n_res - this->hnsw_results.size(), + auto tail = this->hnsw_iterator->getNextResults(n_res - this->backend_results.size(), BY_SCORE_THEN_ID); hnsw_code = tail->code; // Set the hnsw_results code to the last `getNextResults` code. // New batch may contain better results than the previous batch, so we need to merge. // We don't expect duplications (hence the ), as the iterator guarantees that // no result is returned twice. VecSimQueryResultContainer cur_hnsw_results(this->allocator); - merge_results(cur_hnsw_results, this->hnsw_results, tail->results, n_res); + merge_results(cur_hnsw_results, this->backend_results, tail->results, n_res); VecSimQueryReply_Free(tail); - this->hnsw_results.swap(cur_hnsw_results); - filter_irrelevant_results(this->hnsw_results); + this->backend_results.swap(cur_hnsw_results); + filter_irrelevant_results(this->backend_results); if (this->hnsw_iterator->isDepleted()) { delete this->hnsw_iterator; this->hnsw_iterator = DEPLETED; @@ -1913,7 +1749,7 @@ VecSimQueryReply *TieredHNSWIndex::TieredHNSW_BatchIterator: template bool TieredHNSWIndex::TieredHNSW_BatchIterator::isDepleted() { return this->flat_results.empty() && this->flat_iterator->isDepleted() && - this->hnsw_results.empty() && this->hnsw_iterator == DEPLETED; + this->backend_results.empty() && this->hnsw_iterator == DEPLETED; } template @@ -1926,85 +1762,10 @@ void TieredHNSWIndex::TieredHNSW_BatchIterator::reset() { this->flat_iterator->reset(); this->hnsw_iterator = UNINITIALIZED; this->flat_results.clear(); - this->hnsw_results.clear(); + this->backend_results.clear(); returned_results_set.clear(); } -/****************** Helper Functions **************/ - -template -template -VecSimQueryReply * -TieredHNSWIndex::TieredHNSW_BatchIterator::compute_current_batch(size_t n_res) { - // Merge results - // This call will update `hnsw_res` and `bf_res` to point to the end of the merged results. - auto batch_res = new VecSimQueryReply(this->allocator); - std::pair p; - if (isMultiValue) { - p = merge_results(batch_res->results, this->hnsw_results, this->flat_results, n_res); - } else { - p = merge_results(batch_res->results, this->hnsw_results, this->flat_results, n_res); - } - auto [from_hnsw, from_flat] = p; - - if (!isMultiValue) { - // If we're on a single-value index, update the set of results returned from the FLAT index - // before popping them, to prevent them to be returned from the HNSW index in later batches. - for (size_t i = 0; i < from_flat; ++i) { - this->returned_results_set.insert(this->flat_results[i].id); - } - } else { - // If we're on a multi-value index, update the set of results returned (from `batch_res`) - for (size_t i = 0; i < batch_res->results.size(); ++i) { - this->returned_results_set.insert(batch_res->results[i].id); - } - } - - // Update results - this->flat_results.erase(this->flat_results.begin(), this->flat_results.begin() + from_flat); - this->hnsw_results.erase(this->hnsw_results.begin(), this->hnsw_results.begin() + from_hnsw); - - // clean up the results - // On multi-value indexes, one (or both) results lists may contain results that are already - // returned form the other list (with a different score). We need to filter them out. - if (isMultiValue) { - filter_irrelevant_results(this->flat_results); - filter_irrelevant_results(this->hnsw_results); - } - - // Return current batch - return batch_res; -} - -template -void TieredHNSWIndex::TieredHNSW_BatchIterator::filter_irrelevant_results( - VecSimQueryResultContainer &results) { - // Filter out results that were already returned. - auto it = results.begin(); - const auto end = results.end(); - // Skip results that not returned yet - while (it != end && this->returned_results_set.count(it->id) == 0) { - ++it; - } - // If none of the results were returned, return - if (it == end) { - return; - } - // Mark the current result as the first result to be filtered - auto cur_end = it; - ++it; - // "Append" all results that were not returned from the FLAT index - while (it != end) { - if (this->returned_results_set.count(it->id) == 0) { - *cur_end = *it; - ++cur_end; - } - ++it; - } - // Update number of results (pop the tail) - results.resize(cur_end - results.begin()); -} - template VecSimIndexStatsInfo TieredHNSWIndex::statisticInfo() const { auto stats = VecSimTieredIndex::statisticInfo(); diff --git a/src/VecSim/algorithms/svs/svs.h b/src/VecSim/algorithms/svs/svs.h index 33162e4f7..2b998dc11 100644 --- a/src/VecSim/algorithms/svs/svs.h +++ b/src/VecSim/algorithms/svs/svs.h @@ -15,7 +15,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -39,21 +42,32 @@ struct SVSIndexBase SVSIndexBase() : num_marked_deleted{0} {}; virtual ~SVSIndexBase() = default; + virtual int addVector(const void *vector_data, labelType label) = 0; virtual int addVectors(const void *vectors_data, const labelType *labels, size_t n) = 0; + virtual int deleteVector(labelType label) = 0; virtual int deleteVectors(const labelType *labels, size_t n) = 0; + virtual void consolidate(const std::vector &labels) = 0; virtual bool isLabelExists(labelType label) const = 0; virtual size_t indexStorageSize() const = 0; virtual size_t getParallelism() const = 0; virtual void setParallelism(size_t parallelism) = 0; virtual size_t getPoolSize() const = 0; + // Run synchronously using caller-reserved threads, without changing the + // backend's per-operation parallelism. + virtual void runParallel(std::function task, size_t num_threads) = 0; virtual bool isCompressed() const = 0; - size_t getNumMarkedDeleted() const { return num_marked_deleted; } + virtual bool ready() const = 0; + + size_t getNumMarkedDeleted() const { + return num_marked_deleted.load(std::memory_order_relaxed); + } // Abstract handler to manage SVS implementation instance // declared to avoid unsafe unique_ptr usage // Derived SVSIndex class should implement it struct ImplHandler { virtual ~ImplHandler() = default; + virtual void deleteVectors(const labelType *labels, size_t n) = 0; }; virtual std::unique_ptr createImpl(const void *vectors_data, const labelType *labels, size_t n) = 0; @@ -69,7 +83,7 @@ struct SVSIndexBase protected: // Index marked deleted vectors counter to initiate reindexing if it exceeds threshold // markIndexUpdate() manages this counter - size_t num_marked_deleted; + std::atomic num_marked_deleted; }; /** Thread Management Strategy: @@ -93,10 +107,12 @@ class SVSIndex : public VecSimIndexAbstract, fl using graph_builder_t = SVSGraphBuilder; using graph_type = typename graph_builder_t::graph_type; + // Note. svs::concurrent::* classes differ from the svs::index::vamana::* classes + // in their mutation contracts. using impl_type = std::conditional_t< isMulti, - svs::index::vamana::MultiMutableVamanaIndex, - svs::index::vamana::MutableVamanaIndex>; + svs::concurrent::MultiMutableVamanaIndex, + svs::concurrent::MutableVamanaIndex>; bool forcePreprocessing; @@ -118,8 +134,12 @@ class SVSIndex : public VecSimIndexAbstract, fl // SVS thread pool VecSimSVSThreadPool threadpool_; svs::logging::logger_ptr logger_; - // SVS Index implementation instance - std::unique_ptr impl_; + // SVS Index implementation instance. + // initImpl() needs at least one point to compute an entry point. + // so impl_ is null until the first insert + std::atomic impl_{nullptr}; + + impl_type *getImpl() const { return this->impl_.load(std::memory_order_acquire); } static double toVecSimDistance(float v) { return svs_details::toVecSimDistance(v); } @@ -196,6 +216,7 @@ class SVSIndex : public VecSimIndexAbstract, fl // Configure default search parameters auto sp = impl->get_search_parameters(); sp.buffer_config({this->search_window_size, this->search_buffer_capacity}); + sp.search_buffer_visited_set(false); impl->set_search_parameters(sp); impl->reset_performance_parameters(); return impl; @@ -227,7 +248,14 @@ class SVSIndex : public VecSimIndexAbstract, fl // Handler to manage SVS implementation instance struct SVSImplHandler : public SVSIndexBase::ImplHandler { std::unique_ptr impl; + size_t num_deleted{0}; SVSImplHandler(std::unique_ptr impl) : impl{std::move(impl)} {} + + void deleteVectors(const labelType *labels, size_t n) override { + if (impl) { + num_deleted += impl->delete_entries(std::span{labels, n}); + } + } }; std::unique_ptr createImpl(const void *vectors_data, const labelType *labels, @@ -246,16 +274,26 @@ class SVSIndex : public VecSimIndexAbstract, fl return std::make_unique(initImpl(points, ids)); } - void setImpl(std::unique_ptr handler) override { - if (impl_ != nullptr) { - throw std::logic_error("SVSIndex::setImpl called on non-empty impl_"); - } + // if impl_ == nullptr, store impl, return true + // if impl_ != nullptr, do nothing, return false + bool CASImpl(impl_type *impl) { + impl_type *expected = nullptr; + return this->impl_.compare_exchange_strong(expected, impl, std::memory_order_acq_rel); + } + void setImpl(std::unique_ptr handler) override { SVSImplHandler *svs_handler = dynamic_cast(handler.get()); if (!svs_handler) { throw std::logic_error("Failed to cast to SVSImplHandler"); } - this->impl_ = std::move(svs_handler->impl); + + if (!CASImpl(svs_handler->impl.get())) { + throw std::logic_error("SVSIndex::setImpl called on non-empty impl_"); + } + svs_handler->impl.release(); + if (svs_handler->num_deleted > 0) { + this->markIndexUpdate(svs_handler->num_deleted); + } } // Assuming parallelism was updated to reflect the number of available threads before this @@ -289,11 +327,6 @@ class SVSIndex : public VecSimIndexAbstract, fl } int deleted_num = 0; - if constexpr (!isMulti) { - // SVS index does not support overriding vectors with the same label - // so we have to delete them first if needed - deleted_num = deleteVectorsImpl(labels, n); - } std::span ids(labels, n); auto processed_blob = this->preprocessForBatchStorage(vectors_data, n); @@ -301,65 +334,70 @@ class SVSIndex : public VecSimIndexAbstract, fl // Wrap data into SVS SimpleDataView for SVS API auto points = svs::data::SimpleDataView{typed_vectors_data, n, this->dim}; - if (!impl_) { - // SVS index instance cannot be empty, so we have to construct it at first rows - impl_ = initImpl(points, ids); - } else { + if constexpr (!isMulti) { + // SVS index does not support overriding vectors with the same label + // so we have to delete them first if needed + deleted_num = deleteVectorsImpl(labels, n); + } + + if (auto impl = getImpl()) { // Add new points to existing SVS index - impl_->add_points(points, ids); + impl->add_points(points, ids, /*reuse_empty*/ false); + } else { + // SVS index instance cannot be empty, so we have to construct it at first rows + auto built = initImpl(points, ids); + if (CASImpl(built.get())) { + built.release(); + } else { + // Another writer published first. + if constexpr (!isMulti) { + deleted_num = deleteVectorsImpl(labels, n); + } + getImpl()->add_points(points, ids, /*reuse_empty*/ false); + } } return n - deleted_num; } - int deleteVectorImpl(const labelType label) { - if (indexLabelCount() == 0 || !impl_->has_id(label)) { - return 0; + void consolidate(const std::vector &labels) override { + if (auto impl = getImpl()) { + impl->consolidate(labels); } - - const auto deleted_num = impl_->delete_entries(std::span{&label, 1}); - - this->markIndexUpdate(deleted_num); - return deleted_num; } - int deleteVectorsImpl(const labelType *labels, size_t n) { - if (indexLabelCount() == 0) { + int deleteVectorImpl(const labelType label) { + if (auto impl = getImpl()) { + const int deleted_num = impl->delete_entries(std::span{&label, 1}); + this->markIndexUpdate(deleted_num); + return deleted_num; + } else { return 0; } + } - // SVS fails if we try to delete non-existing entries - std::vector entries_to_delete; - entries_to_delete.reserve(n); - for (size_t i = 0; i < n; i++) { - if (impl_->has_id(labels[i])) { - entries_to_delete.push_back(labels[i]); + int deleteVectorsImpl(const labelType *labels, size_t n) { + if (auto impl = getImpl()) { + const int deleted_num = impl->delete_entries(std::span{labels, n}); + if (deleted_num > 0) { + this->markIndexUpdate(deleted_num); } - } - - if (entries_to_delete.size() == 0) { + return deleted_num; + } else { return 0; } - - const auto deleted_num = impl_->delete_entries(entries_to_delete); - - this->markIndexUpdate(deleted_num); - return deleted_num; } - // Count deletions and consolidate index if needed void markIndexUpdate(size_t n = 1) { - if (!impl_) - return; + num_marked_deleted.fetch_add(n, std::memory_order_relaxed); + } - // SVS index instance should not be empty - if (indexLabelCount() == 0) { - this->impl_.reset(); - num_marked_deleted = 0; - return; + static size_t labelCountOf(const impl_type &impl) { + if constexpr (isMulti) { + return impl.labelcount(); + } else { + return impl.size(); } - - num_marked_deleted += n; } bool isTwoLevelLVQ(const VecSimSvsQuantBits &qbits) { @@ -387,7 +425,7 @@ class SVSIndex : public VecSimIndexAbstract, fl svs_details::getOrDefault(params.leanvec_dim, SVS_VAMANA_DEFAULT_LEANVEC_DIM)}, epsilon{svs_details::getOrDefault(params.epsilon, SVS_VAMANA_DEFAULT_EPSILON)}, is_two_level_lvq{isTwoLevelLVQ(params.quantBits)}, - threadpool_{this->allocator, this->logCallbackCtx}, impl_{nullptr} { + threadpool_{this->allocator, this->logCallbackCtx} { logger_ = makeLogger(); if (params.num_threads != 0) { this->log(VecSimCommonStrings::LOG_WARNING_STRING, @@ -396,28 +434,31 @@ class SVSIndex : public VecSimIndexAbstract, fl } } - ~SVSIndex() = default; + ~SVSIndex() { delete this->impl_.load(std::memory_order_relaxed); } + + bool ready() const override { return getImpl() != nullptr; } size_t indexSize() const override { return indexStorageSize(); } - size_t indexStorageSize() const override { return impl_ ? impl_->view_data().size() : 0; } + size_t indexStorageSize() const override { + auto impl = getImpl(); + return impl ? impl->view_data().size() : 0; + } size_t indexCapacity() const override { - return impl_ ? storage_traits_t::storage_capacity(impl_->view_data()) : 0; + auto impl = getImpl(); + return impl ? storage_traits_t::storage_capacity(impl->view_data()) : 0; } size_t indexLabelCount() const override { - if constexpr (isMulti) { - return impl_ ? impl_->labelcount() : 0; - } else { - return impl_ ? impl_->size() : 0; - } + auto impl = getImpl(); + return impl ? labelCountOf(*impl) : 0; } vecsim_stl::set getLabelsSet() const override { vecsim_stl::set labels(this->allocator); - if (impl_) { - impl_->on_ids([&labels](size_t label) { labels.insert(label); }); + if (auto impl = getImpl()) { + impl->on_ids([&labels](size_t label) { labels.insert(label); }); } return labels; } @@ -629,22 +670,31 @@ class SVSIndex : public VecSimIndexAbstract, fl if (old_label == new_label) { return VecSimRelabel_SameLabel; } - // `isLabelExists` also covers the index that never held a vector, where `impl_` has not - // been created yet and so trivially holds nothing. - if (!isLabelExists(old_label)) { + + auto impl = getImpl(); + + // A null impl_ is an index that never held a vector, so it holds neither label. + if (!impl) { return VecSimRelabel_OldLabelMissing; } - if (isLabelExists(new_label)) { + + // SVS checks both ends under its own translator lock + switch (impl->replace_external_id(old_label, new_label)) { + case svs::concurrent::ReplaceExternalIdResult::Ok: + return VecSimRelabel_OK; + case svs::concurrent::ReplaceExternalIdResult::NewIdExists: return VecSimRelabel_NewLabelTaken; + case svs::concurrent::ReplaceExternalIdResult::OldIdMissing: + return VecSimRelabel_OldLabelMissing; + default: + throw std::logic_error("SVSIndex::relabelVector: unreachable SVS result"); } - - impl_->replace_external_id(old_label, new_label); - return VecSimRelabel_OK; } #endif // HAVE_SVS_REPLACE_EXTERNAL_ID bool isLabelExists(labelType label) const override { - return impl_ ? impl_->has_id(label) : false; + auto impl = getImpl(); + return impl ? impl->has_id(label) : false; } size_t getParallelism() const override { return threadpool_.getParallelism(); } @@ -652,6 +702,10 @@ class SVSIndex : public VecSimIndexAbstract, fl size_t getPoolSize() const override { return threadpool_.poolSize(); } + void runParallel(std::function task, size_t num_threads) override { + threadpool_.parallel_for(std::move(task), num_threads); + } + bool isCompressed() const override { return storage_traits_t::is_compressed(); } VecSimSvsQuantBits getCompressionMode() const { @@ -659,12 +713,13 @@ class SVSIndex : public VecSimIndexAbstract, fl } double getDistanceFrom_Unsafe(labelType label, const void *vector_data) const override { - if (!impl_ || !impl_->has_id(label)) { + auto impl = getImpl(); + if (!impl || !impl->has_id(label)) { return std::numeric_limits::quiet_NaN(); - }; + } auto query_datum = std::span{static_cast(vector_data), this->dim}; - auto dist = impl_->get_distance(label, query_datum); + auto dist = impl->get_distance(label, query_datum); return toVecSimDistance(dist); } @@ -672,39 +727,48 @@ class SVSIndex : public VecSimIndexAbstract, fl VecSimQueryParams *queryParams) const override { auto rep = new VecSimQueryReply(this->allocator); this->lastMode = STANDARD_KNN; - if (k == 0 || this->indexLabelCount() == 0) { + auto impl = getImpl(); + if (k == 0 || !impl) { return rep; } + { + const auto label_count = labelCountOf(*impl); + if (label_count == 0) { + return rep; + } + // limit result size to index size + k = std::min(k, label_count); - // limit result size to index size - k = std::min(k, this->indexLabelCount()); - - auto processed_query_ptr = this->preprocessQuery(queryBlob); - const void *processed_query = processed_query_ptr.get(); + auto processed_query_ptr = this->preprocessQuery(queryBlob); + const void *processed_query = processed_query_ptr.get(); - auto query = svs::data::ConstSimpleDataView{ - static_cast(processed_query), 1, this->dim}; - auto result = svs::QueryResult{query.size(), k}; - auto sp = svs_details::joinSearchParams(impl_->get_search_parameters(), queryParams, - is_two_level_lvq); + auto query = svs::data::ConstSimpleDataView{ + static_cast(processed_query), 1, this->dim}; + auto result = svs::QueryResult{query.size(), k}; + auto sp = svs_details::joinSearchParams(impl->get_search_parameters(), queryParams, + is_two_level_lvq); - auto timeoutCtx = queryParams ? queryParams->timeoutCtx : nullptr; - auto cancel = [timeoutCtx]() { return VECSIM_TIMEOUT(timeoutCtx); }; + auto timeoutCtx = queryParams ? queryParams->timeoutCtx : nullptr; + auto cancel = [timeoutCtx]() { return VECSIM_TIMEOUT(timeoutCtx); }; - impl_->search(result.view(), query, sp, cancel); - if (cancel()) { - rep->code = VecSim_QueryReply_TimedOut; - return rep; - } + impl->search(result.view(), query, sp, cancel); + if (cancel()) { + rep->code = VecSim_QueryReply_TimedOut; + return rep; + } - assert(result.n_queries() == 1); + assert(result.n_queries() == 1); - const auto n_neighbors = result.n_neighbors(); - rep->results.reserve(n_neighbors); + const auto n_neighbors = result.n_neighbors(); + rep->results.reserve(n_neighbors); - for (size_t i = 0; i < n_neighbors; i++) { - rep->results.push_back( - VecSimQueryResult{result.index(0, i), toVecSimDistance(result.distance(0, i))}); + for (size_t i = 0; i < n_neighbors; i++) { + if (result.index(0, i) == svs::concurrent::no_external_id) { + continue; + } + rep->results.push_back( + VecSimQueryResult{result.index(0, i), toVecSimDistance(result.distance(0, i))}); + } } // Workaround for VecSim merge_results() that expects results to be sorted // by score, then by id from both indices. @@ -717,61 +781,63 @@ class SVSIndex : public VecSimIndexAbstract, fl VecSimQueryParams *queryParams) const override { auto rep = new VecSimQueryReply(this->allocator); this->lastMode = RANGE_QUERY; - if (radius == 0 || this->indexLabelCount() == 0) { + auto impl = getImpl(); + if (radius == 0 || !impl || labelCountOf(*impl) == 0) { return rep; } + { + auto timeoutCtx = queryParams ? queryParams->timeoutCtx : nullptr; + auto cancel = [timeoutCtx]() { return VECSIM_TIMEOUT(timeoutCtx); }; + + // Prepare query blob for SVS + auto processed_query_ptr = this->preprocessQuery(queryBlob); + const void *processed_query = processed_query_ptr.get(); + std::span query{static_cast(processed_query), + this->dim}; + + // Base search parameters for the SVS iterator schedule. + auto sp = svs_details::joinSearchParams(impl->get_search_parameters(), queryParams, + is_two_level_lvq); + // SVS BatchIterator handles the search in batches + // The batch size is set to the index search window size by default + const size_t batch_size = sp.buffer_config_.get_search_window_size(); + + // Create SVS BatchIterator for range search + // Search result is cached in the iterator and can be accessed by the user + auto svs_it = impl->make_batch_iterator(query); + svs_it.next(batch_size, cancel); + if (cancel()) { + rep->code = VecSim_QueryReply_TimedOut; + return rep; + } - auto timeoutCtx = queryParams ? queryParams->timeoutCtx : nullptr; - auto cancel = [timeoutCtx]() { return VECSIM_TIMEOUT(timeoutCtx); }; - - // Prepare query blob for SVS - auto processed_query_ptr = this->preprocessQuery(queryBlob); - const void *processed_query = processed_query_ptr.get(); - std::span query{static_cast(processed_query), - this->dim}; - - // Base search parameters for the SVS iterator schedule. - auto sp = svs_details::joinSearchParams(impl_->get_search_parameters(), queryParams, - is_two_level_lvq); - // SVS BatchIterator handles the search in batches - // The batch size is set to the index search window size by default - const size_t batch_size = sp.buffer_config_.get_search_window_size(); - - // Create SVS BatchIterator for range search - // Search result is cached in the iterator and can be accessed by the user - auto svs_it = impl_->make_batch_iterator(query); - svs_it.next(batch_size, cancel); - if (cancel()) { - rep->code = VecSim_QueryReply_TimedOut; - return rep; - } - - // range search using epsilon - const auto epsilon = queryParams && queryParams->svsRuntimeParams.epsilon != 0 - ? queryParams->svsRuntimeParams.epsilon - : this->epsilon; - - const auto range_search_boundaries = radius * (1.0 + std::abs(epsilon)); - bool keep_searching = true; - - // Loop while iterator cache is not empty and search radius + epsilon is not exceeded - while (keep_searching && svs_it.size() > 0) { - // Iterate over the cached search results - for (auto &neighbor : svs_it) { - const auto dist = toVecSimDistance(neighbor.distance()); - if (dist <= radius) { - rep->results.push_back(VecSimQueryResult{neighbor.id(), dist}); - } else if (dist > range_search_boundaries) { - keep_searching = false; + // range search using epsilon + const auto epsilon = queryParams && queryParams->svsRuntimeParams.epsilon != 0 + ? queryParams->svsRuntimeParams.epsilon + : this->epsilon; + + const auto range_search_boundaries = radius * (1.0 + std::abs(epsilon)); + bool keep_searching = true; + + // Loop while iterator cache is not empty and search radius + epsilon is not exceeded + while (keep_searching && svs_it.size() > 0) { + // Iterate over the cached search results + for (auto &neighbor : svs_it) { + const auto dist = toVecSimDistance(neighbor.distance()); + if (dist <= radius) { + rep->results.push_back(VecSimQueryResult{neighbor.id(), dist}); + } else if (dist > range_search_boundaries) { + keep_searching = false; + } } - } - // If search radius + epsilon is not exceeded, request SVS BatchIterator for the next - // batch - if (keep_searching) { - svs_it.next(batch_size, cancel); - if (cancel()) { - rep->code = VecSim_QueryReply_TimedOut; - return rep; + // If search radius + epsilon is not exceeded, request SVS BatchIterator for the + // next batch + if (keep_searching) { + svs_it.next(batch_size, cancel); + if (cancel()) { + rep->code = VecSim_QueryReply_TimedOut; + return rep; + } } } } @@ -790,12 +856,13 @@ class SVSIndex : public VecSimIndexAbstract, fl // take ownership of the blob copy and pass it to the batch iterator. auto *queryBlobCopyPtr = queryBlobCopy.release(); // Ownership of queryBlobCopy moves to VecSimBatchIterator that will free it at the end. - if (indexLabelCount() == 0) { + auto impl = getImpl(); + if (!impl || labelCountOf(*impl) == 0) { return new (this->getAllocator()) NullSVS_BatchIterator(queryBlobCopyPtr, queryParams, this->getAllocator()); } else { return new (this->getAllocator()) SVS_BatchIterator( - queryBlobCopyPtr, impl_.get(), queryParams, this->getAllocator(), is_two_level_lvq); + queryBlobCopyPtr, impl, queryParams, this->getAllocator(), is_two_level_lvq); } } @@ -829,15 +896,12 @@ class SVSIndex : public VecSimIndexAbstract, fl } void runGC() override { - if (impl_) { - // There is documentation for consolidate(): - // https://intel.github.io/ScalableVectorSearch/python/dynamic.html#svs.DynamicVamana.consolidate - impl_->consolidate(); - // There is documentation for compact(): - // https://intel.github.io/ScalableVectorSearch/python/dynamic.html#svs.DynamicVamana.compact - impl_->compact(); + if (auto impl = getImpl()) { + // concurrent compact() consolidates every soft-deleted points first + // pending per-label consolidate jobs need not be drained - they become no-ops + impl->compact(); } - num_marked_deleted = 0; + num_marked_deleted.store(0, std::memory_order_relaxed); } private: @@ -847,7 +911,10 @@ class SVSIndex : public VecSimIndexAbstract, fl template void appendStoredDataByLabel(labelType label, std::vector> &vectors_output) const { - if (!impl_) { + // The spans handed to `append_datum` point straight into the dataset, so the instance has + // to stay alive for the whole walk. + auto impl = getImpl(); + if (!impl) { return; } auto append_datum = [&](auto indexed_span) { @@ -867,18 +934,9 @@ class SVSIndex : public VecSimIndexAbstract, fl }; if constexpr (isMulti) { - auto it = impl_->get_label_to_external_lookup().find(label); - if (it == impl_->get_label_to_external_lookup().end()) { - return; - } - for (auto external_id : it->second) { - append_datum(impl_->get_parent_index().get_datum(external_id)); - } + impl->on_data(label, append_datum); } else { - if (!impl_->has_id(label)) { - return; - } - append_datum(impl_->get_datum(label)); + impl->on_datum(label, append_datum); } } diff --git a/src/VecSim/algorithms/svs/svs_extensions.h b/src/VecSim/algorithms/svs/svs_extensions.h index 2e93a6a58..f1734c79f 100644 --- a/src/VecSim/algorithms/svs/svs_extensions.h +++ b/src/VecSim/algorithms/svs/svs_extensions.h @@ -10,10 +10,20 @@ #pragma once #include "VecSim/algorithms/svs/svs_utils.h" #include "svs/extensions/vamana/scalar.h" +// Tells SQDataset that the concurrent SegmentedBlocked tag is a blocked allocator, so +// that resize()/compact() are not constrained away. Without it a scalar-quantized +// concurrent index fails to compile the moment the index grows. +#include "svs/concurrent/extensions/scalar.h" #if HAVE_SVS_LVQ #include SVS_LVQ_HEADER #include SVS_LEANVEC_HEADER +// The same treatment for the LVQ and LeanVec datasets, each of which keeps its own private +// copy of the blocked-allocator trait. LeanVec additionally picks the allocator for its inner +// datasets through that trait, and its blocked branch hardcodes svs::data::Blocked -- so +// without this the inner storage would quietly lose the grow-stable guarantee. +#include "svs/concurrent/extensions/lvq.h" +#include "svs/concurrent/extensions/leanvec.h" #endif // HAVE_SVS_LVQ // Scalar Quantization traits for SVS @@ -21,7 +31,7 @@ template struct SVSStorageTraits { using element_type = std::int8_t; using allocator_type = svs_details::SVSAllocator; - using blocked_type = svs::data::Blocked>; + using blocked_type = svs::concurrent::SegmentedBlocked>; using index_storage_type = svs::quantization::scalar::SQDataset; @@ -32,7 +42,7 @@ struct SVSStorageTraits { // SVS block size is a power of two, so we can use it directly auto svs_bs = svs_details::SVSBlockSize(block_size, element_size(dim)); allocator_type data_allocator{std::move(allocator)}; - return svs::make_blocked_allocator_handle({svs_bs}, data_allocator); + return svs_details::make_segmented_blocked_allocator_handle({svs_bs}, data_allocator); } static constexpr VecSimSvsQuantBits get_compression_mode() { return VecSimSvsQuant_Scalar; } @@ -89,7 +99,7 @@ template struct SVSStorageTraits 1)>> { using allocator_type = svs_details::SVSAllocator; - using blocked_type = svs::data::Blocked>; + using blocked_type = svs::concurrent::SegmentedBlocked>; using strategy_type = typename svs_details::LVQSelector::strategy; using index_storage_type = svs::quantization::lvq::LVQDataset @@ -181,7 +191,7 @@ struct SVSStorageTraits struct SVSStorageTraits { using allocator_type = svs_details::SVSAllocator; - using blocked_type = svs::data::Blocked>; + using blocked_type = svs::concurrent::SegmentedBlocked>; using index_storage_type = svs::leanvec::LeanDataset, svs::leanvec::UsingLVQ, svs::Dynamic, svs::Dynamic, blocked_type>; @@ -212,7 +222,7 @@ struct SVSStorageTraits { auto elem_size = std::max(primary_element_size(dim), secondary_element_size(dim)); auto svs_bs = svs_details::SVSBlockSize(block_size, elem_size); allocator_type data_allocator{std::move(allocator)}; - return svs::make_blocked_allocator_handle({svs_bs}, data_allocator); + return svs_details::make_segmented_blocked_allocator_handle({svs_bs}, data_allocator); } template diff --git a/src/VecSim/algorithms/svs/svs_serializer_impl.h b/src/VecSim/algorithms/svs/svs_serializer_impl.h index 2780d3457..565d904d9 100644 --- a/src/VecSim/algorithms/svs/svs_serializer_impl.h +++ b/src/VecSim/algorithms/svs/svs_serializer_impl.h @@ -76,7 +76,11 @@ template void SVSIndex::impl_save( const std::string &location) { - impl_->save(location + "/config", location + "/graph", location + "/data"); + auto impl = getImpl(); + if (!impl) { + throw std::runtime_error("SVSIndex::impl_save called on an empty index"); + } + impl->save(location + "/config", location + "/graph", location + "/data"); } // This function will load the serialized svs index from the given folder path @@ -92,26 +96,29 @@ void SVSIndex // Verify metadata compatibility, will throw runtime exception if not compatible compareMetadataFile(folder_path + "/metadata"); + std::unique_ptr loaded_impl; if constexpr (isMulti) { - auto loaded = svs::index::vamana::auto_multi_dynamic_assemble( + auto loaded = svs::concurrent::auto_multi_dynamic_assemble( folder_path + "/config", SVS_LAZY(graph_builder_t::load(folder_path + "/graph", this->blockSize, this->buildParams, this->getAllocator())), SVS_LAZY(storage_traits_t::load(folder_path + "/data", this->blockSize, this->dim, this->getAllocator())), distance_f(), std::move(threadpool_handle), - svs::index::vamana::MultiMutableVamanaLoad::FROM_MULTI, logger_); - impl_ = std::make_unique(std::move(loaded)); + svs::concurrent::MultiMutableVamanaLoad::FROM_MULTI, logger_); + loaded_impl = std::make_unique(std::move(loaded)); } else { - auto loaded = svs::index::vamana::auto_dynamic_assemble( + auto loaded = svs::concurrent::auto_dynamic_assemble( folder_path + "/config", SVS_LAZY(graph_builder_t::load(folder_path + "/graph", this->blockSize, this->buildParams, this->getAllocator())), SVS_LAZY(storage_traits_t::load(folder_path + "/data", this->blockSize, this->dim, this->getAllocator())), distance_f(), std::move(threadpool_handle), false, logger_); - impl_ = std::make_unique(std::move(loaded)); + loaded_impl = std::make_unique(std::move(loaded)); } + + delete impl_.exchange(loaded_impl.release(), std::memory_order_release); } template bool SVSIndex::checkIntegrity() const { - if (!impl_) { + auto impl = getImpl(); + if (!impl) { throw std::runtime_error( "SVSIndex integrity check failed: index implementation (impl_) is null."); } @@ -170,9 +178,9 @@ bool SVSIndex try { // SVS internal index integrity validation if constexpr (isMulti) { - impl_->get_parent_index().debug_check_invariants(true); + impl->get_parent_index().debug_check_invariants(true); } else { - impl_->debug_check_invariants(true); + impl->debug_check_invariants(true); } } // debug_check_invariants throws svs::lib::ANNException : public std::runtime_error in case of @@ -182,10 +190,10 @@ bool SVSIndex } try { - size_t index_size = impl_->size(); - size_t storage_size = impl_->view_data().size(); - size_t capacity = storage_traits_t::storage_capacity(impl_->view_data()); - size_t label_count = this->indexLabelCount(); + size_t index_size = impl->size(); + size_t storage_size = impl->view_data().size(); + size_t capacity = storage_traits_t::storage_capacity(impl->view_data()); + size_t label_count = labelCountOf(*impl); // Storage size must match index size if (storage_size != index_size) { @@ -203,7 +211,7 @@ bool SVSIndex bool label_validation_passed = true; try { - impl_->on_ids([&](size_t label) { labels_counted++; }); + impl->on_ids([&](size_t label) { labels_counted++; }); // Validate label count consistency label_validation_passed = (labels_counted == label_count); diff --git a/src/VecSim/algorithms/svs/svs_tiered.h b/src/VecSim/algorithms/svs/svs_tiered.h index e83bccf60..ee66005e0 100644 --- a/src/VecSim/algorithms/svs/svs_tiered.h +++ b/src/VecSim/algorithms/svs/svs_tiered.h @@ -5,14 +5,50 @@ #include "VecSim/algorithms/svs/svs.h" #include "VecSim/index_factories/svs_factory.h" +#include +#include #include #include #include #include #include #include +#include +#include #include +/** + * Definition of a job that inserts a new vector from flat into SVS Index. + */ +struct SVSInsertJob : public TieredInsertJob { + // Pending -> Executing -> Done. Delayed means a relabel is remapping the label in the backend + // right now, and publishing would put the vector under a label that is not settled yet. + // Parked means a worker consumed that submission; relabel must submit it again when settled. + // Written under flatIndexGuard, except Executing/Done which waitForJob() reads unlocked. + enum class Status : uint8_t { Pending, Delayed, Parked, Executing, Done }; + std::atomic status{Status::Pending}; + + SVSInsertJob(std::shared_ptr allocator, labelType label_, idType id_, + JobCallback insertCb, VecSimIndex *index_) + : TieredInsertJob(allocator, SVS_INSERT_VECTOR_JOB, label_, id_, insertCb, index_) {} +}; + +/** + * Definition of a job that launches partial consolidation on SVS Index. + */ +struct SVSConsolidateJob : public AsyncJob { + std::vector labels; + // Pending -> Executing -> Done, or Pending -> TakenOver when relabelVector() consolidates these + // labels itself. + enum class Status : uint8_t { Pending, TakenOver, Executing, Done }; + std::atomic status{Status::Pending}; + + SVSConsolidateJob(std::shared_ptr allocator, + const std::vector &labels_, JobCallback insertCb, + VecSimIndex *index_) + : AsyncJob(allocator, SVS_CONSOLIDATE_JOB, insertCb, index_), labels(labels_) {} +}; + /** * @class SVSMultiThreadJob * @brief Represents a multi-threaded asynchronous job for the SVS algorithm. @@ -219,29 +255,13 @@ class SVSMultiThreadJob : public AsyncJob { template class TieredSVSIndex : public VecSimTieredIndex { + using DistType = float; using Self = TieredSVSIndex; - using Base = VecSimTieredIndex; - using flat_index_t = BruteForceIndex; - using backend_index_t = VecSimIndexAbstract; + using Base = VecSimTieredIndex; + using flat_index_t = BruteForceIndex; + using backend_index_t = VecSimIndexAbstract; using svs_index_t = SVSIndexBase; - // swaps_journal is used by updateSVSIndex() to track vectors swap operations that were done in - // the Flat index during SVS index updating. - // The journal contains tuples of (label, oldId, newId). - // oldId is the index of the label in flat index before the swap. - // newId is the index of the label in flat index after the swap. - // if oldId == newId, it means that the vector was not moved in the Flat index, but was removed - // from the end of the Flat index (no id swaps occurred internally). - // if label == SKIP_LABEL, it means that the vector was not moved in the Flat index, but - // updated in-place (hence was already removed and no need to remove it again) - using swap_record = std::tuple; - constexpr static size_t SKIP_LABEL = std::numeric_limits::max(); - std::vector swaps_journal; - // deleted_labels_journal is used by updateSVSIndex() to track vectors that were deleted from - // Flat index during SVS index updating. The journal contains the deleted labels. These labels - // are used to delete the same vectors from the SVS index at the end of the update. - std::vector deleted_labels_journal; - size_t trainingTriggerThreshold; size_t updateTriggerThreshold; size_t updateJobWaitTime; @@ -255,17 +275,40 @@ class TieredSVSIndex : public VecSimTieredIndex { std::atomic defragTimeNs{0}; // Used to prevent running multiple index update jobs in parallel. // Even if update jobs scheduled sequentially, they can be started in parallel. - mutable std::mutex updateJobMutex; + mutable std::shared_mutex updateJobMutex; // The reason of following container just to properly destroy jobs which not executed yet SVSMultiThreadJob::JobsRegistry uncompletedJobs; + // frontend ids holding training data for the backend initialization + std::unordered_set ids_to_init_; + + vecsim_stl::unordered_map> + labelToConsolidateJobs; + mutable std::shared_mutex consolidateJobsGuard; + + // trade main-thread latency for bounded growth during active init. + // The effective limit is max(this limit, trainingTriggerThreshold). + // So this value should be ~ max(trainingTriggerThreshold) ~ 1e5 + static constexpr size_t INIT_BACKPRESSURE_LIMIT = 100'000; + size_t flat_buffer_bound; + + // flatIndexGuard: init (serialized by updateJobMutex) writes under a shared lock. + // Every other access to these two fields must hold the lock exclusively. + bool initInProgress{false}; + std::vector initDeletedLabels; + + // Deferred execution list; labelToInsertJobs/invalidJobs retain ownership of these jobs. + // Access under flatIndexGuard exclusively. Drain after releasing updateJobMutex so + // inserts arriving during initialization do not occupy workers waiting for the build. + vecsim_stl::vector initInsertJobs; + /// //////////////////////////////////////////////////////////////////////////////////////////////////// // TieredSVS_BatchIterator // //////////////////////////////////////////////////////////////////////////////////////////////////// - class TieredSVS_BatchIterator : public VecSimBatchIterator { + class TieredSVS_BatchIterator : public TieredIndex_BatchIterator { // Defining spacial values for the svs_iterator field, to indicate if the iterator is // uninitialized or depleted when we don't have a valid iterator. static constexpr VecSimBatchIterator *depleted() { @@ -278,78 +321,11 @@ class TieredSVSIndex : public VecSimTieredIndex { const Index *index; VecSimQueryParams *queryParams; - VecSimQueryResultContainer flat_results; - VecSimQueryResultContainer svs_results; - VecSimBatchIterator *flat_iterator; VecSimBatchIterator *svs_iterator; - std::shared_lock svs_lock; - - // On single value indices, this set holds the IDs of the results that were returned from - // the flat buffer. - // On multi value indices, this set holds the IDs of all the results that were returned. - // The difference between the two cases is that on multi value indices, the same ID can - // appear in both indexes and results with different scores, and therefore we can't tell in - // advance when we expect a possibility of a duplicate. - // On single value indices, a duplicate may appear at the same batch (and we will handle it - // when merging the results) Or it may appear in a different batches, first from the flat - // buffer and then from the SVS, in the cases where a better result if found later in SVS - // because of the approximate nature of the algorithm. - vecsim_stl::unordered_set returned_results_set; - - VecSimQueryReply *compute_current_batch(size_t n_res, bool isMultiValue) { - // Merge results - // This call will update `svs_res` and `bf_res` to point to the end of the merged - // results. - auto batch_res = new VecSimQueryReply(allocator); - // VecSim and SVS distance computation is implemented differently, so we always have to - // merge results with set. - auto [from_svs, from_flat] = - merge_results(batch_res->results, svs_results, flat_results, n_res); - - if (!isMultiValue) { - // If we're on a single-value index, update the set of results returned from the - // FLAT index before popping them, to prevent them to be returned from the SVS index - // in later batches. - for (size_t i = 0; i < from_flat; ++i) { - this->returned_results_set.insert(this->flat_results[i].id); - } - } else { - // If we're on a multi-value index, update the set of results returned (from - // `batch_res`) - for (size_t i = 0; i < batch_res->results.size(); ++i) { - this->returned_results_set.insert(batch_res->results[i].id); - } - } - - // Update results - flat_results.erase(flat_results.begin(), flat_results.begin() + from_flat); - svs_results.erase(svs_results.begin(), svs_results.begin() + from_svs); - - // clean up the results - // On multi-value indexes, one (or both) results lists may contain results that are - // already returned form the other list (with a different score). We need to filter them - // out. - if (isMultiValue) { - filter_irrelevant_results(this->flat_results); - filter_irrelevant_results(this->svs_results); - } - - // Return current batch - return batch_res; - } - - void filter_irrelevant_results(VecSimQueryResultContainer &results) { - // Filter out results that were already returned. - const auto it = std::remove_if(results.begin(), results.end(), [this](const auto &r) { - return returned_results_set.count(r.id) != 0; - }); - results.erase(it, results.end()); - } void acquire_svs_iterator() { assert(svs_iterator == nullptr); - this->index->mainIndexGuard.lock_shared(); svs_iterator = index->backendIndex->newBatchIterator( this->flat_iterator->getQueryBlob(), queryParams); } @@ -358,7 +334,6 @@ class TieredSVSIndex : public VecSimTieredIndex { if (svs_iterator != nullptr && svs_iterator != depleted()) { delete svs_iterator; svs_iterator = nullptr; - this->index->mainIndexGuard.unlock_shared(); } } @@ -379,12 +354,11 @@ class TieredSVSIndex : public VecSimTieredIndex { // own copies: flat_iterator copy is created during TieredSVS_BatchIterator // construction When TieredSVS_BatchIterator::getNextResults() is called and // svs_iterator is not initialized, it retrieves the blob from flat_iterator - : VecSimBatchIterator(nullptr, queryParams ? queryParams->timeoutCtx : nullptr, - std::move(allocator)), - index(index), flat_results(this->allocator), svs_results(this->allocator), + : TieredIndex_BatchIterator(nullptr, queryParams ? queryParams->timeoutCtx : nullptr, + std::move(allocator)), + index(index), flat_iterator(index->frontendIndex->newBatchIterator(query_vector, queryParams)), - svs_iterator(nullptr), svs_lock(index->mainIndexGuard, std::defer_lock), - returned_results_set(this->allocator) { + svs_iterator(nullptr) { if (queryParams) { this->queryParams = (VecSimQueryParams *)this->allocator->allocate(sizeof(VecSimQueryParams)); @@ -405,7 +379,6 @@ class TieredSVSIndex : public VecSimTieredIndex { VecSimQueryReply *getNextResults(size_t n_res, VecSimQueryReply_Order order) override { auto svs_code = VecSim_QueryReply_OK; - const bool isMulti = this->index->backendIndex->isMultiValue(); if (svs_iterator == nullptr) { // first call // First call to getNextResults. The call to the BF iterator will include // calculating all the distances and access the BF index. We take the lock on this @@ -420,12 +393,10 @@ class TieredSVSIndex : public VecSimTieredIndex { } flat_results.swap(cur_flat_results->results); VecSimQueryReply_Free(cur_flat_results); - // We also take the lock on the main index on the first call to getNextResults, and - // we hold it until the iterator is depleted or freed. acquire_svs_iterator(); auto cur_svs_results = svs_iterator->getNextResults(n_res, BY_SCORE_THEN_ID); svs_code = cur_svs_results->code; - svs_results.swap(cur_svs_results->results); + backend_results.swap(cur_svs_results->results); VecSimQueryReply_Free(cur_svs_results); handle_svs_depletion(); } else { @@ -436,33 +407,24 @@ class TieredSVSIndex : public VecSimTieredIndex { tail->results.end()); VecSimQueryReply_Free(tail); - if (!isMulti) { - // On single-value indexes, duplicates will never appear in the hnsw results - // before they appear in the flat results (at the same time or later if the - // approximation misses) so we don't need to try and filter the flat results - // (and recheck conditions). - break; - } else { - // On multi-value indexes, the flat results may contain results that are - // already returned from the hnsw index. We need to filter them out. - filter_irrelevant_results(this->flat_results); - } + // The flat results may contain labels already returned from the SVS index. + filter_irrelevant_results(this->flat_results); } - while (svs_results.size() < n_res && svs_iterator != depleted() && + while (backend_results.size() < n_res && svs_iterator != depleted() && svs_code == VecSim_OK) { - auto tail = - svs_iterator->getNextResults(n_res - svs_results.size(), BY_SCORE_THEN_ID); + auto tail = svs_iterator->getNextResults(n_res - backend_results.size(), + BY_SCORE_THEN_ID); svs_code = tail->code; // Set the svs_results code to the last `getNextResults` code. // New batch may contain better results than the previous batch, so we need to // merge. We don't expect duplications (hence the ), as the iterator // guarantees that no result is returned twice. VecSimQueryResultContainer cur_svs_results(this->allocator); - merge_results(cur_svs_results, svs_results, tail->results, n_res); + merge_results(cur_svs_results, backend_results, tail->results, n_res); VecSimQueryReply_Free(tail); - svs_results.swap(cur_svs_results); - filter_irrelevant_results(svs_results); + backend_results.swap(cur_svs_results); + filter_irrelevant_results(backend_results); handle_svs_depletion(); } } @@ -472,7 +434,8 @@ class TieredSVSIndex : public VecSimTieredIndex { } VecSimQueryReply *batch; - batch = compute_current_batch(n_res, isMulti); + // In concurent execution we can observe a vector duplication in backend and frontend + batch = compute_current_batch(n_res); if (order == BY_ID) { sort_results_by_id(batch); @@ -490,7 +453,7 @@ class TieredSVSIndex : public VecSimTieredIndex { // The next call to `getNextResults` will return an empty batch, and then the iterators will // correctly report that they are depleted. bool isDepleted() override { - return flat_results.empty() && flat_iterator->isDepleted() && svs_results.empty() && + return flat_results.empty() && flat_iterator->isDepleted() && backend_results.empty() && svs_iterator == depleted(); } @@ -500,7 +463,7 @@ class TieredSVSIndex : public VecSimTieredIndex { flat_iterator->reset(); svs_iterator = nullptr; flat_results.clear(); - svs_results.clear(); + backend_results.clear(); returned_results_set.clear(); } }; @@ -541,7 +504,6 @@ class TieredSVSIndex : public VecSimTieredIndex { } size_t indexMetaDataCapacity() const override { std::shared_lock flat_lock(this->flatIndexGuard); - std::shared_lock main_lock(this->mainIndexGuard); return this->frontendIndex->indexMetaDataCapacity() + this->backendIndex->indexMetaDataCapacity(); } @@ -551,37 +513,94 @@ class TieredSVSIndex : public VecSimTieredIndex { } #endif +protected: + ScopedLocks lockMainIndexForQuery() const override { +#ifdef BUILD_TESTS + executeTracingCallback("Query::before_backend"); +#endif + // No-op: SVS does its query locking internally. + return ScopedLocks(); + } + + ScopedLocks lockIndexForSize() const override { return ScopedLocks(this->flatIndexLockable); } + + ScopedLocks lockIndexForCapacity() const override { + return ScopedLocks(this->flatIndexLockable); + } + private: /** - * @brief Updates the SVS index in a thread-safe manner. + * @brief Init the SVS index in a thread-safe manner. * * This static wrapper function performs the following actions: * - Acquires a lock on the index's updateJobMutex to prevent concurrent updates. - * - Clears the indexUpdateScheduled flag to allow future scheduling. * - Configures the number of threads for the underlying SVS index update operation. - * - Calls the updateSVSIndex method to perform the actual index update. + * - Calls the initSVSIndex method to perform the actual index update. + * - Clears the indexUpdateScheduled flag to allow future scheduling. * * @param idx Pointer to the VecSimIndex to be updated. * @param availableThreads The number of threads available for the update operation. Current * thread us used as well, so the minimal value is 1. */ - static void updateSVSIndexWrapper(VecSimIndex *idx, size_t availableThreads) { + static void initSVSIndexWrapper(VecSimIndex *idx, size_t availableThreads) { assert(availableThreads > 0); auto index = static_cast *>(idx); assert(index); - // prevent parallel updates - std::lock_guard lock(index->updateJobMutex); - // Release the scheduled flag to allow scheduling again - index->indexUpdateScheduled.clear(); - // Update the SVS index - index->updateSVSIndex(availableThreads); + std::vector deleted_labels; + vecsim_stl::vector insert_jobs(index->allocator); + { + std::lock_guard lock(index->updateJobMutex); + deleted_labels = index->initSVSIndex(availableThreads); + std::lock_guard flat_lock{index->flatIndexGuard}; + insert_jobs.swap(index->initInsertJobs); + } + for (labelType label : deleted_labels) { + index->scheduleSVSIndexConsolidate(label); + } + // Include invalidated jobs: their callbacks remove them from invalidJobs and delete them. + // Catch-up is synchronous: the wrapper returns only after this parallel pass finishes. + // With availableThreads == 1, including in-place initialization, the caller runs it. + // Init resets per-operation parallelism to 1. Each insert's shared updateJobMutex + // excludes multi-threaded GC, so nested single-vector SVS work runs inline without + // renting more workers. runParallel supplies only the outer job-level parallelism. + if (!insert_jobs.empty()) { + const size_t num_jobs = insert_jobs.size(); + std::atomic next_job{0}; + index->GetSVSIndex()->runParallel( + [&](size_t) { + while (true) { + const size_t i = next_job.fetch_add(1, std::memory_order_relaxed); + if (i >= num_jobs) { + return; + } + executeInsertJobWrapper(insert_jobs[i]); + } + }, + std::min(availableThreads, num_jobs)); + } + } + + static void executeInsertJobWrapper(AsyncJob *job) { + auto *insert_job = static_cast(job); + auto *job_index = static_cast *>(insert_job->index); + InsertJobOutcome outcome; + { + // prevent parallel execution with index initilizing job + std::shared_lock lock(job_index->updateJobMutex); + outcome = job_index->executeInsertJob(insert_job); + } + + if (outcome == InsertJobOutcome::Completed) { + delete job; + } } /** * @brief Run SVS index GC in a thread-safe manner. * * This static wrapper function performs the following actions: - * - Acquires a lock on the index's mainIndexGuard to ensure thread safety during the GC + * - Acquires updateJobMutex exclusively to wait for backend update jobs. + * - Acquires mainIndexGuard exclusively to exclude base-class readers during GC. * - Configures the number of threads for the underlying SVS index update operation. * - Calls the SVSIndex::runGC() method to perform the actual index update. * - Clears the indexGCScheduled flag to allow future scheduling. @@ -598,9 +617,6 @@ class TieredSVSIndex : public VecSimTieredIndex { assert(index); const auto start = std::chrono::steady_clock::now(); - std::lock_guard lock{index->mainIndexGuard}; - // Release the scheduled flag to allow scheduling again - index->indexGCScheduled.clear(); // Do SVS index GC index->backendIndex->log(VecSimCommonStrings::LOG_VERBOSE_STRING, @@ -608,17 +624,54 @@ class TieredSVSIndex : public VecSimTieredIndex { auto svs_index = index->GetSVSIndex(); if (index->backendIndex->indexSize() == 0) { // No need to run GC on an empty index. + index->indexGCScheduled.clear(); return; } index->executeTracingCallback("GCJob::before_run_gc"); + // Take updateJobMutex first: workers may need flatIndexGuard while a base-class reader + // holds it and waits for mainIndexGuard. + std::lock_guard lock(index->updateJobMutex); + std::lock_guard main_lock(index->mainIndexGuard); + + // Release the scheduled flag to allow scheduling again + index->indexGCScheduled.clear(); + svs_index->setParallelism(std::min(availableThreads, index->backendIndex->indexSize())); // VecSimIndexAbstract::runGC() is protected static_cast(index->backendIndex)->runGC(); + svs_index->setParallelism(1); index->recordDefrag(start); } - // Account for one backend GC run that started at `start`. The time includes the wait for the - // main index guard, which is part of what GC costs, as in the tiered HNSW defrag(). + static void SVSIndexConsolidateWrapper(AsyncJob *job) { + auto consolidate_job = static_cast(job); + auto index = static_cast *>(consolidate_job->index); + + std::shared_lock lock(index->updateJobMutex); + auto svs_index = index->GetSVSIndex(); + svs_index->setParallelism(1); + + auto expected = SVSConsolidateJob::Status::Pending; + if (consolidate_job->status.compare_exchange_strong( + expected, SVSConsolidateJob::Status::Executing, std::memory_order_acq_rel)) { + svs_index->consolidate(consolidate_job->labels); +#ifdef BUILD_TESTS + index->executeTracingCallback("ConsolidateJob::before_done"); +#endif + // Published before forgetConsolidateJob() needs the registry lock, so a waiter holding + // it shared is released before this job can be deleted. + consolidate_job->status.store(SVSConsolidateJob::Status::Done, + std::memory_order_release); + } + index->forgetConsolidateJob(consolidate_job); + delete job; +#ifdef BUILD_TESTS + index->executeTracingCallback("ConsolidateJob::after_delete"); +#endif + } + + // Account for one backend GC run that started at `start`. The time includes waits for + // mainIndexGuard and, in async mode, updateJobMutex. void recordDefrag(std::chrono::steady_clock::time_point start) { const auto elapsedNs = std::chrono::duration_cast( std::chrono::steady_clock::now() - start) @@ -633,14 +686,15 @@ class TieredSVSIndex : public VecSimTieredIndex { #ifdef BUILD_TESTS public: #endif - void scheduleSVSIndexUpdate() { + + void scheduleSVSIndexInit() { // do not schedule if scheduled already if (indexUpdateScheduled.test_and_set()) { return; } auto jobs = SVSMultiThreadJob::createScheduledJobs( - this->allocator, SVS_BATCH_UPDATE_JOB, updateSVSIndexWrapper, this, + this->allocator, SVS_BATCH_UPDATE_JOB, initSVSIndexWrapper, this, std::chrono::microseconds(updateJobWaitTime), &uncompletedJobs); this->submitJobs(jobs); } @@ -657,128 +711,215 @@ class TieredSVSIndex : public VecSimTieredIndex { this->submitJobs(jobs); } + void scheduleSVSIndexConsolidate(labelType label) { + auto *new_consolidate_job = new (this->allocator) + SVSConsolidateJob(this->allocator, {label}, SVSIndexConsolidateWrapper, this); + + { + std::lock_guard lock(this->consolidateJobsGuard); + auto it = this->labelToConsolidateJobs.find(label); + if (it != this->labelToConsolidateJobs.end()) { + it->second.push_back(new_consolidate_job); + } else { + vecsim_stl::vector jobs(1, new_consolidate_job, + this->allocator); + this->labelToConsolidateJobs.insert({label, std::move(jobs)}); + } + } + // Insert job to the queue. + this->submitSingleJob(new_consolidate_job); + } + private: - static void applySwapsToLabelsArray(std::vector &labels, - const std::vector &swaps) { - // Enumerate journal and reflect swaps in the labels. - // The journal contains tuples of (label, oldId, newId). - // oldId is the index of the label in flat index before the swap. - // newId is the index of the label in flat index after the swap. - for (const auto &p : swaps) { - auto oldId = std::get<1>(p); - auto newId = std::get<2>(p); - - if (oldId == newId || oldId >= labels.size()) { - // If oldId == newId, it means that the vector was not moved in the Flat index, - // but was removed or updated in-place. - // If oldId is out of bounds - new vector was added and swapped meanwhile. - // In both cases, we should not touch flat index at the new position. - if (newId < labels.size()) { - labels[newId] = SKIP_LABEL; - } - continue; // Next swap record. + void forgetConsolidateJob(SVSConsolidateJob *job) { + std::lock_guard lock(this->consolidateJobsGuard); + for (auto label : job->labels) { + auto it = this->labelToConsolidateJobs.find(label); + if (it == this->labelToConsolidateJobs.end()) { + continue; } + auto &jobs = it->second; + jobs.erase(std::remove(jobs.begin(), jobs.end(), job), jobs.end()); + if (jobs.empty()) { + this->labelToConsolidateJobs.erase(it); + } + } + } + std::vector takeOverConsolidateOf(labelType label) { + std::vector taken_over; + std::shared_lock lock(this->consolidateJobsGuard); + auto it = this->labelToConsolidateJobs.find(label); + if (it == this->labelToConsolidateJobs.end()) { + return taken_over; + } + for (auto *job : it->second) { + auto expected = SVSConsolidateJob::Status::Pending; + if (job->status.compare_exchange_strong(expected, SVSConsolidateJob::Status::TakenOver, + std::memory_order_acq_rel)) { + taken_over.insert(taken_over.end(), job->labels.begin(), job->labels.end()); + } else if (expected == SVSConsolidateJob::Status::Executing) { +#ifdef BUILD_TESTS + executeTracingCallback("Relabel::before_wait_for_consolidate"); +#endif + waitForJob(job); // its own consolidate covers these labels + } + // already TakenOver or Done + } + return taken_over; + } + + // Wait until the job leaves its publish window + // Safe to call while holding flatIndexGuard exclusive + static void waitForJob(SVSInsertJob *job) { + while (job->status.load(std::memory_order_acquire) == SVSInsertJob::Status::Executing) { + std::this_thread::yield(); + } + } - // Real swap case. - // If oldId != newId and oldId is in bounds, it means that the tracked vector was moved - // So, move the label to new position and skip the old position from deletions. - // NOTE: labels[oldId] can be SKIP_LABEL, but we still need to move it to newId. - labels[newId] = labels[oldId]; - labels[oldId] = SKIP_LABEL; + // Safe to call while holding consolidateJobsGuard shared + static void waitForJob(SVSConsolidateJob *job) { + while (job->status.load(std::memory_order_acquire) == + SVSConsolidateJob::Status::Executing) { + std::this_thread::yield(); } } - void updateSVSIndex(size_t availableThreads) { + idType setAndSaveInvalidJob(AsyncJob *job) override { + waitForJob(static_cast(job)); + return Base::setAndSaveInvalidJob(job); + } + + enum class InsertJobOutcome { Completed, Parked }; + + InsertJobOutcome executeInsertJob(SVSInsertJob *job) { + auto svs_index = GetSVSIndex(); + + // Note that accessing the job fields should occur with flat index guard held (here and + // later). + this->flatIndexGuard.lock_shared(); + if (!job->isValid) { + this->flatIndexGuard.unlock_shared(); + // Job has been invalidated in the meantime - nothing to execute, and remove it from the + // lookup. + this->invalidJobsLookupGuard.lock(); + this->invalidJobs.erase(job->id); + this->invalidJobsLookupGuard.unlock(); + return InsertJobOutcome::Completed; + } + + if (job->status.load(std::memory_order_relaxed) == SVSInsertJob::Status::Delayed) { + job->status.store(SVSInsertJob::Status::Parked, std::memory_order_relaxed); + this->flatIndexGuard.unlock_shared(); + // Relabel owns resubmission. It may already have resumed the job after this unlock. + return InsertJobOutcome::Parked; + } + + // Copy the vector blob out of the flat buffer while holding flatIndexGuard, so we + // can release the flat lock before indexing into the SVS backend + size_t data_size = this->frontendIndex->getStoredDataSize(); + auto blob_copy = this->getAllocator()->allocate_unique(data_size); + memcpy(blob_copy.get(), this->frontendIndex->getDataByInternalId(job->id), data_size); + + job->status.store(SVSInsertJob::Status::Executing, std::memory_order_release); + this->flatIndexGuard.unlock_shared(); + + svs_index->addVector(blob_copy.get(), job->label); + + job->status.store(SVSInsertJob::Status::Done, std::memory_order_release); + // Remove the vector and the insert job from the flat buffer. + this->removeIngestedVectorFromFlat(job); + return InsertJobOutcome::Completed; + } + + std::vector initSVSIndex(size_t availableThreads) { + std::vector ids_to_move; std::vector labels_to_move; std::vector vectors_to_move; - std::vector deleted_labels_during_update; + std::vector deleted_labels; + std::unique_ptr impl; - { // lock frontendIndex from modifications + executeTracingCallback("UpdateJob::before_add_to_svs"); + { std::shared_lock flat_lock{this->flatIndexGuard}; auto flat_index = this->GetFlatIndex(); - const auto frontend_index_size = this->frontendIndex->indexSize(); + const auto init_batch_size = ids_to_init_.size(); + if (init_batch_size == 0) { + indexUpdateScheduled.clear(); + return {}; + } const size_t dim = flat_index->getDim(); - labels_to_move.reserve(frontend_index_size); - vectors_to_move.reserve(frontend_index_size * dim); - - for (idType i = 0; i < frontend_index_size; i++) { - labels_to_move.push_back(flat_index->getVectorLabel(i)); - auto data = flat_index->getDataByInternalId(i); + ids_to_move.reserve(init_batch_size); + labels_to_move.reserve(init_batch_size); + vectors_to_move.reserve(init_batch_size * dim); + initDeletedLabels.reserve(init_batch_size); + + for (idType id : ids_to_init_) { + ids_to_move.push_back(id); + labels_to_move.push_back(flat_index->getVectorLabel(id)); + auto data = flat_index->getDataByInternalId(id); vectors_to_move.insert(vectors_to_move.end(), data, data + dim); } - // reset journal to the current frontend index state - swaps_journal.clear(); - deleted_labels_journal.clear(); - } // release frontend index - - executeTracingCallback("UpdateJob::before_add_to_svs"); - if (!labels_to_move.empty()) { - // lock backend index for writing and add vectors there - std::shared_lock main_shared_lock(this->mainIndexGuard); - auto svs_index = GetSVSIndex(); - assert(labels_to_move.size() == vectors_to_move.size() / this->frontendIndex->getDim()); - if (this->backendIndex->indexSize() == 0) { - // If backend index is empty, we need to initialize it first. + // Keep tracking surviving snapshot IDs; new arrivals use insert jobs. + initInProgress = true; + } + executeTracingCallback("UpdateJob::after_snapshot"); + { + // Nothing to initialize from an empty batch, and setParallelism(0) is not a + // valid request against the shared pool. + if (!labels_to_move.empty()) { + auto svs_index = GetSVSIndex(); svs_index->setParallelism(std::min(availableThreads, labels_to_move.size())); - auto impl = svs_index->createImpl(vectors_to_move.data(), labels_to_move.data(), - labels_to_move.size()); + assert(labels_to_move.size() == + vectors_to_move.size() / this->frontendIndex->getDim()); + assert(!svs_index->ready()); - // Upgrade to unique lock to set the new impl - main_shared_lock.unlock(); - std::lock_guard lock(this->mainIndexGuard); - svs_index->setImpl(std::move(impl)); - } else { - // Backend index is initialized - just add the vectors. - main_shared_lock.unlock(); - std::lock_guard lock(this->mainIndexGuard); - // Upgrade to unique lock to add vectors - svs_index->setParallelism(std::min(availableThreads, labels_to_move.size())); - svs_index->addVectors(vectors_to_move.data(), labels_to_move.data(), - labels_to_move.size()); + impl = svs_index->createImpl(vectors_to_move.data(), labels_to_move.data(), + labels_to_move.size()); + svs_index->setParallelism(1); } } - executeTracingCallback("UpdateJob::after_add_to_svs"); - // clean-up frontend index - { // lock frontend index for writing and delete moved vectors - std::lock_guard lock(this->flatIndexGuard); - - // swap deleted labels journal with the local variable to track deleted labels during - // update - std::swap(deleted_labels_during_update, deleted_labels_journal); - - // Apply swaps from journal to labels_to_move to reflect changes made in meanwhile. - applySwapsToLabelsArray(labels_to_move, this->swaps_journal); - - // delete vectors from the frontend index in reverse order - // it increases the chance of avoiding swaps in the frontend index and performance - // improvement - int deleted = 0; - idType id = labels_to_move.size(); - while (id-- > 0) { - auto label = labels_to_move[id]; + executeTracingCallback("UpdateJob::before_publish"); + { + std::lock_guard flat_lock{this->flatIndexGuard}; + deleted_labels.swap(initDeletedLabels); + if (!deleted_labels.empty()) { + impl->deleteVectors(deleted_labels.data(), deleted_labels.size()); + } + // Deletions can swap surviving rows; cleanup needs their current IDs. + ids_to_move.assign(ids_to_init_.begin(), ids_to_init_.end()); + ids_to_init_.clear(); + std::sort(ids_to_move.begin(), ids_to_move.end()); + [[maybe_unused]] size_t total_deleted = 0; + for (auto it = ids_to_move.rbegin(); it != ids_to_move.rend(); ++it) { + idType id = *it; + auto label = this->frontendIndex->getVectorLabel(id); // Delete the vector from the frontend index if not in-place updated. - if (label != SKIP_LABEL) { - deleted += this->frontendIndex->deleteVectorById(label, id); + labelType last_vec_label = + this->frontendIndex->getVectorLabel(this->frontendIndex->indexSize() - 1); + int deleted = this->frontendIndex->deleteVectorById(label, id); + if (deleted && id != this->frontendIndex->indexSize()) { + // If the vector removal caused a swap with the last id, update the relevant + // insert job. + this->updateInsertJobInternalId(this->frontendIndex->indexSize(), id, + last_vec_label); } + total_deleted += deleted; } - assert(deleted == std::count_if(labels_to_move.begin(), labels_to_move.end(), - [](labelType label) { return label != SKIP_LABEL; }) && + + assert(total_deleted == ids_to_move.size() && "Deleted vectors count does not match the number of labels to delete"); - } - // delete vectors from backend index that were deleted from the frontend index during - // the update process. - { - std::lock_guard main_lock(this->mainIndexGuard); - std::sort(deleted_labels_during_update.begin(), deleted_labels_during_update.end()); - auto it = std::unique(deleted_labels_during_update.begin(), - deleted_labels_during_update.end()); - deleted_labels_during_update.erase(it, deleted_labels_during_update.end()); - auto svs_index = GetSVSIndex(); - svs_index->deleteVectors(deleted_labels_during_update.data(), - deleted_labels_during_update.size()); - } + // Publish only after applying writes made while the snapshot was built. + if (impl) { + GetSVSIndex()->setImpl(std::move(impl)); + executeTracingCallback("UpdateJob::after_publish"); + } + initInProgress = false; + indexUpdateScheduled.clear(); + } // release frontend index + executeTracingCallback("UpdateJob::after_add_to_svs"); + return deleted_labels; } public: @@ -786,13 +927,14 @@ class TieredSVSIndex : public VecSimTieredIndex { const TieredIndexParams &tiered_index_params, std::shared_ptr allocator) : Base(svs_index, bf_index, tiered_index_params, allocator), - uncompletedJobs(this->allocator) { + uncompletedJobs(this->allocator), labelToConsolidateJobs(this->allocator), + initInsertJobs(this->allocator) { const auto &tiered_svs_params = tiered_index_params.specificParams.tieredSVSParams; // If flatBufferLimit is not initialized (0), use the default update threshold. - const size_t flat_buffer_bound = tiered_index_params.flatBufferLimit == 0 - ? SVS_VAMANA_DEFAULT_UPDATE_THRESHOLD - : tiered_index_params.flatBufferLimit; + flat_buffer_bound = tiered_index_params.flatBufferLimit == 0 + ? SVS_VAMANA_DEFAULT_UPDATE_THRESHOLD + : tiered_index_params.flatBufferLimit; this->updateTriggerThreshold = tiered_svs_params.updateTriggerThreshold == 0 @@ -812,184 +954,273 @@ class TieredSVSIndex : public VecSimTieredIndex { this->updateJobWaitTime = tiered_svs_params.updateJobWaitTime == 0 ? SVS_DEFAULT_UPDATE_JOB_WAIT_TIME : tiered_svs_params.updateJobWaitTime; + } - // Reserve space for the journal to avoid reallocation. - this->swaps_journal.reserve(this->trainingTriggerThreshold); + ~TieredSVSIndex() { + // Delete all the pending consolidate jobs + for (auto &jobs : this->labelToConsolidateJobs) { + for (auto *job : jobs.second) { + delete job; + } + } } int addVector(const void *blob, labelType label) override { int ret = 0; auto svs_index = GetSVSIndex(); - size_t update_threshold = 0; - size_t frontend_index_size = 0; // In-Place mode - add vector syncronously to the backend index. if (this->getWriteMode() == VecSim_WriteInPlace) { - // It is ok to lock everything at once for in-place mode, - // but we will have to unlock averything before calling updateSVSIndexWrapper() - // so make the minimal needed lock here. - std::shared_lock backend_shared_lock(this->mainIndexGuard); // Backend index initialization data have to be buffered for proper // compression/training. - if (this->backendIndex->indexSize() == 0) { - // If backend index size is 0, first collect vectors in frontend index - // lock in scope to ensure that these will be released before - // updateSVSIndexWrapper() is called. - { - std::lock_guard lock(this->flatIndexGuard); - ret = this->frontendIndex->addVector(blob, label); - // If frontend size exceeds the update job threshold, ... - frontend_index_size = this->frontendIndex->indexSize(); - } - // ... move vectors to the backend index. - if (frontend_index_size >= this->trainingTriggerThreshold) { - // updateSVSIndexWrapper() accures it's own locks - backend_shared_lock.unlock(); - // initialize the SVS index synchonously using current thread only - updateSVSIndexWrapper(this, 1); + if (!svs_index->ready()) { + if (auto buffered = tryBufferForTraining(blob, label, /*in_place=*/true)) { + return *buffered; } - return ret; - } else { - // backend index is initialized - we can add the vector directly - backend_shared_lock.unlock(); - auto storage_blob = this->frontendIndex->preprocessForStorage(blob); - // prevent update job from running in parallel and lock any access to the backend - // index - std::scoped_lock lock(this->updateJobMutex, this->mainIndexGuard); - // Defensive: ensure single-threaded operation for write-in-place mode. - // parallelism_ defaults to 1, so this is a no-op in the normal case. - svs_index->setParallelism(1); - return this->backendIndex->addVector(storage_blob.get(), label); } + // backend index is initialized - we can add the vector directly + auto storage_blob = this->frontendIndex->preprocessForStorage(blob); + // Only updateJobMutex is needed here + std::lock_guard lock(this->updateJobMutex); + // Defensive: ensure single-threaded operation for write-in-place mode. + // parallelism_ defaults to 1, so this is a no-op in the normal case. + svs_index->setParallelism(1); + int deleted = 0; + if (!this->backendIndex->isMultiValue()) { + deleted = svs_index->deleteVector(label); + if (deleted > 0) + svs_index->consolidate({label}); + } + return this->backendIndex->addVector(storage_blob.get(), label) - deleted; } assert(this->getWriteMode() != VecSim_WriteInPlace && "InPlace mode returns early"); - // Async mode - add vector to the frontend index and schedule an update job if needed. - if (!this->backendIndex->isMultiValue()) { - { - std::shared_lock lock(this->flatIndexGuard); - // If the label already exists in the frontend index, we should count it - // to prevent the case when existing vector is moved meanwhile by the update job. - if (this->frontendIndex->isLabelExists(label)) { - ret = -1; - } - } - // Remove vector from the backend index if it exists in case of non-MULTI. - auto label_exists = [&]() { - std::shared_lock lock(this->mainIndexGuard); - return svs_index->isLabelExists(label); - }(); - - if (label_exists) { - std::lock_guard lock(this->mainIndexGuard); - ret -= this->backendIndex->deleteVector(label); + // Async mode - buffer training data until the backend can be inited from a batch. + if (!svs_index->ready()) { + if (auto buffered = tryBufferForTraining(blob, label, /*in_place=*/false)) { + return *buffered; } } - { // Add vector to the frontend index. - std::lock_guard lock(this->flatIndexGuard); - const auto ft_ret = this->frontendIndex->addVector(blob, label); - if (ft_ret == 0) { // Vector was overriden - add 'skiping' swap to the journal. - assert(!this->backendIndex->isMultiValue() && - "addVector() may return 0 for single value indices only"); - for (auto id : this->frontendIndex->getElementIds(label)) { - this->swaps_journal.emplace_back(SKIP_LABEL, id, id); + this->flatIndexGuard.lock_shared(); + // Before the backend is ready, a direct insert would wait for the build. + if (svs_index->ready() && this->frontendIndex->indexSize() >= flat_buffer_bound) { + this->flatIndexGuard.unlock_shared(); + auto storage_blob = this->frontendIndex->preprocessForStorage(blob); + + const int overwritten = + this->backendIndex->isMultiValue() ? 0 : this->deleteVector(label); + + std::shared_lock lock(updateJobMutex); + ret = svs_index->addVector(storage_blob.get(), label); + return std::max(ret - overwritten, 0); + } else { + this->flatIndexGuard.unlock_shared(); + this->flatIndexGuard.lock(); + if (initInProgress && this->frontendIndex->indexSize() >= + std::max(INIT_BACKPRESSURE_LIMIT, trainingTriggerThreshold)) { + // Init owns updateJobMutex exclusively and needs flat access to finish. + this->flatIndexGuard.unlock(); + executeTracingCallback("AddVector::before_init_wait"); + { std::shared_lock wait_for_init(updateJobMutex); } + // Release the update lock before retrying normal insertion. + return addVector(blob, label); + } + idType new_flat_id = this->frontendIndex->indexSize(); + if (this->frontendIndex->isLabelExists(label) && !this->frontendIndex->isMultiValue()) { + // We are going to update the internal id that currently holds the vector associated + // with the given label. + if (initInProgress && !this->labelToInsertJobs.count(label)) { + // A captured row has no insert job; cancel it through the journal. + ret = -removeLabelFromFlat(label); + new_flat_id = this->frontendIndex->indexSize(); + } else { + new_flat_id = this->invalidatePendingInsertJob(label); + ret = 0; } - deleted_labels_journal.push_back(label); } - ret = std::max(ret + ft_ret, 0); - // Check frontend index size to determine if an update job schedule is needed. - frontend_index_size = this->frontendIndex->indexSize(); - } - { - // If main index is empty then update_threshold is trainingTriggerThreshold, - // overwise it is updateTriggerThreshold. - std::shared_lock lock(this->mainIndexGuard); - update_threshold = this->backendIndex->indexSize() == 0 ? this->trainingTriggerThreshold - : this->updateTriggerThreshold; - } - if (frontend_index_size >= update_threshold) { - scheduleSVSIndexUpdate(); + // If this label already exists, this will do overwrite. + ret += this->frontendIndex->addVector(blob, label); + + TieredInsertJob *new_insert_job = createInsertJob(label, new_flat_id); + const bool defer_until_initialized = initInProgress; + if (defer_until_initialized) { + initInsertJobs.push_back(new_insert_job); + } + this->flatIndexGuard.unlock(); + + // Here, a worker might ingest the previous vector that was stored under "label" + // (in case of override in non-MULTI index) - so if it's there, we remove it + // we submit the insert job. + if (!this->backendIndex->isMultiValue() && !defer_until_initialized) { + // If we removed the previous vector from both svs and flat in the overwrite + // process, we still return 0 (not -1). + auto deleted = svs_index->deleteVector(label); + if (deleted > 0) + scheduleSVSIndexConsolidate(label); + ret = std::max(ret - deleted, 0); + } + + // Initialization drains its deferred jobs after releasing updateJobMutex. + if (!defer_until_initialized) { + this->submitSingleJob(new_insert_job); + } + return ret; } + } - return ret; + // training data is tracked by frontend id, so it has to follow the same swaps as insert jobs + void updateInsertJobInternalId(idType prev_id, idType new_id, labelType label) override { + if (ids_to_init_.erase(prev_id) > 0) { + ids_to_init_.insert(new_id); + } + Base::updateInsertJobInternalId(prev_id, new_id, label); } - int deleteAndRecordSwaps_Unsafe(labelType label) { + // Returns the number of vectors removed from the frontend index. + int deleteAndUpdateInitIds(labelType label) { auto deleting_ids = this->frontendIndex->getElementIds(label); - // assert if all elements of deleting_ids are unique - assert(std::set(deleting_ids.begin(), deleting_ids.end()).size() == deleting_ids.size() && - "deleting_ids should contain unique ids"); - - // Sort deleting_ids by id descending order - std::sort(deleting_ids.begin(), deleting_ids.end(), - [](const auto &a, const auto &b) { return a > b; }); + std::sort(deleting_ids.begin(), deleting_ids.end()); // Delete vector from the frontend index. auto updated_ids = this->frontendIndex->deleteVectorAndGetUpdatedIds(label); assert(std::all_of(updated_ids.begin(), updated_ids.end(), [&deleting_ids](const auto &pair) { - return std::find(deleting_ids.begin(), deleting_ids.end(), - pair.first) != deleting_ids.end(); + return std::binary_search(deleting_ids.begin(), deleting_ids.end(), + pair.first); }) && "updated_ids should be a subset of deleting_ids"); - // Record swaps in the journal. - for (auto id : deleting_ids) { - auto it = updated_ids.find(id); - if (it != updated_ids.end()) { - assert(id == it->first && "id in updated_ids should match the id in deleting_ids"); - auto newId = id; - auto oldId = it->second.first; - auto oldLabel = it->second.second; - this->swaps_journal.emplace_back(oldLabel, oldId, newId); - } else { - // No swap, just delete is marked by oldId == newId == deleted id - this->swaps_journal.emplace_back(SKIP_LABEL, id, id); + std::vector new_ids_to_init; + for (auto &it : updated_ids) { + idType prev_id = it.second.first; + idType new_id = it.first; + + if (ids_to_init_.count(prev_id) > 0) { + ids_to_init_.erase(prev_id); + new_ids_to_init.push_back(new_id); } + labelType updated_vec_label = it.second.second; + this->updateInsertJobInternalId(prev_id, new_id, updated_vec_label); } - deleted_labels_journal.push_back(label); - return deleting_ids.size(); + for (idType new_id : new_ids_to_init) { + ids_to_init_.insert(new_id); + } + + return static_cast(deleting_ids.size()); } - int deleteVector(labelType label) override { + int removeLabelFromFlat(labelType label) { + if (!this->frontendIndex->isLabelExists(label)) { + return 0; + } + auto deleting_ids = this->frontendIndex->getElementIds(label); + if (deleting_ids.size() == 0) { + return 0; + } + + // assert if all elements of deleting_ids are unique + assert(std::set(deleting_ids.begin(), deleting_ids.end()).size() == deleting_ids.size() && + "deleting_ids should contain unique ids"); + + // If id is deleted, don't use it for initialization + bool captured = false; + if (!ids_to_init_.empty()) { + for (idType id : deleting_ids) { + captured |= ids_to_init_.erase(id) > 0; + } + } + if (captured && initInProgress) { + initDeletedLabels.push_back(label); + } + + if (this->labelToInsertJobs.count(label) > 0) { + // Invalidate the pending insert job(s) into SVS associated with this label + auto &insert_jobs = this->labelToInsertJobs.at(label); + for (auto *job : insert_jobs) { + job->id = this->setAndSaveInvalidJob(job); + } + // Remove the pending insert job(s) from the labelToInsertJobs mapping. + this->labelToInsertJobs.erase(label); + } + + deleteAndUpdateInitIds(label); + return static_cast(deleting_ids.size()); + } + + // Create a pending insert job for a vector already buffered in the frontend index + // Caller must hold flatIndexGuard exclusively. + TieredInsertJob *createInsertJob(labelType label, idType flat_id) { + TieredInsertJob *job = new (this->allocator) + SVSInsertJob(this->allocator, label, flat_id, executeInsertJobWrapper, this); + this->registerInsertJob(label, job); + return job; + } + + // Buffer a vector in the frontend index as training data for the backend initialization. + // Returns 1 for a new label, 0 for an overwrite, or nullopt if the backend turned ready while + // the guard was taken + std::optional tryBufferForTraining(const void *blob, labelType label, bool in_place) { int ret = 0; - auto svs_index = GetSVSIndex(); - // Backend index deletions to be synchronized with the frontend index, - // elsewhere there is the risk of labels duplication in both indices which can lead to wrong - // results of topK queries. In such case we should behave as if InPlace mode is always set. - bool label_exists = [&]() { - std::shared_lock lock(this->flatIndexGuard); - return this->frontendIndex->isLabelExists(label); - }(); + bool run_init = false; + { + std::lock_guard flat_lock{this->flatIndexGuard}; + // Once the snapshot is taken, use the existing insertion path. + if (GetSVSIndex()->ready() || initInProgress) { + return std::nullopt; + } + int deleted = 0; + if (!this->frontendIndex->isMultiValue() && this->frontendIndex->isLabelExists(label)) { + deleted = deleteAndUpdateInitIds(label); + } + ids_to_init_.insert(this->frontendIndex->indexSize()); + ret = std::max(this->frontendIndex->addVector(blob, label) - deleted, 0); + run_init = ids_to_init_.size() >= this->trainingTriggerThreshold; + } + if (run_init && !in_place) { + scheduleSVSIndexInit(); + run_init = false; + } - if (label_exists) { - std::lock_guard lock(this->flatIndexGuard); - ret = this->deleteAndRecordSwaps_Unsafe(label); + if (run_init) { + initSVSIndexWrapper(this, 1); } + return ret; + } - label_exists = [&]() { - std::shared_lock lock(this->mainIndexGuard); - return svs_index->isLabelExists(label); - }(); + int deleteVector(labelType label) override { + int ret = 0; + + this->flatIndexGuard.lock_shared(); + if (this->frontendIndex->isLabelExists(label)) { + this->flatIndexGuard.unlock_shared(); + std::lock_guard flat_lock{this->flatIndexGuard}; + ret += removeLabelFromFlat(label); + } else { + this->flatIndexGuard.unlock_shared(); + } - if (label_exists) { - std::lock_guard lock(this->mainIndexGuard); - ret += this->backendIndex->deleteVector(label); + auto deleted = this->backendIndex->deleteVector(label); + if (deleted > 0) { + if (this->getWriteMode() == VecSim_WriteInPlace) { + GetSVSIndex()->consolidate({label}); + } else { + scheduleSVSIndexConsolidate(label); + } } + ret += deleted; return ret; } /** * Set the vectors stored under `label` to the given ones, in whichever tier holds them. * * Both halves go through this index's own delete and add, which is what makes every tier and - * write mode fall out: the delete removes the label's buffered vectors - journaling the label - * so that an update job running at the same time does not carry them into the backend - and - * deletes its backend vectors, and each insert then lands where a fresh vector would, with the - * training and update thresholds deciding when the batch moves. + * write mode fall out: the delete removes the label's buffered vectors - invalidating their + * pending insert jobs and deletes its backend vectors, and each insert then lands where a fresh + * vector would, with the training and update thresholds deciding when the batch moves. * * Not atomic: a query landing between the two halves sees the label with none of its vectors, * the same window the delete-then-insert it replaces would leave. @@ -1003,9 +1234,11 @@ class TieredSVSIndex : public VecSimTieredIndex { // The locks the paths below take are scoped, so unwinding releases them. try { // One vector replacing a single-value label *is* an overwrite, which `addVector` - // already performs across the tiers: it writes over the buffered copy in place and - // journals the overwrite so an update job in flight does not carry the replaced vector - // away, where a delete followed by an insert would drop that copy and buffer a new one. + // already performs across the tiers: it writes over the buffered copy in place - + // and invalidating the pending insert job, + // a worker in flight does not carry the replaced vector away, + // where a delete followed by an insert would + // discard the buffered copy and buffer a new one. if (n == 1 && !this->backendIndex->isMultiValue()) { this->addVector(new_blobs, label); return VecSimUpdate_OK; @@ -1040,11 +1273,8 @@ class TieredSVSIndex : public VecSimTieredIndex { * case when built against an SVS without `replace_external_id`. All-or-nothing: on any code * other than `VecSimRelabel_OK` both tiers are untouched. * - * `updateJobMutex` is taken first, in the order `updateSVSIndex` takes its own locks. Holding - * it is what makes this correct rather than merely serialised: an update job snapshots the - * buffer's labels *by value* and afterwards reconciles only id swaps and deletions, so a - * rename landing inside its window would be invisible to it and the vector would reach the - * backend under the old label. + * Relabel is not atomic with concurrent queries. An overlapping query may return both the + * old and new labels, using two top-k positions for the moved label. */ VecSimRelabelCode relabelVector(labelType old_label, labelType new_label) override { if (old_label == new_label) { @@ -1052,79 +1282,115 @@ class TieredSVSIndex : public VecSimTieredIndex { } auto *svs_index = GetSVSIndex(); - // Taken together through `std::lock`'s back-off rather than one after another, because - // the orders in this class disagree: nearly everything acquires flat before main, but the - // in-place add path takes `mainIndexGuard` before `flatIndexGuard` while the backend is - // still empty. Acquiring in sequence would hold one guard while blocking on another and - // close a cycle with that path; acquiring them together cannot, whichever order the other - // side uses. Same idiom that path already uses for its own two locks. - std::scoped_lock lock(this->updateJobMutex, this->flatIndexGuard, this->mainIndexGuard); - - const bool in_flat = this->frontendIndex->isLabelExists(old_label); - const bool in_backend = svs_index->isLabelExists(old_label); - if (!in_flat && !in_backend) { - return VecSimRelabel_OldLabelMissing; - } - if (this->frontendIndex->isLabelExists(new_label) || svs_index->isLabelExists(new_label)) { - return VecSimRelabel_NewLabelTaken; - } - - // The backend goes first because it is the tier that can refuse; refusing after the - // buffer had already moved would leave the label half applied. - if (in_backend) { - const VecSimRelabelCode backend_ret = - this->backendIndex->relabelVector(old_label, new_label); - if (backend_ret != VecSimRelabel_OK) { - return backend_ret; - } - } - if (in_flat) { - const VecSimRelabelCode flat_ret = - this->frontendIndex->relabelVector(old_label, new_label); -#ifdef BUILD_TESTS - assert(flat_ret == VecSimRelabel_OK && "the buffer just reported holding this label"); -#endif - UNUSED(flat_ret); - } - return VecSimRelabel_OK; - } + std::shared_lock lock(this->updateJobMutex); -#endif // HAVE_SVS_REPLACE_EXTERNAL_ID + bool flat_holds_old = false; + bool delayed_any = false; + std::vector taken_over; + vecsim_stl::vector parked_jobs(this->allocator); + { + std::lock_guard flat_lock{this->flatIndexGuard}; - size_t getNumMarkedDeleted() const override { - return this->GetSVSIndex()->getNumMarkedDeleted(); - } + flat_holds_old = this->frontendIndex->isLabelExists(old_label); + if (!flat_holds_old && !svs_index->isLabelExists(old_label)) { + return VecSimRelabel_OldLabelMissing; + } + if (this->frontendIndex->isLabelExists(new_label) || + svs_index->isLabelExists(new_label)) { + return VecSimRelabel_NewLabelTaken; + } - size_t indexSize() const override { - std::shared_lock flat_lock(this->flatIndexGuard); - std::shared_lock main_lock(this->mainIndexGuard); - return this->frontendIndex->indexSize() + this->backendIndex->indexSize(); - } + auto pending = this->labelToInsertJobs.find(old_label); + if (pending != this->labelToInsertJobs.end()) { + auto jobs = std::move(pending->second); + this->labelToInsertJobs.erase(pending); + for (auto *job : jobs) { + auto *insert_job = static_cast(job); + // wait it out: a job that published before the remap is remapped with the + // backend, and still drops its own buffer copy under the new label + waitForJob(insert_job); + job->label = new_label; + if (insert_job->status.load(std::memory_order_relaxed) == + SVSInsertJob::Status::Pending) { + // did not publish yet - hold it until the remap settles the label + insert_job->status.store(SVSInsertJob::Status::Delayed, + std::memory_order_relaxed); + delayed_any = true; + } + } + this->labelToInsertJobs.emplace(new_label, std::move(jobs)); + } - size_t indexCapacity() const override { - std::shared_lock flat_lock(this->flatIndexGuard); - std::shared_lock main_lock(this->mainIndexGuard); - return this->frontendIndex->indexCapacity() + this->backendIndex->indexCapacity(); - } + // isLabelExists() said new_label is free, but SVS keeps the translator entry of a + // soft-deleted label until its consolidate job runs, and would refuse the remap. + taken_over = takeOverConsolidateOf(new_label); - double getDistanceFrom_Unsafe(labelType label, const void *blob) const override { - // Try to get the distance from the flat buffer. - // If the label doesn't exist, the distance will be NaN. - auto flat_dist = this->frontendIndex->getDistanceFrom_Unsafe(label, blob); + if (flat_holds_old) { + const VecSimRelabelCode flat_ret = + this->frontendIndex->relabelVector(old_label, new_label); + assert(flat_ret == VecSimRelabel_OK && + "the buffer just reported holding this label"); + UNUSED(flat_ret); + } + } + + executeTracingCallback("Relabel::before_backend_remap"); + // Do the taken-over job's, + // the remap below isn't refused by a leftover translator entry + if (!taken_over.empty()) { + svs_index->consolidate(taken_over); + } - // Optimization. TODO: consider having different implementations for single and multi - // indexes, to avoid checking the index type on every query. - if (!this->backendIndex->isMultiValue() && !std::isnan(flat_dist)) { - // If the index is single value, and we got a valid distance from the flat buffer, - // we can return the distance without querying the Main index. - return flat_dist; + const bool rollback = + this->backendIndex->relabelVector(old_label, new_label) == VecSimRelabel_NewLabelTaken; + const labelType final_label = rollback ? old_label : new_label; + + if (rollback || delayed_any) { + std::lock_guard flat_lock{this->flatIndexGuard}; + if (rollback) { + auto pending = this->labelToInsertJobs.find(new_label); + if (pending != this->labelToInsertJobs.end()) { + auto jobs = std::move(pending->second); + this->labelToInsertJobs.erase(pending); + for (auto *job : jobs) { + job->label = old_label; + } + this->labelToInsertJobs.emplace(old_label, std::move(jobs)); + } + if (flat_holds_old) { + this->frontendIndex->relabelVector(new_label, old_label); + } + } + if (delayed_any) { + auto it = this->labelToInsertJobs.find(final_label); + if (it != this->labelToInsertJobs.end()) { + for (auto *job : it->second) { + auto *insert_job = static_cast(job); + auto status = insert_job->status.load(std::memory_order_relaxed); + if (status == SVSInsertJob::Status::Parked) { + parked_jobs.push_back(job); + } + if (status == SVSInsertJob::Status::Delayed || + status == SVSInsertJob::Status::Parked) { + insert_job->status.store(SVSInsertJob::Status::Pending, + std::memory_order_relaxed); + } + } + } + } } + // The pool's submit callback may wait for workers blocked on either of our locks. + lock.unlock(); + if (!parked_jobs.empty()) { + this->submitJobs(parked_jobs); + } + return rollback ? VecSimRelabel_NewLabelTaken : VecSimRelabel_OK; + } - // Try to get the distance from the Main index. - auto svs_dist = this->backendIndex->getDistanceFrom_Unsafe(label, blob); +#endif // HAVE_SVS_REPLACE_EXTERNAL_ID - // Return the minimum distance that is not NaN. - return std::fmin(flat_dist, svs_dist); + size_t getNumMarkedDeleted() const override { + return this->GetSVSIndex()->getNumMarkedDeleted(); } VecSimIndexDebugInfo debugInfo() const override { @@ -1137,25 +1403,15 @@ class TieredSVSIndex : public VecSimTieredIndex { .defragRuns = defragRuns.load(std::memory_order_relaxed), .defragTimeNs = defragTimeNs.load(std::memory_order_relaxed), }; - { - // Use try_lock to avoid blocking the main thread during long-running - // training operations. updateSVSIndexWrapper holds updateJobMutex for - // the entire training duration (which can take 40-85s on slow machines). - // If the mutex is held, training is actively running, so we report - // indexUpdateScheduled = true (BACKGROUND_INDEXING = 1). - std::unique_lock lock(this->updateJobMutex, std::try_to_lock); - if (lock.owns_lock()) { - svsTieredInfo.indexUpdateScheduled = this->indexUpdateScheduled.test(); - } else { - // Mutex is held by updateSVSIndexWrapper — training is in progress. - svsTieredInfo.indexUpdateScheduled = true; - } - } + + svsTieredInfo.indexUpdateScheduled = + (info.tieredInfo.frontendCommonInfo.indexSize > 0) && + this->indexUpdateScheduled.test(std::memory_order_acquire); info.tieredInfo.specificTieredBackendInfo.svsTieredInfo = svsTieredInfo; + + // Background indexing is in progress whenever the flat buffer is non-empty info.tieredInfo.backgroundIndexing = - svsTieredInfo.indexUpdateScheduled && info.tieredInfo.frontendCommonInfo.indexSize > 0 - ? VecSimBool_TRUE - : VecSimBool_FALSE; + info.tieredInfo.frontendCommonInfo.indexSize > 0 ? VecSimBool_TRUE : VecSimBool_FALSE; return info; } @@ -1224,7 +1480,7 @@ class TieredSVSIndex : public VecSimTieredIndex { "running synchronous GC for tiered SVS index in write-in-place mode"); // In write-in-place mode, we run GC synchronously. const auto start = std::chrono::steady_clock::now(); - std::lock_guard lock{this->mainIndexGuard}; + std::lock_guard main_lock(this->mainIndexGuard); if (this->backendIndex->indexSize() == 0) { // No need to run GC on an empty index. return; @@ -1241,13 +1497,7 @@ class TieredSVSIndex : public VecSimTieredIndex { scheduleSVSIndexGC(); } - void acquireSharedLocks() override { - this->flatIndexGuard.lock_shared(); - this->mainIndexGuard.lock_shared(); - } + void acquireSharedLocks() override { this->flatIndexGuard.lock_shared(); } - void releaseSharedLocks() override { - this->mainIndexGuard.unlock_shared(); - this->flatIndexGuard.unlock_shared(); - } + void releaseSharedLocks() override { this->flatIndexGuard.unlock_shared(); } }; diff --git a/src/VecSim/algorithms/svs/svs_utils.h b/src/VecSim/algorithms/svs/svs_utils.h index 7d90f97c9..232240ae9 100644 --- a/src/VecSim/algorithms/svs/svs_utils.h +++ b/src/VecSim/algorithms/svs/svs_utils.h @@ -16,6 +16,10 @@ #include "svs/core/distance.h" #include "svs/lib/float16.h" #include "svs/index/vamana/dynamic_index.h" +// Pulls in `svs::index::vamana::concurrent` and the `svs::concurrent` alias. The concurrent +// namespace redeclares only the entities it has to replace; every other name (build/search +// parameters, extension points) still resolves to the enclosing `svs::index::vamana`. +#include "svs/concurrent/concurrent.h" #if HAVE_SVS_LVQ #include "svs/cpuid.h" @@ -204,6 +208,21 @@ inline std::pair isSVSQuantBitsSupported(VecSimSvsQuan // unreachable code, but to avoid compiler warning return std::make_pair(VecSimSvsQuant_NONE, false); } + +// The `svs::concurrent` counterpart of `svs::make_blocked_allocator_handle`. +// +// The upstream helper cannot be reused here: it hardcodes `svs::data::Blocked` in its return +// type, so there is no way to ask it for the grow-stable tag. The type-erasing +// `AllocatorHandle` wrapper is what makes the compressed datasets (which take an allocator +// handle rather than a concrete allocator) work with VecSim's own allocator. +template +svs::concurrent::SegmentedBlocked> +make_segmented_blocked_allocator_handle(const svs::data::BlockingParameters ¶meters, + Alloc alloc) { + using handle_type = svs::AllocatorHandle; + return svs::concurrent::SegmentedBlocked{ + parameters, svs::make_allocator_handle(std::move(alloc))}; +} } // namespace svs_details template ; // Used in creating storage + // `SegmentedBlocked` is a drop-in replacement for `svs::data::Blocked` (it derives from + // it and carries the same blocking parameters) that additionally keeps already-published + // elements at a stable address when the dataset grows, which is what allows a concurrent + // reader to keep a pointer into the data while a writer appends. + using blocked_type = svs::concurrent::SegmentedBlocked; // svs::Dynamic means runtime dimensionality in opposite to compile-time dimensionality - using index_storage_type = svs::data::BlockedData; + using index_storage_type = + svs::concurrent::SegmentedBlockedData; static constexpr bool is_compressed() { return false; } @@ -278,9 +302,13 @@ struct SVSStorageTraits { template struct SVSGraphBuilder { using allocator_type = svs_details::SVSAllocator; - using blocked_type = svs::data::Blocked; - using graph_data_type = svs::data::BlockedData; - using graph_type = svs::graphs::SimpleGraph; + using blocked_type = svs::concurrent::SegmentedBlocked; + using graph_data_type = + svs::concurrent::SegmentedBlockedData; + // The concurrent `SimpleGraph` has the same shape as the upstream one but guards each + // adjacency list with a sequence lock, so a reader can traverse an edge list that a + // writer is concurrently rewriting. + using graph_type = svs::concurrent::graphs::SimpleGraph; static blocked_type make_blocked_allocator(size_t block_size, size_t graph_max_degree, std::shared_ptr allocator) { @@ -310,8 +338,8 @@ struct SVSGraphBuilder { // based on the data types, which we found to perform better through heuristic analysis. auto prefetch_parameters = svs::index::vamana::extensions::estimate_prefetch_parameters(data); - auto builder = svs::index::vamana::VamanaBuilder( - graph, data, std::move(distance), parameters, threadpool, prefetch_parameters, logger); + auto builder = svs::concurrent::VamanaBuilder(graph, data, std::move(distance), parameters, + threadpool, prefetch_parameters, logger); // Specific to the Vamana algorithm: // It builds in two rounds, one with alpha=1 and the second time with the user/config @@ -345,6 +373,12 @@ struct SVSGraphBuilder { // graph_max_degree neighbors return sizeof(SVSIdType) * (graph_max_degree + 1); } + + // Per-node size of the concurrent graph's reverse-edge index. + static constexpr size_t reverse_edges_element_size() { + using reverse_edges_type = typename graph_type::reverse_edges_type; + return sizeof(typename reverse_edges_type::list_type) + sizeof(svs::concurrent::SpinLock); + } }; // A slot in the shared SVS thread pool. Wraps an SVS Thread with an occupancy flag diff --git a/src/VecSim/index_factories/svs_factory.cpp b/src/VecSim/index_factories/svs_factory.cpp index 33adf6fb1..bae7c021d 100644 --- a/src/VecSim/index_factories/svs_factory.cpp +++ b/src/VecSim/index_factories/svs_factory.cpp @@ -234,10 +234,12 @@ size_t EstimateElementSize(const SVSParams *params) { // Assuming that the graph_max_degree can be unset in params. const auto graph_max_degree = svs_details::makeVamanaBuildParameters(*params).graph_max_degree; const auto graph_node_size = SVSGraphBuilder::element_size(graph_max_degree); + const auto reverse_edges_node_size = + SVSGraphBuilder::reverse_edges_element_size(); const auto vector_size = QuantizedVectorSize(params->type, params->quantBits, params->dim, 0, params->leanvec_dim); - return vector_size + graph_node_size; + return vector_size + graph_node_size + reverse_edges_node_size; } size_t EstimateInitialSize(const SVSParams *params, bool is_normalized) { diff --git a/src/VecSim/utils/scoped_locks.h b/src/VecSim/utils/scoped_locks.h new file mode 100644 index 000000000..234047a9b --- /dev/null +++ b/src/VecSim/utils/scoped_locks.h @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ + +#pragma once + +#include +#include + +// A lockable resource, so that different mutex-like primitives (a shared_mutex, HNSW's own +// data guard, or nothing at all) can be composed under one RAII type below. +class Lockable { +public: + virtual void lock() const = 0; + virtual void unlock() const = 0; + virtual ~Lockable() = default; +}; + +// Adapts a std::shared_mutex, locked/unlocked for shared (read) access, to Lockable. +class SharedMutexLockable : public Lockable { + std::shared_mutex &guard; + +public: + explicit SharedMutexLockable(std::shared_mutex &guard) : guard(guard) {} + void lock() const override { guard.lock_shared(); } + void unlock() const override { guard.unlock_shared(); } +}; + +// RAII acquisition of zero, one or two Lockables, so callers can plug in whatever locking +// (or none) they need without repeating the lock/unlock bookkeeping. +class ScopedLocks { + std::array locks{}; + size_t count = 0; + +public: + ScopedLocks() = default; + explicit ScopedLocks(const Lockable &a) : locks{&a, nullptr}, count(1) { a.lock(); } + ScopedLocks(const Lockable &a, const Lockable &b) : locks{&a, &b}, count(2) { + a.lock(); + b.lock(); + } + ScopedLocks(ScopedLocks &&other) noexcept : locks(other.locks), count(other.count) { + other.count = 0; + } + ScopedLocks(const ScopedLocks &) = delete; + ScopedLocks &operator=(ScopedLocks &&) = delete; + ~ScopedLocks() { + // Unlock in reverse acquisition order. + for (size_t i = count; i-- > 0;) { + locks[i]->unlock(); + } + } +}; diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h index 7d0676fb8..1da279829 100644 --- a/src/VecSim/vec_sim_common.h +++ b/src/VecSim/vec_sim_common.h @@ -318,7 +318,9 @@ typedef enum { HNSW_SWAP_JOB, HNSW_DISK_JOB, SVS_BATCH_UPDATE_JOB, + SVS_INSERT_VECTOR_JOB, SVS_GC_JOB, + SVS_CONSOLIDATE_JOB, INVALID_JOB // to indicate that finding a JobType >= INVALID_JOB is an error } JobType; @@ -464,7 +466,7 @@ typedef struct SvsTieredInfo { // before executing the scheduled SVS Index update job. bool indexUpdateScheduled; size_t defragRuns; // Number of backend GC runs so far. - size_t defragTimeNs; // Wall time spent in them, including the main index guard wait. + size_t defragTimeNs; // Wall time spent in backend GC, including lock waits. } SvsTieredInfo; typedef struct { diff --git a/src/VecSim/vec_sim_tiered_index.h b/src/VecSim/vec_sim_tiered_index.h index 286d4f057..9cf574dc9 100644 --- a/src/VecSim/vec_sim_tiered_index.h +++ b/src/VecSim/vec_sim_tiered_index.h @@ -11,12 +11,15 @@ #include "vec_sim_index.h" #include "algorithms/brute_force/brute_force.h" +#include "algorithms/brute_force/brute_force_single.h" #include "VecSim/batch_iterator.h" #include "VecSim/tombstone_interface.h" #include "VecSim/utils/query_result_utils.h" #include "VecSim/utils/alignment.h" +#include "VecSim/utils/scoped_locks.h" #include +#include #include #include @@ -42,6 +45,92 @@ struct AsyncJob : public VecsimBaseObject { isValid(true) {} }; +/** + * Definition of a job that inserts a new vector from flat into the backend index. + * Backend specific insert jobs derive from it to set their own job type. + */ +struct TieredInsertJob : public AsyncJob { + labelType label; + idType id; + + TieredInsertJob(std::shared_ptr allocator, JobType type, labelType label_, + idType id_, JobCallback insertCb, VecSimIndex *index_) + : AsyncJob(allocator, type, insertCb, index_), label(label_), id(id_) {} +}; + +class TieredIndex_BatchIterator : public VecSimBatchIterator { +protected: + VecSimQueryResultContainer flat_results; + VecSimQueryResultContainer backend_results; + + // On single value indices, this set holds the IDs of the results that were returned from + // the flat buffer. + // On multi value indices, this set holds the IDs of all the results that were returned. + // The difference between the two cases is that on multi value indices, the same ID can + // appear in both indexes and results with different scores, and therefore we can't tell in + // advance when we expect a possibility of a duplicate. + vecsim_stl::unordered_set returned_results_set; + + TieredIndex_BatchIterator(void *query_vector, void *tctx, + std::shared_ptr allocator) + : VecSimBatchIterator(query_vector, tctx, std::move(allocator)), + flat_results(this->allocator), backend_results(this->allocator), + returned_results_set(this->allocator) {} + + template + VecSimQueryReply *compute_current_batch(size_t n_res) { + // Merge results + auto batch_res = new VecSimQueryReply(this->allocator); + std::pair p; + if (needsDedup) { + p = merge_results(batch_res->results, this->backend_results, this->flat_results, + n_res); + } else { + p = merge_results(batch_res->results, this->backend_results, this->flat_results, + n_res); + } + auto [from_backend, from_flat] = p; + + if (!needsDedup) { + // Update the set of results returned from the FLAT + // index before popping them. + for (size_t i = 0; i < from_flat; ++i) { + this->returned_results_set.insert(this->flat_results[i].id); + } + } else { + // Update the set of results returned (from `batch_res`) + for (size_t i = 0; i < batch_res->results.size(); ++i) { + this->returned_results_set.insert(batch_res->results[i].id); + } + } + + // Update results + this->flat_results.erase(this->flat_results.begin(), + this->flat_results.begin() + from_flat); + this->backend_results.erase(this->backend_results.begin(), + this->backend_results.begin() + from_backend); + + // clean up the results + // One (or both) results lists may contain results that are already + // returned form the other list (with a different score). We need to filter them out. + if (needsDedup) { + this->filter_irrelevant_results(this->flat_results); + this->filter_irrelevant_results(this->backend_results); + } + + // Return current batch + return batch_res; + } + + void filter_irrelevant_results(VecSimQueryResultContainer &results) { + // Filter out results that were already returned. + const auto it = std::remove_if(results.begin(), results.end(), [this](const auto &r) { + return returned_results_set.count(r.id) != 0; + }); + results.erase(it, results.end()); + } +}; + // All read operations (including KNN, range, batch iterators and get-distance-from) are guaranteed // to consider all vectors that were added to the index before the query was submitted. The results // may include vectors that were added after the query was submitted, with no guarantees. @@ -57,6 +146,18 @@ class VecSimTieredIndex : public VecSimIndexInterface { mutable std::shared_mutex flatIndexGuard; mutable std::shared_mutex mainIndexGuard; + SharedMutexLockable flatIndexLockable{flatIndexGuard}; + SharedMutexLockable mainIndexLockable{mainIndexGuard}; + + // Locking behavior for topKQuery/rangeQuery + virtual ScopedLocks lockMainIndexForQuery() const = 0; + + // Locking behavior for indexSize() + virtual ScopedLocks lockIndexForSize() const = 0; + + // Locking behavior for indexCapacity() + virtual ScopedLocks lockIndexForCapacity() const = 0; + void lockMainIndexGuard() const { mainIndexGuard.lock(); #ifdef BUILD_TESTS @@ -121,6 +222,111 @@ class VecSimTieredIndex : public VecSimIndexInterface { VecSimQueryParams *queryParams, VecSimQueryReply_Order order) const; +#ifdef BUILD_TESTS +public: +#endif + /// Mappings from id/label to associated jobs, for invalidating and update ids if necessary. + // In MULTI, we can have more than one insert job pending per label. + // **This map is protected with the flat buffer lock** + vecsim_stl::unordered_map> labelToInsertJobs; + + // Helper function for updating the pending insert job(s) of a label after the flat buffer + // swapped the vector's internal id + virtual void updateInsertJobInternalId(idType prev_id, idType new_id, labelType label) { + // Update the pending job id, due to a swap that was caused after the removal of new_id. + assert(new_id != INVALID_ID && prev_id != INVALID_ID); + auto it = this->labelToInsertJobs.find(label); + if (it != this->labelToInsertJobs.end()) { + // There is a pending job for the label of the swapped last id - update its id. + for (TieredInsertJob *job_it : it->second) { + if (job_it->id == prev_id) { + job_it->id = new_id; + } + } + } + } + + // A mapping to hold invalid jobs, so we can dispose them upon index deletion. + vecsim_stl::unordered_map invalidJobs; + idType currInvalidJobId; // A unique arbitrary identifier for accessing invalid jobs + std::mutex invalidJobsLookupGuard; + + // Set an insert/repair job as invalid, put the job pointer in the invalid jobs lookup under + // the current available id, increase it and return it (while holding invalidJobsLookupGuard). + // Returns the id that the job was stored under (to be set in the job id field). + virtual idType setAndSaveInvalidJob(AsyncJob *job) { + std::lock_guard lock(this->invalidJobsLookupGuard); + job->isValid = false; + idType curInvalidId = currInvalidJobId++; + this->invalidJobs.insert({curInvalidId, job}); + return curInvalidId; + } + + // Remove the job pointer from the labelToInsertJobs mapping. Must hold flatIndexGuard. + void detachInsertJob(TieredInsertJob *job) { + auto it = this->labelToInsertJobs.find(job->label); + if (it == this->labelToInsertJobs.end()) { + return; + } + auto &jobs = it->second; + for (size_t i = 0; i < jobs.size(); i++) { + if (jobs[i] == job) { + jobs.erase(jobs.begin() + (long)i); + break; + } + } + if (jobs.empty()) { + this->labelToInsertJobs.erase(it); + } + } + + // A caller must hold flatIndexGuard. + idType invalidatePendingInsertJob(labelType label) { + auto *old_job = this->labelToInsertJobs.at(label).at(0); + old_job->id = this->setAndSaveInvalidJob(old_job); + this->labelToInsertJobs.erase(label); + return dynamic_cast *>(this->frontendIndex) + ->getIdOfLabel(label); + } + + // A caller must hold flatIndexGuard. + void registerInsertJob(labelType label, TieredInsertJob *job) { + // Construct the job vector only for a new label; multi-value labels append to the existing + // one. + auto [it, inserted] = this->labelToInsertJobs.try_emplace(label, 1, job, this->allocator); + if (!inserted) { + assert(this->backendIndex->isMultiValue()); + it->second.push_back(job); + } + } + + // Remove a vector and its insert job from the flat buffer + void removeIngestedVectorFromFlat(TieredInsertJob *job) { + std::lock_guard flat_lock(this->flatIndexGuard); + // The job might have been invalidated due to overwrite in the meantime. In this case, + // it was already deleted and the job has been evicted. Otherwise, we need to do it now. + if (!job->isValid) { + // Remove the current job from the invalid jobs' lookup, as we are about to delete it + // now. + std::lock_guard invalid_jobs_lock(this->invalidJobsLookupGuard); + this->invalidJobs.erase(job->id); + return; + } + this->detachInsertJob(job); + // Remove the vector from the flat buffer. This may cause the last vector id to swap with + // the deleted id. Hold the label for the last id, so we can later on update its + // corresponding job id. Note that after calling deleteVectorById, the last id's label + // shouldn't be available, since it is removed from the lookup. + labelType last_vec_label = + this->frontendIndex->getVectorLabel(this->frontendIndex->indexSize() - 1); + int deleted = this->frontendIndex->deleteVectorById(job->label, job->id); + if (deleted && job->id != this->frontendIndex->indexSize()) { + // If the vector removal caused a swap with the last id, update the relevant insert job. + this->updateInsertJobInternalId(this->frontendIndex->indexSize(), job->id, + last_vec_label); + } + } + public: /** * @brief Get the vector elements stored under a label, in insertion order. @@ -188,17 +394,72 @@ class VecSimTieredIndex : public VecSimIndexInterface { } } + // `getDistanceFrom` returns the minimum distance between the given blob and the vector with + // the given label. If the label doesn't exist, the distance will be NaN. + // Therefore, it's better to just call `getDistanceFrom` on both indexes and return the minimum + // instead of checking if the label exists in each index. We first try to get the distance from + // the flat buffer, as vectors in the buffer might move to the backend while we're "between" + // the locks. + // Behavior for single (regular) index: + // 1. label doesn't exist in both indexes - return NaN + // 2. label exists in one of the indexes only - return the distance from that index (valid) + // 3. label exists in both indexes - return the value from the flat buffer (valid and equal to + // the value from the backend index), saving us from locking the backend index. + // Behavior for multi index: + // 1. label doesn't exist in both indexes - return NaN + // 2. label exists in one of the indexes only - return the distance from that index (valid) + // 3. label exists in both indexes - we may have some of the vectors with the same label in the + // flat buffer only and some in the backend index only (and maybe temporal duplications). So, + // we get the distance from both indexes and return the minimum. + // + // IMPORTANT: this should be called when the *tiered index locks are locked for shared + // ownership*, along with the backend index's own data guard lock if it has one. That is since + // the internal getDistanceFrom calls access the indexes' data, and it is not safe to run + // insert/delete operations in parallel. Also, we avoid acquiring the locks internally, since + // this is usually called for every vector individually, and the overhead of acquiring and + // releasing the locks is significant in that case. + double getDistanceFrom_Unsafe(labelType label, const void *blob) const override { + // Try to get the distance from the flat buffer. + // If the label doesn't exist, the distance will be NaN. + auto flat_dist = this->frontendIndex->getDistanceFrom_Unsafe(label, blob); + + // Optimization. TODO: consider having different implementations for single and multi + // indexes, to avoid checking the index type on every query. + if (!this->backendIndex->isMultiValue() && !std::isnan(flat_dist)) { + // If the index is single value, and we got a valid distance from the flat buffer, + // we can return the distance without querying the backend index. + return flat_dist; + } + + // Try to get the distance from the backend index. + auto backend_dist = this->backendIndex->getDistanceFrom_Unsafe(label, blob); + + // Return the minimum distance that is not NaN. + return std::fmin(flat_dist, backend_dist); + } + VecSimTieredIndex(VecSimIndexAbstract *backendIndex_, BruteForceIndex *frontendIndex_, TieredIndexParams tieredParams, std::shared_ptr allocator) : VecSimIndexInterface(allocator), backendIndex(backendIndex_), frontendIndex(frontendIndex_), jobQueue(tieredParams.jobQueue), jobQueueCtx(tieredParams.jobQueueCtx), SubmitJobsToQueue(tieredParams.submitCb), - flatBufferLimit(tieredParams.flatBufferLimit) { + flatBufferLimit(tieredParams.flatBufferLimit), labelToInsertJobs(this->allocator), + invalidJobs(this->allocator), currInvalidJobId(0) { assert(backendIndex != nullptr); } virtual ~VecSimTieredIndex() { + // Delete all the pending insert jobs. + for (auto &jobs : this->labelToInsertJobs) { + for (auto *job : jobs.second) { + delete job; + } + } + // Delete all the pending invalid jobs. + for (auto &it : this->invalidJobs) { + delete it.second; + } VecSimIndex_Free(backendIndex); VecSimIndex_Free(frontendIndex); } @@ -210,6 +471,16 @@ class VecSimTieredIndex : public VecSimIndexInterface { VecSimQueryParams *queryParams, VecSimQueryReply_Order order) const override; + size_t indexSize() const override { + auto locks = lockIndexForSize(); + return this->frontendIndex->indexSize() + this->backendIndex->indexSize(); + } + + size_t indexCapacity() const override { + auto locks = lockIndexForCapacity(); + return this->frontendIndex->indexCapacity() + this->backendIndex->indexCapacity(); + } + virtual inline uint64_t getAllocationSize() const override { return this->allocator->getAllocationSize() + this->backendIndex->getAllocationSize() + this->frontendIndex->getAllocationSize(); @@ -255,8 +526,10 @@ VecSimTieredIndex::topKQueryImp(const void *queryBlob, size_ // Simply query the main index and return the results while holding the lock. auto processed_query_ptr = this->frontendIndex->preprocessQuery(queryBlob); const void *processed_query = processed_query_ptr.get(); - std::shared_lock main_lock(this->mainIndexGuard); - return this->backendIndex->topKQuery(processed_query, k, queryParams); + auto mainLock = lockMainIndexForQuery(); + auto res = this->backendIndex->topKQuery(processed_query, k, queryParams); + + return res; } else { // No luck... first query the flat buffer and release the lock. // The query blob is already processed according to the frontend index. @@ -271,10 +544,12 @@ VecSimTieredIndex::topKQueryImp(const void *queryBlob, size_ auto processed_query_ptr = this->frontendIndex->preprocessQuery(queryBlob); const void *processed_query = processed_query_ptr.get(); - // Lock the main index and query it. - std::shared_lock main_lock(this->mainIndexGuard); - auto main_results = this->backendIndex->topKQuery(processed_query, k, queryParams); - main_lock.unlock(); + VecSimQueryReply *main_results; + { + // Lock the main index and query it. + auto mainLock = lockMainIndexForQuery(); + main_results = this->backendIndex->topKQuery(processed_query, k, queryParams); + } // If the query failed (currently only on timeout), return the error code. if (main_results->code != VecSim_QueryReply_OK) { @@ -317,10 +592,12 @@ VecSimTieredIndex::rangeQueryImp(const void *queryBlob, doub auto processed_query_ptr = this->frontendIndex->preprocessQuery(queryBlob); const void *processed_query = processed_query_ptr.get(); - // Simply query the main index and return the results while holding the lock. - this->mainIndexGuard.lock_shared(); - auto res = this->backendIndex->rangeQuery(processed_query, radius, queryParams); - this->mainIndexGuard.unlock_shared(); + VecSimQueryReply *res; + { + auto mainLock = lockMainIndexForQuery(); + // Simply query the main index and return the results while holding the lock. + res = this->backendIndex->rangeQuery(processed_query, radius, queryParams); + } // We could have passed the order to the main index, but we can sort them here after // unlocking it instead. @@ -341,9 +618,12 @@ VecSimTieredIndex::rangeQueryImp(const void *queryBlob, doub auto processed_query_ptr = this->frontendIndex->preprocessQuery(queryBlob); const void *processed_query = processed_query_ptr.get(); // Lock the main index and query it. - this->mainIndexGuard.lock_shared(); - auto main_results = this->backendIndex->rangeQuery(processed_query, radius, queryParams); - this->mainIndexGuard.unlock_shared(); + + VecSimQueryReply *main_results; + { + auto mainLock = lockMainIndexForQuery(); + main_results = this->backendIndex->rangeQuery(processed_query, radius, queryParams); + } // Merge the results and return, avoiding duplicates. // At this point, the return code of the FLAT index is OK, and the return code of the MAIN diff --git a/src/python_bindings/bindings.cpp b/src/python_bindings/bindings.cpp index b157e8b67..ab46ffb14 100644 --- a/src/python_bindings/bindings.cpp +++ b/src/python_bindings/bindings.cpp @@ -266,7 +266,9 @@ class PyVecSimIndex { return PyBatchIterator(index, py_batch_ptr); } - void runGC() { VecSimTieredIndex_GC(index.get()); } + // Not VecSimTieredIndex_GC(): that is a no-op unless the index is tiered, while + // SVSIndex::runGC() compacts and rebuilds reverse edges on a plain index too. + void runGC() { index->runGC(); } VecSimRelabelCode relabelVector(labelType old_label, labelType new_label) { py::gil_scoped_release py_gil; diff --git a/tests/flow/test_svs_tiered.py b/tests/flow/test_svs_tiered.py index 09ab8c93b..876aad8f8 100644 --- a/tests/flow/test_svs_tiered.py +++ b/tests/flow/test_svs_tiered.py @@ -166,6 +166,9 @@ def create_tiered_index(test_logger, is_multi: bool, num_per_label=1, data_type= test_logger.info(f"Total time for inserting vectors to the tiered index and indexing them into SVS using {threads_num}" f" threads took {round_ms(tiered_index_time)} ms") + index.run_gc() + index.wait_for_index() + # Measure total memory of the tiered index. tiered_memory = bytes_to_mega(index.index_memory()) @@ -175,6 +178,8 @@ def create_tiered_index(test_logger, is_multi: bool, num_per_label=1, data_type= _, svs_index_time, _ = indices_ctx.populate_index(svs_index) test_logger.info(f"Insert {num_elements} vectors directly to SVS index (one by one) took {round_(svs_index_time)} s") + + svs_index.run_gc() svs_memory = bytes_to_mega(svs_index.index_memory()) test_logger.info(f"total memory of svs index = {svs_memory} MB") diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index 539f239d2..c25fa1997 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -170,7 +170,7 @@ TYPED_TEST(HNSWTieredIndexTest, CreateIndexInstance) { auto job = new (allocator) HNSWInsertJob(tiered_index->allocator, vector_label, 0, insert_to_index, tiered_index); - auto jobs_vec = vecsim_stl::vector(1, job, allocator); + auto jobs_vec = vecsim_stl::vector(1, job, allocator); tiered_index->labelToInsertJobs.insert({vector_label, jobs_vec}); // Wrap this job with an array and submit the jobs to the queue. @@ -425,10 +425,10 @@ TYPED_TEST(HNSWTieredIndexTest, addVector) { sizeof(void *) + sizeof(size_t); // Account for the memory that was allocated in the labelToInsertJobs map (approx.) expected_mem += - sizeof( - vecsim_stl::unordered_map>::value_type) + + sizeof(vecsim_stl::unordered_map>::value_type) + sizeof(void *) + sizeof(size_t); - // Account for the inner buffer of the std::vector in the map. + // Account for the inner buffer of the std::vector in the map. expected_mem += sizeof(void *) + sizeof(size_t); // Account for the insert job that was created. expected_mem += sizeof(HNSWInsertJob) + sizeof(size_t); @@ -5374,7 +5374,7 @@ TYPED_TEST(HNSWTieredIndexTestBasic, relabelVectorMulti) { } ASSERT_EQ(tiered_index->labelToInsertJobs.count(7), 0); ASSERT_EQ(tiered_index->labelToInsertJobs.at(70).size(), per_label); - for (HNSWInsertJob *job : tiered_index->labelToInsertJobs.at(70)) { + for (TieredInsertJob *job : tiered_index->labelToInsertJobs.at(70)) { ASSERT_EQ(job->label, 70); } diff --git a/tests/unit/test_svs.cpp b/tests/unit/test_svs.cpp index 7684ff017..18803988c 100644 --- a/tests/unit/test_svs.cpp +++ b/tests/unit/test_svs.cpp @@ -35,6 +35,16 @@ } \ } +// Expected allocation for one block of `block_size` element slots, of which only +// `num_elements` are actually populated. +// +static size_t EstimateBlockSize(const SVSParams ¶ms, size_t block_size, size_t num_elements) { + const size_t reverse_edges_per_slot = SVSGraphBuilder::reverse_edges_element_size(); + const size_t reverse_edges_slots = svs::lib::SegmentedVector(num_elements).capacity(); + return (EstimateElementSize(params) - reverse_edges_per_slot) * block_size + + reverse_edges_per_slot * reverse_edges_slots; +} + // Log callback function to print non-debug log messages static void svsTestLogCallBackNoDebug(void *ctx, const char *level, const char *message) { if (level == nullptr || message == nullptr) { @@ -177,11 +187,22 @@ TYPED_TEST(SVSTest, svs_vector_update_test) { // Prepare new vector data and call addVector with the same id, different data. GenerateAndAddVector(index, dim, 1, 2.0); - // Index size shouldn't change. + // Now we have one deleted (not empty) slot + // and one live slot + EXPECT_EQ(VecSimIndex_IndexSize(index), 2); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), 1); + + index->runGC(); + // GC consolidated and compacted deleted slot. EXPECT_EQ(VecSimIndex_IndexSize(index), 1); // Delete the last vector. VecSimIndex_DeleteVector(index, 1); + // Still have a single "deleted" slot + EXPECT_EQ(VecSimIndex_IndexSize(index), 1); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), 1); + + index->runGC(); EXPECT_EQ(VecSimIndex_IndexSize(index), 0); ASSERT_EQ(svs_index->getNumMarkedDeleted(), 0); @@ -272,8 +293,11 @@ TYPED_TEST(SVSTest, svs_bulk_vectors_add_delete_test) { // Delete rest of the vectors // num_marked_deleted should reset. ASSERT_EQ(svs_index->deleteVectors(ids.data() + n - keep_num, keep_num), keep_num); - ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + ASSERT_EQ(VecSimIndex_IndexSize(index), n); ASSERT_EQ(index->indexLabelCount(), 0); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), n); + index->runGC(); + ASSERT_EQ(VecSimIndex_IndexSize(index), 0); ASSERT_EQ(svs_index->getNumMarkedDeleted(), 0); VecSimIndex_Free(index); } @@ -338,14 +362,13 @@ TYPED_TEST(SVSTest, two_stage_initialization_test) { impl = svs_index->createImpl(v.data(), ids.data(), n); EXPECT_THROW(svs_index->setImpl(std::move(impl)), std::logic_error); - // Delete rest of the vectors - index should be empty now and setImpl() should succeed. + // Delete rest of the vectors - index should be empty, but impl still alive. ASSERT_EQ(svs_index->deleteVectors(ids.data() + n - keep_num, keep_num), keep_num); - ASSERT_EQ(VecSimIndex_IndexSize(index), 0); - // Re-initialization should succeed. - impl = svs_index->createImpl(v.data(), ids.data(), n); - svs_index->setImpl(std::move(impl)); ASSERT_EQ(VecSimIndex_IndexSize(index), n); - runTopKSearchTest(index, query, k, verify_res, nullptr, BY_ID); + + // Re-initialization should fail again. + impl = svs_index->createImpl(v.data(), ids.data(), n); + EXPECT_THROW(svs_index->setImpl(std::move(impl)), std::logic_error); VecSimIndex_Free(index); } @@ -954,10 +977,11 @@ TYPED_TEST(SVSTest, svs_empty_index) { VecSimIndex_DeleteVector(index, 1); // Size equals 0. + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); - // The expected capacity should be 0 for empty index. - ASSERT_EQ(index->indexCapacity(), 0); + // Even an empty index holds the first block for nonquantized data + // ASSERT_EQ(index->indexCapacity(), 0); ASSERT_EQ(index->indexMetaDataCapacity(), index->indexCapacity()); // Try to remove it again. @@ -1205,6 +1229,7 @@ TYPED_TEST(SVSTest, test_dynamic_svs_info_iterator) { // Delete vector. VecSimIndex_DeleteVector(index, 0); + index->runGC(); info = VecSimIndex_DebugInfo(index); infoIter = VecSimIndex_DebugInfoIterator(index); ASSERT_EQ(0, info.commonInfo.indexSize); @@ -1416,6 +1441,7 @@ TYPED_TEST(SVSTest, svs_search_empty_index) { for (size_t i = 0; i < n; i++) { VecSimIndex_DeleteVector(index, i); } + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Again - we do not expect any results. @@ -1996,7 +2022,7 @@ TYPED_TEST(SVSTest, testSizeEstimation) { size_t actual = index->getAllocationSize(); ASSERT_EQ(estimation, actual); - estimation = EstimateElementSize(params) * bs; + estimation = EstimateBlockSize(params, bs, 1); GenerateAndAddVector(index, dim, 0); actual = index->getAllocationSize() - actual; // get the delta @@ -2855,7 +2881,8 @@ TEST(SVSTest, quant_modes) { const size_t dim = 4; const size_t n = 100; - const size_t k = 10; + // Include both boundary labels, 45 and 55, which tie for query 50. + const size_t k = 11; for (auto quant_bits : {VecSimSvsQuant_NONE, VecSimSvsQuant_Scalar, VecSimSvsQuant_8, VecSimSvsQuant_4, VecSimSvsQuant_4x4, VecSimSvsQuant_4x8, @@ -2909,7 +2936,7 @@ TEST(SVSTest, quant_modes) { ASSERT_EQ(VecSimIndex_IndexSize(index), n); - estimation = EstimateElementSize(params) * params.blockSize; + estimation = EstimateBlockSize(params, params.blockSize, n); actual = index->getAllocationSize() - actual; // get the delta ASSERT_GT(actual, 0); // LVQ element size estimation accuracy is low @@ -2936,7 +2963,8 @@ TEST(SVSTest, save_load) { const size_t dim = 4; const size_t n = 100; - const size_t k = 10; + // Include both boundary labels, 45 and 55, in the single-value case. + const size_t k = 11; // Helper function to convert quant_bits to string for error messages auto quant_bits_to_string = [](VecSimSvsQuantBits quant_bits) -> std::string { @@ -3029,7 +3057,7 @@ TEST(SVSTest, save_load) { float query[] = {50, 50, 50, 50}; auto verify_res = [&](size_t id, double score, size_t idx) { EXPECT_DOUBLE_EQ(VecSimIndex_GetDistanceFrom_Unsafe(index, id, query), score); - // Both single and multi should return labels starting from 45 + // Single-value labels are 45..55; multi-value labels are 20..30. if (is_multi) { // For multi, that label of {50,50,50,50} is 25 size_t expected_label = (20 + idx); @@ -3271,7 +3299,7 @@ TEST(SVSTest, scalar_quantization_query) { ASSERT_EQ(VecSimIndex_IndexSize(index_sq), n); ASSERT_EQ(index_sq->indexCapacity(), n); - estimation = EstimateElementSize(params) * params.blockSize; + estimation = EstimateBlockSize(params, params.blockSize, n); actual = index_sq->getAllocationSize() - actual; // get the delta ASSERT_GT(actual, 0); ASSERT_GE(estimation * 1.01, actual); diff --git a/tests/unit/test_svs_fp16.cpp b/tests/unit/test_svs_fp16.cpp index fcf2c7dd7..08c989dc1 100644 --- a/tests/unit/test_svs_fp16.cpp +++ b/tests/unit/test_svs_fp16.cpp @@ -40,6 +40,13 @@ using float16 = vecsim_types::float16; } \ } +static size_t EstimateBlockSize(const SVSParams ¶ms, size_t block_size, size_t num_elements) { + const size_t reverse_edges_per_slot = SVSGraphBuilder::reverse_edges_element_size(); + const size_t reverse_edges_slots = svs::lib::SegmentedVector(num_elements).capacity(); + return (EstimateElementSize(params) - reverse_edges_per_slot) * block_size + + reverse_edges_per_slot * reverse_edges_slots; +} + // Log callback function to print non-debug log messages static void svsTestLogCallBackNoDebug(void *ctx, const char *level, const char *message) { if (level == nullptr || message == nullptr) { @@ -192,12 +199,24 @@ TYPED_TEST(FP16SVSTest, svs_vector_update_test) { // Prepare new vector data and call addVector with the same id, different data. this->GenerateAndAddVector(index, dim, 1, 2.0); - // Index size shouldn't change. + // Now we have one deleted (not empty) slot + // and one live slot + EXPECT_EQ(VecSimIndex_IndexSize(index), 2); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), 1); + + index->runGC(); + // GC consolidated and compacted deleted slot. EXPECT_EQ(VecSimIndex_IndexSize(index), 1); // Delete the last vector. VecSimIndex_DeleteVector(index, 1); + // Still have a single "deleted" slot + EXPECT_EQ(VecSimIndex_IndexSize(index), 1); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), 1); + + index->runGC(); EXPECT_EQ(VecSimIndex_IndexSize(index), 0); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), 0); VecSimIndex_Free(index); } @@ -257,8 +276,11 @@ TYPED_TEST(FP16SVSTest, svs_bulk_vectors_add_delete_test) { // Delete rest of the vectors // num_marked_deleted should reset. ASSERT_EQ(svs_index->deleteVectors(ids.data() + n - keep_num, keep_num), keep_num); - ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + ASSERT_EQ(VecSimIndex_IndexSize(index), n); ASSERT_EQ(index->indexLabelCount(), 0); + ASSERT_EQ(svs_index->getNumMarkedDeleted(), n); + index->runGC(); + ASSERT_EQ(VecSimIndex_IndexSize(index), 0); ASSERT_EQ(svs_index->getNumMarkedDeleted(), 0); VecSimIndex_Free(index); } @@ -753,10 +775,11 @@ TYPED_TEST(FP16SVSTest, svs_empty_index) { for (size_t i = 0; i < n; i++) { VecSimIndex_DeleteVector(index, i); } + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); - // The expected capacity should be 0 for empty index. - ASSERT_EQ(index->indexCapacity(), 0); + // Even an empty index holds the first block for nonquantized data + // ASSERT_EQ(index->indexCapacity(), 0); // Try to remove it again. VecSimIndex_DeleteVector(index, 1); @@ -1018,6 +1041,7 @@ TYPED_TEST(FP16SVSTest, test_dynamic_svs_info_iterator) { // Delete vector. VecSimIndex_DeleteVector(index, 0); + index->runGC(); info = VecSimIndex_DebugInfo(index); infoIter = VecSimIndex_DebugInfoIterator(index); ASSERT_EQ(0, info.commonInfo.indexSize); @@ -1388,7 +1412,7 @@ TYPED_TEST(FP16SVSTest, testSizeEstimation) { size_t actual = index->getAllocationSize(); ASSERT_EQ(estimation, actual); - estimation = EstimateElementSize(params) * bs; + estimation = EstimateBlockSize(params, bs, 1); this->GenerateAndAddVector(index, dim, 0); actual = index->getAllocationSize() - actual; // get the delta @@ -1589,7 +1613,7 @@ TYPED_TEST(FP16SVSTest, quant_modes) { ASSERT_EQ(VecSimIndex_IndexSize(index), n); - estimation = EstimateElementSize(params) * DEFAULT_BLOCK_SIZE; + estimation = EstimateBlockSize(params, DEFAULT_BLOCK_SIZE, n); actual = index->getAllocationSize() - actual; // get the delta ASSERT_GT(actual, 0); // LVQ element size estimation accuracy is low @@ -1666,15 +1690,19 @@ TYPED_TEST(FP16SVSTest, scalar_quantization_query) { const size_t k = 10; const double quant_precision = 1.0 / (1 << 7); // int8 quantization precision + std::default_random_engine gen; + std::uniform_real_distribution dist(-1.0, 1.0); std::vector> dataset(n); for (size_t i = 0; i < n; i++) { - test_utils::populate_float16_vec(dataset[i].data(), dim, i, -1.0, 1.0); + for (size_t j = 0; j < dim; j++) { + dataset[i][j] = vecsim_types::FP32_to_FP16(dist(gen)); + } } std::vector ids(n); std::iota(ids.begin(), ids.end(), 0); float16 query[dim]; - test_utils::populate_float16_vec(query, dim, n, -1.0, 1.0); + this->GenerateVector(query, dim, 0.1f); VecSimQueryReply *fp_results = nullptr; auto verify_res = [&](size_t id, double score, size_t result_rank) { @@ -1724,7 +1752,7 @@ TYPED_TEST(FP16SVSTest, scalar_quantization_query) { ASSERT_EQ(VecSimIndex_IndexSize(index_sq), n); ASSERT_EQ(index_sq->indexCapacity(), n); - estimation = EstimateElementSize(params) * params.blockSize; + estimation = EstimateBlockSize(params, params.blockSize, n); actual = index_sq->getAllocationSize() - actual; // get the delta ASSERT_GT(actual, 0); ASSERT_GE(estimation * 1.01, actual); @@ -1794,6 +1822,7 @@ TYPED_TEST(FP16SVSMultiTest, vector_add_multiple_test) { // Deleting the label. All the vectors should be deleted. ASSERT_EQ(VecSimIndex_DeleteVector(index, label), n); + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); ASSERT_EQ(index->indexLabelCount(), 0); @@ -2031,6 +2060,7 @@ TYPED_TEST(FP16SVSMultiTest, reindexing_same_vector_different_id) { for (size_t i = 0; i < n; i++) { VecSimIndex_DeleteVector(index, i); } + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Reinsert the same vectors under different ids than before. @@ -2328,7 +2358,7 @@ class FP16SVSTieredIndexTest : public FP16SVSTest { } } // Submit the index update job. - tiered_index->scheduleSVSIndexUpdate(); + tiered_index->scheduleSVSIndexInit(); ASSERT_EQ(mock_thread_pool.jobQ.size(), mock_thread_pool.thread_pool_size); // Execute the job from the queue and validate that the index was updated properly. @@ -2839,6 +2869,9 @@ TYPED_TEST(FP16SVSTieredIndexTest, deleteVector) { // Remove from main index. ASSERT_EQ(tiered_index->deleteVector(vec_label), 1); ASSERT_EQ(tiered_index->indexLabelCount(), 0); + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->indexSize(), 0); // Re-insert a deleted label with a different vector. @@ -2848,8 +2881,9 @@ TYPED_TEST(FP16SVSTieredIndexTest, deleteVector) { ASSERT_EQ(tiered_index->indexSize(), 1); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); - // Move the vector to SVS by executing the insert job. - mock_thread_pool.thread_iteration(); + // Move the vector to SVS by executing the pending jobs. + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); // Scalar quantization accuracy is insufficient for this check. @@ -2894,9 +2928,15 @@ TYPED_TEST(FP16SVSTieredIndexTest, deleteVectorMulti) { ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); ASSERT_EQ(tiered_index->deleteVector(vec_label), 2); + + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->indexSize(), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 0); - mock_thread_pool.thread_iteration(); + + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); // Test deleting a label for which both of its vector's is in the flat index. @@ -2910,13 +2950,21 @@ TYPED_TEST(FP16SVSTieredIndexTest, deleteVectorMulti) { ASSERT_EQ(tiered_index->deleteVector(vec_label), 2); ASSERT_EQ(tiered_index->indexLabelCount(), 0); + + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); + mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); mock_thread_pool.thread_iteration(); ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); // Test deleting a label for which both of its vector's is in SVS index. this->GenerateAndAddVector(tiered_index, dim, vec_label, vec_label); this->GenerateAndAddVector(tiered_index, dim, vec_label, other_vec_val); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 2); diff --git a/tests/unit/test_svs_multi.cpp b/tests/unit/test_svs_multi.cpp index 010765657..8460e6ac2 100644 --- a/tests/unit/test_svs_multi.cpp +++ b/tests/unit/test_svs_multi.cpp @@ -23,6 +23,13 @@ } \ } +static size_t EstimateBlockSize(const SVSParams ¶ms, size_t block_size, size_t num_elements) { + const size_t reverse_edges_per_slot = SVSGraphBuilder::reverse_edges_element_size(); + const size_t reverse_edges_slots = svs::lib::SegmentedVector(num_elements).capacity(); + return (EstimateElementSize(params) - reverse_edges_per_slot) * block_size + + reverse_edges_per_slot * reverse_edges_slots; +} + // Log callback function to print non-debug log messages static void svsTestLogCallBackNoDebug(void *ctx, const char *level, const char *message) { if (level == nullptr || message == nullptr) { @@ -121,6 +128,7 @@ TYPED_TEST(SVSMultiTest, vector_add_multiple_test) { // Deleting the label. All the vectors should be deleted. ASSERT_EQ(VecSimIndex_DeleteVector(index, id), rep); + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); ASSERT_EQ(index->indexLabelCount(), 0); @@ -157,6 +165,7 @@ TYPED_TEST(SVSMultiTest, empty_index) { ASSERT_EQ(VecSimIndex_DeleteVector(index, 1), 3); // Size equals 0. + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Try to remove it again. @@ -424,6 +433,7 @@ TYPED_TEST(SVSMultiTest, reindexing_same_vector_different_id) { for (size_t i = 0; i < n; i++) { VecSimIndex_DeleteVector(index, i); } + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Reinsert the same vectors under different ids than before. @@ -609,6 +619,7 @@ TYPED_TEST(SVSMultiTest, search_empty_index) { } ASSERT_EQ(VecSimIndex_IndexSize(index), n); VecSimIndex_DeleteVector(index, 46); + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Again - we do not expect any results. @@ -748,7 +759,7 @@ TYPED_TEST(SVSMultiTest, testSizeEstimation) { size_t actual = index->getAllocationSize(); ASSERT_EQ(estimation, actual); - estimation = EstimateElementSize(params) * bs; + estimation = EstimateBlockSize(params, bs, 1); GenerateAndAddVector(index, dim, 0); actual = index->getAllocationSize() - actual; // get the delta @@ -779,16 +790,15 @@ TYPED_TEST(SVSMultiTest, emptyIndex) { // Try to remove it. ASSERT_EQ(VecSimIndex_DeleteVector(index, 1), 1); - // The capacity should change to be zero. - ASSERT_EQ(index->indexCapacity(), 0); + // Even an empty index holds the first block for nonquantized data + // ASSERT_EQ(index->indexCapacity(), 0); // Size equals 0. + index->runGC(); ASSERT_EQ(VecSimIndex_IndexSize(index), 0); // Try to remove it again. - // The capacity should remain unchanged, as we are trying to delete a label that doesn't exist. ASSERT_EQ(VecSimIndex_DeleteVector(index, 1), 0); - ASSERT_EQ(index->indexCapacity(), 0); // Nor the size. ASSERT_EQ(VecSimIndex_IndexSize(index), 0); diff --git a/tests/unit/test_svs_tiered.cpp b/tests/unit/test_svs_tiered.cpp index 4f765453e..6fabc1cf5 100644 --- a/tests/unit/test_svs_tiered.cpp +++ b/tests/unit/test_svs_tiered.cpp @@ -11,8 +11,16 @@ #include "mock_thread_pool.h" #if HAVE_SVS +#include +#include +#include +#include +#include +#include +#include #include #include +#include // For getAvailableCPUs(): #include @@ -32,6 +40,11 @@ } \ } +static size_t EstimateReverseEdgesSize(size_t num_elements) { + return SVSGraphBuilder::reverse_edges_element_size() * + svs::lib::SegmentedVector(num_elements).capacity(); +} + // Get available number of CPUs // Returns the number of logical processors on the process // Returns std::thread::hardware_concurrency() if the number of logical processors is not available @@ -371,6 +384,10 @@ TYPED_TEST(SVSTieredIndexTest, updateVectors) { ASSERT_EQ(tiered_index->updateVectors(7, two, 2), VecSimUpdate_OK); mock_thread_pool.init_threads(); mock_thread_pool.thread_pool_join(); + // The replaced vectors leave tombstones behind, so compact before counting. + VecSimTieredIndex_GC(tiered_index); + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); ASSERT_EQ(tiered_index->indexSize(), 2); ASSERT_EQ(tiered_index->indexLabelCount(), 1); for (size_t i = 0; i < 2; i++) { @@ -457,6 +474,7 @@ TYPED_TEST(SVSTieredIndexTest, updateVectorsDuringUpdateJob) { ASSERT_EQ(tiered_index->updateVectors(i, replacements, new_per_label), VecSimUpdate_OK) << "label " << i; } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); auto *svs_index = tiered_index->GetSVSIndex(); @@ -572,6 +590,7 @@ TYPED_TEST(SVSTieredIndexTest, addRelabelAndUpdateTogetherUnderLoad) { GenerateVector(fresh, dim, 7000.0f + 10.0f * i); ASSERT_EQ(VecSimIndex_AddVector(tiered_index, fresh, fresh_base + i), 1); } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); auto *svs_index = tiered_index->GetSVSIndex(); @@ -682,6 +701,342 @@ TYPED_TEST(SVSTieredIndexTest, relabelVectorMovesTheLabelInBothWriteStates) { runTopKSearchTest(tiered_index, vector, 1, verify); } +// The test above never reaches `labelToInsertJobs`: with a training threshold of 1 the first +// vector is training data, and training data has no insert job. Draining once makes the backend +// ready, so the next add does get one. +TYPED_TEST(SVSTieredIndexTest, relabelVectorRemapsAPendingInsertJob) { + size_t dim = 4; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti(), + .quantBits = TypeParam::get_quant_bits()}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + + mock_thread_pool.init_threads(); + GenerateAndAddVector(tiered_index, dim, 0, 0); + mock_thread_pool.thread_pool_join(); + ASSERT_TRUE(tiered_index->GetSVSIndex()->ready()); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); + + const size_t per_label = TypeParam::isMulti() ? 2 : 1; + for (size_t j = 0; j < per_label; j++) { + GenerateAndAddVector(tiered_index, dim, 7, 7 + j); + } + ASSERT_EQ(mock_thread_pool.jobQ.size(), per_label); + ASSERT_EQ(tiered_index->labelToInsertJobs.at(7).size(), per_label); + + ASSERT_EQ(VecSimIndex_RelabelVector(tiered_index, 7, 70), VecSimRelabel_OK); + + ASSERT_EQ(tiered_index->labelToInsertJobs.count(7), 0); + ASSERT_EQ(tiered_index->labelToInsertJobs.at(70).size(), per_label); + for (auto *job : tiered_index->labelToInsertJobs.at(70)) { + ASSERT_EQ(job->label, 70); + // Left Delayed, the ingest below would defer forever. + ASSERT_EQ(static_cast(job)->status.load(), SVSInsertJob::Status::Pending); + } + ASSERT_TRUE(tiered_index->GetFlatIndex()->isLabelExists(70)); + ASSERT_FALSE(tiered_index->GetFlatIndex()->isLabelExists(7)); + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(70)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(7)); + ASSERT_EQ(tiered_index->indexLabelCount(), 2); +} + +// The worker consumes the delayed submission; relabel must resume it exactly once. +TYPED_TEST(SVSTieredIndexTest, relabelVectorDefersAnInsertJobInsideTheDelayWindow) { + size_t dim = 4; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti(), + .quantBits = TypeParam::get_quant_bits()}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + + mock_thread_pool.init_threads(); + GenerateAndAddVector(tiered_index, dim, 0, 0); + mock_thread_pool.thread_pool_join(); + ASSERT_TRUE(tiered_index->GetSVSIndex()->ready()); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); + + GenerateAndAddVector(tiered_index, dim, 7, 7); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + + auto *job = static_cast(mock_thread_pool.jobQ.front().job); + tiered_index->registerTracingCallback("Relabel::before_backend_remap", [&]() { + std::thread worker([&]() { mock_thread_pool.thread_iteration(); }); + worker.join(); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + EXPECT_EQ(job->status.load(), SVSInsertJob::Status::Parked); + EXPECT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(70)); + }); + ASSERT_EQ(VecSimIndex_RelabelVector(tiered_index, 7, 70), VecSimRelabel_OK); + + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); + ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(7)); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + ASSERT_EQ(mock_thread_pool.jobQ.front().job, job); + + ASSERT_EQ(job->status.load(), SVSInsertJob::Status::Pending); + ASSERT_EQ(job->label, 70); + mock_thread_pool.thread_iteration(); + + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(70)); + ASSERT_TRUE(tiered_index->labelToInsertJobs.empty()); + ASSERT_EQ(tiered_index->indexLabelCount(), 2); +} + +// isLabelExists() reports a soft-deleted label as free, but SVS keeps its translator entry +// until the label's consolidate job runs, and would refuse the remap. So relabelVector takes +// those jobs over and consolidates them itself. +TYPED_TEST(SVSTieredIndexTest, relabelVectorTakesOverPendingConsolidatesOfTheTarget) { + size_t dim = 4; + const size_t n = 3; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti(), + .quantBits = TypeParam::get_quant_bits()}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + + mock_thread_pool.init_threads(); + for (size_t i = 0; i < n; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + for (size_t i = 0; i < n; i++) { + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(i)); + } + + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 1), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + ASSERT_EQ(mock_thread_pool.jobQ.front().job->jobType, SVS_CONSOLIDATE_JOB); + auto *first_job = static_cast(mock_thread_pool.jobQ.front().job); + + ASSERT_EQ(VecSimIndex_RelabelVector(tiered_index, 2, 1), VecSimRelabel_OK); + ASSERT_EQ(first_job->status.load(), SVSConsolidateJob::Status::TakenOver); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(1)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(2)); + + // The first job is still registered under label 1, so this one is appended to it. + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 1), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); + + // The takeover now has to skip the job it already invalidated. + ASSERT_EQ(VecSimIndex_RelabelVector(tiered_index, 0, 1), VecSimRelabel_OK); + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(1)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(0)); + ASSERT_EQ(tiered_index->indexLabelCount(), 1); +} + +// A consolidate job that is already inside consolidate() cannot be taken over, so the takeover +// waits for it to leave. The test holds the flag and does the consolidation the wrapper would, +// so the remap that follows really is unblocked. +TYPED_TEST(SVSTieredIndexTest, relabelVectorWaitsForAnExecutingConsolidate) { + size_t dim = 4; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti(), + .quantBits = TypeParam::get_quant_bits()}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + + mock_thread_pool.init_threads(); + for (size_t i = 0; i < 2; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(1)); + + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 1), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + auto *consolidate_job = static_cast(mock_thread_pool.jobQ.front().job); + consolidate_job->status.store(SVSConsolidateJob::Status::Executing, std::memory_order_release); + + std::atomic relabel_code{-1}; + std::thread relabeler( + [&]() { relabel_code = static_cast(VecSimIndex_RelabelVector(tiered_index, 0, 1)); }); + + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + // EXPECT, so the thread is always joined below. + EXPECT_EQ(relabel_code.load(), -1) << "the takeover did not wait for the running job"; + + tiered_index->GetSVSIndex()->consolidate({1}); + consolidate_job->status.store(SVSConsolidateJob::Status::Done, std::memory_order_release); + relabeler.join(); + + ASSERT_EQ(relabel_code.load(), static_cast(VecSimRelabel_OK)); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(1)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(0)); + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + ASSERT_EQ(tiered_index->indexLabelCount(), 1); +} + +// Use the real worker's Done -> unregister -> delete path while relabel holds its job pointer. +// Without the shared registry guard, deletion can precede waitForJob's next load (an ASan UAF). +TYPED_TEST(SVSTieredIndexTest, relabelVectorProtectsExecutingConsolidateJobLifetime) { + const size_t dim = 4; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti(), + .quantBits = TypeParam::get_quant_bits()}; + VecSimParams svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 2); + ASSERT_INDEX(tiered_index); + + GenerateAndAddVector(tiered_index, dim, 0, 0); + GenerateAndAddVector(tiered_index, dim, 1, 1); + mock_thread_pool.thread_iteration(); // initialize before starting the consolidate worker + ASSERT_TRUE(tiered_index->GetSVSIndex()->ready()); + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 1), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + ASSERT_EQ(mock_thread_pool.jobQ.front().job->jobType, SVS_CONSOLIDATE_JOB); + + std::promise worker_paused, relabel_waiting, job_deleted; + auto paused = worker_paused.get_future(); + auto waiting = relabel_waiting.get_future(); + auto deleted = job_deleted.get_future(); + const auto timeout = std::chrono::seconds(5); + bool worker_saw_waiter = false; + bool waiter_entered = false; + bool deleted_while_waiting = false; + tiered_index->registerTracingCallback("ConsolidateJob::before_done", [&] { + worker_paused.set_value(); + worker_saw_waiter = waiting.wait_for(timeout) == std::future_status::ready; + }); + tiered_index->registerTracingCallback("Relabel::before_wait_for_consolidate", [&] { + waiter_entered = true; + relabel_waiting.set_value(); + // Let the worker publish Done and attempt cleanup before waitForJob reads the job. + // The registry guard must prevent deletion until this callback and the wait return. + deleted_while_waiting = + deleted.wait_for(std::chrono::milliseconds(100)) == std::future_status::ready; + }); + tiered_index->registerTracingCallback("ConsolidateJob::after_delete", + [&] { job_deleted.set_value(); }); + + mock_thread_pool.init_threads(); + const bool reached_worker = paused.wait_for(timeout) == std::future_status::ready; + auto relabel_code = VecSimRelabel_OldLabelMissing; + if (reached_worker) { + // Public writes stay on the test thread; the worker alone changes the job's status. + relabel_code = VecSimIndex_RelabelVector(tiered_index, 0, 1); + } + if (!waiter_entered) { + relabel_waiting.set_value(); // release the worker if relabel never reached its wait + } + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + ASSERT_TRUE(reached_worker); + ASSERT_TRUE(worker_saw_waiter); + ASSERT_TRUE(waiter_entered); + EXPECT_FALSE(deleted_while_waiting); + EXPECT_EQ(deleted.wait_for(std::chrono::seconds(0)), std::future_status::ready); + EXPECT_EQ(relabel_code, VecSimRelabel_OK); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + EXPECT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(1)); + EXPECT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(0)); + EXPECT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + EXPECT_EQ(tiered_index->indexLabelCount(), 1); +} + +// Rollback needs the backend to refuse the target, and to hold old_label live while the buffer +// still holds it with a registered job. Both are built by hand: +// - the target keeps a stale translator entry, because a backend-side delete leaves no +// consolidate job for the takeover to run. Single-value only - the multi backend checks live +// labels only, so a soft-deleted target never blocks it. +// - old_label is live in both tiers, the state an insert job is in after it published and +// before it dropped its buffer copy. Published here directly, then the job marked Done. +TYPED_TEST(SVSTieredIndexTestBasic, relabelVectorRollsBackWhenTheBackendRefusesTheTarget) { + size_t dim = 4; + SVSParams params = { + .type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_L2, .multi = false}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + + mock_thread_pool.init_threads(); + GenerateAndAddVector(tiered_index, dim, 0, 0); + mock_thread_pool.thread_pool_join(); + ASSERT_TRUE(tiered_index->GetSVSIndex()->ready()); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); + + GenerateAndAddVector(tiered_index->GetBackendIndex(), dim, 5, 5); + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index->GetBackendIndex(), 5), 1); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(5)); + ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 1); + + GenerateAndAddVector(tiered_index, dim, 7, 7); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + auto *job = static_cast(mock_thread_pool.jobQ.front().job); + + GenerateAndAddVector(tiered_index->GetBackendIndex(), dim, 7, 7); + job->status.store(SVSInsertJob::Status::Done, std::memory_order_relaxed); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(7)); + ASSERT_TRUE(tiered_index->GetFlatIndex()->isLabelExists(7)); + + // Keep the original completed-job case, then exercise rollback with a parked submission. + for (auto status : {SVSInsertJob::Status::Done, SVSInsertJob::Status::Pending}) { + job->status.store(status, std::memory_order_relaxed); + tiered_index->registerTracingCallback("Relabel::before_backend_remap", [&, status]() { + if (status == SVSInsertJob::Status::Pending) { + std::thread worker([&]() { mock_thread_pool.thread_iteration(); }); + worker.join(); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + EXPECT_EQ(job->status.load(), SVSInsertJob::Status::Parked); + } + }); + ASSERT_EQ(VecSimIndex_RelabelVector(tiered_index, 7, 5), VecSimRelabel_NewLabelTaken); + + ASSERT_TRUE(tiered_index->GetFlatIndex()->isLabelExists(7)); + ASSERT_FALSE(tiered_index->GetFlatIndex()->isLabelExists(5)); + ASSERT_EQ(tiered_index->labelToInsertJobs.count(5), 0); + ASSERT_EQ(tiered_index->labelToInsertJobs.at(7).size(), 1); + ASSERT_EQ(job->label, 7); + ASSERT_EQ(job->status.load(), status); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + ASSERT_EQ(mock_thread_pool.jobQ.front().job, job); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(7)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(5)); + } + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(7)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(5)); + ASSERT_EQ(tiered_index->indexLabelCount(), 2); +} + TYPED_TEST(SVSTieredIndexTest, relabelVectorRejectsOnATier) { size_t dim = 4; SVSParams params = {.type = TypeParam::get_index_type(), @@ -933,6 +1288,387 @@ TYPED_TEST(SVSTieredIndexTest, relabelVectorCannotLandInsideAnUpdateJobsWindow) #endif // HAVE_SVS_REPLACE_EXTERNAL_ID +TYPED_TEST(SVSTieredIndexTestBasic, addVectorDuringInitializationWithFullFlatBuffer) { + constexpr size_t dim = 4; + constexpr size_t n = 8; + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = multi}; + auto svs_params = CreateParams(params); + std::atomic submitted_insert_jobs{0}; + tieredIndexMock mock_thread_pool(1); + auto tiered_params = this->CreateTieredSVSParams(svs_params, mock_thread_pool, n); + tiered_params.flatBufferLimit = n; + // The mock resolves its queue through index_ctx; use jobQueue to count submissions. + tiered_params.jobQueue = &submitted_insert_jobs; + tiered_params.submitCb = [](void *counter, void *ctx, AsyncJob **jobs, + JobCallback *callbacks, size_t count) { + for (size_t i = 0; i < count; ++i) { + if (jobs[i]->jobType == SVS_INSERT_VECTOR_JOB) { + static_cast *>(counter)->fetch_add( + 1, std::memory_order_relaxed); + } + } + return tieredIndexMock::submit_callback(nullptr, ctx, jobs, callbacks, count); + }; + auto *index = this->CreateTieredSVSIndex(tiered_params, mock_thread_pool); + ASSERT_INDEX(index); + + for (size_t i = 0; i < n; ++i) { + GenerateAndAddVector(index, dim, i, i); + } + ASSERT_EQ(index->GetFlatIndex()->indexSize(), n); + ASSERT_FALSE(index->GetSVSIndex()->ready()); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + + std::promise snapshot_taken; + auto snapshot_ready = snapshot_taken.get_future(); + std::promise additions_finished; + auto additions_ready = additions_finished.get_future(); + bool completed_during_init = false; + size_t buffered_during_init = 0; + // Pause with the snapshot fixed, flat unlocked, and updateJobMutex held exclusively. + index->registerTracingCallback("UpdateJob::after_snapshot", [&] { + snapshot_taken.set_value(); + completed_during_init = + additions_ready.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + if (completed_during_init) { + buffered_during_init = index->GetFlatIndex()->indexSize(); + } + // On timeout, let init finish so a blocked add can return without hanging the test. + }); + mock_thread_pool.init_threads(); + const bool reached_snapshot = + snapshot_ready.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + if (reached_snapshot) { + // Public writes stay on the test thread. The first add sees a full buffer; + // the second sees a buffer above the limit. Both must return before init resumes. + EXPECT_EQ(GenerateAndAddVector(index, dim, n, n), 1); + EXPECT_EQ(GenerateAndAddVector(index, dim, n + 1, n + 1), 1); + // The only worker is paused in init. Neither insert may be submitted yet, + // where another worker could pick it up and block behind updateJobMutex. + std::lock_guard queue_lock{mock_thread_pool.queue_guard}; + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } + additions_finished.set_value(); + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + ASSERT_TRUE(reached_snapshot) << "Initialization did not reach the snapshot hook"; + ASSERT_TRUE(completed_during_init) << "addVector waited for initialization"; + EXPECT_EQ(buffered_during_init, n + 2); + EXPECT_TRUE(index->GetSVSIndex()->ready()); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), 0); + EXPECT_EQ(index->GetBackendIndex()->indexSize(), n + 2); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), n + 2); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + // Both deferred inserts were consumed by init, not queued and later drained by the mock. + EXPECT_EQ(submitted_insert_jobs.load(std::memory_order_relaxed), 0); + EXPECT_EQ(index->GetSVSIndex()->getParallelism(), 1); + for (size_t i = 0; i < n + 2; ++i) { + std::vector> stored; + index->getDataByLabel(i, stored); + const std::vector> expected{ + std::vector(dim, static_cast(i))}; + EXPECT_EQ(stored, expected) << "label=" << i; + } + } +} + +TYPED_TEST(SVSTieredIndexTestBasic, mainLockLimitWaitsForInitialization) { + constexpr size_t dim = 4; + constexpr size_t training_rows = 8; + constexpr size_t main_lock_limit = 100'000; + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = multi}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto tiered_params = + this->CreateTieredSVSParams(svs_params, mock_thread_pool, training_rows); + tiered_params.flatBufferLimit = training_rows; + auto *index = this->CreateTieredSVSIndex(tiered_params, mock_thread_pool); + ASSERT_INDEX(index); + + std::promise before_publish, resume_init, before_wait; + auto publish_ready = before_publish.get_future(); + auto resume_ready = resume_init.get_future(); + auto wait_ready = before_wait.get_future(); + std::atomic pause_timed_out{false}; + std::atomic wait_seen{false}; + index->registerTracingCallback("UpdateJob::before_publish", [&] { + before_publish.set_value(); + pause_timed_out.store(resume_ready.wait_for(std::chrono::seconds(60)) != + std::future_status::ready); + }); + index->registerTracingCallback("AddVector::before_init_wait", [&] { + if (!wait_seen.exchange(true)) { + before_wait.set_value(); + } + }); + + for (size_t i = 0; i < training_rows; ++i) { + EXPECT_EQ(GenerateAndAddVector(index, dim, i, i), 1); + } + mock_thread_pool.init_threads(); + const bool reached_publish = + publish_ready.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + + size_t added = training_rows; + if (reached_publish) { + // Exercise the actual policy floor, not a test-only smaller limit. + // Each addition below the limit must return while init is held open. + for (; added < main_lock_limit; ++added) { + EXPECT_EQ(GenerateAndAddVector(index, dim, added, added), 1); + if (pause_timed_out.load()) { + break; + } + } + } + const bool filled_during_init = + reached_publish && added == main_lock_limit && !pause_timed_out.load(); + std::future writer; + bool reached_wait = false; + bool blocked_at_limit = false; + bool still_waiting_after_delete = false; + if (filled_during_init) { + writer = std::async(std::launch::async, [&] { + return GenerateAndAddVector(index, dim, main_lock_limit, + main_lock_limit); + }); + reached_wait = + wait_ready.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + blocked_at_limit = reached_wait && writer.wait_for(std::chrono::milliseconds(50)) == + std::future_status::timeout; + if (blocked_at_limit && !pause_timed_out.load()) { + // Acquiring flat shared here also checks that the waiting add released it. + index->acquireSharedLocks(); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), main_lock_limit); + index->releaseSharedLocks(); + EXPECT_FALSE(index->GetSVSIndex()->ready()); + + // Keep the snapshot, but invalidate all deferred filler jobs. This avoids + // constructing a large graph; the wait must still last until init completes. + for (size_t i = training_rows; i < main_lock_limit; ++i) { + EXPECT_EQ(index->deleteVector(i), 1); + } + still_waiting_after_delete = + writer.wait_for(std::chrono::milliseconds(0)) == std::future_status::timeout; + } + } + + // Release init before any fatal assertion or future destruction. + resume_init.set_value(); + int add_result = -1; + if (writer.valid()) { + EXPECT_NO_THROW(add_result = writer.get()); + } + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + ASSERT_TRUE(reached_publish); + ASSERT_TRUE(filled_during_init); + ASSERT_TRUE(reached_wait); + ASSERT_TRUE(blocked_at_limit); + ASSERT_TRUE(still_waiting_after_delete); + ASSERT_FALSE(pause_timed_out.load()); + EXPECT_EQ(add_result, 1); + EXPECT_TRUE(index->GetSVSIndex()->ready()); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), 0); + EXPECT_EQ(index->GetBackendIndex()->indexSize(), training_rows + 1); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), training_rows + 1); + EXPECT_TRUE(index->GetSVSIndex()->isLabelExists(main_lock_limit)); + EXPECT_TRUE(index->labelToInsertJobs.empty()); + EXPECT_TRUE(index->invalidJobs.empty()); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } +} + +TYPED_TEST(SVSTieredIndexTestBasic, overwriteSnapshotLabelDuringInitialization) { + constexpr size_t dim = 4; + constexpr size_t n = 8; + constexpr labelType overwritten_label = 0; + for (size_t overwrite_count : {size_t{1}, size_t{2}}) { + SCOPED_TRACE(::testing::Message() << "overwrite_count=" << overwrite_count); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = false}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto tiered_params = this->CreateTieredSVSParams(svs_params, mock_thread_pool, n); + tiered_params.flatBufferLimit = n; + auto *index = this->CreateTieredSVSIndex(tiered_params, mock_thread_pool); + ASSERT_INDEX(index); + + for (size_t i = 0; i < n; ++i) { + GenerateAndAddVector(index, dim, i, i); + } + ASSERT_FALSE(index->GetSVSIndex()->ready()); + ASSERT_EQ(index->GetFlatIndex()->indexSize(), n); + ASSERT_EQ(index->labelToInsertJobs.count(overwritten_label), 0); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + + std::promise snapshot_taken, overwrites_finished; + auto snapshot_ready = snapshot_taken.get_future(); + auto overwrites_ready = overwrites_finished.get_future(); + const auto timeout = std::chrono::seconds(5); + bool completed_during_init = false; + index->registerTracingCallback("UpdateJob::after_snapshot", [&] { + snapshot_taken.set_value(); + completed_during_init = overwrites_ready.wait_for(timeout) == std::future_status::ready; + // Release init on timeout as well, so a regressed blocking add can return. + }); + + mock_thread_pool.init_threads(); + const bool reached_snapshot = snapshot_ready.wait_for(timeout) == std::future_status::ready; + if (reached_snapshot) { + // The first overwrite cancels a captured row with no insert job. Removing + // label 0 swaps the last captured row down; init cleanup must later preserve + // the replacement and update its insert job's id as the buffer shrinks. + // A second overwrite also exercises invalidation of an unsubmitted job. + for (size_t i = 0; i < overwrite_count; ++i) { + EXPECT_EQ( + GenerateAndAddVector(index, dim, overwritten_label, n + i + 1), 0); + } + } + overwrites_finished.set_value(); + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + ASSERT_TRUE(reached_snapshot) << "Initialization did not reach the snapshot hook"; + ASSERT_TRUE(completed_during_init) << "Overwrite waited for initialization"; + EXPECT_TRUE(index->GetSVSIndex()->ready()); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), 0); + // Per-label consolidation does not compact the deleted snapshot row. + EXPECT_EQ(index->GetBackendIndex()->indexSize(), n + 1); + EXPECT_EQ(index->getNumMarkedDeleted(), 1); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), n); + EXPECT_TRUE(index->labelToInsertJobs.empty()); + EXPECT_TRUE(index->invalidJobs.empty()); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + for (size_t i = 0; i < n; ++i) { + // Read the backend directly so a surviving flat row cannot hide a stale + // snapshot copy. Exact cardinality and values detect duplicates as well. + std::vector> stored; + index->GetBackendIndex()->getDataByLabel(i, stored); + const auto value = i == overwritten_label ? n + overwrite_count : i; + const std::vector> expected{ + std::vector(dim, static_cast(value))}; + EXPECT_EQ(stored, expected) << "label=" << i; + } + } +} + +// A search already occupies one worker, so init proceeds after its reservation timeout. +// The search reaches the backend while init is still paused immediately after publication. +TYPED_TEST(SVSTieredIndexTestBasic, deleteVectorBeforeInitializationPublishes) { + if (std::thread::hardware_concurrency() < 2) { + GTEST_SKIP() << "Requires two workers"; + } + constexpr size_t dim = 4; + constexpr size_t n = 8; + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = multi}; + params.search_window_size = n; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(2); + auto *index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, n + (multi ? 1 : 0)); + ASSERT_INDEX(index); + for (size_t i = 0; i < n - 1; ++i) { + GenerateAndAddVector(index, dim, i, i); + } + if (multi) { + GenerateAndAddVector(index, dim, 0, 20); + } + + std::promise worker_started, start_search, build_finished; + std::promise query_between_tiers, published, search_finished; + auto worker_ready = worker_started.get_future(); + auto start_ready = start_search.get_future(); + auto build_ready = build_finished.get_future(); + auto query_ready = query_between_tiers.get_future(); + auto publish_ready = published.get_future(); + auto search_ready = search_finished.get_future(); + const auto timeout = std::chrono::seconds(5); + bool query_reached_hook = false; + bool query_saw_publication = false; + bool search_completed_during_init = false; + index->registerTracingCallback("UpdateJob::before_publish", [&] { + EXPECT_FALSE(index->GetSVSIndex()->ready()); + build_finished.set_value(); + query_reached_hook = query_ready.wait_for(timeout) == std::future_status::ready; + }); + index->registerTracingCallback("Query::before_backend", [&] { + query_between_tiers.set_value(); + query_saw_publication = publish_ready.wait_for(timeout) == std::future_status::ready; + }); + index->registerTracingCallback("UpdateJob::after_publish", [&] { + published.set_value(); + search_completed_during_init = + search_ready.wait_for(timeout) == std::future_status::ready; + }); + + VecSimQueryReply *reply = nullptr; + auto search = [&](VecSimIndex *, size_t) { + worker_started.set_value(); + if (start_ready.wait_for(timeout) == std::future_status::ready) { + TEST_DATA_T query[dim] = {}; + reply = VecSimIndex_TopKQuery(index, query, n, nullptr, BY_ID); + } + search_finished.set_value(); + }; + SVSMultiThreadJob::JobsRegistry registry(index->getAllocator()); + auto jobs = SVSMultiThreadJob::createJobs(index->getAllocator(), HNSW_SEARCH_JOB, search, + index, 1, timeout, ®istry); + index->submitJobs(jobs); + mock_thread_pool.init_threads(); + const bool search_occupies_worker = + worker_ready.wait_for(timeout) == std::future_status::ready; + // Cross the training threshold only after the search worker has been occupied. + GenerateAndAddVector(index, dim, n - 1, n - 1); + const bool reached_build_hook = build_ready.wait_for(timeout) == std::future_status::ready; + if (reached_build_hook) { + EXPECT_EQ(VecSimIndex_DeleteVector(index, 0), multi ? 2 : 1); + EXPECT_EQ(VecSimIndex_DeleteVector(index, 0), 0); + } + start_search.set_value(); + // Init can still enqueue consolidation after the queue first becomes empty. + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + std::unique_ptr result( + reply, VecSimQueryReply_Free); + ASSERT_TRUE(search_occupies_worker); + ASSERT_TRUE(reached_build_hook); + ASSERT_TRUE(query_reached_hook); + ASSERT_TRUE(query_saw_publication); + ASSERT_TRUE(search_completed_during_init); + ASSERT_NE(reply, nullptr); + EXPECT_EQ(VecSimQueryReply_GetCode(reply), VecSim_QueryReply_OK); + ASSERT_EQ(VecSimQueryReply_Len(reply), n - 1); + for (size_t i = 0; i < n - 1; ++i) { + EXPECT_EQ(VecSimQueryResult_GetId(reply->results.data() + i), i + 1); + } + EXPECT_TRUE(index->GetSVSIndex()->ready()); + EXPECT_FALSE(index->GetSVSIndex()->isLabelExists(0)); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), 0); + EXPECT_EQ(index->indexLabelCount(), n - 1); + EXPECT_EQ(index->getNumMarkedDeleted(), multi ? 2 : 1); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } +} + TYPED_TEST(SVSTieredIndexTest, TestDebugInfoThreadCount) { // Set thread_pool_size to 4 or actual number of available CPUs const auto num_threads = std::min(4U, getAvailableCPUs()); @@ -1066,7 +1802,7 @@ TYPED_TEST(SVSTieredIndexTest, TestDebugInfoThreadCountWriteInPlace) { std::this_thread::sleep_for(std::chrono::milliseconds(100)); } // Verify: numThreads (pool size) unchanged, lastReservedThreads (parallelism) reflects - // we only used one thread (write-in-place calls updateSVSIndexWrapper with availableThreads=1). + // we only used one thread (write-in-place calls initSVSIndexWrapper with availableThreads=1). backendIndexInfo = tiered_index->GetBackendIndex()->debugInfo(); ASSERT_EQ(backendIndexInfo.svsInfo.numThreads, num_threads); ASSERT_EQ(backendIndexInfo.svsInfo.lastReservedThreads, 1); @@ -1151,7 +1887,7 @@ TYPED_TEST(SVSTieredIndexTest, CreateIndexInstance) { ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 0); // Submit the index update job. - tiered_index->scheduleSVSIndexUpdate(); + tiered_index->scheduleSVSIndexInit(); ASSERT_EQ(mock_thread_pool.jobQ.size(), mock_thread_pool.thread_pool_size); // Execute the job from the queue and validate that the index was updated properly. @@ -1208,7 +1944,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, getDataByLabelReadsSvsBackendWhenUncompresse ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); // Ingest it into the SVS backend. - tiered_index->scheduleSVSIndexUpdate(); + tiered_index->scheduleSVSIndexInit(); mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); @@ -1290,7 +2026,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, ShrinkDuringScheduledUpdateIsDeferred) { shrink_callback_ran = true; }); - tiered_index->scheduleSVSIndexUpdate(); + tiered_index->scheduleSVSIndexInit(); ASSERT_EQ(mock_thread_pool.jobQ.size(), num_threads); mock_thread_pool.init_threads(); @@ -1374,7 +2110,9 @@ TYPED_TEST(SVSTieredIndexTestBasic, ShrinkDuringScheduledGCIsDeferred) { }); VecSimTieredIndex_GC(tiered_index); - ASSERT_EQ(mock_thread_pool.jobQ.size(), num_threads); + // Each soft-deleted slot has corresponding consolidate() job + size_t num_consolidate = tiered_index->GetSVSIndex()->getNumMarkedDeleted(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), num_threads + num_consolidate); mock_thread_pool.init_threads(); mock_thread_pool.thread_pool_join(); @@ -1442,13 +2180,13 @@ TYPED_TEST(SVSTieredIndexTest, addVector) { ASSERT_LE(expected_mem, tiered_index->getAllocationSize()); if constexpr (TypeParam::isMulti()) { - // Add another vector under the same label + // Add another vector under the same label. VecSimIndex_AddVector(tiered_index, vector, vec_label); ASSERT_EQ(tiered_index->indexSize(), 2); ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 2); - // Validate that there still 1 update jobs set + // The batch-init jobs. ASSERT_EQ(mock_thread_pool.jobQ.size(), mock_thread_pool.thread_pool_size); } } @@ -1499,6 +2237,12 @@ TYPED_TEST(SVSTieredIndexTest, background_indexing_check) { std::this_thread::sleep_for(std::chrono::milliseconds(100)); } + // The backgroundIndexing flag tracks the batch update job only; individual async + // SVS_INSERT_VECTOR_JOBs may still be in flight, each transiently holding its vector in + // both the flat and backend indices. Drain the queue so the total size is stable before + // asserting the exact count. + mock_thread_pool.thread_pool_join(); + ASSERT_GT(tiered_index->GetBackendIndex()->indexSize(), training_th + second_batch / update_th); ASSERT_LT(tiered_index->GetFlatIndex()->indexSize(), update_th); ASSERT_EQ(tiered_index->indexSize(), second_batch + training_th); @@ -1569,6 +2313,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, insertJobAsync) { GenerateAndAddVector(tiered_index, dim, i, i / (TEST_DATA_T)n); } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); ASSERT_EQ(tiered_index->indexSize(), n); ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); @@ -1621,6 +2366,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, insertJobAsyncMulti) { } } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); EXPECT_EQ(mock_thread_pool.jobQ.size(), 0); EXPECT_EQ(tiered_index->indexLabelCount(), n / per_label); @@ -1956,6 +2702,10 @@ TYPED_TEST(SVSTieredIndexTest, deleteVector) { // Remove from main index. ASSERT_EQ(tiered_index->deleteVector(vec_label), 1); ASSERT_EQ(tiered_index->indexLabelCount(), 0); + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); + ASSERT_EQ(tiered_index->indexSize(), 0); // Re-insert a deleted label with a different vector. @@ -1965,8 +2715,9 @@ TYPED_TEST(SVSTieredIndexTest, deleteVector) { ASSERT_EQ(tiered_index->indexSize(), 1); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); - // Move the vector to SVS by executing the insert job. - mock_thread_pool.thread_iteration(); + // Move the vector to SVS by executing the pending jobs. + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); // Scalar quantization accuracy is insufficient for this check. @@ -2007,7 +2758,8 @@ TYPED_TEST(SVSTieredIndexTestBasic, markedDeleted) { ASSERT_EQ(tiered_index->getNumMarkedDeleted(), 0); // Move vectors to the backend - mock_thread_pool.thread_iteration(); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), n); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->indexSize(), n); @@ -2049,7 +2801,10 @@ TYPED_TEST(SVSTieredIndexTestBasic, markedDeleted) { VecSimIndex_DeleteVector(tiered_index, i); } - // Consolidate should be triggered and mark deleted count should be zeroed. + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); + ASSERT_EQ(tiered_index->indexSize(), 0); ASSERT_EQ(tiered_index->getNumMarkedDeleted(), 0); ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 0); @@ -2087,10 +2842,12 @@ TYPED_TEST(SVSTieredIndexTestBasic, deleteVectorMulti) { ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->indexSize(), 2); ASSERT_EQ(tiered_index->deleteVector(vec_label), 2); + + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->indexSize(), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 0); - mock_thread_pool.thread_iteration(); - ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); // Test deleting a label for which both of its vector's is in the flat index. GenerateAndAddVector(tiered_index, dim, vec_label, vec_label); @@ -2101,13 +2858,20 @@ TYPED_TEST(SVSTieredIndexTestBasic, deleteVectorMulti) { ASSERT_EQ(tiered_index->indexSize(), 2); ASSERT_EQ(tiered_index->deleteVector(vec_label), 2); ASSERT_EQ(tiered_index->indexLabelCount(), 0); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); + mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); mock_thread_pool.thread_iteration(); ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); // Test deleting a label for which both of its vector's is in SVS index. GenerateAndAddVector(tiered_index, dim, vec_label, vec_label); GenerateAndAddVector(tiered_index, dim, vec_label, other_vec_val); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); + mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); mock_thread_pool.thread_iteration(); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 1); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 2); @@ -2323,6 +3087,7 @@ TYPED_TEST(SVSTieredIndexTest, parallelInsertSearch) { tiered_index->submitSingleJob(search_job); } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); // Verify that vectors were moved to SVS as expected @@ -2427,8 +3192,9 @@ TYPED_TEST(SVSTieredIndexTest, testSizeEstimation) { } mock_thread_pool.thread_pool_wait(); - // Estimate memory delta for filling up the first block and adding another block. - size_t estimation = VecSimIndex_EstimateElementSize(¶ms) * bs; + const size_t reverse_edges_per_slot = SVSGraphBuilder::reverse_edges_element_size(); + size_t estimation = (VecSimIndex_EstimateElementSize(¶ms) - reverse_edges_per_slot) * bs + + EstimateReverseEdgesSize(n + 1) - EstimateReverseEdgesSize(n); size_t before = index->getAllocationSize(); GenerateAndAddVector(index, dim, bs + n, bs + n); @@ -2501,7 +3267,7 @@ TYPED_TEST(SVSTieredIndexTest, parallelInsertAdHoc) { tiered_index->submitSingleJob(search_job); } - tiered_index->scheduleSVSIndexUpdate(); + tiered_index->scheduleSVSIndexInit(); mock_thread_pool.thread_pool_join(); EXPECT_EQ(successful_searches, n); @@ -3518,12 +4284,14 @@ TYPED_TEST(SVSTieredIndexTestBasic, overwriteVectorBasic) { ASSERT_EQ(tiered_index->addVector(overwritten_vec, 0), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 1); // Overriding vector in tiered index should remove the vector from SVS to avoid duplicates. - ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 0); + ASSERT_EQ(tiered_index->GetBackendIndex()->indexLabelCount(), 0); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); ASSERT_EQ(tiered_index->getDistanceFrom_Unsafe(0, overwritten_vec), 0); // Ingest the updated vector to SVS. - mock_thread_pool.thread_iteration(); + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 1); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->indexLabelCount(), 1); @@ -3573,6 +4341,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, overwriteVectorAsync) { EXPECT_EQ(tiered_index->addVector(vector, label_to_overwrite), 0); } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); EXPECT_LE(tiered_index->GetFlatIndex()->indexSize(), updateThreshold); @@ -3706,6 +4475,10 @@ TYPED_TEST(SVSTieredIndexTest, testInfo) { } VecSimIndex_DeleteVector(tiered_index, 1); + VecSimTieredIndex_GC(tiered_index); + while (mock_thread_pool.jobQ.size() > 0) + mock_thread_pool.thread_iteration(); + info = tiered_index->debugInfo(); EXPECT_EQ(info.commonInfo.indexSize, 0); @@ -3802,8 +4575,10 @@ TYPED_TEST(SVSTieredIndexTest, writeInPlaceMode) { ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 2); ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); - // Overwrite inplace - only in single-value mode + // In write-in-place mode every mutation runs synchronously, including consolidation + // of soft-deleted slots. size_t expected_marked_deleted = 0; + // Overwrite inplace - only in single-value mode if (!TypeParam::isMulti()) { TEST_DATA_T overwritten_vec[] = {1, 1, 1, 1}; tiered_index->addVector(overwritten_vec, vec_label); @@ -3815,7 +4590,8 @@ TYPED_TEST(SVSTieredIndexTest, writeInPlaceMode) { ASSERT_EQ(tiered_index->getDistanceFrom_Unsafe(vec_label, overwritten_vec), 0); ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), expected_marked_deleted); } - // Validate that the vector is marked as deleted. + // Validate that the vector is deleted and consolidated synchronously (in-place mode), + // so no marked-deleted entry remains. tiered_index->deleteVector(vec_label); expected_marked_deleted++; ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), TypeParam::isMulti() ? 2 : 3); @@ -3867,6 +4643,7 @@ TYPED_TEST(SVSTieredIndexTest, switchWriteModes) { } VecSimIndex_AddVector(tiered_index, vector, i % n_labels + n_labels); } + mock_thread_pool.thread_pool_wait(); mock_thread_pool.thread_pool_join(); EXPECT_GE(tiered_index->GetBackendIndex()->indexSize(), 2 * n - this->getUpdateThreshold()); @@ -3965,6 +4742,531 @@ TYPED_TEST(SVSTieredIndexTestBasic, preferAdHocOptimization) { ASSERT_NO_THROW(tiered_index->preferAdHocSearch(10, 5, false)); } +TYPED_TEST(SVSTieredIndexTest, runGCBetweenIteratorBatches) { + if (this->isFallbackToSQ()) { + GTEST_SKIP() << "Scalar fallback does not preserve the ordering checked by this test."; + } + constexpr size_t dim = 4; + constexpr size_t n_labels = 6; + for (auto mode : {VecSim_WriteAsync, VecSim_WriteInPlace}) { + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .multi = TypeParam::isMulti()}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(1); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool); + ASSERT_INDEX(tiered_index); + auto *backend = tiered_index->GetBackendIndex(); + for (size_t i = 0; i < n_labels; ++i) { + GenerateAndAddVector(backend, dim, i + 1000, i + 1); + } + + TEST_DATA_T query[dim] = {}; + VecSimQueryParams query_params{}; + query_params.batchSize = 3; + auto *iterator = VecSimBatchIterator_New(tiered_index, query, &query_params); + auto verify_first = [](size_t id, double score, size_t i) { + EXPECT_EQ(id, i + 1000); + if constexpr (TypeParam::get_quant_bits() == VecSimSvsQuant_NONE) { + EXPECT_DOUBLE_EQ(score, dim * double(i + 1) * double(i + 1)); + } + }; + runBatchIteratorSearchTest(iterator, 3, verify_first); + for (size_t i = 0; i < 3; ++i) { + backend->deleteVector(i + 1000); + } + + VecSim_SetWriteMode(mode); + std::promise gc_finished; + auto completion = gc_finished.get_future(); + std::thread gc([&]() { + VecSimTieredIndex_GC(tiered_index); + if (mode == VecSim_WriteAsync) { + mock_thread_pool.thread_iteration(); + } + gc_finished.set_value(); + }); + const bool completed = + completion.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + // If the old lifetime lock is still present, release it on its owning thread + // so GC can finish and the regression fails without hanging the suite. + if (!completed) { + VecSimBatchIterator_Free(iterator); + iterator = nullptr; + } + gc.join(); + VecSim_SetWriteMode(VecSim_WriteAsync); + ASSERT_TRUE(completed) << "GC waited for an idle iterator to be freed"; + EXPECT_EQ(backend->indexSize(), 3); + EXPECT_TRUE(VecSimBatchIterator_HasNext(iterator)); + auto verify_remaining = [](size_t id, double score, size_t i) { + EXPECT_EQ(id, i + 1003); + if constexpr (TypeParam::get_quant_bits() == VecSimSvsQuant_NONE) { + EXPECT_DOUBLE_EQ(score, dim * double(i + 4) * double(i + 4)); + } + }; + runBatchIteratorSearchTest(iterator, 3, verify_remaining); + EXPECT_FALSE(VecSimBatchIterator_HasNext(iterator)); + VecSimBatchIterator_Free(iterator); + } +} + +namespace { + +constexpr size_t svs_gc_test_dim = 4; +constexpr size_t svs_gc_label_base = 1'000'000; +constexpr size_t svs_gc_label_stride = 17; +using SVSReferenceVectors = std::vector>>; +using SVSTestReply = std::unique_ptr; + +size_t svsGCLabel(size_t i) { return svs_gc_label_base + svs_gc_label_stride * i; } + +SVSReferenceVectors makeSVSReferenceVectors(size_t n, bool multi, bool distinct = false) { + SVSReferenceVectors reference(n); + for (size_t i = 0; i < n; ++i) { + for (size_t copy = 0; copy < (multi ? 2U : 1U); ++copy) { + std::vector vector(svs_gc_test_dim); + for (size_t d = 0; d < vector.size(); ++d) { + vector[d] = float(i + 1) + float(d) / 4 + (distinct ? float(copy) / 8 : 0); + } + reference[i].push_back(std::move(vector)); + } + } + return reference; +} + +void addSVSReferenceVectors(VecSimIndex *index, const SVSReferenceVectors &reference, size_t begin, + size_t end) { + for (size_t i = begin; i < end; ++i) { + for (const auto &vector : reference[i]) { + VecSimIndex_AddVector(index, vector.data(), svsGCLabel(i)); + } + } +} + +double svsReferenceDistance(const std::vector &vector, const std::vector &query) { + double distance = 0; + for (size_t d = 0; d < vector.size(); ++d) { + const double delta = double(vector[d]) - double(query[d]); + distance += delta * delta; + } + return distance; +} + +// All references are immutable. Concurrently deleted labels need only remain valid external +// labels; scores are checked for permanent survivors, whose vectors never change. +std::string checkSVSReply(const VecSimQueryReply &reply, const SVSReferenceVectors &reference, + const std::vector &query, size_t survivors_begin, + size_t &checked_survivors) { + if (reply.code != VecSim_QueryReply_OK) { + return "Query did not complete successfully"; + } + std::unordered_set labels; + for (const auto &result : reply.results) { + if (result.id < svs_gc_label_base || + (result.id - svs_gc_label_base) % svs_gc_label_stride != 0 || + (result.id - svs_gc_label_base) / svs_gc_label_stride >= reference.size()) { + return "Unknown external label: " + std::to_string(result.id); + } + if (!std::isfinite(result.score) || !labels.insert(result.id).second) { + return "Non-finite score or duplicate label: " + std::to_string(result.id); + } + const size_t i = (result.id - svs_gc_label_base) / svs_gc_label_stride; + if (i >= survivors_begin) { + const double expected = svsReferenceDistance(reference[i][0], query); + if (std::abs(result.score - expected) > 1e-5 * std::max(1.0, expected)) { + return "Wrong score for label " + std::to_string(result.id) + ": got " + + std::to_string(result.score) + ", expected " + std::to_string(expected); + } + ++checked_survivors; + } + } + return {}; +} + +std::string checkSVSCopies(TieredSVSIndex *index, const SVSReferenceVectors &reference, + size_t begin, size_t end) { + for (size_t i = begin; i < end; ++i) { + std::vector> copied; + index->getDataByLabel(svsGCLabel(i), copied); + std::sort(copied.begin(), copied.end()); + // Reference copies are generated in sorted order. Equality also checks multiplicity. + if (copied != reference[i]) { + return "Wrong vectors copied for label " + std::to_string(svsGCLabel(i)); + } + } + return {}; +} + +class SVSTestWriteMode { + VecSimWriteMode previous = TieredSVSIndex::getWriteMode(); + +public: + explicit SVSTestWriteMode(VecSimWriteMode mode) { VecSim_SetWriteMode(mode); } + ~SVSTestWriteMode() { VecSim_SetWriteMode(previous); } +}; + +// Reader checks run outside the synchronization mutex. Failures are reported after joining, +// and each mutation round waits for every reader to make progress without using sleeps. +class SVSTestReaders { + std::vector> checks; + std::vector threads; + std::vector completed; + std::mutex mutex; + std::condition_variable changed; + std::atomic stopped{false}; + bool started = false; + std::string failure; + +public: + explicit SVSTestReaders(std::vector> checks) + : checks(std::move(checks)), completed(this->checks.size(), 0) {} + + void start() { + for (size_t i = 0; i < checks.size(); ++i) { + threads.emplace_back([this, i] { + { + std::unique_lock lock(mutex); + changed.wait(lock, [&] { return started || stopped.load(); }); + } + while (!stopped.load()) { + auto error = checks[i](); + { + std::lock_guard lock(mutex); + ++completed[i]; + if (!error.empty()) { + if (failure.empty()) { + failure = "Reader " + std::to_string(i) + ": " + error; + } + stopped.store(true); + } + } + changed.notify_all(); + } + }); + } + { + std::lock_guard lock(mutex); + started = true; + } + changed.notify_all(); + } + + bool waitForProgress() { + std::unique_lock lock(mutex); + const auto before = completed; + const bool progressed = changed.wait_for(lock, std::chrono::seconds(5), [&] { + if (!failure.empty()) { + return true; + } + for (size_t i = 0; i < completed.size(); ++i) { + if (completed[i] == before[i]) { + return false; + } + } + return true; + }); + if (!progressed && failure.empty()) { + failure = "Timed out waiting for reader progress"; + stopped.store(true); + } + return progressed && failure.empty(); + } + + void finish() { + { + std::lock_guard lock(mutex); + stopped.store(true); + } + changed.notify_all(); + for (auto &thread : threads) { + if (thread.joinable()) { + thread.join(); + } + } + } + + const std::string &error() const { return failure; } // Read only after finish(). + ~SVSTestReaders() { finish(); } +}; + +} // namespace + +TYPED_TEST(SVSTieredIndexTestBasic, runGCBetweenIteratorBatchesWithStorageShrink) { + constexpr size_t n_labels = 96; + constexpr size_t batch_size = 16; + const auto reference = makeSVSReferenceVectors(n_labels, false); + const std::vector query(svs_gc_test_dim, 0); + for (auto mode : {VecSim_WriteAsync, VecSim_WriteInPlace}) { + for (bool delete_returned : {true, false}) { + SCOPED_TRACE(::testing::Message() + << "mode=" << mode << ", delete_returned=" << delete_returned); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = svs_gc_test_dim, + .metric = VecSimMetric_L2, + .blockSize = 8, + .search_window_size = n_labels, + .search_buffer_capacity = n_labels}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto *index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool); + ASSERT_INDEX(index); + // Index creation resizes the thread pool, which switches write mode to async. + SVSTestWriteMode write_mode(mode); + auto *backend = index->GetBackendIndex(); + addSVSReferenceVectors(backend, reference, 0, n_labels); + ASSERT_EQ(index->GetFlatIndex()->indexSize(), 0); + + VecSimQueryParams query_params{}; + // Consume the entire native batch, leaving no prefetched external labels. + query_params.batchSize = batch_size; + std::unique_ptr iterator( + VecSimBatchIterator_New(index, query.data(), &query_params), + VecSimBatchIterator_Free); + std::unordered_set seen, deleted; + auto read_batch = [&] { + SVSTestReply reply(VecSimBatchIterator_Next(iterator.get(), batch_size, BY_SCORE), + VecSimQueryReply_Free); + size_t checked = 0; + EXPECT_EQ(checkSVSReply(*reply, reference, query, 0, checked), ""); + for (const auto &result : reply->results) { + EXPECT_EQ(deleted.count(result.id), 0) << result.id; + EXPECT_TRUE(seen.insert(result.id).second) << result.id; + } + return reply->results.size(); + }; + ASSERT_EQ(read_batch(), batch_size); + for (size_t round = 0; round < 2; ++round) { + const size_t begin = (delete_returned ? round : 2 * round + 1) * batch_size; + for (size_t i = begin; i < begin + batch_size; ++i) { + ASSERT_EQ(seen.count(svsGCLabel(i)), delete_returned ? 1 : 0); + // As in runGCBetweenIteratorBatches, defer consolidation to GC so the + // timeout below can unwind an old iterator holding a lifetime lock. + ASSERT_EQ(backend->deleteVector(svsGCLabel(i)), 1); + deleted.insert(svsGCLabel(i)); + } + const auto capacity_before = backend->indexCapacity(); + std::promise gc_finished; + auto completion = gc_finished.get_future(); + std::thread gc([&] { + VecSimTieredIndex_GC(index); + if (mode == VecSim_WriteAsync) { + mock_thread_pool.thread_iteration(); + } + gc_finished.set_value(); + }); + const bool completed = + completion.wait_for(std::chrono::seconds(5)) == std::future_status::ready; + if (!completed) { + iterator.reset(); + } + gc.join(); + ASSERT_TRUE(completed) << "GC waited for an idle iterator to be freed"; + EXPECT_EQ(index->GetSVSIndex()->indexStorageSize(), + n_labels - (round + 1) * batch_size); + EXPECT_LT(backend->indexCapacity(), capacity_before); + ASSERT_EQ(read_batch(), batch_size); + } + + auto drain_iterator = [&] { + // Bound the loop so failure to reach exhaustion cannot hang the test. + for (size_t batch = 0; + batch <= n_labels / batch_size && VecSimBatchIterator_HasNext(iterator.get()); + ++batch) { + const auto count = read_batch(); + EXPECT_TRUE(count > 0 || !VecSimBatchIterator_HasNext(iterator.get())); + } + EXPECT_FALSE(VecSimBatchIterator_HasNext(iterator.get())); + }; + drain_iterator(); + for (size_t i = 0; i < n_labels; ++i) { + EXPECT_EQ(seen.count(svsGCLabel(i)), + delete_returned || !deleted.count(svsGCLabel(i)) ? 1 : 0); + } + + VecSimBatchIterator_Reset(iterator.get()); + seen.clear(); + drain_iterator(); + EXPECT_EQ(seen.size(), n_labels - deleted.size()); + for (size_t i = 0; i < n_labels; ++i) { + EXPECT_EQ(seen.count(svsGCLabel(i)), deleted.count(svsGCLabel(i)) ? 0 : 1); + } + } + } +} + +TYPED_TEST(SVSTieredIndexTestBasic, parallelQueriesDuringGC) { + constexpr size_t n_labels = 256; + constexpr size_t rounds = 3; + constexpr size_t delete_per_round = 64; + constexpr size_t survivors_begin = rounds * delete_per_round; + SVSTestWriteMode write_mode(VecSim_WriteAsync); + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + const auto reference = makeSVSReferenceVectors(n_labels, multi); + const auto &query = reference.back()[0]; + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = svs_gc_test_dim, + .metric = VecSimMetric_L2, + .multi = multi, + .blockSize = 16, + .search_window_size = n_labels}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto *index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool); + ASSERT_INDEX(index); + addSVSReferenceVectors(index->GetBackendIndex(), reference, 0, n_labels); + ASSERT_EQ(index->GetFlatIndex()->indexSize(), 0); + + std::array checked_survivors{}; + auto query_once = [&](bool range) -> std::string { + SVSTestReply reply( + range ? VecSimIndex_RangeQuery(index, query.data(), 4096, nullptr, BY_SCORE) + : VecSimIndex_TopKQuery(index, query.data(), 10, nullptr, BY_SCORE), + VecSimQueryReply_Free); + if (reply->results.empty()) { + return "Empty reply for a query near permanent survivors"; + } + return checkSVSReply(*reply, reference, query, survivors_begin, + checked_survivors[range ? 1 : 0]); + }; + SVSTestReaders readers( + {[&] { return query_once(false); }, [&] { return query_once(true); }}); + bool gc_reader_progress = false; + index->registerTracingCallback("GCJob::before_run_gc", + [&] { gc_reader_progress = readers.waitForProgress(); }); + readers.start(); + bool progress = readers.waitForProgress(); + for (size_t round = 0; round < rounds && progress; ++round) { + const auto storage_before = index->GetSVSIndex()->indexStorageSize(); + const auto capacity_before = index->GetBackendIndex()->indexCapacity(); + for (size_t i = round * delete_per_round; i < (round + 1) * delete_per_round; ++i) { + EXPECT_EQ(index->deleteVector(svsGCLabel(i)), multi ? 2 : 1); + } + gc_reader_progress = false; + VecSimTieredIndex_GC(index); + // Also drain the consolidation jobs queued by deleteVector. + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + progress = gc_reader_progress; + EXPECT_LT(index->GetSVSIndex()->indexStorageSize(), storage_before); + EXPECT_LT(index->GetBackendIndex()->indexCapacity(), capacity_before); + } + readers.finish(); + ASSERT_TRUE(progress) << readers.error(); + ASSERT_EQ(readers.error(), ""); + EXPECT_GT(checked_survivors[0], 0); + EXPECT_GT(checked_survivors[1], 0); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), n_labels - survivors_begin); + EXPECT_EQ(index->GetSVSIndex()->getNumMarkedDeleted(), 0); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } +} + +TYPED_TEST(SVSTieredIndexTestBasic, getDataByLabelDuringGC) { + constexpr size_t n_labels = 256; + constexpr size_t rounds = 3; + constexpr size_t delete_per_round = 64; + constexpr size_t monitored_begin = n_labels - 8; + SVSTestWriteMode write_mode(VecSim_WriteAsync); + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + const auto reference = makeSVSReferenceVectors(n_labels, multi, true); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = svs_gc_test_dim, + .metric = VecSimMetric_L2, + .multi = multi, + .blockSize = 16}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto *index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool); + ASSERT_INDEX(index); + addSVSReferenceVectors(index->GetBackendIndex(), reference, 0, n_labels); + ASSERT_EQ(index->GetFlatIndex()->indexSize(), 0); + + SVSTestReaders readers( + {[&] { return checkSVSCopies(index, reference, monitored_begin, n_labels); }}); + bool gc_reader_progress = false; + index->registerTracingCallback("GCJob::before_run_gc", + [&] { gc_reader_progress = readers.waitForProgress(); }); + readers.start(); + bool progress = readers.waitForProgress(); + for (size_t round = 0; round < rounds && progress; ++round) { + const auto capacity_before = index->GetBackendIndex()->indexCapacity(); + for (size_t i = round * delete_per_round; i < (round + 1) * delete_per_round; ++i) { + EXPECT_EQ(index->deleteVector(svsGCLabel(i)), multi ? 2 : 1); + } + gc_reader_progress = false; + VecSimTieredIndex_GC(index); + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + progress = gc_reader_progress; + EXPECT_EQ(index->GetSVSIndex()->indexStorageSize(), + (n_labels - (round + 1) * delete_per_round) * (multi ? 2 : 1)); + EXPECT_LT(index->GetBackendIndex()->indexCapacity(), capacity_before); + } + readers.finish(); + ASSERT_TRUE(progress) << readers.error(); + ASSERT_EQ(readers.error(), ""); + EXPECT_EQ(checkSVSCopies(index, reference, monitored_begin, n_labels), ""); + EXPECT_EQ(index->GetSVSIndex()->getNumMarkedDeleted(), 0); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } +} + +TYPED_TEST(SVSTieredIndexTestBasic, getDataByLabelDuringIngestion) { + constexpr size_t initial_labels = 64; + constexpr size_t rounds = 4; + constexpr size_t labels_per_round = 64; + constexpr size_t n_labels = initial_labels + rounds * labels_per_round; + SVSTestWriteMode write_mode(VecSim_WriteAsync); + for (bool multi : {false, true}) { + SCOPED_TRACE(::testing::Message() << "multi=" << multi); + const auto reference = makeSVSReferenceVectors(n_labels, multi, true); + SVSParams params = {.type = TypeParam::get_index_type(), + .dim = svs_gc_test_dim, + .metric = VecSimMetric_L2, + .multi = multi, + .blockSize = 16}; + auto svs_params = CreateParams(params); + tieredIndexMock mock_thread_pool(1); + auto tiered_params = this->CreateTieredSVSParams(svs_params, mock_thread_pool); + // Keep each batch on the worker ingestion path instead of bypassing a full buffer. + tiered_params.flatBufferLimit = 2 * labels_per_round; + auto *index = this->CreateTieredSVSIndex(tiered_params, mock_thread_pool); + ASSERT_INDEX(index); + addSVSReferenceVectors(index->GetBackendIndex(), reference, 0, initial_labels); + ASSERT_TRUE(index->GetSVSIndex()->ready()); + + SVSTestReaders readers( + {[&] { return checkSVSCopies(index, reference, initial_labels - 8, initial_labels); }}); + readers.start(); + bool progress = readers.waitForProgress(); + for (size_t round = 0; round < rounds && progress; ++round) { + const size_t begin = initial_labels + round * labels_per_round; + const auto capacity_before = index->GetBackendIndex()->indexCapacity(); + addSVSReferenceVectors(index, reference, begin, begin + labels_per_round); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), labels_per_round * (multi ? 2 : 1)); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), begin); + mock_thread_pool.init_threads(); + progress = readers.waitForProgress(); + mock_thread_pool.thread_pool_join(); + EXPECT_EQ(index->GetFlatIndex()->indexSize(), 0); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), begin + labels_per_round); + EXPECT_EQ(index->GetSVSIndex()->indexStorageSize(), + (begin + labels_per_round) * (multi ? 2 : 1)); + EXPECT_GT(index->GetBackendIndex()->indexCapacity(), capacity_before); + EXPECT_TRUE(mock_thread_pool.jobQ.empty()); + } + readers.finish(); + ASSERT_TRUE(progress) << readers.error(); + ASSERT_EQ(readers.error(), ""); + EXPECT_EQ(checkSVSCopies(index, reference, initial_labels - 8, initial_labels), ""); + EXPECT_EQ(index->GetBackendIndex()->indexLabelCount(), n_labels); + } +} + TYPED_TEST(SVSTieredIndexTestBasic, runGCAPI) { // Create TieredSVS index instance with a mock queue. size_t dim = 4; @@ -4013,11 +5315,12 @@ TYPED_TEST(SVSTieredIndexTestBasic, runGCAPI) { auto jobs_before_gc = mock_thread_pool.jobQ.size(); // Run the GC API call, expect that we will clean up the SVS index. VecSimTieredIndex_GC(tiered_index); - // Expected that GC jobs were added to the queue. + // Expected that a multi-thread GC job was added to the queue. ASSERT_EQ(mock_thread_pool.jobQ.size(), jobs_before_gc + mock_thread_pool.thread_pool_size); // Run GC twice. VecSimTieredIndex_GC(tiered_index); - // Expected that no new GC jobs were added to the queue. + // Expected that no new GC jobs were added to the queue (indexGCScheduled is still set until the + // pending jobs run). ASSERT_EQ(mock_thread_pool.jobQ.size(), jobs_before_gc + mock_thread_pool.thread_pool_size); // Wait for any pending jobs to complete. As far as SVS GC is done via a job. mock_thread_pool.init_threads(); @@ -4297,7 +5600,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalSingle) { // update job paused, we have vectors 0-(n-1) in the index, let's do index modifications // Remove vector label=n-2, it is copied to backend index. - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 2), 2); + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 2), 1); // Update vector label=1. EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, 1, 10), 0); // Add a new vector @@ -4305,9 +5608,9 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalSingle) { // Add another one EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, n + 1, n + 1), 1); // Remove vector label=0, it is copied to backend index. - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 2); + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 1); // Remove the last vector copied to backend index - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 1), 2); + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 1), 1); // Update vector label=2. EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, 2, 20), 0); // Remove vector label=2. @@ -4329,16 +5632,15 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalSingle) { // 0:deleted, 1: 10, 2: deleted, 3:3, ..., n-2:deleted n-1: 10(n-1), n+1: n+1; // total: n-2 vectors and labels ASSERT_EQ(tiered_index->indexLabelCount(), n - 2); - EXPECT_EQ(tiered_index->GetBackendIndex()->indexLabelCount(), n - 5); + // Nothing remains in flat. + EXPECT_EQ(tiered_index->GetBackendIndex()->indexLabelCount(), n - 2); - // We added 3 vectors to the flat index and removed 5 vectors from the backend index. - // Backend index: 0:deleted, 1:deleted, 2:deleted, 3:3, ..., n-2:deleted, n-1:deleted; - // total: n-5 - EXPECT_EQ(tiered_index->GetBackendIndex()->indexSize(), n); + EXPECT_EQ(tiered_index->GetBackendIndex()->indexSize(), n + 3); ASSERT_EQ(tiered_index->getNumMarkedDeleted(), 5); - // Frontend index: 1:10, n-1:10(n-1), n+1:n+1 - ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 3); - ASSERT_EQ(tiered_index->indexSize(), n + tiered_index->GetFlatIndex()->indexSize()); + // Flat buffer fully drained by the insert jobs. + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), tiered_index->GetBackendIndex()->indexSize() + + tiered_index->GetFlatIndex()->indexSize()); double abs_err = 1e-2; // Allow a larger relative error for quantization. TEST_DATA_T expected_vector[dim]; @@ -4351,7 +5653,7 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalSingle) { GenerateVector(expected_vector, dim, 0); ASSERT_TRUE(std::isnan(tiered_index->getDistanceFrom_Unsafe(n - 2, expected_vector))); - // Vector label=1, with value 10 should be in the flat index. + // Vector label=1 was updated to value 10; its latest value is retrievable from the index. GenerateVector(expected_vector, dim, 10); ASSERT_NEAR(tiered_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); @@ -4430,22 +5732,24 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalMulti) { // update job paused, we have vectors 0-(n-1) in the index, let's do index modifications - // Remove vector label=n-2, it is copied to backend index. - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 2), 2); + // Remove vector label=n-2. At the pause it lives only in the backend + //, so a single vector is removed. + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 2), 1); // Add one more vector label=1. EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, 1, 10), 1); // Add a new vector EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, n, n), 1); // Add another one EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, n + 1, n + 1), 1); - // Remove vector label=0, it is copied to backend index. - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 2); - // Remove the last vector copied to backend index - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 1), 2); + // Remove vector label=0, only in the backend at the pause. + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 1); + // Remove the last vector, only in the backend at the pause. + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n - 1), 1); // Add one more vector label=2. EXPECT_EQ(GenerateAndAddVector(tiered_index, dim, 2, 20), 1); - // Remove vector label=2: for multi: old is copied to backend , old + new are in flat - EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 2), 3); + // Remove vector label=2: for multi both the original (2) and the new (20) copies sit in + // the flat buffer at this point, so 2 vectors are removed. + EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, 2), 2); // Remove vector label=n - in flat only EXPECT_EQ(VecSimIndex_DeleteVector(tiered_index, n), 1); // Add vector (n-1) again @@ -4463,15 +5767,16 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalMulti) { // 0: deleted, 1: (1,10), 2: deleted, 3:3, ..., n-2: deleted n-1: 10(n-1), n+1: n+1; // total: n-2 labels, n-1 vectors ASSERT_EQ(tiered_index->indexLabelCount(), n - 2); - EXPECT_EQ(tiered_index->GetBackendIndex()->indexLabelCount(), n - 4); + // In the async-insert design the per-vector insert jobs drain every surviving vector from + // the flat buffer into the backend, so nothing remains in flat. + EXPECT_EQ(tiered_index->GetBackendIndex()->indexLabelCount(), n - 2); - // We added 3 vectors to the flat index and removed 4 vectors from the backend index. - // Backend index: 0:deleted, 1:1, 2:deleted, 3:3, ..., n-2:deleted, n-1:deleted; total: n-4 - EXPECT_EQ(tiered_index->GetBackendIndex()->indexSize(), n); + EXPECT_EQ(tiered_index->GetBackendIndex()->indexSize(), n + 3); ASSERT_EQ(tiered_index->getNumMarkedDeleted(), 4); - // Frontend index: 1:10, n-1:10(n-1), n+1:n+1 - ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 3); - ASSERT_EQ(tiered_index->indexSize(), n + tiered_index->GetFlatIndex()->indexSize()); + // Flat buffer fully drained by the insert jobs. + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), tiered_index->GetBackendIndex()->indexSize() + + tiered_index->GetFlatIndex()->indexSize()); double abs_err = 1e-2; // Allow a larger relative error for quantization. TEST_DATA_T expected_vector[dim]; @@ -4484,11 +5789,11 @@ TYPED_TEST(SVSTieredIndexTestBasic, testSwapJournalMulti) { GenerateVector(expected_vector, dim, 0); ASSERT_TRUE(std::isnan(tiered_index->getDistanceFrom_Unsafe(n - 2, expected_vector))); - // There are 2 vectors labeled "1" with values 1 in backend and 10 in flat. - // We expect the minimal distance for the query 10 to be taken from flat index. + // Label "1" is multi-valued and keeps both of its vectors (values 1 and 10). + // The minimal distance for the query 10 should match the second value. GenerateVector(expected_vector, dim, 10); ASSERT_NEAR(tiered_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); - // And the minimal distance for the query 1.0 to be taken from backend + // And the minimal distance for the query 1.0 should match the first value. GenerateVector(expected_vector, dim, 1); ASSERT_NEAR(tiered_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); @@ -4588,8 +5893,9 @@ TYPED_TEST(SVSTieredIndexTestBasic, testDeletedJournalSingle) { mock_thread_pool.thread_pool_join(); - // Verify that vectors labels: {0, 1, 2, n-1} are marked as deleted in the SVS index. - ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 4); + // The update job is paused before it transfers vectors to the backend, so at the pause the + // backend is still empty. + ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 0); // Verify that the deleted vectors are not accessible. double abs_err = 1e-2; // Allow a larger relative error for quantization. @@ -4606,9 +5912,10 @@ TYPED_TEST(SVSTieredIndexTestBasic, testDeletedJournalSingle) { ASSERT_TRUE(std::isnan(tiered_index->getDistanceFrom_Unsafe(label, expected_vector))); } - // label 1 - updated to 10 but deleted in the SVS index - ASSERT_TRUE(flat_index->isLabelExists(1)); - ASSERT_FALSE(svs_index->isLabelExists(1)); + // label 1 - updated to 10 during the pause. The insert job drains it from the flat buffer + // into the backend, so it now lives in the SVS index (not flat) with its updated value. + ASSERT_FALSE(flat_index->isLabelExists(1)); + ASSERT_TRUE(svs_index->isLabelExists(1)); GenerateVector(expected_vector, dim, 10); ASSERT_NEAR(tiered_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); } @@ -4690,8 +5997,9 @@ TYPED_TEST(SVSTieredIndexTestBasic, testDeletedJournalMulti) { mock_thread_pool.thread_pool_join(); - // Verify that vectors labels: {0, 2, n-1} are marked as deleted in the SVS index. - ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 3); + // The update job is paused before it transfers vectors to the backend, so at the pause the + // backend is still empty. + ASSERT_EQ(tiered_index->GetSVSIndex()->getNumMarkedDeleted(), 0); // Verify that the deleted vectors are not accessible. double abs_err = 1e-2; // Allow a larger relative error for quantization. @@ -4708,13 +6016,149 @@ TYPED_TEST(SVSTieredIndexTestBasic, testDeletedJournalMulti) { ASSERT_TRUE(std::isnan(tiered_index->getDistanceFrom_Unsafe(label, expected_vector))); } - // label 1 - multi-value 1 (in SVS) and 10 (in flat) - ASSERT_TRUE(flat_index->isLabelExists(1)); + // label 1 - multi-value with both vectors (1 and 10). The insert jobs drain the flat buffer + // into the backend, so both values now live in the SVS index and none remain in flat. + ASSERT_FALSE(flat_index->isLabelExists(1)); ASSERT_TRUE(svs_index->isLabelExists(1)); GenerateVector(expected_vector, dim, 1); ASSERT_NEAR(backend_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); GenerateVector(expected_vector, dim, 10); - ASSERT_NEAR(flat_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); + ASSERT_NEAR(backend_index->getDistanceFrom_Unsafe(1, expected_vector), 0, abs_err); +} + +// scheduleSVSIndexConsolidate() always registers the label, so a job with an unregistered one +// can only be built by hand. The callback is taken from a real job because the wrapper is +// private. +TYPED_TEST(SVSTieredIndexTestBasic, forgetConsolidateJobIgnoresAnUnregisteredLabel) { + const size_t dim = 4; + SVSParams params = { + .type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_L2, .multi = false}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 1, 1); + ASSERT_INDEX(tiered_index); + auto allocator = tiered_index->getAllocator(); + + mock_thread_pool.init_threads(); + for (size_t i = 0; i < 2; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + mock_thread_pool.thread_pool_wait(); + mock_thread_pool.thread_pool_join(); + + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 1), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + auto *registered_job = mock_thread_pool.jobQ.front().job; + ASSERT_EQ(registered_job->jobType, SVS_CONSOLIDATE_JOB); + + auto *orphan_job = + new (allocator) SVSConsolidateJob(allocator, {999}, registered_job->Execute, tiered_index); + tiered_index->submitSingleJob(orphan_job); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 2); + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(0)); + ASSERT_FALSE(tiered_index->GetSVSIndex()->isLabelExists(1)); + ASSERT_EQ(tiered_index->indexLabelCount(), 1); +} + +// Every add made while the backend is not ready becomes training data, so the init normally +// deletes the whole buffer and never swaps. Adding straight to the frontend leaves a vector the +// batch does not own, which the init has to swap down and keep. +TYPED_TEST(SVSTieredIndexTestBasic, initSVSIndexKeepsANonTrainingFlatVector) { + const size_t dim = 4; + const size_t n = 5; + const labelType kept_label = 99; + SVSParams params = { + .type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_L2, .multi = false}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 100, 100); + ASSERT_INDEX(tiered_index); + + for (size_t i = 0; i < n; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); + ASSERT_FALSE(tiered_index->GetSVSIndex()->ready()); + + GenerateAndAddVector(tiered_index->GetFlatIndex(), dim, kept_label, kept_label); + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), n + 1); + + tiered_index->scheduleSVSIndexInit(); + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 1); + ASSERT_EQ(tiered_index->GetFlatIndex()->getVectorLabel(0), kept_label); + TEST_DATA_T expected_vector[dim]; + GenerateVector(expected_vector, dim, kept_label); + ASSERT_EQ(tiered_index->GetFlatIndex()->getDistanceFrom_Unsafe(kept_label, expected_vector), 0); + for (size_t i = 0; i < n; i++) { + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(i)); + } + ASSERT_EQ(tiered_index->indexLabelCount(), n + 1); +} + +// ids_to_init_ is only non-empty while the backend is not ready, and insert jobs only exist once +// it is, so the swap bookkeeping of the two together is out of reach of addVector. The hand-made +// job below builds that state. Do not trigger an init after it: initSVSIndex asserts on a ready +// backend. +TYPED_TEST(SVSTieredIndexTestBasic, insertJobFollowsTrainingIdsThroughAFlatSwap) { + const size_t dim = 4; + const size_t n = 5; + const labelType ingested_label = 2; + SVSParams params = { + .type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_L2, .multi = false}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 100, 100); + ASSERT_INDEX(tiered_index); + + for (size_t i = 0; i < n; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); + ASSERT_FALSE(tiered_index->GetSVSIndex()->ready()); + + // Label i sits at flat id i, so the last id gets swapped onto the ingested one. + tiered_index->submitSingleJob(tiered_index->createInsertJob(ingested_label, ingested_label)); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 1); + mock_thread_pool.thread_iteration(); + + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), n - 1); + ASSERT_TRUE(tiered_index->GetSVSIndex()->isLabelExists(ingested_label)); + ASSERT_TRUE(tiered_index->labelToInsertJobs.empty()); + + TEST_DATA_T expected_vector[dim]; + for (size_t i : {size_t{0}, size_t{1}, size_t{3}, size_t{4}}) { + GenerateVector(expected_vector, dim, i); + ASSERT_EQ(tiered_index->GetFlatIndex()->getDistanceFrom_Unsafe(i, expected_vector), 0) + << "label " << i; + } + ASSERT_EQ(tiered_index->indexLabelCount(), n); +} + +// deleteVector() drops its shared lock before taking the exclusive one, so the label can be gone +// by the time removeLabelFromFlat() runs. +TYPED_TEST(SVSTieredIndexTestBasic, removeLabelFromFlatIgnoresAnAbsentLabel) { + const size_t dim = 4; + SVSParams params = { + .type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_L2, .multi = false}; + VecSimParams svs_params = CreateParams(params); + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateTieredSVSIndex(svs_params, mock_thread_pool, 100, 100); + ASSERT_INDEX(tiered_index); + + for (size_t i = 0; i < 2; i++) { + GenerateAndAddVector(tiered_index, dim, i, i); + } + + ASSERT_EQ(tiered_index->removeLabelFromFlat(999), 0); + ASSERT_EQ(tiered_index->GetFlatIndex()->indexSize(), 2); + ASSERT_EQ(tiered_index->indexLabelCount(), 2); } TEST(SVSTieredIndexTest, testThreadPool) {