/* * Copyright (c) 2006-Present, Redis Ltd. * All rights reserved. * * Licensed under your choice of the Redis Source Available License 2.0 * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the * GNU Affero General Public License v3 (AGPLv3). */ #include "gtest/gtest.h" #include "VecSim/vec_sim.h" #include "VecSim/memory/vecsim_malloc.h" #include "VecSim/memory/vecsim_base.h" #include "VecSim/algorithms/brute_force/brute_force_single.h" #include "VecSim/algorithms/hnsw/hnsw_single.h" #include "unit_test_utils.h" #include "VecSim/utils/serializer.h" #include "VecSim/index_factories/hnsw_factory.h" const size_t vecsimAllocationOverhead = VecSimAllocator::getAllocationOverheadSize(); const size_t hashTableNodeSize = getLabelsLookupNodeSize(); class AllocatorTest : public ::testing::Test {}; struct SimpleObject : public VecsimBaseObject { public: SimpleObject(std::shared_ptr allocator) : VecsimBaseObject(allocator) {} int x; }; struct ObjectWithSTL : public VecsimBaseObject { std::vector> test_vec; public: ObjectWithSTL(std::shared_ptr allocator) : VecsimBaseObject(allocator), test_vec(allocator) {}; }; struct NestedObject : public VecsimBaseObject { ObjectWithSTL stl_object; SimpleObject simpleObject; public: NestedObject(std::shared_ptr allocator) : VecsimBaseObject(allocator), stl_object(allocator), simpleObject(allocator) {}; }; TEST_F(AllocatorTest, test_simple_object) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); size_t expectedAllocationSize = sizeof(VecSimAllocator); ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); SimpleObject *obj = new (allocator) SimpleObject(allocator); expectedAllocationSize += sizeof(SimpleObject) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); delete obj; expectedAllocationSize -= sizeof(SimpleObject) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), sizeof(VecSimAllocator)); } TEST_F(AllocatorTest, test_object_with_stl) { std::shared_ptr allocator(VecSimAllocator::newVecsimAllocator()); size_t expectedAllocationSize = sizeof(VecSimAllocator); ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); ObjectWithSTL *obj = new (allocator) ObjectWithSTL(allocator); expectedAllocationSize += sizeof(ObjectWithSTL) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); obj->test_vec.push_back(1); expectedAllocationSize += sizeof(int) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); delete obj; } TEST_F(AllocatorTest, test_nested_object) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); size_t expectedAllocationSize = sizeof(VecSimAllocator); ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); NestedObject *obj = new (allocator) NestedObject(allocator); expectedAllocationSize += sizeof(NestedObject) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); obj->stl_object.test_vec.push_back(1); expectedAllocationSize += sizeof(int) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); delete obj; } template class IndexAllocatorTest : public ::testing::Test {}; // DataTypeSet, TEST_DATA_T and TEST_DIST_T are defined in unit_test_utils.h 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, .blockSize = blockSize}; auto *bfIndex = dynamic_cast *>( BruteForceFactory::NewIndex(¶ms)); auto allocator = bfIndex->getAllocator(); TEST_DATA_T vec[128] = {}; uint64_t expectedAllocationSize = sizeof(VecSimAllocator); expectedAllocationSize += sizeof(BruteForceIndex_Single) + vecsimAllocationOverhead; expectedAllocationSize += sizeof(DataBlocksContainer) + vecsimAllocationOverhead; expectedAllocationSize += sizeof(DistanceCalculatorCommon) + vecsimAllocationOverhead; expectedAllocationSize += sizeof(PreprocessorsContainerAbstract) + vecsimAllocationOverhead; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize); 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(dynamic_cast(bfIndex->vectors)->numBlocks(), expected_data_container_blocks); ASSERT_EQ(bfIndex->vectors->size(), expected_size); 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); }; // =========== Add label 1 =========== int before = allocator->getAllocationSize(); size_t buckets_num_before = bfIndex->labelToIdLookup.bucket_count(); auto vectors_blocks = dynamic_cast(bfIndex->vectors); size_t vectors_blocks_capacity = vectors_blocks->capacity(); VecSimIndex_AddVector(bfIndex, vec, 1); int addCommandAllocationDelta = allocator->getAllocationSize() - before; int64_t expectedAllocationDelta = sizeof(labelType) + vecsimAllocationOverhead; // resize idToLabelMapping expectedAllocationDelta += (vectors_blocks->capacity() - vectors_blocks_capacity) * sizeof(DataBlock) + vecsimAllocationOverhead; // New vectors blocks expectedAllocationDelta += blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + bfIndex->getStorageAlignment(); // 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 += (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. { 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 += (vectors_blocks->capacity() - vectors_blocks_capacity) * sizeof(DataBlock); // New vector blocks expectedAllocationDelta += 2 * sizeof(labelType); // resize idToLabelMapping expectedAllocationDelta += 2 * (blockSize * sizeof(TEST_DATA_T) * dim + vecsimAllocationOverhead + bfIndex->getStorageAlignment()); // Two block vectors buffer expectedAllocationDelta += 2 * hashTableNodeSize; // New nodes in the label lookup expectedAllocationDelta += (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(); 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->getStorageAlignment(); // 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); } // =========== labels = [2, 3], vector blocks = 2, maps capacity = 3. Add label 4 =========== // 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->getStorageAlignment(); // 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(); 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->getStorageAlignment()); // 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; verify_containers_size(0, 0, 0); // 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->getStorageAlignment()); // 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 -= 2 * sizeof(labelType) + vecsimAllocationOverhead; // remove two idToLabelMapping and free the container // resizing labelToIdLookup to 0 size_t buckets_after = bfIndex->labelToIdLookup.bucket_count(); ASSERT_EQ(bfIndex->labelToIdLookup.size(), 0); ASSERT_LE(buckets_after, buckets_num_before); 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); } VecSimIndex_Free(bfIndex); } TYPED_TEST(IndexAllocatorTest, test_hnsw) { size_t d = 128; // Build with default args HNSWParams params = {.type = TypeParam::get_index_type(), .dim = d, .metric = VecSimMetric_L2}; TEST_DATA_T vec[128] = {}; auto *hnswIndex = dynamic_cast *>(HNSWFactory::NewIndex(¶ms)); auto allocator = hnswIndex->getAllocator(); uint64_t expectedAllocationSize = sizeof(VecSimAllocator); expectedAllocationSize += sizeof(HNSWIndex_Single) + vecsimAllocationOverhead; ASSERT_GE(allocator->getAllocationSize(), expectedAllocationSize); size_t memory = VecSimIndex_StatsInfo(hnswIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); expectedAllocationSize = memory; int before = allocator->getAllocationSize(); VecSimIndex_AddVector(hnswIndex, vec, 1); int addCommandAllocationDelta = allocator->getAllocationSize() - before; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize + addCommandAllocationDelta); memory = VecSimIndex_StatsInfo(hnswIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); expectedAllocationSize = memory; before = allocator->getAllocationSize(); VecSimIndex_AddVector(hnswIndex, vec, 2); addCommandAllocationDelta = allocator->getAllocationSize() - before; ASSERT_EQ(allocator->getAllocationSize(), expectedAllocationSize + addCommandAllocationDelta); memory = VecSimIndex_StatsInfo(hnswIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); expectedAllocationSize = memory; before = allocator->getAllocationSize(); VecSimIndex_DeleteVector(hnswIndex, 2); int deleteCommandAllocationDelta = allocator->getAllocationSize() - before; ASSERT_EQ(expectedAllocationSize + deleteCommandAllocationDelta, allocator->getAllocationSize()); memory = VecSimIndex_StatsInfo(hnswIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); expectedAllocationSize = memory; before = allocator->getAllocationSize(); VecSimIndex_DeleteVector(hnswIndex, 1); deleteCommandAllocationDelta = allocator->getAllocationSize() - before; ASSERT_EQ(expectedAllocationSize + deleteCommandAllocationDelta, allocator->getAllocationSize()); memory = VecSimIndex_StatsInfo(hnswIndex).memory; ASSERT_EQ(allocator->getAllocationSize(), memory); VecSimIndex_Free(hnswIndex); } TYPED_TEST(IndexAllocatorTest, testIncomingEdgesSet) { size_t d = 2; // Build index, use small M to simplify the scenario. HNSWParams params = { .type = TypeParam::get_index_type(), .dim = d, .metric = VecSimMetric_L2, .M = 2}; auto *hnswIndex = dynamic_cast *>(HNSWFactory::NewIndex(¶ms)); auto allocator = hnswIndex->getAllocator(); // Add a "dummy" vector - labels_lookup hash table will allocate initial size of buckets here. GenerateAndAddVector(hnswIndex, d, 0, 0.0); // Add another vector and validate it's exact memory allocation delta. TEST_DATA_T vec1[] = {1.0, 0.0}; int before = allocator->getAllocationSize(); VecSimIndex_AddVector(hnswIndex, vec1, 1); // Since the memory before did not account for the visited nodes handler pool (it was created // lazily), we need to clear it before calculating the delta. hnswIndex->visitedNodesHandlerPool.clearPool(); int allocation_delta = allocator->getAllocationSize() - before; size_t vec_max_level = hnswIndex->getGraphDataByInternalId(1)->toplevel; // Expect the creation of an empty incoming edges set in every level (+ the allocator header // overhead), and a single node in the labels' lookup hash table. size_t expected_allocation_delta = (vec_max_level + 1) * (sizeof(vecsim_stl::vector) + vecsimAllocationOverhead); expected_allocation_delta += hashTableNodeSize; // Account for allocating link lists for levels higher than 0, if exists. if (vec_max_level > 0) { expected_allocation_delta += hnswIndex->levelDataSize * vec_max_level + vecsimAllocationOverhead; } ASSERT_EQ(allocation_delta, expected_allocation_delta); // Add three more vectors, all should have a connections to vec1. TEST_DATA_T vec2[] = {2.0f, 0.0f}; VecSimIndex_AddVector(hnswIndex, vec2, 2); TEST_DATA_T vec3[] = {1.0f, 1.0f}; VecSimIndex_AddVector(hnswIndex, vec3, 3); TEST_DATA_T vec4[] = {1.0f, -1.0f}; VecSimIndex_AddVector(hnswIndex, vec4, 4); // Layer 0 should look like this (all edges bidirectional): // 3 3 // | | // 0--1--2 => 0--5--1--2 // | |----^| // 4 4 // Next, insertion of vec5 should make 0->1 unidirectional, thus adding 0 to 1's incoming edges // set. TEST_DATA_T vec5[] = {0.5f, 0.0f}; size_t buckets_num_before = hnswIndex->labelLookup.bucket_count(); before = allocator->getAllocationSize(); VecSimIndex_AddVector(hnswIndex, vec5, 5); allocation_delta = allocator->getAllocationSize() - before; vec_max_level = hnswIndex->getGraphDataByInternalId(5)->toplevel; /* Compute the expected allocation delta: * 1. empty incoming edges set in every level (+ allocator's header). * 2. A node in the labels_lookup has table (+ allocator's header). If rehashing occurred, we * account also for the diff in the buckets size (each bucket has sizeof(size_t) overhead). * 3. Account for allocating link lists for levels higher than 0, if exists. * 4. Finally, expect an allocation of the data buffer in the incoming edges vector of vec1 due * to the insertion, and the fact that vec1 will re-select its neighbours. */ expected_allocation_delta = (vec_max_level + 1) * (sizeof(vecsim_stl::vector) + vecsimAllocationOverhead) + hashTableNodeSize; size_t buckets_diff = hnswIndex->labelLookup.bucket_count() - buckets_num_before; expected_allocation_delta += buckets_diff * sizeof(size_t); if (vec_max_level > 0) { expected_allocation_delta += hnswIndex->levelDataSize * vec_max_level + vecsimAllocationOverhead; } // Expect that the first element is pushed to the incoming edges vector of element 1 in level 0. // Then, we account for the capacity of the buffer that is allocated for the vector data. expected_allocation_delta += hnswIndex->getElementLevelData(1, 0).incomingUnidirectionalEdges->capacity() * sizeof(idType) + vecsimAllocationOverhead; ASSERT_EQ(allocation_delta, expected_allocation_delta); VecSimIndex_Free(hnswIndex); } TYPED_TEST(IndexAllocatorTest, test_hnsw_reclaim_memory) { size_t d = 128; VecSimType type = TypeParam::get_index_type(); // Build HNSW index with default args and initial capacity of zero. HNSWParams params = {.type = type, .dim = d, .metric = VecSimMetric_L2}; auto *hnswIndex = dynamic_cast *>(HNSWFactory::NewIndex(¶ms)); auto allocator = hnswIndex->getAllocator(); ASSERT_EQ(hnswIndex->indexCapacity(), 0); size_t initial_memory_size = allocator->getAllocationSize(); // labels_lookup and element_levels containers are not allocated at all in some platforms, // when initial capacity is zero, while in other platforms labels_lookup is allocated with a // single bucket. This, we get the following range in which we expect the initial memory to be // in. ASSERT_GE(initial_memory_size, HNSWFactory::EstimateInitialSize(¶ms)); ASSERT_LE(initial_memory_size, HNSWFactory::EstimateInitialSize(¶ms) + sizeof(size_t) + 2 * vecsimAllocationOverhead); // 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 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. 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(dynamic_cast(hnswIndex->vectors)->numBlocks(), expected_data_container_blocks); 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 // incoming edges sets). ASSERT_EQ(hnswIndex->checkIntegrity().unidirectional_connections, 0); }; // Validate that a single block exists. 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. GenerateAndAddVector(hnswIndex, d, block_size, block_size); 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 last_vec_graph_data_mem = (sizeof(vecsim_stl::vector) + vecsimAllocationOverhead) + hashTableNodeSize; if (vec_max_level > 0) { 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->getStorageAlignment(); size_t size_total_data_per_element = hnswIndex->elementGraphDataSize + hnswIndex->getStoredDataSize(); data_containers_block_mem += size_total_data_per_element * block_size; // account for idToMetaData and visitedNodesHandlerPool entries. expected_mem_delta += (sizeof(tag_t) + sizeof(ElementMetaData) + sizeof(vecsim_stl::one_byte_mutex)) * 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() - one_block_buckets) * sizeof(size_t); // New blocks allocated - 1 aligned block for vectors and 1 unaligned block for graph data. ASSERT_EQ(expected_mem_delta, mem_delta); // 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 = dynamic_cast(hnswIndex->vectors); size_t vectors_blocks_capacity = vectors_blocks->capacity(); VecSimIndex_DeleteVector(hnswIndex, block_size); verify_containers_size(block_size, 1, 2 * block_size); size_t expected_allocation_size = before_delete_mem - last_vec_graph_data_mem - hnswIndex->getStorageAlignment(); // 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. for (size_t i = 0; i < block_size; i++) { VecSimIndex_DeleteVector(hnswIndex, i); } ASSERT_EQ(hnswIndex->indexSize(), 0); 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. 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() + dynamic_cast(hnswIndex->vectors)->capacity()) + 2 * vecsimAllocationOverhead; // Current memory should be back as it was initially. The label_lookup hash table is an // exception, since in some platforms, empty buckets remain even when the capacity is set to // zero, while in others the entire capacity reduced to zero (including the header). // Also, visitedNodesHandlerPool that was created lazily is not freed, but it should not be // accounted when comparing to the initial memory size estimation of the index. hnswIndex->visitedNodesHandlerPool.clearPool(); ASSERT_LE(allocator->getAllocationSize(), HNSWFactory::EstimateInitialSize(¶ms) + block_vectors_memory + hash_table_memory + 2 * vecsimAllocationOverhead); ASSERT_GE(allocator->getAllocationSize(), HNSWFactory::EstimateInitialSize(¶ms) + block_vectors_memory + hash_table_memory); VecSimIndex_Free(hnswIndex); }