Skip to content

Commit bcc4d67

Browse files
meiravgrimeiravgri
authored
[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
1 parent 3a7ec14 commit bcc4d67

File tree

10 files changed

+780
-237
lines changed

10 files changed

+780
-237
lines changed

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

Lines changed: 60 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,15 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
8787
idToLabelMapping.shrink_to_fit();
8888
resizeLabelLookup(idToLabelMapping.size());
8989
}
90+
91+
size_t getStoredVectorsCount() const {
92+
size_t actual_stored_vec = 0;
93+
for (auto &block : vectorBlocks) {
94+
actual_stored_vec += block.getLength();
95+
}
96+
97+
return actual_stored_vec;
98+
}
9099
#endif
91100

92101
protected:
@@ -96,36 +105,63 @@ class BruteForceIndex : public VecSimIndexAbstract<DataType, DistType> {
96105
// Private internal function that implements generic single vector deletion.
97106
virtual void removeVector(idType id);
98107

99-
inline void growByBlock() {
108+
void resizeIndexCommon(size_t new_max_elements) {
109+
assert(new_max_elements % this->blockSize == 0 &&
110+
"new_max_elements must be a multiple of blockSize");
111+
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing FLAT index from %zu to %zu",
112+
idToLabelMapping.capacity(), new_max_elements);
113+
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
114+
idToLabelMapping.resize(new_max_elements);
115+
idToLabelMapping.shrink_to_fit();
116+
assert(idToLabelMapping.capacity() == idToLabelMapping.size());
117+
resizeLabelLookup(new_max_elements);
118+
}
119+
120+
void growByBlock() {
121+
assert(indexCapacity() == idToLabelMapping.capacity());
122+
assert(indexCapacity() % this->blockSize == 0);
123+
assert(indexCapacity() == indexSize());
124+
100125
assert(vectorBlocks.size() == 0 || vectorBlocks.back().getLength() == this->blockSize);
101126
vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
102127
this->alignment);
103-
idToLabelMapping.resize(idToLabelMapping.size() + this->blockSize);
104-
idToLabelMapping.shrink_to_fit();
105-
resizeLabelLookup(idToLabelMapping.size());
128+
resizeIndexCommon(indexCapacity() + this->blockSize);
106129
}
107130

108-
inline void shrinkByBlock() {
109-
assert(indexCapacity() > 0); // should not be called when index is empty
110-
131+
void shrinkByBlock() {
132+
assert(indexCapacity() >= this->blockSize);
133+
assert(indexCapacity() % this->blockSize == 0);
111134
// remove last block (should be empty)
112135
assert(vectorBlocks.size() > 0 && vectorBlocks.back().getLength() == 0);
113136
vectorBlocks.pop_back();
114-
115-
// remove a block size of labels.
116-
assert(idToLabelMapping.size() >= this->blockSize);
117-
idToLabelMapping.resize(idToLabelMapping.size() - this->blockSize);
118-
idToLabelMapping.shrink_to_fit();
119-
resizeLabelLookup(idToLabelMapping.size());
137+
assert(vectorBlocks.size() * this->blockSize == indexSize());
138+
139+
if (indexCapacity() >= (indexSize() + 2 * this->blockSize)) {
140+
assert(indexCapacity() == idToLabelMapping.capacity());
141+
assert(idToLabelMapping.size() == idToLabelMapping.capacity());
142+
// There are at least two free blocks.
143+
assert(vectorBlocks.size() * this->blockSize + 2 * this->blockSize <=
144+
idToLabelMapping.capacity());
145+
resizeIndexCommon(indexCapacity() - this->blockSize);
146+
} else if (indexCapacity() == this->blockSize) {
147+
// Special case to handle last block.
148+
// This special condition resolves the ambiguity: when capacity==blockSize, we can't
149+
// tell if this block came from growth (should shrink to 0) or initial capacity (should
150+
// keep it). We choose to always shrink to 0 to maintain the one-block removal
151+
// guarantee. In contrast, newer branches without initial capacity support use simpler
152+
// logic: immediately shrink to 0 whenever index size becomes 0.
153+
assert(vectorBlocks.empty());
154+
assert(indexSize() == 0);
155+
resizeIndexCommon(0);
156+
return;
157+
}
120158
}
121159

122160
inline DataBlock &getVectorVectorBlock(idType id) {
123161
return vectorBlocks.at(id / this->blockSize);
124162
}
125163
inline size_t getVectorRelativeIndex(idType id) const { return id % this->blockSize; }
126-
inline void setVectorLabel(idType id, labelType new_label) {
127-
idToLabelMapping.at(id) = new_label;
128-
}
164+
void setVectorLabel(idType id, labelType new_label) { idToLabelMapping.at(id) = new_label; }
129165
// inline priority queue getter that need to be implemented by derived class
130166
virtual inline vecsim_stl::abstract_priority_queue<DistType, labelType> *
131167
getNewMaxPriorityQueue() const = 0;
@@ -166,19 +202,19 @@ BruteForceIndex<DataType, DistType>::BruteForceIndex(
166202

167203
template <typename DataType, typename DistType>
168204
void BruteForceIndex<DataType, DistType>::appendVector(const void *vector_data, labelType label) {
169-
// Give the vector new id and increase count.
170-
idType id = this->count++;
171-
172-
// Resize the index if needed.
173-
if (indexSize() > indexCapacity()) {
205+
// Resize the index meta data structures if needed
206+
if (indexSize() >= indexCapacity()) {
174207
growByBlock();
175-
} else if (id % this->blockSize == 0) {
176-
// If we didn't reach the initial capacity but the last block is full, add a new block
177-
// only.
208+
} else if (this->count % this->blockSize == 0) {
209+
// If we didn't reach the initial capacity but the last block is full, initialize a new
210+
// block only.
178211
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
179212
this->alignment);
180213
}
181214

215+
// Give the vector new id and increase count.
216+
idType id = this->count++;
217+
182218
// Get the last vectors block to store the vector in.
183219
DataBlock &vectorBlock = this->vectorBlocks.back();
184220
assert(&vectorBlock == &getVectorVectorBlock(id));

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

Lines changed: 74 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,11 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
233233
double getEpsilon() const;
234234
size_t indexSize() const override;
235235
size_t indexCapacity() const override;
236+
/**
237+
* Checks if the index capacity is full to hint the caller a resize is needed.
238+
* @note Must be called with indexDataGuard locked.
239+
*/
240+
size_t isCapacityFull() const;
236241
size_t getEfConstruction() const;
237242
size_t getM() const;
238243
size_t getMaxLevel() const;
@@ -312,6 +317,15 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
312317
idToMetaData.shrink_to_fit();
313318
resizeLabelLookup(idToMetaData.size());
314319
}
320+
321+
size_t getStoredVectorsCount() const {
322+
size_t actual_stored_vec = 0;
323+
for (auto &block : vectorBlocks) {
324+
actual_stored_vec += block.getLength();
325+
}
326+
327+
return actual_stored_vec;
328+
}
315329
#endif
316330

317331
protected:
@@ -357,6 +371,11 @@ size_t HNSWIndex<DataType, DistType>::indexCapacity() const {
357371
return this->maxElements;
358372
}
359373

374+
template <typename DataType, typename DistType>
375+
size_t HNSWIndex<DataType, DistType>::isCapacityFull() const {
376+
return indexSize() == this->maxElements;
377+
}
378+
360379
template <typename DataType, typename DistType>
361380
size_t HNSWIndex<DataType, DistType>::getEfConstruction() const {
362381
return this->efConstruction;
@@ -1288,44 +1307,69 @@ template <typename DataType, typename DistType>
12881307
void HNSWIndex<DataType, DistType>::resizeIndexCommon(size_t new_max_elements) {
12891308
assert(new_max_elements % this->blockSize == 0 &&
12901309
"new_max_elements must be a multiple of blockSize");
1291-
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING,
1292-
"Updating HNSW index capacity from %zu to %zu", this->maxElements, new_max_elements);
1310+
this->log(VecSimCommonStrings::LOG_VERBOSE_STRING, "Resizing HNSW index from %zu to %zu",
1311+
idToMetaData.capacity(), new_max_elements);
12931312
resizeLabelLookup(new_max_elements);
12941313
visitedNodesHandlerPool.resize(new_max_elements);
1314+
assert(idToMetaData.capacity() == idToMetaData.size());
12951315
idToMetaData.resize(new_max_elements);
12961316
idToMetaData.shrink_to_fit();
1297-
1298-
maxElements = new_max_elements;
1317+
assert(idToMetaData.capacity() == idToMetaData.size());
12991318
}
13001319

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

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

1312-
resizeIndexCommon(new_max_elements);
1338+
if (idToMetaData.capacity() == indexSize()) {
1339+
resizeIndexCommon(maxElements);
1340+
}
13131341
}
13141342

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

13311375
template <typename DataType, typename DistType>
@@ -1685,9 +1729,7 @@ void HNSWIndex<DataType, DistType>::removeAndSwap(idType internalId) {
16851729

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

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

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

1793-
if (indexSize() > indexCapacity()) {
1794-
growByBlock();
1795-
} else if (state.newElementId % this->blockSize == 0) {
1796-
// If we had an initial capacity, we might have to allocate new blocks for the data and
1797-
// meta-data.
1798-
this->vectorBlocks.emplace_back(this->blockSize, this->dataSize, this->allocator,
1799-
this->alignment);
1800-
this->graphDataBlocks.emplace_back(this->blockSize, this->elementGraphDataSize,
1801-
this->allocator);
1802-
}
1803-
18041845
// Insert the new element to the data block
18051846
this->vectorBlocks.back().addElement(vector_data);
18061847
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
@@ -308,7 +308,7 @@ template <typename DataType, typename DistType>
308308
void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToRun) {
309309

310310
// Execute swap jobs - acquire hnsw write lock.
311-
this->mainIndexGuard.lock();
311+
this->lockMainIndexGuard();
312312
TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
313313
"Tiered HNSW index GC: there are %zu ready swap jobs. Start executing %zu swap jobs",
314314
readySwapJobs, std::min(readySwapJobs, maxJobsToRun));
@@ -333,7 +333,7 @@ void TieredHNSWIndex<DataType, DistType>::executeReadySwapJobs(size_t maxJobsToR
333333
readySwapJobs -= idsToRemove.size();
334334
TIERED_LOG(VecSimCommonStrings::LOG_VERBOSE_STRING,
335335
"Tiered HNSW index GC: done executing %zu swap jobs", idsToRemove.size());
336-
this->mainIndexGuard.unlock();
336+
this->unlockMainIndexGuard();
337337
}
338338

339339
template <typename DataType, typename DistType>
@@ -426,11 +426,11 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
426426
this->mainIndexGuard.lock_shared();
427427
hnsw_index->lockIndexDataGuard();
428428
// Check if resizing is needed for HNSW index (requires write lock).
429-
if (hnsw_index->indexCapacity() == hnsw_index->indexSize()) {
429+
if (hnsw_index->isCapacityFull()) {
430430
// Release the inner HNSW data lock before we re-acquire the global HNSW lock.
431431
this->mainIndexGuard.unlock_shared();
432432
hnsw_index->unlockIndexDataGuard();
433-
this->mainIndexGuard.lock();
433+
this->lockMainIndexGuard();
434434
hnsw_index->lockIndexDataGuard();
435435

436436
// Hold the index data lock while we store the new element. If the new node's max level is
@@ -455,7 +455,7 @@ void TieredHNSWIndex<DataType, DistType>::insertVectorToHNSW(
455455
if (state.elementMaxLevel > state.currMaxLevel) {
456456
hnsw_index->unlockIndexDataGuard();
457457
}
458-
this->mainIndexGuard.unlock();
458+
this->unlockMainIndexGuard();
459459
} else {
460460
// Do the same as above except for changing the capacity, but with *shared* lock held:
461461
// Hold the index data lock while we store the new element. If the new node's max level is
@@ -709,9 +709,9 @@ int TieredHNSWIndex<DataType, DistType>::addVector(const void *blob, labelType l
709709
}
710710
// Insert the vector to the HNSW index. Internally, we will never have to overwrite the
711711
// label since we already checked it outside.
712-
this->mainIndexGuard.lock();
712+
this->lockMainIndexGuard();
713713
hnsw_index->addVector(blob, label);
714-
this->mainIndexGuard.unlock();
714+
this->unlockMainIndexGuard();
715715
return ret;
716716
}
717717
if (this->frontendIndex->indexSize() >= this->flatBufferLimit) {
@@ -834,9 +834,9 @@ int TieredHNSWIndex<DataType, DistType>::deleteVector(labelType label) {
834834
}
835835
} else {
836836
// delete in place.
837-
this->mainIndexGuard.lock();
837+
this->lockMainIndexGuard();
838838
num_deleted_vectors += this->deleteLabelFromHNSWInplace(label);
839-
this->mainIndexGuard.unlock();
839+
this->unlockMainIndexGuard();
840840
}
841841

842842
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
friend class BF16TieredTest;
5960
friend class FP16TieredTest;

‎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<DataType, DistType> *backendIndex_,
6275
BruteForceIndex<DataType, DistType> *frontendIndex_,
6376
TieredIndexParams tieredParams, std::shared_ptr<VecSimAllocator> allocator)

0 commit comments

Comments
 (0)