Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
meiravgri merged 1 commit into from
Sep 16, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 59 additions & 21 deletions src/VecSim/algorithms/brute_force/brute_force.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,15 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
*/
virtual void getDataByLabel(labelType label,
std::vector<std::vector<DataType>> &vectors_output) const = 0;

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:
Expand All @@ -92,27 +101,56 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
// Private internal function that implements generic single vector deletion.
virtual void removeVector(idType id);

inline void growByBlock() {
void resizeIndexCommon(size_t new_max_elements) {
assert(new_max_elements % this->blockSize == 0 &&
"new_max_elements must be a multiple of blockSize");
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing FLAT index from %zu to %zu",
idToLabelMapping.capacity(), new_max_elements);
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
idToLabelMapping.resize(new_max_elements);
idToLabelMapping.shrink_to_fit();
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
resizeLabelLookup(new_max_elements);
}

void growByBlock() {
assert(indexCapacity() == idToLabelMapping.capacity());
assert(indexCapacity() % this->blockSize == 0);
assert(indexCapacity() == indexSize());

assert(vectorBlocks.size() == 0 || vectorBlocks.back().getLength() == this->blockSize);
vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
this->alignment);
idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
idToLabelMapping.shrink_to_fit();
resizeLabelLookup(idToLabelMapping.size());
resizeIndexCommon(indexCapacity() + this->blockSize);
}

inline void shrinkByBlock() {
assert(indexCapacity() > 0); // should not be called when index is empty

void shrinkByBlock() {
assert(indexCapacity() >= this->blockSize);
assert(indexCapacity() % this->blockSize == 0);
// remove last block (should be empty)
assert(vectorBlocks.size() > 0 && vectorBlocks.back().getLength() == 0);
vectorBlocks.pop_back();

// remove a block size of labels.
assert(idToLabelMapping.size() >= this->blockSize);
idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
idToLabelMapping.shrink_to_fit();
resizeLabelLookup(idToLabelMapping.size());
assert(vectorBlocks.size() * this->blockSize == indexSize());

if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {
assert(indexCapacity() == idToLabelMapping.capacity());
assert(idToLabelMapping.size() == idToLabelMapping.capacity());
// There are at least two free blocks.
assert(vectorBlocks.size() * this->blockSize + 2 * this->blockSize <=
idToLabelMapping.capacity());
resizeIndexCommon(indexCapacity() - this->blockSize);
} else if (indexCapacity() == this->blockSize) {
// 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) {
Expand Down Expand Up @@ -162,19 +200,19 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(

template <typename DataType, typename DistType>
void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
// Give the vector new id and increase count.
idType id = this->count++;

// Resize the index if needed.
if (indexSize() > indexCapacity()) {
// Resize the index meta data structures if needed
if (indexSize() >= indexCapacity()) {
growByBlock();
} else if (id % this->blockSize == 0) {
// If we 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));
Expand Down
107 changes: 74 additions & 33 deletions src/VecSim/algorithms/hnsw/hnsw.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,11 @@ class HNSWIndex : public VecSimIndexAbstract<DistType>,
double getEpsilon() const;
size_t indexSize() const override;
size_t indexCapacity() const override;
/**
* Checks if the index capacity is full to hint the caller a resize is needed.
* @note Must be called with indexDataGuard locked.
*/
size_t isCapacityFull() const;
size_t getEfConstruction() const;
size_t getM() const;
size_t getMaxLevel() const;
Expand Down Expand Up @@ -313,6 +318,15 @@ class HNSWIndex : public VecSimIndexAbstract<DistType>,
*/
virtual void getDataByLabel(labelType label,
std::vector<std::vector<DataType>> &vectors_output) const = 0;

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:
Expand Down Expand Up @@ -358,6 +372,11 @@ size_t HNSWIndex<DataType, DistType>::indexCapacity() const {
return this->maxElements;
}

template <typename DataType, typename DistType>
size_t HNSWIndex<DataType, DistType>::isCapacityFull() const {
return indexSize() == this->maxElements;
}

template <typename DataType, typename DistType>
size_t HNSWIndex<DataType, DistType>::getEfConstruction() const {
return this->efConstruction;
Expand Down Expand Up @@ -1289,44 +1308,69 @@ template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::resizeIndexCommon(size_t new_max_elements) {
assert(new_max_elements % this->blockSize == 0 &&
"new_max_elements must be a multiple of blockSize");
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
"Updating HNSW index capacity from %zu to %zu", this->maxElements, new_max_elements);
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing HNSW index from %zu to %zu",
idToMetaData.capacity(), new_max_elements);
resizeLabelLookup(new_max_elements);
visitedNodesHandlerPool.resize(new_max_elements);
assert(idToMetaData.capacity() == idToMetaData.size());
idToMetaData.resize(new_max_elements);
idToMetaData.shrink_to_fit();

maxElements = new_max_elements;
assert(idToMetaData.capacity() == idToMetaData.size());
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::growByBlock() {
size_t new_max_elements = maxElements + this->blockSize;

// Validations
assert(vectorBlocks.size() == graphDataBlocks.size());
assert(vectorBlocks.empty() || vectorBlocks.back().getLength() == this->blockSize);
assert(this->maxElements % this->blockSize == 0);
assert(this->maxElements == indexSize());
assert(graphDataBlocks.size() == this->maxElements / this->blockSize);
assert(idToMetaData.capacity() == maxElements ||
idToMetaData.capacity() == maxElements + this->blockSize);

this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
"Updating HNSW index capacity from %zu to %zu", maxElements,
maxElements + this->blockSize);
maxElements += this->blockSize;
vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator, this->alignment);
graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize, this->allocator);

resizeIndexCommon(new_max_elements);
if (idToMetaData.capacity() == indexSize()) {
resizeIndexCommon(maxElements);
}
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::shrinkByBlock() {
assert(maxElements >= this->blockSize);
size_t new_max_elements = maxElements - this->blockSize;

// Validations
assert(vectorBlocks.size() == graphDataBlocks.size());
assert(!vectorBlocks.empty());
assert(vectorBlocks.back().getLength() == 0);

vectorBlocks.pop_back();
graphDataBlocks.pop_back();

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 <typename DataType, typename DistType>
Expand Down Expand Up @@ -1686,9 +1730,7 @@ void HNSWIndex<DataType, DistType>::removeAndSwap(idType internalId) {

// If we need to free a complete block and there is at least one block between the
// capacity and the size.
if (curElementCount % this->blockSize == 0) {
shrinkByBlock();
}
shrinkByBlock();
}

template <typename DataType, typename DistType>
Expand Down Expand Up @@ -1764,6 +1806,16 @@ void HNSWIndex<DataType, DistType>::removeVectorInPlace(const idType element_int
template <typename DataType, typename DistType>
AddVectorCtx HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
const void *vector_data) {
if (isCapacityFull()) {
growByBlock();
} else if (curElementCount % this->blockSize == 0) {
// If we had an initial capacity, we might have to initialize new blocks for the data and
// meta-data.
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
this->alignment);
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
this->allocator);
}
AddVectorCtx state{};

// Choose randomly the maximum level in which the new element will be in the index.
Expand Down Expand Up @@ -1791,17 +1843,6 @@ AddVectorCtx HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
throw e;
}

if (indexSize() > indexCapacity()) {
growByBlock();
} else if (state.newElementId % this->blockSize == 0) {
// If we had an initial capacity, we might have to allocate new blocks for the data and
// meta-data.
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
this->alignment);
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
this->allocator);
}

// Insert the new element to the data block
this->vectorBlocks.back().addElement(vector_data);
this->graphDataBlocks.back().addElement(cur_egd);
Expand Down
18 changes: 9 additions & 9 deletions src/VecSim/algorithms/hnsw/hnsw_tiered.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,7 +301,7 @@ template <typename DataType, typename DistType>
void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToRun) {

// Execute swap jobs - acquire hnsw write lock.
this->mainIndexGuard.lock();
this->lockMainIndexGuard();
TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
"Tiered HNSW index GC: there are %zu ready swap jobs. Start executing %zu swap jobs",
readySwapJobs, std::min(readySwapJobs, maxJobsToRun));
Expand All @@ -326,7 +326,7 @@ void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToR
readySwapJobs -= idsToRemove.size();
TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
"Tiered HNSW index GC: done executing %zu swap jobs", idsToRemove.size());
this->mainIndexGuard.unlock();
this->unlockMainIndexGuard();
}

template <typename DataType, typename DistType>
Expand Down Expand Up @@ -419,11 +419,11 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
this->mainIndexGuard.lock_shared();
hnsw_index->lockIndexDataGuard();
// Check if resizing is needed for HNSW index (requires write lock).
if (hnsw_index->indexCapacity() == hnsw_index->indexSize()) {
if (hnsw_index->isCapacityFull()) {
// Release the inner HNSW data lock before we re-acquire the global HNSW lock.
this->mainIndexGuard.unlock_shared();
hnsw_index->unlockIndexDataGuard();
this->mainIndexGuard.lock();
this->lockMainIndexGuard();
hnsw_index->lockIndexDataGuard();

// Hold the index data lock while we store the new element. If the new node's max level is
Expand All @@ -448,7 +448,7 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
if (state.elementMaxLevel > state.currMaxLevel) {
hnsw_index->unlockIndexDataGuard();
}
this->mainIndexGuard.unlock();
this->unlockMainIndexGuard();
} else {
// Do the same as above except for changing the capacity, but with *shared* lock held:
// Hold the index data lock while we store the new element. If the new node's max level is
Expand Down Expand Up @@ -701,9 +701,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
}
// Insert the vector to the HNSW index. Internally, we will never have to overwrite the
// label since we already checked it outside.
this->mainIndexGuard.lock();
this->lockMainIndexGuard();
hnsw_index->addVector(blob, label);
this->mainIndexGuard.unlock();
this->unlockMainIndexGuard();
return ret;
}
if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
Expand Down Expand Up @@ -826,9 +826,9 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
}
} else {
// delete in place.
this->mainIndexGuard.lock();
this->lockMainIndexGuard();
num_deleted_vectors += this->deleteLabelFromHNSWInplace(label);
this->mainIndexGuard.unlock();
this->unlockMainIndexGuard();
}

return num_deleted_vectors;
Expand Down
1 change: 1 addition & 0 deletions src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

INDEX_TEST_FRIEND_CLASS(BM_VecSimBasics)
INDEX_TEST_FRIEND_CLASS(BM_VecSimCommon)
13 changes: 13 additions & 0 deletions src/VecSim/vec_sim_tiered_index.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -58,6 +68,9 @@ class VecSimTieredIndex : public VecSimIndexInterface {
}

public:
#ifdef BUILD_TESTS
int getMainIndexGuardWriteLockCount() const { return mainIndexGuard_write_lock_count; }
#endif
VecSimTieredIndex(VecSimIndexAbstract<DistType> *backendIndex_,
BruteForceIndex<DataType, DistType> *frontendIndex_,
TieredIndexParams tieredParams, std::shared_ptr<VecSimAllocator> allocator)
Expand Down
Loading
Loading