Skip to content

Commit 942e8f2

Browse files
meiravgrimeiravgri
authored
[0.7] [MOD-10559] Decouple the shrinking and growing logic of large containers in Flat and HNSW (#780)
[0.8] [MOD-10559] Decouple the shrinking and growing logic of large containers in Flat and HNSW (#777) * backport #753 shrink by blocksize shrink to zero only if capcity is 1 blocksize. * move public outside * revert size (cherry picked from commit bcc4d67)
1 parent 40aa744 commit 942e8f2

File tree

10 files changed

+779
-234
lines changed

10 files changed

+779
-234
lines changed

‎src/VecSim/algorithms/brute_force/brute_force.h‎

Lines changed: 59 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,15 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
8383
*/
8484
virtual void getDataByLabel(labelType label,
8585
std::vector<std::vector<DataType>> &vectors_output) const = 0;
86+
87+
size_t getStoredVectorsCount() const {
88+
size_t actual_stored_vec = 0;
89+
for (auto &block : vectorBlocks) {
90+
actual_stored_vec += block.getLength();
91+
}
92+
93+
return actual_stored_vec;
94+
}
8695
#endif
8796

8897
protected:
@@ -92,27 +101,56 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
92101
// Private internal function that implements generic single vector deletion.
93102
virtual void removeVector(idType id);
94103

95-
inline void growByBlock() {
104+
void resizeIndexCommon(size_t new_max_elements) {
105+
assert(new_max_elements % this->blockSize == 0 &&
106+
"new_max_elements must be a multiple of blockSize");
107+
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing FLAT index from %zu to %zu",
108+
idToLabelMapping.capacity(), new_max_elements);
109+
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
110+
idToLabelMapping.resize(new_max_elements);
111+
idToLabelMapping.shrink_to_fit();
112+
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
113+
resizeLabelLookup(new_max_elements);
114+
}
115+
116+
void growByBlock() {
117+
assert(indexCapacity() == idToLabelMapping.capacity());
118+
assert(indexCapacity() % this->blockSize == 0);
119+
assert(indexCapacity() == indexSize());
120+
96121
assert(vectorBlocks.size() == 0 || vectorBlocks.back().getLength() == this->blockSize);
97122
vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
98123
this->alignment);
99-
idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
100-
idToLabelMapping.shrink_to_fit();
101-
resizeLabelLookup(idToLabelMapping.size());
124+
resizeIndexCommon(indexCapacity() + this->blockSize);
102125
}
103126

104-
inline void shrinkByBlock() {
105-
assert(indexCapacity() > 0); // should not be called when index is empty
106-
127+
void shrinkByBlock() {
128+
assert(indexCapacity() >= this->blockSize);
129+
assert(indexCapacity() % this->blockSize == 0);
107130
// remove last block (should be empty)
108131
assert(vectorBlocks.size() > 0 && vectorBlocks.back().getLength() == 0);
109132
vectorBlocks.pop_back();
110-
111-
// remove a block size of labels.
112-
assert(idToLabelMapping.size() >= this->blockSize);
113-
idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
114-
idToLabelMapping.shrink_to_fit();
115-
resizeLabelLookup(idToLabelMapping.size());
133+
assert(vectorBlocks.size() * this->blockSize == indexSize());
134+
135+
if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {
136+
assert(indexCapacity() == idToLabelMapping.capacity());
137+
assert(idToLabelMapping.size() == idToLabelMapping.capacity());
138+
// There are at least two free blocks.
139+
assert(vectorBlocks.size() * this->blockSize + 2 * this->blockSize <=
140+
idToLabelMapping.capacity());
141+
resizeIndexCommon(indexCapacity() - this->blockSize);
142+
} else if (indexCapacity() == this->blockSize) {
143+
// Special case to handle last block.
144+
// This special condition resolves the ambiguity: when capacity==blockSize, we can't
145+
// tell if this block came from growth (should shrink to 0) or initial capacity (should
146+
// keep it). We choose to always shrink to 0 to maintain the one-block removal
147+
// guarantee. In contrast, newer branches without initial capacity support use simpler
148+
// logic: immediately shrink to 0 whenever index size becomes 0.
149+
assert(vectorBlocks.empty());
150+
assert(indexSize() == 0);
151+
resizeIndexCommon(0);
152+
return;
153+
}
116154
}
117155

118156
inline DataBlock &getVectorVectorBlock(idType id) {
@@ -162,19 +200,19 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(
162200

163201
template <typename DataType, typename DistType>
164202
void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
165-
// Give the vector new id and increase count.
166-
idType id = this->count++;
167-
168-
// Resize the index if needed.
169-
if (indexSize() > indexCapacity()) {
203+
// Resize the index meta data structures if needed
204+
if (indexSize() >= indexCapacity()) {
170205
growByBlock();
171-
} else if (id % this->blockSize == 0) {
172-
// If we we didn't reach the initial capacity but the last block is full, add a new block
173-
// only.
206+
} else if (this->count % this->blockSize == 0) {
207+
// If we didn't reach the initial capacity but the last block is full, initialize a new
208+
// block only.
174209
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
175210
this->alignment);
176211
}
177212

213+
// Give the vector new id and increase count.
214+
idType id = this->count++;
215+
178216
// Get the last vectors block to store the vector in.
179217
DataBlock &vectorBlock = this->vectorBlocks.back();
180218
assert(&vectorBlock == &getVectorVectorBlock(id));

‎src/VecSim/algorithms/hnsw/hnsw.h‎

Lines changed: 74 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,11 @@ class HNSWIndex : public VecSimIndexAbstract<DistType>,
234234
double getEpsilon() const;
235235
size_t indexSize() const override;
236236
size_t indexCapacity() const override;
237+
/**
238+
* Checks if the index capacity is full to hint the caller a resize is needed.
239+
* @note Must be called with indexDataGuard locked.
240+
*/
241+
size_t isCapacityFull() const;
237242
size_t getEfConstruction() const;
238243
size_t getM() const;
239244
size_t getMaxLevel() const;
@@ -313,6 +318,15 @@ class HNSWIndex : public VecSimIndexAbstract<DistType>,
313318
*/
314319
virtual void getDataByLabel(labelType label,
315320
std::vector<std::vector<DataType>> &vectors_output) const = 0;
321+
322+
size_t getStoredVectorsCount() const {
323+
size_t actual_stored_vec = 0;
324+
for (auto &block : vectorBlocks) {
325+
actual_stored_vec += block.getLength();
326+
}
327+
328+
return actual_stored_vec;
329+
}
316330
#endif
317331

318332
protected:
@@ -358,6 +372,11 @@ size_t HNSWIndex<DataType, DistType>::indexCapacity() const {
358372
return this->maxElements;
359373
}
360374

375+
template <typename DataType, typename DistType>
376+
size_t HNSWIndex<DataType, DistType>::isCapacityFull() const {
377+
return indexSize() == this->maxElements;
378+
}
379+
361380
template <typename DataType, typename DistType>
362381
size_t HNSWIndex<DataType, DistType>::getEfConstruction() const {
363382
return this->efConstruction;
@@ -1289,44 +1308,69 @@ template <typename DataType, typename DistType>
12891308
void HNSWIndex<DataType, DistType>::resizeIndexCommon(size_t new_max_elements) {
12901309
assert(new_max_elements % this->blockSize == 0 &&
12911310
"new_max_elements must be a multiple of blockSize");
1292-
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
1293-
"Updating HNSW index capacity from %zu to %zu", this->maxElements, new_max_elements);
1311+
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing HNSW index from %zu to %zu",
1312+
idToMetaData.capacity(), new_max_elements);
12941313
resizeLabelLookup(new_max_elements);
12951314
visitedNodesHandlerPool.resize(new_max_elements);
1315+
assert(idToMetaData.capacity() == idToMetaData.size());
12961316
idToMetaData.resize(new_max_elements);
12971317
idToMetaData.shrink_to_fit();
1298-
1299-
maxElements = new_max_elements;
1318+
assert(idToMetaData.capacity() == idToMetaData.size());
13001319
}
13011320

13021321
template <typename DataType, typename DistType>
13031322
void HNSWIndex<DataType, DistType>::growByBlock() {
1304-
size_t new_max_elements = maxElements + this->blockSize;
1305-
13061323
// Validations
13071324
assert(vectorBlocks.size() == graphDataBlocks.size());
13081325
assert(vectorBlocks.empty() || vectorBlocks.back().getLength() == this->blockSize);
1326+
assert(this->maxElements % this->blockSize == 0);
1327+
assert(this->maxElements == indexSize());
1328+
assert(graphDataBlocks.size() == this->maxElements / this->blockSize);
1329+
assert(idToMetaData.capacity() == maxElements ||
1330+
idToMetaData.capacity() == maxElements + this->blockSize);
13091331

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

1313-
resizeIndexCommon(new_max_elements);
1339+
if (idToMetaData.capacity() == indexSize()) {
1340+
resizeIndexCommon(maxElements);
1341+
}
13141342
}
13151343

13161344
template <typename DataType, typename DistType>
13171345
void HNSWIndex<DataType, DistType>::shrinkByBlock() {
1318-
assert(maxElements >= this->blockSize);
1319-
size_t new_max_elements = maxElements - this->blockSize;
1320-
1321-
// Validations
1322-
assert(vectorBlocks.size() == graphDataBlocks.size());
1323-
assert(!vectorBlocks.empty());
1324-
assert(vectorBlocks.back().getLength() == 0);
1325-
1326-
vectorBlocks.pop_back();
1327-
graphDataBlocks.pop_back();
1328-
1329-
resizeIndexCommon(new_max_elements);
1346+
assert(this->maxElements >= this->blockSize);
1347+
assert(this->maxElements % this->blockSize == 0);
1348+
if (indexSize() % this->blockSize == 0) {
1349+
assert(vectorBlocks.back().getLength() == 0);
1350+
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
1351+
"Updating HNSW index capacity from %zu to %zu", maxElements,
1352+
maxElements - this->blockSize);
1353+
vectorBlocks.pop_back();
1354+
graphDataBlocks.pop_back();
1355+
assert(graphDataBlocks.size() * this->blockSize == indexSize());
1356+
1357+
if (idToMetaData.capacity() >= (indexSize() + 2 * this->blockSize)) {
1358+
resizeIndexCommon(idToMetaData.capacity() - this->blockSize);
1359+
} else if (idToMetaData.capacity() == this->blockSize) {
1360+
// Special case to handle last block.
1361+
// This special condition resolves the ambiguity: when capacity==blockSize, we can't
1362+
// tell if this block came from growth (should shrink to 0) or initial capacity (should
1363+
// keep it). We choose to always shrink to 0 to maintain the one-block removal
1364+
// guarantee. In contrast, newer branches without initial capacity support use simpler
1365+
// logic: immediately shrink to 0 whenever index size becomes 0.
1366+
assert(vectorBlocks.empty());
1367+
assert(indexSize() == 0);
1368+
assert(maxElements == this->blockSize);
1369+
resizeIndexCommon(0);
1370+
}
1371+
// Take the lower bound into account.
1372+
maxElements -= this->blockSize;
1373+
}
13301374
}
13311375

13321376
template <typename DataType, typename DistType>
@@ -1686,9 +1730,7 @@ void HNSWIndex<DataType, DistType>::removeAndSwap(idType internalId) {
16861730

16871731
// If we need to free a complete block and there is at least one block between the
16881732
// capacity and the size.
1689-
if (curElementCount % this->blockSize == 0) {
1690-
shrinkByBlock();
1691-
}
1733+
shrinkByBlock();
16921734
}
16931735

16941736
template <typename DataType, typename DistType>
@@ -1764,6 +1806,16 @@ void HNSWIndex<DataType, DistType>::removeVectorInPlace(const idType element_int
17641806
template <typename DataType, typename DistType>
17651807
AddVectorCtx HNSWIndex<DataType, DistType>::storeNewElement(labelType label,
17661808
const void *vector_data) {
1809+
if (isCapacityFull()) {
1810+
growByBlock();
1811+
} else if (curElementCount % this->blockSize == 0) {
1812+
// If we had an initial capacity, we might have to initialize new blocks for the data and
1813+
// meta-data.
1814+
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
1815+
this->alignment);
1816+
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
1817+
this->allocator);
1818+
}
17671819
AddVectorCtx state{};
17681820

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

1794-
if (indexSize() > indexCapacity()) {
1795-
growByBlock();
1796-
} else if (state.newElementId % this->blockSize == 0) {
1797-
// If we had an initial capacity, we might have to allocate new blocks for the data and
1798-
// meta-data.
1799-
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
1800-
this->alignment);
1801-
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
1802-
this->allocator);
1803-
}
1804-
18051846
// Insert the new element to the data block
18061847
this->vectorBlocks.back().addElement(vector_data);
18071848
this->graphDataBlocks.back().addElement(cur_egd);

‎src/VecSim/algorithms/hnsw/hnsw_tiered.h‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ template <typename DataType, typename DistType>
301301
void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToRun) {
302302

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

332332
template <typename DataType, typename DistType>
@@ -419,11 +419,11 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
419419
this->mainIndexGuard.lock_shared();
420420
hnsw_index->lockIndexDataGuard();
421421
// Check if resizing is needed for HNSW index (requires write lock).
422-
if (hnsw_index->indexCapacity() == hnsw_index->indexSize()) {
422+
if (hnsw_index->isCapacityFull()) {
423423
// Release the inner HNSW data lock before we re-acquire the global HNSW lock.
424424
this->mainIndexGuard.unlock_shared();
425425
hnsw_index->unlockIndexDataGuard();
426-
this->mainIndexGuard.lock();
426+
this->lockMainIndexGuard();
427427
hnsw_index->lockIndexDataGuard();
428428

429429
// Hold the index data lock while we store the new element. If the new node's max level is
@@ -448,7 +448,7 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
448448
if (state.elementMaxLevel > state.currMaxLevel) {
449449
hnsw_index->unlockIndexDataGuard();
450450
}
451-
this->mainIndexGuard.unlock();
451+
this->unlockMainIndexGuard();
452452
} else {
453453
// Do the same as above except for changing the capacity, but with *shared* lock held:
454454
// Hold the index data lock while we store the new element. If the new node's max level is
@@ -701,9 +701,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
701701
}
702702
// Insert the vector to the HNSW index. Internally, we will never have to overwrite the
703703
// label since we already checked it outside.
704-
this->mainIndexGuard.lock();
704+
this->lockMainIndexGuard();
705705
hnsw_index->addVector(blob, label);
706-
this->mainIndexGuard.unlock();
706+
this->unlockMainIndexGuard();
707707
return ret;
708708
}
709709
if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
@@ -826,9 +826,9 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
826826
}
827827
} else {
828828
// delete in place.
829-
this->mainIndexGuard.lock();
829+
this->lockMainIndexGuard();
830830
num_deleted_vectors += this->deleteLabelFromHNSWInplace(label);
831-
this->mainIndexGuard.unlock();
831+
this->unlockMainIndexGuard();
832832
}
833833

834834
return num_deleted_vectors;

‎src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteBothAsyncAndInplaceMulti_
5454
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceMultiSwapId_Test)
5555
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarkedDeleted_Test)
5656
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test)
57+
INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test)
5758

5859
INDEX_TEST_FRIEND_CLASS(BM_VecSimBasics)
5960
INDEX_TEST_FRIEND_CLASS(BM_VecSimCommon)

‎src/VecSim/vec_sim_tiered_index.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,17 @@ class VecSimTieredIndex : public VecSimIndexInterface {
4141

4242
mutable std::shared_mutex flatIndexGuard;
4343
mutable std::shared_mutex mainIndexGuard;
44+
void lockMainIndexGuard() const {
45+
mainIndexGuard.lock();
46+
#ifdef BUILD_TESTS
47+
mainIndexGuard_write_lock_count++;
48+
#endif
49+
}
4450

51+
void unlockMainIndexGuard() const { mainIndexGuard.unlock(); }
52+
#ifdef BUILD_TESTS
53+
mutable std::atomic_int mainIndexGuard_write_lock_count = 0;
54+
#endif
4555
size_t flatBufferLimit;
4656

4757
void submitSingleJob(AsyncJob *job) {
@@ -58,6 +68,9 @@ class VecSimTieredIndex : public VecSimIndexInterface {
5868
}
5969

6070
public:
71+
#ifdef BUILD_TESTS
72+
int getMainIndexGuardWriteLockCount() const { return mainIndexGuard_write_lock_count; }
73+
#endif
6174
VecSimTieredIndex(VecSimIndexAbstract<DistType> *backendIndex_,
6275
BruteForceIndex<DataType, DistType> *frontendIndex_,
6376
TieredIndexParams tieredParams, std::shared_ptr<VecSimAllocator> allocator)

0 commit comments

Comments
 (0)