diff --git a/src/VecSim/algorithms/brute_force/brute_force.h b/src/VecSim/algorithms/brute_force/brute_force.h index 23bd142de..3b989ceeb 100644 --- a/src/VecSim/algorithms/brute_force/brute_force.h +++ b/src/VecSim/algorithms/brute_force/brute_force.h @@ -80,6 +80,8 @@ class BruteForceIndex : public VecSimIndexAbstract { idToLabelMapping.shrink_to_fit(); resizeLabelLookup(idToLabelMapping.size()); } + + size_t indexMetaDataCapacity() const override { return idToLabelMapping.capacity(); } #endif protected: diff --git a/src/VecSim/algorithms/hnsw/hnsw.h b/src/VecSim/algorithms/hnsw/hnsw.h index 2c629afa4..c3c71c22b 100644 --- a/src/VecSim/algorithms/hnsw/hnsw.h +++ b/src/VecSim/algorithms/hnsw/hnsw.h @@ -310,6 +310,8 @@ class HNSWIndex : public VecSimIndexAbstract, resizeLabelLookup(idToMetaData.size()); } } + + size_t indexMetaDataCapacity() const override { return idToMetaData.capacity(); } #endif protected: diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered.h b/src/VecSim/algorithms/hnsw/hnsw_tiered.h index 764b9ad35..96b8314e4 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered.h @@ -248,6 +248,10 @@ class TieredHNSWIndex : public VecSimTieredIndex { #ifdef BUILD_TESTS void getDataByLabel(labelType label, std::vector> &vectors_output) const; + size_t indexMetaDataCapacity() const override { + return this->backendIndex->indexMetaDataCapacity() + + this->frontendIndex->indexMetaDataCapacity(); + } #endif }; diff --git a/src/VecSim/algorithms/svs/svs.h b/src/VecSim/algorithms/svs/svs.h index e407642f6..dcec0f184 100644 --- a/src/VecSim/algorithms/svs/svs.h +++ b/src/VecSim/algorithms/svs/svs.h @@ -705,6 +705,7 @@ class SVSIndex : public VecSimIndexAbstract, fl public: void fitMemory() override {} + size_t indexMetaDataCapacity() const override { return this->indexCapacity(); } std::vector> getStoredVectorDataByLabel(labelType label) const override { // For compressed/quantized indices, this function is not meaningful diff --git a/src/VecSim/algorithms/svs/svs_tiered.h b/src/VecSim/algorithms/svs/svs_tiered.h index d4d0796bc..bdaea67d2 100644 --- a/src/VecSim/algorithms/svs/svs_tiered.h +++ b/src/VecSim/algorithms/svs/svs_tiered.h @@ -457,6 +457,12 @@ class TieredSVSIndex : public VecSimTieredIndex { backend_index_t *GetBackendIndex() { return this->backendIndex; } void submitSingleJob(AsyncJob *job) { Base::submitSingleJob(job); } void submitJobs(vecsim_stl::vector &jobs) { Base::submitJobs(jobs); } + 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(); + } #endif private: diff --git a/src/VecSim/vec_sim_interface.h b/src/VecSim/vec_sim_interface.h index 4066e3e6c..0b627bc57 100644 --- a/src/VecSim/vec_sim_interface.h +++ b/src/VecSim/vec_sim_interface.h @@ -217,5 +217,15 @@ struct VecSimIndexInterface : public VecsimBaseObject { } #ifdef BUILD_TESTS virtual void fitMemory() = 0; + /** + * @brief get the capacity of the meta data containers. + * + * @return The capacity of the meta data containers in number of elements. + * The value returned from this function may differ from the indexCapacity() function. For + * example, in HNSW, the capacity of the meta data containers is the capacity of the labels + * lookup table, while the capacity of the data containers is the capacity of the vectors + * container. + */ + virtual size_t indexMetaDataCapacity() const = 0; #endif }; diff --git a/tests/benchmark/bm_vecsim_basics.h b/tests/benchmark/bm_vecsim_basics.h index b7edd3904..b81a6b0d5 100644 --- a/tests/benchmark/bm_vecsim_basics.h +++ b/tests/benchmark/bm_vecsim_basics.h @@ -318,6 +318,9 @@ void BM_VecSimBasics::UpdateAtBlockSize(benchmark::State &st) { // Calculate vectors needed to reach next block boundary size_t vecs_to_blocksize = BM_VecSimGeneral::block_size - (initial_index_size % BM_VecSimGeneral::block_size); + size_t initial_index_cap = index->indexMetaDataCapacity(); + assert(initial_index_cap == N_VECTORS + vecs_to_blocksize); + assert(vecs_to_blocksize < BM_VecSimGeneral::block_size); labelType initial_label_count = index->indexLabelCount(); labelType curr_label = initial_label_count; @@ -342,15 +345,20 @@ void BM_VecSimBasics::UpdateAtBlockSize(benchmark::State &st) { // Benchmark loop: repeatedly delete/add same vector to trigger grow-shrink cycles labelType label_to_update = curr_label - 1; - size_t index_cap = index->indexCapacity(); + size_t index_cap = index->indexMetaDataCapacity(); + std::cout << "index_cap after adding vectors " << index_cap << std::endl; + assert(index_cap == initial_index_cap + BM_VecSimGeneral::block_size); + for (auto _ : st) { // Remove the vector directly from hnsw size_t ret = VecSimIndex_DeleteVector( GET_INDEX(st.range(0) == INDEX_TIERED_HNSW ? INDEX_HNSW : st.range(0)), label_to_update); assert(ret == 1); - assert(index->indexCapacity() == index_cap - BM_VecSimGeneral::block_size); - // Capacity should shrink by one block after deletion + + // Capacity should not change + size_t curr_cap = index->indexMetaDataCapacity(); + assert(curr_cap == index_cap); ret = VecSimIndex_AddVector(index, QUERIES[(added_vec_count - 1) % N_QUERIES].data(), label_to_update); assert(ret == 1); @@ -358,8 +366,8 @@ void BM_VecSimBasics::UpdateAtBlockSize(benchmark::State &st) { assert(VecSimIndex_IndexSize( GET_INDEX(st.range(0) == INDEX_TIERED_HNSW ? INDEX_HNSW : st.range(0))) == N_VECTORS + added_vec_count); - // Capacity should grow back to original size after addition - assert(index->indexCapacity() == index_cap); + // Capacity should not change + assert(index->indexMetaDataCapacity() == index_cap); } assert(VecSimIndex_IndexSize(index) == N_VECTORS + added_vec_count); diff --git a/tests/unit/test_allocator.cpp b/tests/unit/test_allocator.cpp index 0b897d3b9..525f8f4b7 100644 --- a/tests/unit/test_allocator.cpp +++ b/tests/unit/test_allocator.cpp @@ -126,6 +126,7 @@ TYPED_TEST(IndexAllocatorTest, test_bf_index_block_size_1) { ASSERT_EQ(bfIndex->indexCapacity(), expected_map_containers_size); ASSERT_EQ(bfIndex->idToLabelMapping.capacity(), expected_map_containers_size); + ASSERT_EQ(bfIndex->indexMetaDataCapacity(), expected_map_containers_size); ASSERT_EQ(bfIndex->idToLabelMapping.size(), expected_map_containers_size); ASSERT_GE(bfIndex->labelToIdLookup.bucket_count(), expected_map_containers_size); }; @@ -536,6 +537,7 @@ TYPED_TEST(IndexAllocatorTest, test_hnsw_reclaim_memory) { ASSERT_EQ(hnswIndex->vectors->size(), expected_size); ASSERT_EQ(hnswIndex->idToMetaData.capacity(), expected_map_containers_size); + ASSERT_EQ(hnswIndex->indexMetaDataCapacity(), 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 diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index c4c790726..77d16dd50 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -4364,6 +4364,11 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), 1); ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), blockSize); + ASSERT_EQ(tiered_index->frontendIndex->indexMetaDataCapacity(), 0); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), + hnsw_index->indexMetaDataCapacity() + + tiered_index->frontendIndex->indexMetaDataCapacity()); // add up to block size for (size_t i = 1; i < blockSize; i++) { @@ -4374,6 +4379,11 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), blockSize); ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), blockSize); + ASSERT_EQ(tiered_index->frontendIndex->indexMetaDataCapacity(), 0); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), + hnsw_index->indexMetaDataCapacity() + + tiered_index->frontendIndex->indexMetaDataCapacity()); // add one more vector to trigger another resize GenerateAndAddVector(tiered_index, dim, blockSize); @@ -4383,6 +4393,11 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), blockSize + 1); ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), 2 * blockSize); + ASSERT_EQ(tiered_index->frontendIndex->indexMetaDataCapacity(), 0); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), + hnsw_index->indexMetaDataCapacity() + + tiered_index->frontendIndex->indexMetaDataCapacity()); // delete a vector to shrink data blocks ASSERT_EQ(VecSimIndex_DeleteVector(tiered_index, 0), 1) << "Failed to delete vector 0"; @@ -4394,6 +4409,8 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), blockSize); ASSERT_EQ(hnsw_index->indexCapacity(), blockSize); + // meta data capacity should not shrink + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), 2 * blockSize); // add this vector again and verify lock was acquired to resize GenerateAndAddVector(tiered_index, dim, 0); @@ -4402,6 +4419,11 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), blockSize + 1); ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), 2 * blockSize); + ASSERT_EQ(tiered_index->frontendIndex->indexMetaDataCapacity(), 0); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), + hnsw_index->indexMetaDataCapacity() + + tiered_index->frontendIndex->indexMetaDataCapacity()); // add up to block size (count = 2 blockSize), the lock shouldn't be acquired because no resize // is required @@ -4412,4 +4434,9 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { ASSERT_EQ(tiered_index->getMainIndexGuardWriteLockCount(), resize_operations); ASSERT_EQ(hnsw_index->indexSize(), 2 * blockSize); ASSERT_EQ(hnsw_index->indexCapacity(), 2 * blockSize); + ASSERT_EQ(hnsw_index->indexMetaDataCapacity(), 2 * blockSize); + ASSERT_EQ(tiered_index->frontendIndex->indexMetaDataCapacity(), 0); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), + hnsw_index->indexMetaDataCapacity() + + tiered_index->frontendIndex->indexMetaDataCapacity()); } diff --git a/tests/unit/test_svs.cpp b/tests/unit/test_svs.cpp index 95338dc83..3c29dcbcb 100644 --- a/tests/unit/test_svs.cpp +++ b/tests/unit/test_svs.cpp @@ -833,6 +833,7 @@ TYPED_TEST(SVSTest, resizeIndex) { } // The size (+extra) and the capacity should be equal. ASSERT_EQ(index->indexCapacity(), VecSimIndex_IndexSize(index) + extra_cap); + ASSERT_EQ(index->indexMetaDataCapacity(), index->indexCapacity()); // The capacity shouldn't be changed. ASSERT_EQ(index->indexCapacity(), n + extra_cap); @@ -878,6 +879,7 @@ TYPED_TEST(SVSTest, svs_empty_index) { // The expected capacity should be 0 for empty index. ASSERT_EQ(index->indexCapacity(), 0); + ASSERT_EQ(index->indexMetaDataCapacity(), index->indexCapacity()); // Try to remove it again. VecSimIndex_DeleteVector(index, 1); diff --git a/tests/unit/test_svs_tiered.cpp b/tests/unit/test_svs_tiered.cpp index 8a1d6c2a3..c154292d7 100644 --- a/tests/unit/test_svs_tiered.cpp +++ b/tests/unit/test_svs_tiered.cpp @@ -523,6 +523,7 @@ TYPED_TEST(SVSTieredIndexTest, addVector) { ASSERT_EQ(tiered_index->GetBackendIndex()->indexSize(), 0); ASSERT_EQ(tiered_index->GetFlatIndex()->indexCapacity(), DEFAULT_BLOCK_SIZE); ASSERT_EQ(tiered_index->indexCapacity(), DEFAULT_BLOCK_SIZE); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), tiered_index->indexCapacity()); ASSERT_EQ(tiered_index->GetFlatIndex()->getDistanceFrom_Unsafe(vec_label, vector), 0); ASSERT_EQ(mock_thread_pool.jobQ.size(), mock_thread_pool.thread_pool_size); @@ -624,6 +625,7 @@ TYPED_TEST(SVSTieredIndexTest, insertJob) { ? DEFAULT_BLOCK_SIZE : tiered_index->GetBackendIndex()->indexCapacity(); ASSERT_EQ(tiered_index->indexCapacity(), expected_capacity); + ASSERT_EQ(tiered_index->indexMetaDataCapacity(), tiered_index->indexCapacity()); ASSERT_EQ(tiered_index->GetFlatIndex()->indexCapacity(), 0); ASSERT_EQ(tiered_index->getDistanceFrom_Unsafe(vec_label, vector), 0); }