[0.8] [MOD-10559] Decouple the shrinking and growing logic of large containers in Flat and HNSW

src/VecSim/algorithms/brute_force/brute_force.h

@@ -87,6 +87,15 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
         idToLabelMapping.shrink_to_fit();
         resizeLabelLookup(idToLabelMapping.size());
     }
+
+    size_t getStoredVectorsCount() const {
+        size_t actual_stored_vec = 0;
+        for (auto &block : vectorBlocks) {
+            actual_stored_vec += block.getLength();
+        }
+
+        return actual_stored_vec;
+    }
 #endif
 
 protected:
@@ -96,36 +105,61 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
     // Private internal function that implements generic single vector deletion.
     virtual void removeVector(idType id);
 
-    inline void growByBlock() {
+    void resizeIndexCommon(size_t new_max_elements) {
+        assert(new_max_elements % this->blockSize == 0 &&
+               "new_max_elements must be a multiple of blockSize");
+        this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing FLAT index from %zu to %zu",
+                  idToLabelMapping.capacity(), new_max_elements);
+        assert(idToLabelMapping.capacity() == idToLabelMapping.size());
+        idToLabelMapping.resize(new_max_elements);
+        idToLabelMapping.shrink_to_fit();
+        assert(idToLabelMapping.capacity() == idToLabelMapping.size());
+        resizeLabelLookup(new_max_elements);
+    }
+
+    void growByBlock() {
+        assert(indexCapacity() == idToLabelMapping.capacity());
+        assert(indexCapacity() % this->blockSize == 0);
+        assert(indexCapacity() == indexSize());
+
         assert(vectorBlocks.size() == 0 || vectorBlocks.back().getLength() == this->blockSize);
         vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
                                   this->alignment);
-        idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
-        idToLabelMapping.shrink_to_fit();
-        resizeLabelLookup(idToLabelMapping.size());
+        resizeIndexCommon(indexCapacity() + this->blockSize);
     }
 
-    inline void shrinkByBlock() {
-        assert(indexCapacity() > 0); // should not be called when index is empty
-
+    void shrinkByBlock() {
+        assert(indexCapacity() >= this->blockSize);
+        assert(indexCapacity() % this->blockSize == 0);
         // remove last block (should be empty)
         assert(vectorBlocks.size() > 0 && vectorBlocks.back().getLength() == 0);
         vectorBlocks.pop_back();
-
-        // remove a block size of labels.
-        assert(idToLabelMapping.size() >= this->blockSize);
-        idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
-        idToLabelMapping.shrink_to_fit();
-        resizeLabelLookup(idToLabelMapping.size());
+        assert(vectorBlocks.size() * this->blockSize == indexSize());
+
+        if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {
+            assert(indexCapacity() == idToLabelMapping.capacity());
+            assert(idToLabelMapping.size() == idToLabelMapping.capacity());
+            // There are at least two free blocks.
+            assert(vectorBlocks.size() * this->blockSize + 2 * this->blockSize <=
+                   idToLabelMapping.capacity());
+            resizeIndexCommon(indexCapacity() - this->blockSize);
+        } else if (indexCapacity() == this->blockSize) {
+            // Make sure the last block is removed.
+            // To align with the initial capacity behaviour, where we always remove one block, we
+            // have this special condition. On newer version branches, where the initial capacity is
+            // deprecated, we immediately shrunk to 0 when the index size is 0.
+            assert(vectorBlocks.empty());
+            assert(indexSize() == 0);
+            resizeIndexCommon(0);
+            return;
+        }
     }
 
     inline DataBlock &getVectorVectorBlock(idType id) {
         return vectorBlocks.at(id / this->blockSize);
     }
     inline size_t getVectorRelativeIndex(idType id) const { return id % this->blockSize; }
-    inline void setVectorLabel(idType id, labelType new_label) {
-        idToLabelMapping.at(id) = new_label;
-    }
+    void setVectorLabel(idType id, labelType new_label) { idToLabelMapping.at(id) = new_label; }
     // inline priority queue getter that need to be implemented by derived class
     virtual inline vecsim_stl::abstract_priority_queue<DistType, labelType> *
     getNewMaxPriorityQueue() const = 0;
@@ -166,19 +200,19 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(
 
 template <typename DataType, typename DistType>
 void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
-    // Give the vector new id and increase count.
-    idType id = this->count++;
-
-    // Resize the index if needed.
-    if (indexSize() > indexCapacity()) {
+    // Resize the index meta data structures if needed
+    if (indexSize() >= indexCapacity()) {
         growByBlock();
-    } else if (id % this->blockSize == 0) {
-        // If we didn't reach the initial capacity but the last block is full, add a new block
-        // only.
+    } else if (this->count % this->blockSize == 0) {
+        // If we didn't reach the initial capacity but the last block is full, initialize a new
+        // block only.
         this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
                                         this->alignment);
     }
 
+    // Give the vector new id and increase count.
+    idType id = this->count++;
+
     // Get the last vectors block to store the vector in.
     DataBlock &vectorBlock = this->vectorBlocks.back();
     assert(&vectorBlock == &getVectorVectorBlock(id));
@@ -239,7 +273,7 @@ size_t BruteForceIndex<DataType, DistType>::indexSize() const {
 
 template <typename DataType, typename DistType>
 size_t BruteForceIndex<DataType, DistType>::indexCapacity() const {
-    return this->idToLabelMapping.size();
+    return this->idToLabelMapping.capacity();
 }
 
 // Compute the score for every vector in the block by using the given distance function.

src/VecSim/algorithms/hnsw/hnsw.h

@@ -233,6 +233,11 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
     double getEpsilon() const;
     size_t indexSize() const override;
     size_t indexCapacity() const override;
+    /**
+     * Checks if the index capacity is full to hint the caller a resize is needed.
+     * @note Must be called with indexDataGuard locked.
+     */
+    size_t isCapacityFull() const;
     size_t getEfConstruction() const;
     size_t getM() const;
     size_t getMaxLevel() const;
@@ -312,6 +317,15 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
         idToMetaData.shrink_to_fit();
         resizeLabelLookup(idToMetaData.size());
     }
+
+    size_t getStoredVectorsCount() const {
+        size_t actual_stored_vec = 0;
+        for (auto &block : vectorBlocks) {
+            actual_stored_vec += block.getLength();
+        }
+
+        return actual_stored_vec;
+    }
 #endif
 
 protected:
@@ -357,6 +371,11 @@ size_t HNSWIndex<DataType, DistType>::indexCapacity() const {
     return this->maxElements;
 }
 
+template <typename DataType, typename DistType>
+size_t HNSWIndex<DataType, DistType>::isCapacityFull() const {
+    return indexSize() == this->maxElements;
+}
+
 template <typename DataType, typename DistType>
 size_t HNSWIndex<DataType, DistType>::getEfConstruction() const {
     return this->efConstruction;
@@ -1288,44 +1307,64 @@ template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::resizeIndexCommon(size_t new_max_elements) {
     assert(new_max_elements % this->blockSize == 0 &&
            "new_max_elements must be a multiple of blockSize");
-    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
-              "Updating HNSW index capacity from %zu to %zu", this->maxElements, new_max_elements);
+    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing HNSW index from %zu to %zu",
+              idToMetaData.capacity(), new_max_elements);
     resizeLabelLookup(new_max_elements);
     visitedNodesHandlerPool.resize(new_max_elements);
+    assert(idToMetaData.capacity() == idToMetaData.size());
     idToMetaData.resize(new_max_elements);
     idToMetaData.shrink_to_fit();
-
-    maxElements = new_max_elements;
+    assert(idToMetaData.capacity() == idToMetaData.size());
 }
 
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::growByBlock() {
-    size_t new_max_elements = maxElements + this->blockSize;
-
     // Validations
     assert(vectorBlocks.size() == graphDataBlocks.size());
     assert(vectorBlocks.empty() || vectorBlocks.back().getLength() == this->blockSize);
+    assert(this->maxElements % this->blockSize == 0);
+    assert(this->maxElements == indexSize());
+    assert(graphDataBlocks.size() == this->maxElements / this->blockSize);
+    assert(idToMetaData.capacity() == maxElements ||
+           idToMetaData.capacity() == maxElements + this->blockSize);
 
+    this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
+              "Updating HNSW index capacity from %zu to %zu", maxElements,
+              maxElements + this->blockSize);
+    maxElements += this->blockSize;
     vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator, this->alignment);
     graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize, this->allocator);
 
-    resizeIndexCommon(new_max_elements);
+    if (idToMetaData.capacity() == indexSize()) {
+        resizeIndexCommon(maxElements);
+    }
 }
 
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::shrinkByBlock() {
-    assert(maxElements >= this->blockSize);
-    size_t new_max_elements = maxElements - this->blockSize;
-
-    // Validations
-    assert(vectorBlocks.size() == graphDataBlocks.size());
-    assert(!vectorBlocks.empty());
-    assert(vectorBlocks.back().getLength() == 0);
-
-    vectorBlocks.pop_back();
-    graphDataBlocks.pop_back();
+    assert(this->maxElements >= this->blockSize);
+    assert(this->maxElements % this->blockSize == 0);
+    if (indexSize() % this->blockSize == 0) {
+        assert(vectorBlocks.back().getLength() == 0);
+        this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
+                  "Updating HNSW index capacity from %zu to %zu", maxElements,
+                  maxElements - this->blockSize);
+        vectorBlocks.pop_back();
+        graphDataBlocks.pop_back();
+        assert(graphDataBlocks.size() * this->blockSize == indexSize());
+
+        if (idToMetaData.capacity() >= (indexSize() + 2 * this->blockSize)) {
+            resizeIndexCommon(idToMetaData.capacity() - this->blockSize);
+        } else if (idToMetaData.capacity() == this->blockSize) {
+            assert(vectorBlocks.empty());
+            assert(indexSize() == 0);
+            assert(maxElements == this->blockSize);
+            resizeIndexCommon(0);
+        }
 
-    resizeIndexCommon(new_max_elements);
+        // Take the lower bound into account.
+        maxElements -= this->blockSize;
+    }
 }
 
 template <typename DataType, typename DistType>
@@ -1685,9 +1724,7 @@ void HNSWIndex<DataType, DistType>::removeAndSwap(idType internalId) {
 
     // If we need to free a complete block and there is at least one block between the
     // capacity and the size.
-    if (curElementCount % this->blockSize == 0) {
-        shrinkByBlock();
-    }
+    shrinkByBlock();
 }
 
 template <typename DataType, typename DistType>
@@ -1763,6 +1800,16 @@ void HNSWIndex<DataType, DistType>::removeVectorInPlace(const idType element_int
 template <typename DataType, typename DistType>
 AddVectorCtx HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
                                                             const void *vector_data) {
+    if (isCapacityFull()) {
+        growByBlock();
+    } else if (curElementCount % this->blockSize == 0) {
+        // If we had an initial capacity, we might have to initialize new blocks for the data and
+        // meta-data.
+        this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
+                                        this->alignment);
+        this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
+                                           this->allocator);
+    }
     AddVectorCtx state{};
 
     // Choose randomly the maximum level in which the new element will be in the index.
@@ -1790,17 +1837,6 @@ AddVectorCtx HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
         throw e;
     }
 
-    if (indexSize() > indexCapacity()) {
-        growByBlock();
-    } else if (state.newElementId % this->blockSize == 0) {
-        // If we had an initial capacity, we might have to allocate new blocks for the data and
-        // meta-data.
-        this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
-                                        this->alignment);
-        this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
-                                           this->allocator);
-    }
-
     // Insert the new element to the data block
     this->vectorBlocks.back().addElement(vector_data);
     this->graphDataBlocks.back().addElement(cur_egd);

src/VecSim/algorithms/hnsw/hnsw_tiered.h

@@ -308,7 +308,7 @@ template <typename DataType, typename DistType>
 void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToRun) {
 
     // Execute swap jobs - acquire hnsw write lock.
-    this->mainIndexGuard.lock();
+    this->lockMainIndexGuard();
     TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
                "Tiered HNSW index GC: there are %zu ready swap jobs. Start executing %zu swap jobs",
                readySwapJobs, std::min(readySwapJobs, maxJobsToRun));
@@ -333,7 +333,7 @@ void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToR
     readySwapJobs -= idsToRemove.size();
     TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
                "Tiered HNSW index GC: done executing %zu swap jobs", idsToRemove.size());
-    this->mainIndexGuard.unlock();
+    this->unlockMainIndexGuard();
 }
 
 template <typename DataType, typename DistType>
@@ -426,11 +426,11 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
     this->mainIndexGuard.lock_shared();
     hnsw_index->lockIndexDataGuard();
     // Check if resizing is needed for HNSW index (requires write lock).
-    if (hnsw_index->indexCapacity() == hnsw_index->indexSize()) {
+    if (hnsw_index->isCapacityFull()) {
         // Release the inner HNSW data lock before we re-acquire the global HNSW lock.
         this->mainIndexGuard.unlock_shared();
         hnsw_index->unlockIndexDataGuard();
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         hnsw_index->lockIndexDataGuard();
 
         // Hold the index data lock while we store the new element. If the new node's max level is
@@ -455,7 +455,7 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
         if (state.elementMaxLevel > state.currMaxLevel) {
             hnsw_index->unlockIndexDataGuard();
         }
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
     } else {
         // Do the same as above except for changing the capacity, but with *shared* lock held:
         // Hold the index data lock while we store the new element. If the new node's max level is
@@ -709,9 +709,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
         }
         // Insert the vector to the HNSW index. Internally, we will never have to overwrite the
         // label since we already checked it outside.
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         hnsw_index->addVector(blob, label);
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
         return ret;
     }
     if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
@@ -834,9 +834,9 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
         }
     } else {
         // delete in place.
-        this->mainIndexGuard.lock();
+        this->lockMainIndexGuard();
         num_deleted_vectors += this->deleteLabelFromHNSWInplace(label);
-        this->mainIndexGuard.unlock();
+        this->unlockMainIndexGuard();
     }
 
     return num_deleted_vectors;

src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h

@@ -54,6 +54,7 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteBothAsyncAndInplaceMulti_
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceMultiSwapId_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarkedDeleted_Test)
 INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
+INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test)
 
 friend class BF16TieredTest;
 friend class FP16TieredTest;