diff --git a/src/VecSim/algorithms/brute_force/brute_force.h b/src/VecSim/algorithms/brute_force/brute_force.h index cb0eb366c..338b63e80 100644 --- a/src/VecSim/algorithms/brute_force/brute_force.h +++ b/src/VecSim/algorithms/brute_force/brute_force.h @@ -87,6 +87,15 @@ class BruteForceIndex : public VecSimIndexAbstract { 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,63 @@ class BruteForceIndex : public VecSimIndexAbstract { // 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) { + // Special case to handle last block. + // This special condition resolves the ambiguity: when capacity==blockSize, we can't + // tell if this block came from growth (should shrink to 0) or initial capacity (should + // keep it). We choose to always shrink to 0 to maintain the one-block removal + // guarantee. In contrast, newer branches without initial capacity support use simpler + // logic: immediately shrink to 0 whenever index size becomes 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 * getNewMaxPriorityQueue() const = 0; @@ -166,19 +202,19 @@ BruteForceIndex::BruteForceIndex( template void BruteForceIndex::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)); diff --git a/src/VecSim/algorithms/hnsw/hnsw.h b/src/VecSim/algorithms/hnsw/hnsw.h index 7b37aa38f..eaa92199b 100644 --- a/src/VecSim/algorithms/hnsw/hnsw.h +++ b/src/VecSim/algorithms/hnsw/hnsw.h @@ -233,6 +233,11 @@ class HNSWIndex : public VecSimIndexAbstract, 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, 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::indexCapacity() const { return this->maxElements; } +template +size_t HNSWIndex::isCapacityFull() const { + return indexSize() == this->maxElements; +} + template size_t HNSWIndex::getEfConstruction() const { return this->efConstruction; @@ -1288,44 +1307,69 @@ template void HNSWIndex::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 void HNSWIndex::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 void HNSWIndex::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(); - - resizeIndexCommon(new_max_elements); + 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) { + // Special case to handle last block. + // This special condition resolves the ambiguity: when capacity==blockSize, we can't + // tell if this block came from growth (should shrink to 0) or initial capacity (should + // keep it). We choose to always shrink to 0 to maintain the one-block removal + // guarantee. In contrast, newer branches without initial capacity support use simpler + // logic: immediately shrink to 0 whenever index size becomes 0. + assert(vectorBlocks.empty()); + assert(indexSize() == 0); + assert(maxElements == this->blockSize); + resizeIndexCommon(0); + } + // Take the lower bound into account. + maxElements -= this->blockSize; + } } template @@ -1685,9 +1729,7 @@ void HNSWIndex::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 @@ -1763,6 +1805,16 @@ void HNSWIndex::removeVectorInPlace(const idType element_int template AddVectorCtx HNSWIndex::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 +1842,6 @@ AddVectorCtx HNSWIndex::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); diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered.h b/src/VecSim/algorithms/hnsw/hnsw_tiered.h index cb5c88183..1d3c78766 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered.h @@ -308,7 +308,7 @@ template void TieredHNSWIndex::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::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 @@ -426,11 +426,11 @@ void TieredHNSWIndex::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::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::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::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; diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h b/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h index b4cec5fef..abe4c2054 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h +++ b/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; diff --git a/src/VecSim/vec_sim_tiered_index.h b/src/VecSim/vec_sim_tiered_index.h index fb6535c99..60d82e551 100644 --- a/src/VecSim/vec_sim_tiered_index.h +++ b/src/VecSim/vec_sim_tiered_index.h @@ -41,7 +41,17 @@ class VecSimTieredIndex : public VecSimIndexInterface { mutable std::shared_mutex flatIndexGuard; mutable std::shared_mutex mainIndexGuard; + void lockMainIndexGuard() const { + mainIndexGuard.lock(); +#ifdef BUILD_TESTS + mainIndexGuard_write_lock_count++; +#endif + } + void unlockMainIndexGuard() const { mainIndexGuard.unlock(); } +#ifdef BUILD_TESTS + mutable std::atomic_int mainIndexGuard_write_lock_count = 0; +#endif size_t flatBufferLimit; void submitSingleJob(AsyncJob *job) { @@ -58,6 +68,9 @@ class VecSimTieredIndex : public VecSimIndexInterface { } public: +#ifdef BUILD_TESTS + int getMainIndexGuardWriteLockCount() const { return mainIndexGuard_write_lock_count; } +#endif VecSimTieredIndex(VecSimIndexAbstract *backendIndex_, BruteForceIndex *frontendIndex_, TieredIndexParams tieredParams, std::shared_ptr allocator) diff --git a/tests/unit/test_allocator.cpp b/tests/unit/test_allocator.cpp index 24edd9ab1..70f79511d 100644 --- a/tests/unit/test_allocator.cpp +++ b/tests/unit/test_allocator.cpp @@ -90,11 +90,12 @@ TYPED_TEST_SUITE(IndexAllocatorTest, DataTypeSet); TYPED_TEST(IndexAllocatorTest, test_bf_index_block_size_1) { // Create only the minimal struct. size_t dim = 128; + size_t blockSize = 1; BFParams params = {.type = TypeParam::get_index_type(), .dim = dim, .metric = VecSimMetric_IP, .initialCapacity = 0, - .blockSize = 1}; + .blockSize = blockSize}; auto *bfIndex = dynamic_cast *>( BruteForceFactory::NewIndex(¶ms)); bfIndex->alignment = 0; // Disable alignment for testing purposes. @@ -107,95 +108,228 @@ TYPED_TEST(IndexAllocatorTest, test_bf_index_block_size_1) { size_t memory = VecSimIndex_StatsInfo(bfIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); + // @param expected_size - The expected number of elements in the index. + // @param expected_data_container_blocks - The expected number of blocks in the data containers. + // @param expected_map_containers_capacity - The expected capacity of the map containers in + // number of elements. + auto verify_containers_size = [&](size_t expected_size, size_t expected_data_container_blocks, + size_t expected_map_containers_size) { + ASSERT_EQ(bfIndex->indexSize(), expected_size); + ASSERT_EQ(bfIndex->vectorBlocks.size(), expected_data_container_blocks); + ASSERT_EQ(bfIndex->getStoredVectorsCount(), expected_size); + + ASSERT_EQ(bfIndex->indexCapacity(), expected_map_containers_size); + ASSERT_EQ(bfIndex->idToLabelMapping.capacity(), expected_map_containers_size); + ASSERT_EQ(bfIndex->idToLabelMapping.size(), expected_map_containers_size); + ASSERT_GE(bfIndex->labelToIdLookup.bucket_count(), expected_map_containers_size); + }; + // =========== Add label 1 =========== int before = allocator->getAllocationSize(); + size_t buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + auto &vectors_blocks = bfIndex->vectorBlocks; + size_t vectors_blocks_capacity = vectors_blocks.capacity(); + VecSimIndex_AddVector(bfIndex, vec, 1); int addCommandAllocationDelta = allocator->getAllocationSize() - before; - int64_t expectedAllocationDelta = 0; - expectedAllocationDelta += + int64_t expectedAllocationDelta = sizeof(labelType) + vecsimAllocationOverhead; // resize idToLabelMapping - expectedAllocationDelta += sizeof(DataBlock) + vecsimAllocationOverhead; // New vector block expectedAllocationDelta += - sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead; // keep the vector in the vector block + (vectors_blocks.capacity() - vectors_blocks_capacity) * sizeof(DataBlock) + + vecsimAllocationOverhead; // New vectors blocks + expectedAllocationDelta += blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment(); // block vectors buffer + expectedAllocationDelta += hashTableNodeSize; // New node in the label lookup + // Account for the allocation of a new buckets in the labels_lookup hash table. expectedAllocationDelta += - sizeof(std::pair) + vecsimAllocationOverhead; // keep the mapping + (bfIndex->labelToIdLookup.bucket_count() - buckets_num_before) * sizeof(size_t); // Assert that the additional allocated delta did occur, and it is limited, as some STL // collection allocate additional structures for their internal implementation. - ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize + addCommandAllocationDelta); - ASSERT_LE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); - ASSERT_LE(expectedAllocationDelta, addCommandAllocationDelta); - memory = VecSimIndex_StatsInfo(bfIndex).memory; - ASSERT_EQ(allocator->getAllocationSize(), memory); + { + SCOPED_TRACE("Verifying allocation delta for adding first vector"); + verify_containers_size(1, 1, 1); + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + addCommandAllocationDelta); + ASSERT_LE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + ASSERT_LE(expectedAllocationDelta, addCommandAllocationDelta); + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + } + + // =========== labels = [1], vector blocks = 1, maps capacity = 1. Add label 2 + 3 =========== // Prepare for next assertion test expectedAllocationSize = memory; expectedAllocationDelta = 0; before = allocator->getAllocationSize(); + vectors_blocks_capacity = vectors_blocks.capacity(); + buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + VecSimIndex_AddVector(bfIndex, vec, 2); + VecSimIndex_AddVector(bfIndex, vec, 3); addCommandAllocationDelta = allocator->getAllocationSize() - before; - expectedAllocationDelta += sizeof(DataBlock) + vecsimAllocationOverhead; // New vector block - expectedAllocationDelta += sizeof(labelType); // resize idToLabelMapping + expectedAllocationDelta += (vectors_blocks.capacity() - vectors_blocks_capacity) * + sizeof(DataBlock); // New vector blocks + expectedAllocationDelta += 2 * sizeof(labelType); // resize idToLabelMapping expectedAllocationDelta += - sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead; // keep the vector in the vector block + 2 * (blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment()); // Two block vectors buffer + expectedAllocationDelta += 2 * hashTableNodeSize; // New nodes in the label lookup expectedAllocationDelta += - sizeof(std::pair) + vecsimAllocationOverhead; // keep the mapping - // Assert that the additional allocated delta did occur, and it is limited, as some STL - // collection allocate additional structures for their internal implementation. - ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize + addCommandAllocationDelta); - ASSERT_LE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); - ASSERT_LE(expectedAllocationDelta, addCommandAllocationDelta); - memory = VecSimIndex_StatsInfo(bfIndex).memory; - ASSERT_EQ(allocator->getAllocationSize(), memory); + (bfIndex->labelToIdLookup.bucket_count() - buckets_num_before) * sizeof(size_t); + { + SCOPED_TRACE("Index size = 1Verifying allocation delta for adding two more vectors"); + verify_containers_size(3, 3, 3); + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + addCommandAllocationDelta); + ASSERT_EQ(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + ASSERT_EQ(expectedAllocationDelta, addCommandAllocationDelta); + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + } + + // =========== labels = [1, 2, 3], vector blocks = 3, maps capacity = 3. Delete label 1 + // =========== // Prepare for next assertion test expectedAllocationSize = memory; expectedAllocationDelta = 0; before = allocator->getAllocationSize(); - VecSimIndex_DeleteVector(bfIndex, 2); - int deleteCommandAllocationDelta = allocator->getAllocationSize() - before; - expectedAllocationDelta -= - sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead; // Free the vector in the vector block - expectedAllocationDelta -= sizeof(labelType); // resize idToLabelMapping - expectedAllocationDelta -= - sizeof(std::pair) + vecsimAllocationOverhead; // remove one label:id pair + vectors_blocks_capacity = vectors_blocks.capacity(); + buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + { + SCOPED_TRACE("Verifying allocation delta for deleting a vector from index size 3"); + ASSERT_EQ(VecSimIndex_DeleteVector(bfIndex, 1), 1); + int deleteCommandAllocationDelta = allocator->getAllocationSize() - before; + verify_containers_size(2, 2, 3); + // Removing blocks doesn't change vectors_blocks.capacity(), but the block buffer is freed. + ASSERT_EQ(vectors_blocks.capacity(), vectors_blocks_capacity); + expectedAllocationDelta -= + blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment(); // Free the vector buffer in the vector block + expectedAllocationDelta -= hashTableNodeSize; // Remove node from the label lookup + // idToLabelMapping and label:id should not change since count > capacity - 2 * blockSize + ASSERT_EQ(bfIndex->labelToIdLookup.bucket_count(), buckets_num_before); + + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + deleteCommandAllocationDelta); + ASSERT_EQ(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + ASSERT_EQ(expectedAllocationDelta, deleteCommandAllocationDelta); + + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + } - // Assert that the reclaiming of memory did occur, and it is limited, as some STL - // collection allocate additional structures for their internal implementation. - ASSERT_EQ(allocator->getAllocationSize(), - expectedAllocationSize + deleteCommandAllocationDelta); - ASSERT_GE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); - ASSERT_GE(expectedAllocationDelta, deleteCommandAllocationDelta); + // =========== labels = [2, 3], vector blocks = 2, maps capacity = 3. Add label 4 =========== - memory = VecSimIndex_StatsInfo(bfIndex).memory; - ASSERT_EQ(allocator->getAllocationSize(), memory); + // Prepare for next assertion test + expectedAllocationSize = memory; + expectedAllocationDelta = 0; + + before = allocator->getAllocationSize(); + vectors_blocks_capacity = vectors_blocks.capacity(); + buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + size_t idToLabel_size_before = bfIndex->idToLabelMapping.size(); + + VecSimIndex_AddVector(bfIndex, vec, 4); + addCommandAllocationDelta = allocator->getAllocationSize() - before; + expectedAllocationDelta += (vectors_blocks.capacity() - vectors_blocks_capacity) * + sizeof(DataBlock); // New vector block + expectedAllocationDelta += blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment(); // block vectors buffer + expectedAllocationDelta += hashTableNodeSize; // New node in the label lookup + { + SCOPED_TRACE( + "Verifying allocation delta for adding a vector to index size 2 with capacity 3"); + verify_containers_size(3, 3, 3); + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + addCommandAllocationDelta); + ASSERT_EQ(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + ASSERT_EQ(expectedAllocationDelta, addCommandAllocationDelta); + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + + // idToLabelMapping and label:id should not change since if we one free block + ASSERT_EQ(bfIndex->labelToIdLookup.bucket_count(), buckets_num_before); + ASSERT_EQ(bfIndex->idToLabelMapping.size(), idToLabel_size_before); + } + + // =========== labels = [2, 3, 4], vector blocks = 3, maps capacity = 3. Delete label 2 + 3 + // =========== // Prepare for next assertion test expectedAllocationSize = memory; expectedAllocationDelta = 0; before = allocator->getAllocationSize(); - VecSimIndex_DeleteVector(bfIndex, 1); - deleteCommandAllocationDelta = allocator->getAllocationSize() - before; - expectedAllocationDelta -= - (sizeof(DataBlock) + vecsimAllocationOverhead); // Free the vector block - expectedAllocationDelta -= - sizeof(DataBlock *) + vecsimAllocationOverhead; // remove from vectorBlocks vector - expectedAllocationDelta -= - sizeof(labelType) + vecsimAllocationOverhead; // resize idToLabelMapping - expectedAllocationDelta -= (sizeof(TEST_DATA_T) * dim + - vecsimAllocationOverhead); // Free the vector in the vector block - expectedAllocationDelta -= - sizeof(std::pair) + vecsimAllocationOverhead; // remove one label:id pair + vectors_blocks_capacity = vectors_blocks.capacity(); + buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + { + SCOPED_TRACE("Verifying allocation delta for deleting two vectors from index size 3"); + ASSERT_EQ(VecSimIndex_DeleteVector(bfIndex, 2), 1); + ASSERT_EQ(VecSimIndex_DeleteVector(bfIndex, 3), 1); + + int deleteCommandAllocationDelta = allocator->getAllocationSize() - before; + verify_containers_size(1, 1, 2); + // Removing blocks doesn't change vectors_blocks.capacity(), but the block buffer is freed. + ASSERT_EQ(vectors_blocks.capacity(), vectors_blocks_capacity); + expectedAllocationDelta -= + 2 * (blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment()); // Free the vector buffer in the vector block + expectedAllocationDelta -= 2 * hashTableNodeSize; // Remove nodes from the label lookup + // idToLabelMapping and label:id should shrink by block since count >= capacity - 2 * + // blockSize + expectedAllocationDelta -= sizeof(labelType); // remove one idToLabelMapping + expectedAllocationDelta -= + (buckets_num_before - bfIndex->labelToIdLookup.bucket_count()) * sizeof(size_t); + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + deleteCommandAllocationDelta); + ASSERT_EQ(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + ASSERT_EQ(expectedAllocationDelta, deleteCommandAllocationDelta); + + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + } + + // =========== labels = [4], vector blocks = 1, maps capacity = 2. Delete last label =========== + + // Prepare for next assertion test + expectedAllocationSize = memory; + expectedAllocationDelta = 0; + + before = allocator->getAllocationSize(); + vectors_blocks_capacity = vectors_blocks.capacity(); + buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); + { + SCOPED_TRACE("Verifying allocation delta for emptying the index"); + ASSERT_EQ(VecSimIndex_DeleteVector(bfIndex, 4), 1); + + int deleteCommandAllocationDelta = allocator->getAllocationSize() - before; + // We decrease meta data containers size by one block + verify_containers_size(0, 0, 1); + // Removing blocks doesn't change vectors_blocks.capacity(), but the block buffer is freed. + ASSERT_EQ(vectors_blocks.capacity(), vectors_blocks_capacity); + expectedAllocationDelta -= + (blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + + bfIndex->getAlignment()); // Free the vector buffer in the vector block + expectedAllocationDelta -= hashTableNodeSize; // Remove nodes from the label lookup + // idToLabelMapping and label:id should shrink by block since count >= capacity - 2 * + // blockSize + expectedAllocationDelta -= + sizeof(labelType); // remove one idToLabelMapping and free the container + // resizing labelToIdLookup + size_t buckets_after = bfIndex->labelToIdLookup.bucket_count(); + expectedAllocationDelta -= (buckets_num_before - buckets_after) * sizeof(size_t); + ASSERT_EQ(allocator->getAllocationSize(), + expectedAllocationSize + deleteCommandAllocationDelta); + ASSERT_LE(abs(expectedAllocationDelta), abs(deleteCommandAllocationDelta)); + ASSERT_GE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); + + memory = VecSimIndex_StatsInfo(bfIndex).memory; + ASSERT_EQ(allocator->getAllocationSize(), memory); + } - // Assert that the reclaiming of memory did occur, and it is limited, as some STL - // collection allocate additional structures for their internal implementation. - ASSERT_EQ(allocator->getAllocationSize(), - expectedAllocationSize + deleteCommandAllocationDelta); - ASSERT_LE(expectedAllocationSize + expectedAllocationDelta, allocator->getAllocationSize()); - ASSERT_LE(expectedAllocationDelta, deleteCommandAllocationDelta); - memory = VecSimIndex_StatsInfo(bfIndex).memory; - ASSERT_EQ(allocator->getAllocationSize(), memory); VecSimIndex_Free(bfIndex); } @@ -373,80 +507,106 @@ TYPED_TEST(IndexAllocatorTest, test_hnsw_reclaim_memory) { // Add vectors up to the size of a whole block, and calculate the total memory delta. size_t block_size = hnswIndex->basicInfo().blockSize; - size_t accumulated_mem_delta = allocator->getAllocationSize(); + size_t prev_bucket_count = hnswIndex->labelLookup.bucket_count(); for (size_t i = 0; i < block_size; i++) { GenerateAndAddVector(hnswIndex, d, i, i); } // Get the memory delta after adding the block. - accumulated_mem_delta = allocator->getAllocationSize() - accumulated_mem_delta; + size_t one_block_mem_delta = allocator->getAllocationSize() - initial_memory_size; + + size_t one_block_buckets = hnswIndex->labelLookup.bucket_count(); + // @param expected_size - The expected number of elements in the index. + // @param expected_data_container_blocks - The expected number of blocks in the data containers. + // @param expected_map_containers_capacity - The expected capacity of the map containers in + // number of elements. + auto verify_containers_size = [&](size_t expected_size, size_t expected_data_container_blocks, + size_t expected_map_containers_size) { + SCOPED_TRACE("Verifying containers size for size " + std::to_string(expected_size)); + ASSERT_EQ(hnswIndex->indexSize(), expected_size); + ASSERT_EQ(hnswIndex->indexCapacity(), expected_data_container_blocks * block_size); + ASSERT_EQ(hnswIndex->indexCapacity(), hnswIndex->maxElements); + ASSERT_EQ(hnswIndex->graphDataBlocks.size(), expected_data_container_blocks); + ASSERT_EQ(hnswIndex->vectorBlocks.size(), expected_data_container_blocks); + ASSERT_EQ(hnswIndex->getStoredVectorsCount(), expected_size); + + ASSERT_EQ(hnswIndex->idToMetaData.capacity(), expected_map_containers_size); + ASSERT_EQ(hnswIndex->idToMetaData.size(), expected_map_containers_size); + ASSERT_GE(hnswIndex->labelLookup.bucket_count(), expected_map_containers_size); + // Also validate that there are no unidirectional connections (these add memory to the + // incoming edges sets). + ASSERT_EQ(hnswIndex->checkIntegrity().unidirectional_connections, 0); + }; // Validate that a single block exists. - ASSERT_EQ(hnswIndex->indexSize(), block_size); - ASSERT_EQ(hnswIndex->indexCapacity(), block_size); - ASSERT_EQ(allocator->getAllocationSize(), initial_memory_size + accumulated_mem_delta); - // Also validate that there are no unidirectional connections (these add memory to the incoming - // edges sets). - ASSERT_EQ(hnswIndex->checkIntegrity().unidirectional_connections, 0); + verify_containers_size(block_size, 1, block_size); + size_t one_block_mem = allocator->getAllocationSize(); // Add another vector, expect resizing of the index to contain two blocks. - size_t prev_bucket_count = hnswIndex->labelLookup.bucket_count(); - size_t mem_delta = allocator->getAllocationSize(); GenerateAndAddVector(hnswIndex, d, block_size, block_size); - mem_delta = allocator->getAllocationSize() - mem_delta; - - ASSERT_EQ(hnswIndex->indexSize(), block_size + 1); - ASSERT_EQ(hnswIndex->indexCapacity(), 2 * block_size); - ASSERT_EQ(hnswIndex->checkIntegrity().unidirectional_connections, 0); + verify_containers_size(block_size + 1, 2, 2 * block_size); + size_t mem_delta = allocator->getAllocationSize() - one_block_mem; // Compute the expected memory allocation due to the last vector insertion. size_t vec_max_level = hnswIndex->getGraphDataByInternalId(block_size)->toplevel; - size_t expected_mem_delta = - (vec_max_level + 1) * (sizeof(vecsim_stl::vector) + vecsimAllocationOverhead) + - hashTableNodeSize; + size_t last_vec_graph_data_mem = + (sizeof(vecsim_stl::vector) + vecsimAllocationOverhead) + hashTableNodeSize; if (vec_max_level > 0) { - expected_mem_delta += hnswIndex->levelDataSize * vec_max_level + vecsimAllocationOverhead; + last_vec_graph_data_mem += + hnswIndex->levelDataSize * vec_max_level + vecsimAllocationOverhead; } + size_t expected_mem_delta = last_vec_graph_data_mem; // Also account for all the memory allocation caused by the resizing that this vector triggered // except for the bucket count of the labels_lookup hash table that is calculated separately. + // Calculate the expected memory delta for adding a block. + size_t data_containers_block_mem = + 2 * (sizeof(DataBlock) + vecsimAllocationOverhead) + hnswIndex->getAlignment(); size_t size_total_data_per_element = hnswIndex->elementGraphDataSize + hnswIndex->dataSize; + data_containers_block_mem += size_total_data_per_element * block_size; + // account for idToMetaData and visitedNodesHandlerPool entries. expected_mem_delta += - (sizeof(tag_t) + sizeof(labelType) + sizeof(elementFlags) + size_total_data_per_element) * - block_size; + (sizeof(tag_t) + sizeof(ElementMetaData)) * block_size + data_containers_block_mem; + // Account for the allocation of a new bucket in the labels_lookup hash table. expected_mem_delta += - (hnswIndex->labelLookup.bucket_count() - prev_bucket_count) * sizeof(size_t); + (hnswIndex->labelLookup.bucket_count() - one_block_buckets) * sizeof(size_t); // New blocks allocated - 1 aligned block for vectors and 1 unaligned block for graph data. - expected_mem_delta += 2 * (sizeof(DataBlock) + vecsimAllocationOverhead) + hnswIndex->alignment; - expected_mem_delta += - (hnswIndex->vectorBlocks.capacity() - hnswIndex->vectorBlocks.size()) * sizeof(DataBlock); - expected_mem_delta += - (hnswIndex->graphDataBlocks.capacity() - hnswIndex->graphDataBlocks.size()) * - sizeof(DataBlock); ASSERT_EQ(expected_mem_delta, mem_delta); - // Remove the last vector, expect resizing back to a single block, and return to the previous - // memory consumption. + // Remove the last vector, expect datablocks containers (vectors buffer and graph data) resizing + // back to a single block. Index-size container such as id to label mapping, are only freed when + // there two empty blocks. + size_t before_delete_mem = allocator->getAllocationSize(); + size_t graph_data_blocks_capacity = hnswIndex->graphDataBlocks.capacity(); + auto &vectors_blocks = hnswIndex->vectorBlocks; + size_t vectors_blocks_capacity = vectors_blocks.capacity(); VecSimIndex_DeleteVector(hnswIndex, block_size); - ASSERT_EQ(hnswIndex->indexSize(), block_size); - ASSERT_EQ(hnswIndex->indexCapacity(), block_size); - ASSERT_EQ(hnswIndex->checkIntegrity().unidirectional_connections, 0); - size_t expected_allocation_size = initial_memory_size + accumulated_mem_delta; - expected_allocation_size += - (hnswIndex->vectorBlocks.capacity() - hnswIndex->vectorBlocks.size()) * sizeof(DataBlock); - expected_allocation_size += - (hnswIndex->graphDataBlocks.capacity() - hnswIndex->graphDataBlocks.size()) * - sizeof(DataBlock); + verify_containers_size(block_size, 1, 2 * block_size); + + size_t expected_allocation_size = + before_delete_mem - last_vec_graph_data_mem - hnswIndex->getAlignment(); + // Free the buffer of the last block in both data containers. + expected_allocation_size -= + size_total_data_per_element * block_size + 2 * vecsimAllocationOverhead; + expected_allocation_size -= + (graph_data_blocks_capacity - hnswIndex->graphDataBlocks.capacity()) * + (sizeof(DataBlock) + vecsimAllocationOverhead); + expected_allocation_size -= (vectors_blocks_capacity - vectors_blocks.capacity()) * + (sizeof(DataBlock) + vecsimAllocationOverhead); ASSERT_EQ(allocator->getAllocationSize(), expected_allocation_size); - // Remove the rest of the vectors, and validate that the memory returns to its initial state. + // Remove the rest of the vectors, meta data containers size should decrease by one block. for (size_t i = 0; i < block_size; i++) { VecSimIndex_DeleteVector(hnswIndex, i); } ASSERT_EQ(hnswIndex->indexSize(), 0); + // Capacity reelects the lower bound, which is the data containers size ASSERT_EQ(hnswIndex->indexCapacity(), 0); - // All data structures' memory returns to as it was, with the exceptional of the labels_lookup - // (STL unordered_map with hash table implementation), that leaves some empty buckets. + + // Add and remove a vector to trigger another resize + GenerateAndAddVector(hnswIndex, d, 0); + VecSimIndex_DeleteVector(hnswIndex, 0); + size_t hash_table_memory = hnswIndex->labelLookup.bucket_count() * sizeof(size_t); // Data block vectors do not shrink on resize so extra memory is expected. size_t block_vectors_memory = sizeof(DataBlock) * (hnswIndex->graphDataBlocks.capacity() + diff --git a/tests/unit/test_bruteforce.cpp b/tests/unit/test_bruteforce.cpp index 48c3d2ce2..6c6b75c15 100644 --- a/tests/unit/test_bruteforce.cpp +++ b/tests/unit/test_bruteforce.cpp @@ -108,52 +108,175 @@ TYPED_TEST(BruteForceTest, brute_force_vector_update_test) { TYPED_TEST(BruteForceTest, resize_and_align_index) { size_t dim = 4; - size_t n = 14; size_t blockSize = 10; + size_t curr_label = 0; - BFParams params = { - .dim = dim, .metric = VecSimMetric_L2, .initialCapacity = n, .blockSize = blockSize}; + BFParams params = {.dim = dim, .metric = VecSimMetric_L2, .blockSize = blockSize}; VecSimIndex *index = this->CreateNewIndex(params); BruteForceIndex *bf_index = this->CastToBF(index); - ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + auto verify_index_size = [&](size_t expected_size, size_t expected_capacity, std::string msg) { + SCOPED_TRACE("verify_index_size: " + msg); + ASSERT_EQ(VecSimIndex_IndexSize(index), expected_size); + ASSERT_EQ(bf_index->idToLabelMapping.size(), expected_capacity); + ASSERT_EQ(bf_index->indexCapacity(), expected_capacity); + ASSERT_EQ(bf_index->getVectorBlocks().size(), + expected_size / blockSize + (expected_size % blockSize != 0)) + << "expected_size: " << expected_size << " expected_capacity: " << expected_capacity; + }; + // Empty index with no initial capacity + verify_index_size(0, 0, "empty index"); - for (size_t i = 0; i < n; i++) { - GenerateAndAddVector(index, dim, i, i); + // Add one vector, index capacity should grow to blockSize. + GenerateAndAddVector(index, dim, curr_label++); + verify_index_size(1, blockSize, "add one vector"); + + // Add vector up to blocksize, index capacity should remain the same. + while (curr_label < blockSize) { + GenerateAndAddVector(index, dim, curr_label++); } - ASSERT_EQ(bf_index->idToLabelMapping.size(), 2 * blockSize); - ASSERT_EQ(VecSimIndex_IndexSize(index), n); + verify_index_size(blockSize, blockSize, "add up to blocksize"); // remove invalid id VecSimIndex_DeleteVector(index, 3459); // This should do nothing - ASSERT_EQ(VecSimIndex_IndexSize(index), n); - ASSERT_EQ(bf_index->idToLabelMapping.size(), 2 * blockSize); + verify_index_size(blockSize, blockSize, "remove invalid id"); - // Add another vector, since index size equals to the capacity, this should cause resizing - // (to fit a multiplication of block_size). - GenerateAndAddVector(index, dim, n); - ASSERT_EQ(VecSimIndex_IndexSize(index), n + 1); - // Capacity and size should remain blockSize * 2. - ASSERT_EQ(bf_index->idToLabelMapping.size(), 2 * blockSize); - ASSERT_EQ(bf_index->idToLabelMapping.capacity(), 2 * blockSize); + // Add another vector, since index size equals to block size, this should cause resizing + // of the capacity by one blocksize + GenerateAndAddVector(index, dim, curr_label++); + verify_index_size(blockSize + 1, 2 * blockSize, + "add one more vector after reaching block size"); - // Now size = n + 1 (= 15), capacity = 2 * bs (= 20). Test capacity overflow again + // Now size = blocksize + 1 (= 11), capacity = 2 * bs (= 20). Test capacity overflow again // to check that it stays aligned with block size. + size_t add_vectors_count = blockSize + 2; // 12 + while (curr_label < add_vectors_count + blockSize + 1) { + GenerateAndAddVector(index, dim, curr_label++); + } + + // Size should be blocksize + 1 + add_vectors_count (= 23). + verify_index_size( + blockSize + 1 + add_vectors_count, 3 * blockSize, + "add more vectors after reaching 2 * blocksize capacity to trigger another resize"); + + // Delete vectors so that indexsize % blocksize == 0 (and then delete one more) + size_t num_deleted = 0; + auto remove_to_one_below_blocksize = [&](size_t initial_label_to_remove) { + while (VecSimIndex_IndexSize(index) % blockSize != 0) { + ASSERT_EQ(VecSimIndex_DeleteVector(index, initial_label_to_remove++), 1) + << "tried to remove label: " << initial_label_to_remove - 1; + num_deleted++; + } + VecSimIndex_DeleteVector(index, initial_label_to_remove); + num_deleted++; + }; + + // First trigger of remove_to_one_below_blocksize will result in one free block. + // This should not trigger shrinking of metadata containers. + remove_to_one_below_blocksize(0); // remove first block labels. + verify_index_size( + blockSize * 2 - 1, 3 * blockSize, + "delete vectors so that indexsize % blocksize == 0, but there is only one free block"); + + // Second trigger of remove_to_one_below_blocksize will result in two free blocks. + // This should trigger shrinking of metadata containers by one block. + remove_to_one_below_blocksize(num_deleted); + verify_index_size( + blockSize - 1, 2 * blockSize, + "delete vectors so that indexsize % blocksize == 0 and there are two free blocks"); + + // Now there is one block in use and one free. adding vectors up to blocksize should not trigger + // another resize. + GenerateAndAddVector(index, dim, curr_label++); + verify_index_size(blockSize, 2 * blockSize, + "add vectors up to blocksize after deleting two blocks"); + + // Delete all vectors. + while (VecSimIndex_IndexSize(index) > 0) { + VecSimIndex_DeleteVector(index, num_deleted++); + } + + // We shrink the capcity by blocksize. + verify_index_size(0, blockSize, "delete all vectors"); + + // Add one vector and delete it to verify we shrink to zero. + GenerateAndAddVector(index, dim, 0); + verify_index_size(1, blockSize, "add one vector after deleting all"); + VecSimIndex_DeleteVector(index, 0); + verify_index_size(0, 0, "delete the only vector to verify shrinking to zero"); + + VecSimIndex_Free(index); +} + +TYPED_TEST(BruteForceTest, brute_force_no_oscillation_test) { + size_t dim = 4; + size_t blockSize = 2; + size_t cycles = 5; // Number of add/delete cycles to test + + BFParams params = {.dim = dim, .metric = VecSimMetric_L2, .blockSize = blockSize}; + VecSimIndex *index = this->CreateNewIndex(params); + BruteForceIndex *bf_index = this->CastToBF(index); + + auto verify_no_oscillation = [&](size_t expected_size, size_t expected_capacity, + const std::string &msg) { + SCOPED_TRACE("verify_no_oscillation: " + msg); + ASSERT_EQ(VecSimIndex_IndexSize(index), expected_size); + ASSERT_EQ(bf_index->indexCapacity(), expected_capacity); + }; + + // Initial state: empty index + verify_no_oscillation(0, 0, "initial empty state"); + + size_t current_label = 0; - size_t add_vectors_count = 8; - for (size_t i = 0; i < add_vectors_count; i++) { - GenerateAndAddVector(index, dim, n + 2 + i, i); + // Add initial 3 blocks + size_t initial_num_blocks = 3; + for (size_t i = 0; i < initial_num_blocks * blockSize; i++) { + GenerateAndAddVector(index, dim, current_label++); + } + verify_no_oscillation(initial_num_blocks * blockSize, initial_num_blocks * blockSize, + "initial " + std::to_string(initial_num_blocks) + + " blocks vectors added"); + + // Perform oscillation cycles: delete block, add block, delete block, add block... + for (size_t cycle = 0; cycle < cycles; cycle++) { + // Delete blockSize vectors (size becomes blockSize, but capacity should remain 2 * + // blockSize due to buffer zone) + for (size_t i = 0; i < blockSize; i++) { + VecSimIndex_DeleteVector(index, cycle * blockSize + i); + } + verify_no_oscillation((initial_num_blocks - 1) * blockSize, initial_num_blocks * blockSize, + "cycle " + std::to_string(cycle) + + " - after deleting block of vectors"); + + // Add blockSize vectors back (size becomes 2 * blockSize, capacity should remain 2 * + // blockSize) + for (size_t i = 0; i < blockSize; i++) { + GenerateAndAddVector(index, dim, current_label++); + } + verify_no_oscillation(initial_num_blocks * blockSize, initial_num_blocks * blockSize, + "cycle " + std::to_string(cycle) + + " - after adding blockSize vectors back"); } - // Size should be n + 1 + 8 (= 25). - ASSERT_EQ(VecSimIndex_IndexSize(index), n + 1 + add_vectors_count); + // Final verification: delete enough vectors to trigger actual shrinking + // Delete blocksize vectors to have only one block of vectors (2 free blocks = shrinking + // condition) + size_t vectors_to_delete = 2 * blockSize; + for (size_t i = 0; i < vectors_to_delete; i++) { + VecSimIndex_DeleteVector(index, cycles * blockSize + i); + } + verify_no_oscillation(blockSize, 2 * blockSize, + "final shrinking to trigger shrinking by one block"); - // Check new capacity size, should be blockSize * 3. - ASSERT_EQ(bf_index->idToLabelMapping.size(), 3 * blockSize); - ASSERT_EQ(bf_index->idToLabelMapping.capacity(), 3 * blockSize); + // Verify we can still grow normally after the oscillation test + for (size_t i = 0; i < blockSize; i++) { + GenerateAndAddVector(index, dim, current_label++); + } + verify_no_oscillation(2 * blockSize, 2 * blockSize, "growth after oscillation test"); VecSimIndex_Free(index); } @@ -868,8 +991,10 @@ TYPED_TEST(BruteForceTest, brute_force_remove_vector_after_replacing_block) { TYPED_TEST(BruteForceTest, brute_force_zero_minimal_capacity) { size_t dim = 4; size_t n = 2; + size_t bs = 1; - BFParams params = {.dim = dim, .metric = VecSimMetric_L2, .initialCapacity = 0, .blockSize = 1}; + BFParams params = { + .dim = dim, .metric = VecSimMetric_L2, .initialCapacity = 0, .blockSize = bs}; VecSimIndex *index = this->CreateNewIndex(params); @@ -891,8 +1016,8 @@ TYPED_TEST(BruteForceTest, brute_force_zero_minimal_capacity) { VecSimIndex_DeleteVector(index, i); } ASSERT_EQ(VecSimIndex_IndexSize(index), 0); - // id2label size should be the same as index size - ASSERT_EQ(bf_index->idToLabelMapping.size(), 0); + // id2label size should decrease by one block. + ASSERT_EQ(bf_index->idToLabelMapping.size(), n - bs); VecSimIndex_Free(index); } diff --git a/tests/unit/test_bruteforce_multi.cpp b/tests/unit/test_bruteforce_multi.cpp index 921218d14..7639b494e 100644 --- a/tests/unit/test_bruteforce_multi.cpp +++ b/tests/unit/test_bruteforce_multi.cpp @@ -68,9 +68,11 @@ TYPED_TEST(BruteForceMultiTest, vector_add_multiple_test) { TYPED_TEST(BruteForceMultiTest, resize_and_align_index) { size_t dim = 4; - size_t n = 15; - size_t blockSize = 10; - size_t n_labels = 3; + constexpr size_t blockSize = 10; + constexpr size_t per_label = 3; + constexpr size_t n_labels = 4; + constexpr size_t n = n_labels * per_label; + constexpr size_t expected_cap = n - n % blockSize + blockSize; BFParams params = { .dim = dim, .metric = VecSimMetric_L2, .initialCapacity = n, .blockSize = blockSize}; @@ -78,61 +80,109 @@ TYPED_TEST(BruteForceMultiTest, resize_and_align_index) { VecSimIndex *index = this->CreateNewIndex(params); auto bf_index = this->CastToBF_Multi(index); - ASSERT_EQ(VecSimIndex_IndexSize(index), 0); + auto verify_index_size = [&](size_t expected_num_vectors, size_t expected_labels, + size_t expected_capacity, std::string msg) { + SCOPED_TRACE("verify_index_size: " + msg); + VecSimIndexDebugInfo info = VecSimIndex_DebugInfo(index); + ASSERT_EQ(info.commonInfo.indexSize, expected_num_vectors); + ASSERT_EQ(info.commonInfo.indexLabelCount, expected_labels); + ASSERT_EQ(bf_index->idToLabelMapping.size(), expected_capacity); + ASSERT_EQ(bf_index->getStoredVectorsCount(), expected_num_vectors); + ASSERT_EQ(bf_index->indexCapacity(), expected_capacity); + ASSERT_EQ(bf_index->getVectorBlocks().size(), + expected_num_vectors / blockSize + (expected_num_vectors % blockSize != 0)) + << "expected_size: " << expected_num_vectors + << " expected_capacity: " << expected_capacity; + }; + // Empty index with initial capacity + verify_index_size(0, 0, expected_cap, "empty index"); for (size_t i = 0; i < n; i++) { - GenerateAndAddVector(index, dim, i % n_labels, i); + GenerateAndAddVector(index, dim, i % n_labels); } - - VecSimIndexDebugInfo info = VecSimIndex_DebugInfo(index); - ASSERT_EQ(info.commonInfo.indexSize, n); - ASSERT_EQ(info.commonInfo.indexLabelCount, n_labels); - ASSERT_EQ(bf_index->idToLabelMapping.size(), n - n % blockSize + blockSize); - ASSERT_EQ(bf_index->getVectorBlocks().size(), n / blockSize + 1); + verify_index_size(n, n_labels, expected_cap, "add" + std::to_string(n) + " vectors"); // remove invalid id ASSERT_EQ(bf_index->deleteVector(3459), 0); // This should do nothing - info = VecSimIndex_DebugInfo(index); - ASSERT_EQ(info.commonInfo.indexSize, n); - ASSERT_EQ(info.commonInfo.indexLabelCount, n_labels); - ASSERT_EQ(bf_index->idToLabelMapping.size(), n - n % blockSize + blockSize); - ASSERT_EQ(bf_index->getVectorBlocks().size(), n / blockSize + 1); + verify_index_size(n, n_labels, expected_cap, "remove invalid id"); // Add another vector (index capacity should remain the same). + // We add to an existing label - number of labels should not change. GenerateAndAddVector(index, dim, 0); - info = VecSimIndex_DebugInfo(index); - ASSERT_EQ(info.commonInfo.indexSize, n + 1); - // Label count doesn't increase because label 0 already exists - ASSERT_EQ(info.commonInfo.indexLabelCount, n_labels); - // Check new capacity size, should be blockSize * 2. - ASSERT_EQ(bf_index->idToLabelMapping.size(), 2 * blockSize); + verify_index_size(n + 1, n_labels, expected_cap, "add one more vector"); - // Now size = n + 1 = 16, capacity = 2 * bs = 20. Test capacity overflow again + // Now size = n + 1 = 13, capacity = 2 * bs = 20. Test capacity overflow again // to check that it stays aligned with blocksize. - - size_t add_vectors_count = 8; - for (size_t i = 0; i < add_vectors_count; i++) { + for (size_t i = 0; i < blockSize; i++) { GenerateAndAddVector(index, dim, i % n_labels, i); } - // Size should be n + 1 + 8 = 24. - size_t new_n = n + 1 + add_vectors_count; - info = VecSimIndex_DebugInfo(index); + // Size should be n + 1 + blockSize = 23. + // We add to existing labels only - number of labels doesn't change. + // The new capacity size should be increased by one block (blockSize * 3). + size_t new_n = n + 1 + blockSize; + ASSERT_EQ(expected_cap + blockSize, 3 * blockSize); + verify_index_size(new_n, n_labels, expected_cap + blockSize, + "add vectors to trigger another resize"); - ASSERT_EQ(info.commonInfo.indexSize, new_n); - // Label count doesn't increase because label 0 already exists - ASSERT_EQ(info.commonInfo.indexLabelCount, n_labels); size_t total_vectors = 0; for (auto label_ids : bf_index->labelToIdsLookup) { total_vectors += label_ids.second.size(); } ASSERT_EQ(total_vectors, new_n); - // Check new capacity size, should be blockSize * 3. - ASSERT_EQ(bf_index->idToLabelMapping.size(), 3 * blockSize); + // Delete vectors until one block plus some vectors are removed. + size_t current_vectors_block_count = (new_n + blockSize - 1) / blockSize; + ASSERT_EQ(current_vectors_block_count, 3); + size_t num_deleted = 0; + size_t deleted_labels = 0; + size_t label_to_delete = 0; + auto remove_to_one_below_blocksize = [&]() { + while (bf_index->getVectorBlocks().size() == current_vectors_block_count) { + int ret = VecSimIndex_DeleteVector(index, label_to_delete++); + ASSERT_GE(ret, 1) << "Failed to remove label: " << label_to_delete - 1; + num_deleted += ret; + deleted_labels++; + } + if (VecSimIndex_IndexSize(bf_index) % blockSize == 0) { + int ret = VecSimIndex_DeleteVector(index, label_to_delete++); + ASSERT_GE(ret, 1) << "Failed to remove label: " << label_to_delete - 1; + num_deleted += ret; + deleted_labels++; + } + }; + // First trigger of remove_to_one_below_blocksize will result in one free block. + // This should not trigger shrinking of metadata containers. + remove_to_one_below_blocksize(); // remove first block labels. + ASSERT_LT(new_n - num_deleted, 2 * blockSize); + ASSERT_GT(new_n - num_deleted, blockSize); + verify_index_size( + new_n - num_deleted, n_labels - deleted_labels, 3 * blockSize, + "delete vectors so that indexsize < 2 * blocksize, but there is only one free block"); + + // Second trigger of remove_to_one_below_blocksize will result in two free blocks. + // This should trigger shrinking of metadata containers by one block. + current_vectors_block_count--; + remove_to_one_below_blocksize(); + ASSERT_LT(new_n - num_deleted, blockSize); + verify_index_size( + new_n - num_deleted, n_labels - deleted_labels, 2 * blockSize, + "delete vectors so that indexsize % blocksize == 0 and there are two free blocks"); + + new_n -= num_deleted; + + // Delete all vectors. + while (VecSimIndex_IndexSize(index) > 0) { + int ret = VecSimIndex_DeleteVector(index, label_to_delete++); + ASSERT_GE(ret, 1) << "Failed to remove label: " << label_to_delete - 1; + num_deleted += ret; + } + ASSERT_EQ(num_deleted, total_vectors); + // We shrink capacity by one blocksize. + verify_index_size(0, 0, blockSize, "delete all vectors"); VecSimIndex_Free(index); } diff --git a/tests/unit/test_common.cpp b/tests/unit/test_common.cpp index bfae332ba..7beb800c2 100644 --- a/tests/unit/test_common.cpp +++ b/tests/unit/test_common.cpp @@ -539,16 +539,50 @@ TEST(CommonAPITest, testlogTieredIndex) { GenerateAndAddVector(tiered_index, 4, 1); mock_thread_pool.thread_iteration(); tiered_index->deleteVector(1); - ASSERT_EQ(log.logBuffer.size(), 4); - ASSERT_EQ(log.logBuffer[0], - "verbose: " + log.prefix + "Updating HNSW index capacity from 0 to 1024"); - ASSERT_EQ(log.logBuffer[1], + auto buffer_as_string = [&]() { + std::string buffer; + for (size_t i = 0; i < log.logBuffer.size(); i++) { + buffer += log.logBuffer[i] + "\n"; + } + return buffer; + }; + ASSERT_EQ(log.logBuffer.size(), 8) << buffer_as_string(); + size_t log_iter = 0; + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Resizing FLAT index from 0 to 1024") + << "failed at log index:" << log_iter - 1 << "." << std::endl + << "expected log: " << buffer_as_string(); + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Updating HNSW index capacity from 0 to 1024") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Resizing HNSW index from 0 to 1024") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Resizing FLAT index from 1024 to 0") + << "failed at log index:" << log_iter - 1 << "." << std::endl + << "expected log: " << buffer_as_string(); + + ASSERT_EQ(log.logBuffer[log_iter++], "verbose: " + log.prefix + - "Tiered HNSW index GC: there are 1 ready swap jobs. Start executing 1 swap jobs"); - ASSERT_EQ(log.logBuffer[2], - "verbose: " + log.prefix + "Updating HNSW index capacity from 1024 to 0"); - ASSERT_EQ(log.logBuffer[3], - "verbose: " + log.prefix + "Tiered HNSW index GC: done executing 1 swap jobs"); + "Tiered HNSW index GC: there are 1 ready swap jobs. Start executing 1 swap jobs") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Updating HNSW index capacity from 1024 to 0") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Resizing HNSW index from 1024 to 0") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); + + ASSERT_EQ(log.logBuffer[log_iter++], + "verbose: " + log.prefix + "Tiered HNSW index GC: done executing 1 swap jobs") + << "failed at log index:" << log_iter - 1 << std::endl + << "expected log: " << buffer_as_string(); } TEST(CommonAPITest, NormalizeBfloat16) { diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index 9b9dd5b46..111dc696e 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -3826,3 +3826,86 @@ TYPED_TEST(HNSWTieredIndexTestBasic, switchDeleteModes) { ASSERT_EQ(state.valid_state, 1); ASSERT_EQ(state.connections_to_repair, 0); } + +TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { + size_t dim = 4; + constexpr size_t blockSize = 10; + size_t resize_operations = 0; + HNSWParams params = {.type = TypeParam::get_index_type(), + .dim = dim, + .metric = VecSimMetric_L2, + .blockSize = blockSize}; + VecSimParams hnsw_params = CreateParams(params); + + auto mock_thread_pool = tieredIndexMock(); + TieredIndexParams tiered_params = {.jobQueue = &mock_thread_pool.jobQ, + .jobQueueCtx = mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .flatBufferLimit = SIZE_MAX, + .primaryIndexParams = &hnsw_params, + .specificParams = {TieredHNSWParams{.swapJobThreshold = 1}}}; + auto *tiered_index = reinterpret_cast *>( + TieredFactory::NewIndex(&tiered_params)); + mock_thread_pool.ctx->index_strong_ref.reset(tiered_index); + + auto hnsw_index = this->CastToHNSW(tiered_index); + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), 0); + + // add a vector + GenerateAndAddVector(tiered_index, dim, 0); + mock_thread_pool.thread_iteration(); + resize_operations++; + + // first vector should trigger resize + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), 1); + ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + + // add up to block size + for (size_t i = 1; i < blockSize; i++) { + GenerateAndAddVector(tiered_index, dim, i); + mock_thread_pool.thread_iteration(); + } + // should not trigger resize + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), blockSize); + ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + + // add one more vector to trigger another resize + GenerateAndAddVector(tiered_index, dim, blockSize); + mock_thread_pool.thread_iteration(); + resize_operations++; + + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), blockSize + 1); + ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); + + // delete a vector to shrink data blocks + ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 1) << "Failed to delete vector 0"; + + mock_thread_pool.init_threads(); + mock_thread_pool.thread_pool_join(); + tiered_index->executeReadySwapJobs(); + resize_operations++; + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), blockSize); + ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + + // add this vector again and verify lock was acquired to resize + GenerateAndAddVector(tiered_index, dim, 0); + mock_thread_pool.thread_iteration(); + resize_operations++; + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), blockSize + 1); + ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); + + // add up to block size (count = 2 blockSize), the lock shouldn't be acquired because no resize + // is required + for (size_t i = blockSize + 1; i < 2 * blockSize; i++) { + GenerateAndAddVector(tiered_index, dim, i); + mock_thread_pool.thread_iteration(); + } + ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); + ASSERT_EQ(hnsw_index->indexSize(), 2 * blockSize); + ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); +}