Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
56 changes: 41 additions & 15 deletions src/VecSim/algorithms/brute_force/brute_force.h
Original file line number Diff line number Diff line change
Expand Up @@ -89,20 +89,45 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
// Private internal function that implements generic single vector deletion.
virtual void removeVector(idType id);

void growByBlock() {
idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
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();
resizeLabelLookup(idToLabelMapping.size());
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
resizeLabelLookup(new_max_elements);
}

void shrinkByBlock() {
assert(indexCapacity() > 0); // should not be called when index is empty
void growByBlock() {
assert(indexCapacity() == idToLabelMapping.capacity());
assert(indexCapacity() % this->blockSize == 0);
assert(indexCapacity() == indexSize());
assert((dynamic_cast<DataBlocksContainer *>(this->vectors)->numBlocks() ==
(indexSize()) / this->blockSize));

// remove a block size of labels.
assert(idToLabelMapping.size() >= this->blockSize);
idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
idToLabelMapping.shrink_to_fit();
resizeLabelLookup(idToLabelMapping.size());
resizeIndexCommon(indexCapacity() + this->blockSize);
}

void shrinkByBlock() {
assert(indexCapacity() >= this->blockSize);
assert(indexCapacity() % this->blockSize == 0);
assert(dynamic_cast<DataBlocksContainer *>(this->vectors)->numBlocks() ==
indexSize() / this->blockSize);

if (indexSize() == 0) {
resizeIndexCommon(0);
} else if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {

assert(indexCapacity() == idToLabelMapping.capacity());
assert(idToLabelMapping.size() == idToLabelMapping.capacity());
assert(dynamic_cast<DataBlocksContainer *>(this->vectors)->size() +
2 * this->blockSize ==
idToLabelMapping.capacity());
resizeIndexCommon(indexCapacity() - this->blockSize);
}
}

void setVectorLabel(idType id, labelType new_label) { idToLabelMapping.at(id) = new_label; }
Expand Down Expand Up @@ -143,14 +168,15 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(

template <typename DataType, typename DistType>
void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
// Resize the index meta data structures if needed
if (indexSize() >= indexCapacity()) {
growByBlock();
}

auto processed_blob = this->preprocessForStorage(vector_data);
// Give the vector new id and increase count.
idType id = this->count++;

// Resize the index meta data structures if needed
if (indexSize() > indexCapacity()) {
growByBlock();
}
// add vector data to vector raw data container
this->vectors->addElement(processed_blob.get(), id);

Expand Down Expand Up @@ -199,7 +225,7 @@ size_t BruteForceIndex<DataType, DistType>::indexSize() const {

template <typename DataType, typename DistType>
size_t BruteForceIndex<DataType, DistType>::indexCapacity() const {
return this->idToLabelMapping.size();
return this->idToLabelMapping.capacity();
}

template <typename DataType, typename DistType>
Expand Down
73 changes: 52 additions & 21 deletions src/VecSim/algorithms/hnsw/hnsw.h
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,11 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
double getEpsilon() const;
size_t indexSize() const override;
size_t indexCapacity() const override;
/**
* Checks if the index capacity is full to hint the caller a resize is needed.
* @note Must be called with indexDataGuard locked.
*/
size_t isCapacityFull() const;
size_t getEfConstruction() const;
size_t getM() const;
size_t getMaxLevel() const;
Expand Down Expand Up @@ -349,6 +354,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 @@ -1281,31 +1291,59 @@ 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;
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;

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;
graphDataBlocks.pop_back();
assert(this->maxElements >= this->blockSize);
assert(this->maxElements % this->blockSize == 0);

if (indexSize() % this->blockSize == 0) {
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
"Updating HNSW index capacity from %zu to %zu", maxElements,
maxElements - this->blockSize);
graphDataBlocks.pop_back();
assert(graphDataBlocks.size() == indexSize() / this->blockSize);

// assuming idToMetaData reflects the capacity of the heavy reallocation containers.
if (indexSize() == 0) {
resizeIndexCommon(0);
} else if (idToMetaData.capacity() >= (indexSize() + 2 * this->blockSize)) {
assert(this->maxElements + this->blockSize == idToMetaData.capacity());
resizeIndexCommon(idToMetaData.capacity() - this->blockSize);
}

resizeIndexCommon(new_max_elements);
// Take the lower bound into account.
maxElements -= this->blockSize;
}
}

template <typename DataType, typename DistType>
Expand Down Expand Up @@ -1660,9 +1698,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.
this->vectors->removeElement(curElementCount);
if (curElementCount % this->blockSize == 0) {
shrinkByBlock();
}
shrinkByBlock();
}

template <typename DataType, typename DistType>
Expand Down Expand Up @@ -1738,6 +1774,9 @@ void HNSWIndex<DataType, DistType>::removeVectorInPlace(const idType element_int
template <typename DataType, typename DistType>
HNSWAddVectorState HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
const void *vector_data) {
if (isCapacityFull()) {
growByBlock();
}
HNSWAddVectorState state{};

// Choose randomly the maximum level in which the new element will be in the index.
Expand Down Expand Up @@ -1765,14 +1804,6 @@ HNSWAddVectorState HNSWIndex<DataType, DistType>::storeNewElement(labelType labe
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 graph data.
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
this->allocator);
}

// Insert the new element to the data block
this->vectors->addElement(vector_data, state.newElementId);
this->graphDataBlocks.back().addElement(cur_egd);
Expand Down
18 changes: 9 additions & 9 deletions src/VecSim/algorithms/hnsw/hnsw_tiered.h
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,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 @@ -339,7 +339,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 @@ -437,11 +437,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 @@ -466,7 +466,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 @@ -713,9 +713,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
auto storage_blob = this->frontendIndex->preprocessForStorage(blob);
// 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(storage_blob.get(), label);
this->mainIndexGuard.unlock();
this->unlockMainIndexGuard();
return ret;
}
if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
Expand Down Expand Up @@ -841,9 +841,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
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,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 CommonAPITest_SearchDifferentScores_Test;
friend class BF16TieredTest;
Expand Down
4 changes: 2 additions & 2 deletions src/VecSim/containers/data_blocks_container.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,8 @@ std::unique_ptr<RawDataContainer::Iterator> DataBlocksContainer::getIterator() c
return std::make_unique<DataBlocksContainer::Iterator>(*this);
}

size_t DataBlocksContainer::numBlocks() const { return this->blocks.size(); }

#ifdef BUILD_TESTS
void DataBlocksContainer::saveVectorsData(std::ostream &output) const {
// Save data blocks
Expand Down Expand Up @@ -114,8 +116,6 @@ void DataBlocksContainer::restoreBlocks(std::istream &input, size_t num_vectors,

void DataBlocksContainer::shrinkToFit() { this->blocks.shrink_to_fit(); }

size_t DataBlocksContainer::numBlocks() const { return this->blocks.size(); }

#endif
/********************************** Iterator API ************************************************/

Expand Down
4 changes: 3 additions & 1 deletion src/VecSim/containers/data_blocks_container.h
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,10 @@ class DataBlocksContainer : public VecsimBaseObject, public RawDataContainer {
std::shared_ptr<VecSimAllocator> allocator, unsigned char alignment = 0);
~DataBlocksContainer();

// Number of elements in the container.
size_t size() const override;

// Number of blocks allocated.
size_t capacity() const;

size_t blockSize() const;
Expand All @@ -46,13 +48,13 @@ class DataBlocksContainer : public VecsimBaseObject, public RawDataContainer {

std::unique_ptr<RawDataContainer::Iterator> getIterator() const override;

size_t numBlocks() const;
#ifdef BUILD_TESTS
void saveVectorsData(std::ostream &output) const override;
// Use that in deserialization when file was created with old version (v3) that serialized
// the blocks themselves and not just thw raw vector data.
void restoreBlocks(std::istream &input, size_t num_vectors, Serializer::EncodingVersion);
void shrinkToFit();
size_t numBlocks() const;
#endif

class Iterator : public RawDataContainer::Iterator {
Expand Down
11 changes: 11 additions & 0 deletions src/VecSim/vec_sim_tiered_index.h
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,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 Down Expand Up @@ -89,6 +99,7 @@ class VecSimTieredIndex : public VecSimIndexInterface {

#ifdef BUILD_TESTS
public:
int getMainIndexGuardWriteLockCount() const { return mainIndexGuard_write_lock_count; }
#endif
// For both topK and range, Use withSet=false if you can guarantee that shared ids between the
// two lists will also have identical scores. In this case, any duplicates will naturally align
Expand Down
Loading
Loading