diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered.h b/src/VecSim/algorithms/hnsw/hnsw_tiered.h index a4d5e08e4..2ff56795a 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered.h @@ -1,6 +1,7 @@ /* * Copyright (c) 2006-Present, Redis Ltd. * All rights reserved. + * SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates * * 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 @@ -46,6 +47,7 @@ struct HNSWSwapJob : public VecsimBaseObject { static const size_t DEFAULT_PENDING_SWAP_JOBS_THRESHOLD = DEFAULT_BLOCK_SIZE; static const size_t MAX_PENDING_SWAP_JOBS_THRESHOLD = 100000; +static const size_t MAX_QUANT_NORMALIZATION_SET_SIZE = 100 * DEFAULT_BLOCK_SIZE; /** * Definition of a job that repairs a certain node's connection in HNSW Index after delete @@ -99,6 +101,16 @@ class TieredHNSWIndex : public VecSimTieredIndex { // Not atomic since it's only accessed from the main thread. size_t directHNSWInsertions{0}; + bool isQuantized{false}; + + // SQ accumulation phase members, only used for quantized tiered index + vecsim_stl::vector runningSumVec; + size_t quantNormalizationSetSize; + bool isInAccumulationPhase; + VecSimParams backendIndexParams; // Saved params for creating the SQ backend + + void initializeQuantizedBackend(); + void executeInsertJob(HNSWInsertJob *job); void executeRepairJob(HNSWRepairJob *job); @@ -143,6 +155,26 @@ class TieredHNSWIndex : public VecSimTieredIndex { // Handle deletion of vector inplace considering that async deletion might occurred beforehand. int deleteLabelFromHNSWInplace(labelType label); + void addToSum(const DataType *vector) { + for (size_t i = 0; i < this->runningSumVec.size(); i++) { + if constexpr (std::is_same_v) { + this->runningSumVec[i] += vector[i]; + } else if constexpr (std::is_same_v) { + this->runningSumVec[i] += FP16_to_FP32(vector[i]); + } + } + } + + void subtractFromSum(const DataType *vector) { + for (size_t i = 0; i < this->runningSumVec.size(); i++) { + if constexpr (std::is_same_v) { + this->runningSumVec[i] -= vector[i]; + } else if constexpr (std::is_same_v) { + this->runningSumVec[i] -= FP16_to_FP32(vector[i]); + } + } + } + #ifdef BUILD_TESTS #include "VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h" #endif @@ -204,10 +236,23 @@ class TieredHNSWIndex : public VecSimTieredIndex { std::shared_ptr allocator); virtual ~TieredHNSWIndex(); + // Override query routing so that for quantized tiered indices the flat+backend + // merge always uses withSet=true (the flat and quantized backend scores are not + // directly comparable, so the withSet=false optimization is unsound). + // Also short-circuits to a flat-only query while in accumulation phase. + VecSimQueryReply *topKQuery(const void *queryBlob, size_t k, + VecSimQueryParams *queryParams) const override; + VecSimQueryReply *rangeQuery(const void *queryBlob, double radius, + VecSimQueryParams *queryParams, + VecSimQueryReply_Order order) const override; + int addVector(const void *blob, labelType label) override; int deleteVector(labelType label) override; size_t getNumMarkedDeleted() const override { - return this->getHNSWIndex()->getNumMarkedDeleted(); + this->mainIndexGuard.lock_shared(); + size_t num_marked_deleted = this->getHNSWIndex()->getNumMarkedDeleted(); + this->mainIndexGuard.unlock_shared(); + return num_marked_deleted; } size_t indexSize() const override; size_t indexCapacity() const override; @@ -220,12 +265,17 @@ class TieredHNSWIndex : public VecSimTieredIndex { VecSimDebugInfoIterator *debugInfoIterator() const override; VecSimBatchIterator *newBatchIterator(const void *queryBlob, VecSimQueryParams *queryParams) const override { + if (isInAccumulationPhase) { + return this->frontendIndex->newBatchIterator(queryBlob, queryParams); + } // The query blob will be processed and copied by the internal indexes's batch iterator. return new (this->allocator) TieredHNSW_BatchIterator(queryBlob, this, queryParams, this->allocator); } inline void setLastSearchMode(VecSearchMode mode) override { - return this->backendIndex->setLastSearchMode(mode); + this->mainIndexGuard.lock_shared(); + this->backendIndex->setLastSearchMode(mode); + this->mainIndexGuard.unlock_shared(); } void runGC() override { // Run no more than pendingSwapJobsThreshold value jobs. @@ -255,8 +305,10 @@ class TieredHNSWIndex : public VecSimTieredIndex { #ifdef BUILD_TESTS void getDataByLabel(labelType label, std::vector> &vectors_output) const; size_t indexMetaDataCapacity() const override { - return this->backendIndex->indexMetaDataCapacity() + - this->frontendIndex->indexMetaDataCapacity(); + this->mainIndexGuard.lock_shared(); + size_t capacity = this->backendIndex->indexMetaDataCapacity(); + this->mainIndexGuard.unlock_shared(); + return capacity + this->frontendIndex->indexMetaDataCapacity(); } #endif }; @@ -661,7 +713,8 @@ TieredHNSWIndex::TieredHNSWIndex(HNSWIndex(hnsw_index, bf_index, tiered_index_params, allocator), labelToInsertJobs(this->allocator), idToRepairJobs(this->allocator), idToSwapJob(this->allocator), invalidJobs(this->allocator), currInvalidJobId(0), - readySwapJobs(0) { + readySwapJobs(0), runningSumVec(this->allocator), quantNormalizationSetSize(0), + isInAccumulationPhase(false) { // If the param for swapJobThreshold is 0 use the default value, if it exceeds the maximum // allowed, use the maximum value. this->pendingSwapJobsThreshold = @@ -669,6 +722,21 @@ TieredHNSWIndex::TieredHNSWIndex(HNSWIndexalgoParams.hnswParams; + if (hnswParams.quantType != VecSimQuant_NONE) { + isQuantized = true; + size_t normSize = + tiered_index_params.specificParams.tieredHnswParams.QuantNormalizationSetSize; + if (normSize > 0) { + this->quantNormalizationSetSize = std::min(normSize, MAX_QUANT_NORMALIZATION_SET_SIZE); + this->isInAccumulationPhase = true; + this->runningSumVec.resize(hnswParams.dim, 0.0f); + } + // Save tiered params for creating the SQ backend + this->backendIndexParams = *(tiered_index_params.primaryIndexParams); + } } template @@ -695,19 +763,88 @@ TieredHNSWIndex::~TieredHNSWIndex() { } } +template +VecSimQueryReply * +TieredHNSWIndex::topKQuery(const void *queryBlob, size_t k, + VecSimQueryParams *queryParams) const { + // Accumulation phase: backend is an empty placeholder. We can short-circuit + // to a flat-only query. + if (this->isInAccumulationPhase) { + this->flatIndexGuard.lock_shared(); + auto *res = this->frontendIndex->topKQuery(queryBlob, k, queryParams); + this->flatIndexGuard.unlock_shared(); + return res; + } + // For quantized tiered, flat/backend scores are not directly comparable, so we + // must use withSet=true even for single-value indexes. Multi-value already + // uses withSet=true in the base. + if (isQuantized && !this->backendIndex->isMultiValue()) { + return this->template topKQueryImp(queryBlob, k, queryParams); + } + return VecSimTieredIndex::topKQuery(queryBlob, k, queryParams); +} + +template +VecSimQueryReply * +TieredHNSWIndex::rangeQuery(const void *queryBlob, double radius, + VecSimQueryParams *queryParams, + VecSimQueryReply_Order order) const { + if (this->isInAccumulationPhase) { + this->flatIndexGuard.lock_shared(); + auto *res = this->frontendIndex->rangeQuery(queryBlob, radius, queryParams); + this->flatIndexGuard.unlock_shared(); + if (res) { + sort_results(res, order); + } + return res; + } + if (isQuantized && !this->backendIndex->isMultiValue()) { + return this->template rangeQueryImp(queryBlob, radius, queryParams, order); + } + return VecSimTieredIndex::rangeQuery(queryBlob, radius, queryParams, order); +} + +template +void TieredHNSWIndex::initializeQuantizedBackend() { + // Compute mean and create the SQ backend. + size_t dim = this->backendIndexParams.algoParams.hnswParams.dim; + vecsim_stl::vector mean(dim, this->allocator); + for (size_t i = 0; i < dim; i++) { + mean[i] = this->runningSumVec[i] / this->quantNormalizationSetSize; + } + + // Build HNSWParams with SQ quantization and computed mean. + VecSimParams newParams = this->backendIndexParams; + newParams.algoParams.hnswParams.quantParams = mean.data(); + + // Exclude readers while replacing the empty backend with its SQ8 counterpart. + this->lockMainIndexGuard(); + VecSimIndex_Free(this->backendIndex); + + // Create new SQ HNSW backend. Normalization is done by the frontend index. + this->backendIndex = reinterpret_cast *>( + HNSWFactory::NewIndex(&newParams, true)); + this->unlockMainIndexGuard(); +} + template size_t TieredHNSWIndex::indexSize() const { this->flatIndexGuard.lock_shared(); + this->mainIndexGuard.lock_shared(); this->getHNSWIndex()->lockSharedIndexDataGuard(); size_t res = this->backendIndex->indexSize() + this->frontendIndex->indexSize(); this->getHNSWIndex()->unlockSharedIndexDataGuard(); + this->mainIndexGuard.unlock_shared(); this->flatIndexGuard.unlock_shared(); return res; } template size_t TieredHNSWIndex::indexCapacity() const { - return this->backendIndex->indexCapacity() + this->frontendIndex->indexCapacity(); + this->mainIndexGuard.lock_shared(); + size_t capacity = this->backendIndex->indexCapacity(); + this->mainIndexGuard.unlock_shared(); + return capacity + this->frontendIndex->indexCapacity(); } // In the tiered index, we assume that the blobs are processed by the flat buffer @@ -721,7 +858,8 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l auto hnsw_index = this->getHNSWIndex(); // writeMode is not protected since it is assumed to be called only from the "main thread" // (that is the thread that is exclusively calling add/delete vector). - if (this->getWriteMode() == VecSim_WriteInPlace) { + // VecSim_WriteInPlace is ignored during the accumulation phase + if (!this->isInAccumulationPhase && this->getWriteMode() == VecSim_WriteInPlace) { // First, check if we need to overwrite the vector in-place for single (from both indexes). if (!this->backendIndex->isMultiValue()) { ret -= this->deleteVector(label); @@ -739,7 +877,11 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l ++this->directHNSWInsertions; return ret; } - if (this->frontendIndex->indexSize() >= this->flatBufferLimit) { + if (this->isInAccumulationPhase) { + // Accumulate the vectors in the running sum vector + auto storage_blob = this->frontendIndex->preprocessForStorage(blob); + this->addToSum(reinterpret_cast(storage_blob.get())); + } else if (this->frontendIndex->indexSize() >= this->flatBufferLimit) { // Handle overwrite situation. if (!this->backendIndex->isMultiValue()) { // This will do nothing (and return 0) if this label doesn't exist. Otherwise, it may @@ -766,6 +908,10 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l if (this->frontendIndex->isLabelExists(label) && !this->frontendIndex->isMultiValue()) { // Overwrite the vector and invalidate its only pending job (since we are not in MULTI). auto *old_job = this->labelToInsertJobs.at(label).at(0); + if (this->isInAccumulationPhase) { + const DataType *vector_data = this->frontendIndex->getDataByInternalId(old_job->id); + this->subtractFromSum(vector_data); + } old_job->id = this->setAndSaveInvalidJob(old_job); this->labelToInsertJobs.erase(label); ret = 0; @@ -798,7 +944,7 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l // Here, a worker might ingest the previous vector that was stored under "label" // (in case of override in non-MULTI index) - so if it's there, we remove it (and create the // required repair jobs), *before* we submit the insert job. - if (!this->backendIndex->isMultiValue()) { + if (!this->isInAccumulationPhase && !this->backendIndex->isMultiValue()) { // If we removed the previous vector from both HNSW and flat in the overwrite process, // we still return 0 (not -1). ret = std::max(ret - this->deleteLabelFromHNSW(label), 0); @@ -812,8 +958,27 @@ int TieredHNSWIndex::addVector(const void *blob, labelType l this->executeReadySwapJobs(this->pendingSwapJobsThreshold); } - // Insert job to the queue and signal the workers' updater. - this->submitSingleJob(new_insert_job); + if (!this->isInAccumulationPhase) { + // Insert job to the queue and signal the workers' updater. + this->submitSingleJob(new_insert_job); + } else if (this->frontendIndex->indexSize() >= this->quantNormalizationSetSize) { + // If we are in the accumulation phase and we just reached the quantization set size, we + // can initialize the backend index and transition to the regular mode. + this->initializeQuantizedBackend(); + + // Submit all pending insert jobs to the job queue. + vecsim_stl::vector jobs(this->allocator); + jobs.reserve(this->labelToInsertJobs.size()); + for (auto &entry : this->labelToInsertJobs) { + for (auto *job : entry.second) { + jobs.push_back(job); + } + } + this->submitJobs(jobs); + + // Transition complete. + this->isInAccumulationPhase = false; + } return ret; } @@ -829,6 +994,10 @@ int TieredHNSWIndex::deleteVector(labelType label) { // Invalidate the pending insert job(s) into HNSW associated with this label auto &insert_jobs = this->labelToInsertJobs.at(label); for (auto *job : insert_jobs) { + if (this->isInAccumulationPhase) { + const DataType *vector_data = this->frontendIndex->getDataByInternalId(job->id); + this->subtractFromSum(vector_data); + } job->id = this->setAndSaveInvalidJob(job); } num_deleted_vectors += insert_jobs.size(); @@ -850,6 +1019,10 @@ int TieredHNSWIndex::deleteVector(labelType label) { this->flatIndexGuard.unlock_shared(); } + if (this->isInAccumulationPhase) { + return num_deleted_vectors; + } + // Next, check if there vector(s) stored under the given label in HNSW and delete them as well. // Note that we may remove the same vector that has been removed from the flat index, if it was // being ingested at that time. @@ -901,6 +1074,10 @@ double TieredHNSWIndex::getDistanceFrom_Unsafe(labelType lab // If the label doesn't exist, the distance will be NaN. auto flat_dist = this->frontendIndex->getDistanceFrom_Unsafe(label, blob); + if (this->isInAccumulationPhase) { + return flat_dist; + } + // Optimization. TODO: consider having different implementations for single and multi indexes, // to avoid checking the index type on every query. if (!this->backendIndex->isMultiValue() && !std::isnan(flat_dist)) { diff --git a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h b/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h index 21f99f8f5..e66811f4d 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h +++ b/src/VecSim/algorithms/hnsw/hnsw_tiered_tests_friends.h @@ -1,6 +1,7 @@ /* * Copyright (c) 2006-Present, Redis Ltd. * All rights reserved. + * SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates * * 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 @@ -65,6 +66,9 @@ INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_deleteInplaceAvoidUpdatedMarked INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_switchDeleteModes_Test) INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestBasic_HNSWResize_Test) +INDEX_TEST_FRIEND_CLASS(HNSWTieredIndexTestSQ8) +INDEX_TEST_FRIEND_CLASS(SQ8TieredHNSWTest) + friend class CommonAPITest_SearchDifferentScores_Test; friend class BF16TieredTest; friend class FP16TieredTest; diff --git a/src/VecSim/index_factories/tiered_factory.cpp b/src/VecSim/index_factories/tiered_factory.cpp index 337db6cc3..a82f2163b 100644 --- a/src/VecSim/index_factories/tiered_factory.cpp +++ b/src/VecSim/index_factories/tiered_factory.cpp @@ -1,6 +1,7 @@ /* * Copyright (c) 2006-Present, Redis Ltd. * All rights reserved. + * SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates * * 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 @@ -91,12 +92,23 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) { throw std::invalid_argument("Invalid hnsw_params.type"); } + VecSimQuantType quantType = hnsw_params.quantType; + size_t normSize = params->specificParams.tieredHnswParams.QuantNormalizationSetSize; + if (quantType != VecSimQuant_NONE && normSize > 0) { + // Add size of SQ accumulation buffer + est += allocations_overhead + hnsw_params.dim * sizeof(float); + } + return est; } VecSimIndex *NewIndex(const TieredIndexParams *params) { // Tiered index that contains HNSW index as primary index VecSimType type = params->primaryIndexParams->algoParams.hnswParams.type; + VecSimQuantType quantType = params->primaryIndexParams->algoParams.hnswParams.quantType; + if (quantType != VecSimQuant_NONE && type != VecSimType_FLOAT32 && type != VecSimType_FLOAT16) { + return nullptr; + } if (type == VecSimType_FLOAT32) { return TieredHNSWFactory::NewIndex(params); } else if (type == VecSimType_FLOAT64) { diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h index 828f5084b..3605d9fff 100644 --- a/src/VecSim/vec_sim_common.h +++ b/src/VecSim/vec_sim_common.h @@ -216,6 +216,9 @@ typedef struct { typedef struct { size_t swapJobThreshold; // The minimum number of swap jobs to accumulate before applying // all the ready swap jobs in a batch. + size_t QuantNormalizationSetSize; // Number of vectors to accumulate before SQ initialization. + // 0 = skip phase 0 (naive SQ8, no mean). + // Max: 100 * DEFAULT_BLOCK_SIZE (102400). } TieredHNSWParams; // A struct that contains HNSW Disk tiered index specific params. diff --git a/tests/unit/test_hnsw_sq8.cpp b/tests/unit/test_hnsw_sq8.cpp index f20a3f5d7..bdef3f344 100644 --- a/tests/unit/test_hnsw_sq8.cpp +++ b/tests/unit/test_hnsw_sq8.cpp @@ -8,6 +8,7 @@ #include "VecSim/types/float16.h" #include "VecSim/types/sq8.h" #include "VecSim/vec_sim.h" +#include "mock_thread_pool.h" #include "unit_test_utils.h" #include @@ -42,7 +43,7 @@ class HNSWSQ8Test : public ::testing::Test { } } - void SetUp(HNSWParams ¶ms) { + virtual void SetUp(HNSWParams ¶ms) { params.type = index_type_t::get_index_type(); params.quantType = VecSimQuant_SQ8; if constexpr (index_type_t::with_quant_params) { @@ -61,7 +62,7 @@ class HNSWSQ8Test : public ::testing::Test { } } - HNSWIndex *CastToHNSW() { + virtual HNSWIndex *CastToHNSW() { return dynamic_cast *>(index); } @@ -244,7 +245,6 @@ void HNSWSQ8Test::search_empty_index_test() { for (size_t i = 0; i < 100; i++) { VecSimIndex_DeleteVector(index, i); } - ASSERT_EQ(VecSimIndex_IndexSize(index), 0u); // Again, we do not expect any results. reply = VecSimIndex_TopKQuery(index, query, 11, nullptr, BY_SCORE); @@ -260,29 +260,43 @@ TYPED_TEST(HNSWSQ8Test, SearchEmptyIndex) { this->search_empty_index_test(); } template void HNSWSQ8Test::test_override() { - constexpr size_t count = 250; + constexpr size_t n = 100; + constexpr size_t new_n = 250; + // Scale factor to avoid FP16 overflow. FP16 max value is 65504, and L2² = dim × diff². + // With scale=0.1 and max diff=250: L2² = 4 × (250×0.1)² = 10000 < 65504. + constexpr float scale = 0.1f; HNSWParams params = { - .dim = 4, .initialCapacity = 100, .M = 8, .efConstruction = 20, .efRuntime = count}; + .dim = 4, .initialCapacity = n, .M = 8, .efConstruction = 20, .efRuntime = new_n}; SetUp(params); - // Insert 100 vectors and then overwrite each one with the same value. - for (size_t i = 0; i < 100; i++) { - ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 1); - ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 0); + // Insert n vectors. + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(GenerateAndAddVector(i, i * scale), 1); } - // Add vectors up to count. - for (size_t i = 100; i < count; i++) { - ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 1); + + // Override n vectors, the first 100 will be overwritten (deleted first). + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(GenerateAndAddVector(i, i * scale), 0); + } + + // Add up to new_n vectors. + for (size_t i = n; i < new_n; i++) { + ASSERT_EQ(GenerateAndAddVector(i, i * scale), 1); } data_t query[4]; - GenerateVector(query, static_cast(count)); - // The largest label is closest to the query, so labels are returned in descending order. + GenerateVector(query, new_n * scale); + + // Vectors values equal their labels (scaled), so the larger the label, the closer it is to + // the query. auto verify = [&](size_t id, double score, size_t result_index) { - EXPECT_EQ(id, count - result_index - 1); - EXPECT_FLOAT_EQ(score, 4.0f * (count - id) * (count - id)); + EXPECT_EQ(id, new_n - result_index - 1); + const float diff = new_n * scale - id * scale; + const float expected_score = 4 * diff * diff; + const float tolerance = std::max(1.0f, std::abs(expected_score) * 0.002f); + EXPECT_NEAR(score, expected_score, tolerance); }; - runTopKSearchTest(index, query, count, verify); + runTopKSearchTest(index, query, 300, verify); } TYPED_TEST(HNSWSQ8Test, Override) { this->test_override(); } @@ -343,17 +357,20 @@ template void HNSWSQ8Test::test_batch_iterator_basic() { constexpr size_t count = 250; constexpr size_t batch_size = 5; + // Scale factor to avoid FP16 overflow. FP16 max value is 65504, and L2² = dim × diff². + // With scale=0.1 and max diff=250: L2² = 4 × (250×0.1)² = 10000 < 65504. + constexpr float scale = 0.1f; HNSWParams params = { .dim = 4, .initialCapacity = count, .M = 8, .efConstruction = 20, .efRuntime = count}; SetUp(params); // For every i, add the vector (i, i, i, i) under label i. for (size_t i = 0; i < count; i++) { - ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 1); + ASSERT_EQ(GenerateAndAddVector(i, i * scale), 1); } data_t query[4]; - GenerateVector(query, static_cast(count)); + GenerateVector(query, count * scale); VecSimBatchIterator *iterator = VecSimBatchIterator_New(index, query, nullptr); ASSERT_NE(iterator, nullptr); @@ -372,3 +389,119 @@ void HNSWSQ8Test::test_batch_iterator_basic() { } TYPED_TEST(HNSWSQ8Test, BatchIteratorBasic) { this->test_batch_iterator_basic(); } + +/* ---------------------------- Tiered HNSW tests ---------------------------- */ + +using HNSWSQ8TieredDataTypeSet = + ::testing::Types, + HNSWSQ8IndexType>; + +template +class SQ8TieredHNSWTest : public HNSWSQ8Test { +public: + using data_t = typename index_type_t::data_t; + + void create_index_test(); + +protected: + static constexpr size_t normalization_set_size = 10; + + void SetUp(HNSWParams &hnsw_params) override { + hnsw_params.type = index_type_t::get_index_type(); + hnsw_params.quantType = VecSimQuant_SQ8; + VecSimParams vecsim_hnsw_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = { + .jobQueue = &mock_thread_pool.jobQ, + .jobQueueCtx = mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .primaryIndexParams = &vecsim_hnsw_params, + .specificParams = { + TieredHNSWParams{.QuantNormalizationSetSize = normalization_set_size}}}; + VecSimParams vecsim_params = CreateParams(tiered_params); + this->index = VecSimIndex_New(&vecsim_params); + ASSERT_NE(this->index, nullptr); + this->dim = hnsw_params.dim; + mock_thread_pool.ctx->index_strong_ref.reset(this->index); + } + + void TearDown() override {} + + HNSWIndex *CastToHNSW() override { + auto *tiered_index = dynamic_cast *>(this->index); + return tiered_index ? tiered_index->getHNSWIndex() : nullptr; + } + + tieredIndexMock mock_thread_pool; +}; + +template +void SQ8TieredHNSWTest::create_index_test() { + HNSWParams params = {.dim = 40, .M = 16, .efConstruction = 200}; + SetUp(params); + + ASSERT_EQ(VecSimIndex_IndexSize(this->index), 0u); + for (size_t label = 0; label < 100; label++) { + ASSERT_EQ(this->GenerateAndAddVector(label, static_cast(label), 1.0f), 1); + ASSERT_EQ(VecSimIndex_IndexSize(this->index), label + 1); + } + EXPECT_EQ(this->index->basicInfo().type, index_type_t::get_index_type()); + EXPECT_TRUE(this->index->basicInfo().isTiered); +} + +TYPED_TEST_SUITE(SQ8TieredHNSWTest, HNSWSQ8TieredDataTypeSet); + +TYPED_TEST(SQ8TieredHNSWTest, CreateIndex) { this->create_index_test(); } + +TYPED_TEST(SQ8TieredHNSWTest, SizeEstimation) { + constexpr size_t block_size = DEFAULT_BLOCK_SIZE; + HNSWParams hnsw_params = {.dim = 16, .initialCapacity = block_size, .M = 32}; + this->SetUp(hnsw_params); + + VecSimParams vecsim_hnsw_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = { + .jobQueue = &this->mock_thread_pool.jobQ, + .jobQueueCtx = this->mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .primaryIndexParams = &vecsim_hnsw_params, + .specificParams = { + TieredHNSWParams{.QuantNormalizationSetSize = this->normalization_set_size}}}; + VecSimParams params = CreateParams(tiered_params); + + EXPECT_EQ(VecSimIndex_EstimateInitialSize(¶ms), this->index->getAllocationSize()); + + for (size_t label = 0; label < block_size; label++) { + ASSERT_EQ(this->GenerateAndAddVector(label, static_cast(label)), 1); + } + while (!this->mock_thread_pool.jobQ.empty()) { + this->mock_thread_pool.thread_iteration(); + } + + const size_t estimation = VecSimIndex_EstimateElementSize(¶ms) * block_size; + const size_t before = this->index->getAllocationSize(); + ASSERT_EQ(this->GenerateAndAddVector(block_size, static_cast(block_size)), 1); + while (!this->mock_thread_pool.jobQ.empty()) { + this->mock_thread_pool.thread_iteration(); + } + const size_t actual = this->index->getAllocationSize() - before; + + EXPECT_EQ(this->index->indexSize(), block_size + 1); + EXPECT_EQ(this->index->indexCapacity(), 2 * block_size); + EXPECT_GE(estimation, actual * 0.99); + EXPECT_LE(estimation, actual * 1.01); +} + +TYPED_TEST(SQ8TieredHNSWTest, SearchByID) { this->search_by_id_test(); } + +TYPED_TEST(SQ8TieredHNSWTest, SearchByScore) { this->search_by_score_test(); } + +TYPED_TEST(SQ8TieredHNSWTest, SearchEmptyIndex) { this->search_empty_index_test(); } + +TYPED_TEST(SQ8TieredHNSWTest, Override) { this->test_override(); } + +TYPED_TEST(SQ8TieredHNSWTest, RangeQuery) { this->test_range_query(); } + +TYPED_TEST(SQ8TieredHNSWTest, GetDistanceL2) { this->test_get_distance(VecSimMetric_L2); } + +TYPED_TEST(SQ8TieredHNSWTest, GetDistanceIP) { this->test_get_distance(VecSimMetric_IP); } + +TYPED_TEST(SQ8TieredHNSWTest, BatchIteratorBasic) { this->test_batch_iterator_basic(); } diff --git a/tests/unit/test_hnsw_tiered.cpp b/tests/unit/test_hnsw_tiered.cpp index 21f504177..22e7c3391 100644 --- a/tests/unit/test_hnsw_tiered.cpp +++ b/tests/unit/test_hnsw_tiered.cpp @@ -1,6 +1,7 @@ /* * Copyright (c) 2006-Present, Redis Ltd. * All rights reserved. + * SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates * * 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 @@ -11,9 +12,12 @@ #include "VecSim/algorithms/hnsw/hnsw_tiered.h" #include "VecSim/algorithms/hnsw/hnsw_single.h" #include "VecSim/algorithms/hnsw/hnsw_multi.h" +#include "VecSim/types/float16.h" #include "VecSim/vec_sim_debug.h" #include #include +#include +#include #include "unit_test_utils.h" #include "mock_thread_pool.h" @@ -4532,3 +4536,1482 @@ TYPED_TEST(HNSWTieredIndexTestBasic, HNSWResize) { hnsw_index->indexMetaDataCapacity() + tiered_index->frontendIndex->indexMetaDataCapacity()); } + +using float16 = vecsim_types::float16; + +// ------------------------------------------------------------------- +// Type definitions for parameterized tests (float32 and float16) +// ------------------------------------------------------------------- + +template +struct SQ8IndexType { + static VecSimType get_index_type() { return type; } + static bool isMulti() { return IsMulti; } + typedef DataType data_t; + typedef DistType dist_t; +}; + +// ------------------------------------------------------------------- +// Test fixture +// ------------------------------------------------------------------- + +template +class HNSWTieredIndexTestSQ8 : public ::testing::Test { +public: + using data_t = typename index_type_t::data_t; + using dist_t = typename index_type_t::dist_t; + +protected: + VecSimWriteMode original_mode; + + void SetUp() override { original_mode = VecSimIndexInterface::asyncWriteMode; } + void TearDown() override { VecSimIndexInterface::asyncWriteMode = original_mode; } + + // Create a tiered HNSW index with SQ8 quantization and accumulation phase. + TieredHNSWIndex * + CreateSQ8TieredIndex(tieredIndexMock &mock_thread_pool, size_t dim = 16, + VecSimMetric metric = VecSimMetric_L2, size_t normSetSize = 100, + size_t flat_buffer_limit = SIZE_MAX, size_t M = 16, + size_t efConstruction = 200) { + HNSWParams hnsw_params = {.type = index_type_t::get_index_type(), + .dim = dim, + .metric = metric, + .multi = index_type_t::isMulti(), + .M = M, + .efConstruction = efConstruction, + .quantType = VecSimQuant_SQ8}; + VecSimParams vecsim_params = CreateParams(hnsw_params); + TieredIndexParams tiered_params = { + .jobQueue = &mock_thread_pool.jobQ, + .jobQueueCtx = mock_thread_pool.ctx, + .submitCb = tieredIndexMock::submit_callback, + .flatBufferLimit = flat_buffer_limit, + .primaryIndexParams = &vecsim_params, + .specificParams = { + TieredHNSWParams{.swapJobThreshold = 0, .QuantNormalizationSetSize = normSetSize}}}; + auto *tiered_index = reinterpret_cast *>( + TieredFactory::NewIndex(&tiered_params)); + mock_thread_pool.ctx->index_strong_ref.reset(tiered_index); + return tiered_index; + } + + HNSWIndex *CastToHNSW(VecSimIndex *index) { + auto tiered_index = reinterpret_cast *>(index); + return tiered_index->getHNSWIndex(); + } + + // --- Accessor helpers (HNSWTieredIndexTestSQ8 is a friend of TieredHNSWIndex) --- + + bool getIsInAccumulationPhase(TieredHNSWIndex *idx) { + return idx->isInAccumulationPhase; + } + + const vecsim_stl::vector &getRunningSumVec(TieredHNSWIndex *idx) { + return idx->runningSumVec; + } + + size_t getQuantNormalizationSetSize(TieredHNSWIndex *idx) { + return idx->quantNormalizationSetSize; + } + + BruteForceIndex *getFrontendIndex(TieredHNSWIndex *idx) { + return idx->frontendIndex; + } + + VecSimIndexAbstract *getBackendIndex(TieredHNSWIndex *idx) { + return idx->backendIndex; + } + + auto &getLabelToInsertJobs(TieredHNSWIndex *idx) { + return idx->labelToInsertJobs; + } + + void callExecuteReadySwapJobs(TieredHNSWIndex *idx) { + idx->executeReadySwapJobs(); + } + + // Generate a vector with a pattern based on label. + void GenerateVectorData(data_t *output, size_t dim, float base_value) { + for (size_t i = 0; i < dim; i++) { + float val = base_value + static_cast(i) * 0.1f; + if constexpr (std::is_same_v) { + output[i] = val; + } else if constexpr (std::is_same_v) { + output[i] = vecsim_types::FP32_to_FP16(val); + } + } + } + + // Get value as float from data type. + float ToFloat(data_t val) { + if constexpr (std::is_same_v) { + return val; + } else { + return vecsim_types::FP16_to_FP32(val); + } + } +}; + +using SQ8FP32Single = SQ8IndexType; +using SQ8FP32Multi = SQ8IndexType; +using SQ8FP16Single = SQ8IndexType; +using SQ8FP16Multi = SQ8IndexType; + +using SQ8DataTypeSet = ::testing::Types; +using SQ8SingleDataTypeSet = ::testing::Types; +using SQ8MultiDataTypeSet = ::testing::Types; + +template +class HNSWTieredIndexTestSQ8Single : public HNSWTieredIndexTestSQ8 {}; + +template +class HNSWTieredIndexTestSQ8Multi : public HNSWTieredIndexTestSQ8 {}; + +TYPED_TEST_SUITE(HNSWTieredIndexTestSQ8, SQ8DataTypeSet); +TYPED_TEST_SUITE(HNSWTieredIndexTestSQ8Single, SQ8SingleDataTypeSet); +TYPED_TEST_SUITE(HNSWTieredIndexTestSQ8Multi, SQ8MultiDataTypeSet); + +// ------------------------------------------------------------------- +// Accumulation Phase Core Tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, AccumulationPhaseInitialization) { + // Verify that creating an SQ8 tiered index enters accumulation phase. + size_t dim = 16; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Verify accumulation phase state. + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getRunningSumVec(tiered_index).size(), dim); + ASSERT_EQ(this->getQuantNormalizationSetSize(tiered_index), normSetSize); + + // Verify running sum is zero-initialized. + for (size_t i = 0; i < dim; i++) { + ASSERT_FLOAT_EQ(this->getRunningSumVec(tiered_index)[i], 0.0f); + } + + // Backend index exists but should be empty. + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), 0); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, RunningSumAccuracy) { + // Verify that runningSumVec correctly accumulates vector values. + size_t dim = 8; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add vectors and verify running sum. + std::vector expected_sum(dim, 0.0f); + for (size_t i = 0; i < 5; i++) { + TEST_DATA_T vec[dim]; + float base = static_cast(i + 1); + this->GenerateVectorData(vec, dim, base); + VecSimIndex_AddVector(tiered_index, vec, i); + + // Update expected sum. + for (size_t d = 0; d < dim; d++) { + expected_sum[d] += this->ToFloat(vec[d]); + } + } + + // Verify running sum matches expected. + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + for (size_t d = 0; d < dim; d++) { + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], expected_sum[d], 1e-3f) + << "Mismatch at dimension " << d; + } +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, QueryDuringAccumulation) { + // Search should only return flat buffer results during accumulation. + size_t dim = 8; + size_t normSetSize = 100; // High threshold so we stay in accumulation. + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add some vectors. + size_t n = 10; + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), n); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); + + // Run TopK query. + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, 0.0f); + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 5, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + size_t res_count = VecSimQueryReply_Len(results); + ASSERT_GT(res_count, 0); + ASSERT_LE(res_count, 5); + VecSimQueryReply_Free(results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, RangeQueryDuringAccumulation) { + // Range queries should only use flat buffer during accumulation. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add identical vectors (distance 0 from each other). + size_t n = 5; + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, 1.0f); // Same vector + VecSimIndex_AddVector(tiered_index, vec, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + + // Range query with large radius should find all vectors. + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, 1.0f); + auto *results = VecSimIndex_RangeQuery(tiered_index, query, 0.01, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + size_t res_count = VecSimQueryReply_Len(results); + ASSERT_EQ(res_count, n); + VecSimQueryReply_Free(results); +} + +// ------------------------------------------------------------------- +// Accumulation Phase Insert/Delete Tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, AddVectorDuringAccumulation) { + // Vectors added during accumulation go to flat buffer; no jobs submitted to queue. + size_t dim = 8; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, 1.5f); + VecSimIndex_AddVector(tiered_index, vec, 42); + + // Vector should be in flat buffer. + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), 1); + + // Job should be created in labelToInsertJobs but NOT submitted to queue. + ASSERT_EQ(this->getLabelToInsertJobs(tiered_index).size(), 1); + ASSERT_EQ(mock_thread_pool.jobQ.size(), 0); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, DeleteVectorDuringAccumulation) { + // Deletion from flat buffer during accumulation subtracts from running sum. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add two vectors. + TEST_DATA_T vec1[dim], vec2[dim]; + this->GenerateVectorData(vec1, dim, 1.0f); + this->GenerateVectorData(vec2, dim, 2.0f); + VecSimIndex_AddVector(tiered_index, vec1, 1); + VecSimIndex_AddVector(tiered_index, vec2, 2); + + // Record sum before deletion. + std::vector sum_before(this->getRunningSumVec(tiered_index).begin(), + this->getRunningSumVec(tiered_index).end()); + + // Delete label 1. + VecSimIndex_DeleteVector(tiered_index, 1); + + // Verify running sum was updated (subtracted vec1's values). + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + for (size_t d = 0; d < dim; d++) { + float expected = sum_before[d] - this->ToFloat(vec1[d]); + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], expected, 1e-3f); + } + + // Verify index size. + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + ASSERT_EQ(tiered_index->indexSize(), 1); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8Single, OverwriteDuringAccumulation) { + // Vector overwrite should update running sum correctly (only for single-label). + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add vector with label 1. + TEST_DATA_T vec1[dim]; + this->GenerateVectorData(vec1, dim, 1.0f); + VecSimIndex_AddVector(tiered_index, vec1, 1); + + std::vector sum_after_first(this->getRunningSumVec(tiered_index).begin(), + this->getRunningSumVec(tiered_index).end()); + + // Overwrite with different vector. + TEST_DATA_T vec2[dim]; + this->GenerateVectorData(vec2, dim, 3.0f); + VecSimIndex_AddVector(tiered_index, vec2, 1); + + // Running sum should reflect: sum - vec1 + vec2 (overwrite subtracts old + adds new). + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + for (size_t d = 0; d < dim; d++) { + float expected = sum_after_first[d] - this->ToFloat(vec1[d]) + this->ToFloat(vec2[d]); + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], expected, 1e-3f); + } + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); +} + +// ------------------------------------------------------------------- +// Backend Index Initialization Tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, BackendCreatedAtThreshold) { + // When accumulation reaches quantNormalizationSetSize, backend is initialized. + size_t dim = 4; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add vectors up to threshold - 1. + for (size_t i = 0; i < normSetSize - 1; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), normSetSize - 1); + + // Add the threshold-triggering vector. + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(normSetSize - 1)); + VecSimIndex_AddVector(tiered_index, vec, normSetSize - 1); + + // Accumulation phase should be over. + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + // Backend should be initialized (still empty since jobs haven't run). + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); + // Flat buffer should hold all vectors. + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), normSetSize); + // All vectors should have associated insert jobs. + ASSERT_EQ(this->getLabelToInsertJobs(tiered_index).size(), normSetSize); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, MeanComputedCorrectly) { + // Verify the mean vector computed during initializeQuantizedBackend. + size_t dim = 4; + size_t normSetSize = 5; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Track expected sum. + std::vector expected_sum(dim, 0.0f); + for (size_t i = 0; i < normSetSize - 1; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i + 1)); + VecSimIndex_AddVector(tiered_index, vec, i); + for (size_t d = 0; d < dim; d++) { + expected_sum[d] += this->ToFloat(vec[d]); + } + } + + // Verify running sum before triggering. + for (size_t d = 0; d < dim; d++) { + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], expected_sum[d], 1e-2f); + } + + // Add final vector to trigger transition. + TEST_DATA_T last_vec[dim]; + this->GenerateVectorData(last_vec, dim, static_cast(normSetSize)); + VecSimIndex_AddVector(tiered_index, last_vec, normSetSize - 1); + for (size_t d = 0; d < dim; d++) { + expected_sum[d] += this->ToFloat(last_vec[d]); + } + + // After transition, verify the mean was computed correctly. + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + // The running sum should still hold the accumulated values. + for (size_t d = 0; d < dim; d++) { + float expected_mean = expected_sum[d] / normSetSize; + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d] / normSetSize, expected_mean, 1e-2f); + } +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, QueryDuringPartialMigration) { + // SQ8 scores from the flat and backend indexes are not directly comparable. Verify both + // query types find an exact-match vector that is still in the flat index while migration is + // in progress. For multi-value indexes, the vector shares a label with the migrated vector + // to verify duplicate labels are merged into one result. + + size_t dim = 8; + size_t normSetSize = 3; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i * 10)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Migrate only one threshold vector, leaving the remaining vectors in the flat index. + mock_thread_pool.thread_iteration(); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 1); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), normSetSize - 1); + + // This post-transition vector remains in the flat index while the backend has SQ8 data. + // In multi-value indexes, reuse the migrated label to exercise deduplication across indexes. + TEST_DATA_T flat_vec[dim]; + labelType flat_label = TypeParam::isMulti() ? 0 : 100; + this->GenerateVectorData(flat_vec, dim, static_cast(flat_label)); + VecSimIndex_AddVector(tiered_index, flat_vec, flat_label); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 1); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), normSetSize); + + auto *topk_results = VecSimIndex_TopKQuery(tiered_index, flat_vec, 1, nullptr, BY_SCORE); + ASSERT_NE(topk_results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(topk_results), 1); + auto *topk_iterator = VecSimQueryReply_GetIterator(topk_results); + auto *topk_result = VecSimQueryReply_IteratorNext(topk_iterator); + ASSERT_EQ(VecSimQueryResult_GetId(topk_result), flat_label); + VecSimQueryReply_IteratorFree(topk_iterator); + VecSimQueryReply_Free(topk_results); + + auto *range_results = VecSimIndex_RangeQuery(tiered_index, flat_vec, 0.01, nullptr, BY_SCORE); + ASSERT_NE(range_results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(range_results), 1); + auto *range_iterator = VecSimQueryReply_GetIterator(range_results); + auto *range_result = VecSimQueryReply_IteratorNext(range_iterator); + ASSERT_EQ(VecSimQueryResult_GetId(range_result), flat_label); + VecSimQueryReply_IteratorFree(range_iterator); + VecSimQueryReply_Free(range_results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, NewVectorsAfterAccumulation) { + // Vectors added after accumulation are submitted to job queue. + size_t dim = 4; + size_t normSetSize = 5; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Trigger transition. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + size_t queue_size_before = mock_thread_pool.jobQ.size(); + + // Add a new vector after accumulation. + TEST_DATA_T new_vec[dim]; + this->GenerateVectorData(new_vec, dim, 99.0f); + VecSimIndex_AddVector(tiered_index, new_vec, 99); + + // New job should be submitted to queue. + ASSERT_GT(mock_thread_pool.jobQ.size(), queue_size_before); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, DeleteFromBackendAfterAccumulation) { + // Delete operations work on SQ backend after accumulation. + size_t dim = 4; + size_t normSetSize = 5; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Trigger transition. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Execute all jobs to move vectors to HNSW backend. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + ASSERT_EQ(tiered_index->indexSize(), normSetSize); + + // Delete a vector (marks it for deletion in HNSW). + int deleted = VecSimIndex_DeleteVector(tiered_index, 0); + ASSERT_EQ(deleted, 1); + + // Execute repair jobs, then run swap jobs to physically remove the vector. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + this->callExecuteReadySwapJobs(tiered_index); + + ASSERT_EQ(tiered_index->indexSize(), normSetSize - 1); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, ConcurrentSearchDuringAccumulation) { + // Parallel searches should work correctly during accumulation. + size_t dim = 8; + size_t normSetSize = 1000; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add some vectors. + size_t n = 50; + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + + // Launch parallel searches. + std::atomic_int successful_searches(0); + size_t n_threads = 4; + auto search_fn = [&](size_t thread_id) { + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, static_cast(thread_id)); + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 5, nullptr, BY_SCORE); + if (results && VecSimQueryReply_Len(results) > 0) { + successful_searches++; + } + VecSimQueryReply_Free(results); + }; + + std::vector threads; + for (size_t t = 0; t < n_threads; t++) { + threads.emplace_back(search_fn, t); + } + for (auto &t : threads) { + t.join(); + } + + ASSERT_EQ(successful_searches, (int)n_threads); +} + +// ------------------------------------------------------------------- +// Memory & Size Tracking Tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, IndexSizeDuringAccumulation) { + // indexSize() returns flat buffer size during accumulation (backend is empty). + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + for (size_t i = 0; i < 10; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(tiered_index->indexSize(), 10); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 10); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CapacityDuringAccumulation) { + // indexCapacity() reflects flat buffer capacity during accumulation. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, 1.0f); + VecSimIndex_AddVector(tiered_index, vec, 0); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + // Capacity should come from flat buffer (at least DEFAULT_BLOCK_SIZE after first insert). + ASSERT_GE(tiered_index->indexCapacity(), 1); + ASSERT_EQ(tiered_index->indexCapacity(), this->getFrontendIndex(tiered_index)->indexCapacity()); +} + +// ------------------------------------------------------------------- +// Edge Cases +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, ZeroAccumulationThreshold) { + // QuantNormalizationSetSize=0 should skip accumulation phase entirely. + size_t dim = 4; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, 0 /* normSetSize=0 */); + + // Should NOT be in accumulation phase. + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Adding a vector should immediately submit job to queue. + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, 1.0f); + VecSimIndex_AddVector(tiered_index, vec, 0); + ASSERT_GT(mock_thread_pool.jobQ.size(), 0); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, AllVectorsDeletedBeforeThreshold) { + // All vectors deleted before reaching threshold - should remain in accumulation. + size_t dim = 4; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add and delete vectors. + for (size_t i = 0; i < 5; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + for (size_t i = 0; i < 5; i++) { + VecSimIndex_DeleteVector(tiered_index, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), 0); + + // Running sum should be approximately zero. + for (size_t d = 0; d < dim; d++) { + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], 0.0f, 1e-3f); + } +} + +// ------------------------------------------------------------------- +// Multi-Label Accumulation +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8Multi, AccumulationMultiLabel) { + // Multi-label: multiple vectors per label during accumulation. + size_t dim = 4; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add multiple vectors with the same label. + labelType shared_label = 42; + for (size_t i = 0; i < 3; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i + 1)); + VecSimIndex_AddVector(tiered_index, vec, shared_label); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 3); + // labelToInsertJobs should have 3 jobs for the same label. + ASSERT_EQ(this->getLabelToInsertJobs(tiered_index).at(shared_label).size(), 3); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8Multi, DeleteMultiLabelDuringAccumulation) { + // Multi-label: deleting one label removes all its vectors and adjusts the sum. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add vectors with different labels. + TEST_DATA_T vec1[dim], vec2[dim], vec3[dim]; + this->GenerateVectorData(vec1, dim, 1.0f); + this->GenerateVectorData(vec2, dim, 2.0f); + this->GenerateVectorData(vec3, dim, 3.0f); + + VecSimIndex_AddVector(tiered_index, vec1, 10); + VecSimIndex_AddVector(tiered_index, vec2, 10); // Same label + VecSimIndex_AddVector(tiered_index, vec3, 20); // Different label + + std::vector sum_before(this->getRunningSumVec(tiered_index).begin(), + this->getRunningSumVec(tiered_index).end()); + + // Delete label 10 (should remove both vectors). + VecSimIndex_DeleteVector(tiered_index, 10); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + + // Running sum should be adjusted by subtracting both vec1 and vec2. + for (size_t d = 0; d < dim; d++) { + float expected = sum_before[d] - this->ToFloat(vec1[d]) - this->ToFloat(vec2[d]); + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], expected, 1e-2f); + } +} + +// ------------------------------------------------------------------- +// Batch Iterator During Accumulation +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, BatchIteratorDuringAccumulation) { + // Batch iterator works during accumulation (only flat buffer results). + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + size_t n = 20; + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, 0.0f); + auto *batch_iterator = VecSimBatchIterator_New(tiered_index, query, nullptr); + ASSERT_NE(batch_iterator, nullptr); + + // Get first batch. + auto *batch = VecSimBatchIterator_Next(batch_iterator, 5, BY_SCORE); + ASSERT_NE(batch, nullptr); + size_t count = VecSimQueryReply_Len(batch); + ASSERT_GT(count, 0); + ASSERT_LE(count, 5); + VecSimQueryReply_Free(batch); + + VecSimBatchIterator_Free(batch_iterator); +} + +// ------------------------------------------------------------------- +// Index Statistics During Accumulation +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, DebugInfoDuringAccumulation) { + // Debug info during accumulation reflects flat-only state. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + for (size_t i = 0; i < 5; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + VecSimIndexDebugInfo info = tiered_index->debugInfo(); + ASSERT_EQ(info.commonInfo.indexSize, 5); +} + +// ------------------------------------------------------------------- +// Write Mode Interactions with Accumulation +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, WriteInPlaceDuringAccumulation) { + // WriteInPlace mode is ignored during accumulation phase. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Switch to write-in-place mode. + VecSimIndexInterface::asyncWriteMode = VecSim_WriteInPlace; + + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, 1.0f); + VecSimIndex_AddVector(tiered_index, vec, 0); + + // During accumulation, WriteInPlace is ignored - vector goes to flat buffer. + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), 0); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, WriteInPlaceAfterAccumulation) { + // After accumulation, WriteInPlace inserts directly to SQ backend. + size_t dim = 4; + size_t normSetSize = 5; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Trigger transition. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Execute all pending jobs (submitted by initializeQuantizedBackend). + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + // Switch to write-in-place mode. + VecSimIndexInterface::asyncWriteMode = VecSim_WriteInPlace; + + // Add vector - should go directly to HNSW backend. + TEST_DATA_T new_vec[dim]; + this->GenerateVectorData(new_vec, dim, 99.0f); + VecSimIndex_AddVector(tiered_index, new_vec, 99); + + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), normSetSize + 1); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 0); +} + +// ------------------------------------------------------------------- +// Buffer Limit Interactions +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, BufferLimitWithAccumulation) { + // Flat buffer limit is respected during accumulation. + size_t dim = 4; + size_t normSetSize = 100; // High normalization set size. + size_t buffer_limit = 10; // Small buffer limit. + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, + normSetSize, buffer_limit); + + // During accumulation, buffer limit should not trigger direct insert to backend + // (since backend is not ready). The addVector code checks isInAccumulationPhase + // before checking flatBufferLimit. + for (size_t i = 0; i < 15; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + // Should still be in accumulation phase. + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + // All vectors should be in flat buffer (accumulation overrides buffer limit behavior). + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 15); +} + +// ------------------------------------------------------------------- +// Quantization Quality Validation +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, QuantizedSearchQuality) { + // After accumulation, SQ8 backend should produce reasonable search results. + size_t dim = 16; + size_t normSetSize = 50; + size_t n = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add enough vectors to trigger transition and more. + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Execute all jobs. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + // Query with the same vector as label 0 - should find label 0 as nearest. + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, 0.0f); + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 1, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(results), 1); + + auto it = VecSimQueryReply_GetIterator(results); + auto *entry = VecSimQueryReply_IteratorNext(it); + // The closest vector should be label 0 (same as query). + ASSERT_EQ(VecSimQueryResult_GetId(entry), 0); + // Distance should be very small (quantization introduces some error). + ASSERT_LT(VecSimQueryResult_GetScore(entry), 1.0); + + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(results); +} + +// ------------------------------------------------------------------- +// End-to-end flow tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, FullFlowAsyncInsertAndSearch) { + // Full end-to-end test: accumulation -> transition -> async insert -> search. + size_t dim = 8; + size_t normSetSize = 20; + size_t total_vectors = 50; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Phase 1: Accumulation. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Phase 2: Post-accumulation inserts. + for (size_t i = normSetSize; i < total_vectors; i++) { + TEST_DATA_T vec[dim]; + this->GenerateVectorData(vec, dim, static_cast(i)); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + // Phase 3: Execute all jobs (including accumulation-phase jobs submitted during transition). + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + ASSERT_EQ(this->getBackendIndex(tiered_index)->indexSize(), total_vectors); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 0); + ASSERT_EQ(tiered_index->indexSize(), total_vectors); + + // Phase 4: Search. + TEST_DATA_T query[dim]; + this->GenerateVectorData(query, dim, 0.0f); + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 10, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(results), 10); + VecSimQueryReply_Free(results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, DeleteThenReinsertDuringAccumulation) { + // Delete and re-insert a vector during accumulation. + size_t dim = 4; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + // Add vector. + TEST_DATA_T vec1[dim]; + this->GenerateVectorData(vec1, dim, 1.0f); + VecSimIndex_AddVector(tiered_index, vec1, 0); + + // Delete it. + VecSimIndex_DeleteVector(tiered_index, 0); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 0); + + // Re-insert with different data. + TEST_DATA_T vec2[dim]; + this->GenerateVectorData(vec2, dim, 5.0f); + VecSimIndex_AddVector(tiered_index, vec2, 0); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + + // Running sum should only contain vec2's values (vec1 was subtracted, vec2 was added). + for (size_t d = 0; d < dim; d++) { + ASSERT_NEAR(this->getRunningSumVec(tiered_index)[d], this->ToFloat(vec2[d]), 1e-3f); + } +} + +// ------------------------------------------------------------------- +// FP16-specific precision tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, RunningSumPrecision) { + // Verify running sum precision over many insertions. + size_t dim = 4; + size_t normSetSize = 1000; + size_t n = 500; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_L2, normSetSize); + + std::vector reference_sum(dim, 0.0f); + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + float base = static_cast(i) * 0.01f; + this->GenerateVectorData(vec, dim, base); + VecSimIndex_AddVector(tiered_index, vec, i); + + for (size_t d = 0; d < dim; d++) { + reference_sum[d] += this->ToFloat(vec[d]); + } + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + + // The running sum should be close to our reference computation. + // Allow for floating point accumulation differences. + for (size_t d = 0; d < dim; d++) { + float relative_error = + std::abs(this->getRunningSumVec(tiered_index)[d] - reference_sum[d]) / + (std::abs(reference_sum[d]) + 1e-10f); + ASSERT_LT(relative_error, 0.01f) << "Precision loss at dim " << d; + } +} + +// ------------------------------------------------------------------- +// Cosine Metric Tests +// ------------------------------------------------------------------- + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineAccumulationPhase) { + // Verify accumulation phase works correctly with Cosine metric. + // Cosine normalizes vectors before storage, so addToSum must use stored data. + size_t dim = 8; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + + // Add vectors with varying magnitudes (normalization will make them unit vectors). + for (size_t i = 0; i < normSetSize - 1; i++) { + TEST_DATA_T vec[dim]; + float scale = static_cast(i + 1); // Different magnitudes + this->GenerateVectorData(vec, dim, scale); + VecSimIndex_AddVector(tiered_index, vec, i); + } + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), normSetSize - 1); + + // Verify the running sum is computed from STORED (normalized) vectors. + // After normalization, each stored vector has unit length, so each component + // should be bounded by [-1, 1]. The sum of N unit vectors has bounded magnitude. + const auto &running_sum = this->getRunningSumVec(tiered_index); + float sum_norm_sq = 0.0f; + for (size_t d = 0; d < dim; d++) { + sum_norm_sq += running_sum[d] * running_sum[d]; + } + // The magnitude of the sum of (normSetSize-1) unit vectors is at most (normSetSize-1). + float sum_norm = std::sqrt(sum_norm_sq); + ASSERT_LE(sum_norm, static_cast(normSetSize - 1) + 0.1f); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineTransitionAndQuery) { + // Full flow: accumulation -> transition -> query with Cosine metric. + size_t dim = 16; + size_t normSetSize = 20; + size_t total_vectors = 50; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Phase 1: Fill up to threshold to trigger transition. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + // Create vectors with distinct directions by varying the first component. + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? static_cast(i + 1) : 1.0f; + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Phase 2: Add more vectors after transition. + for (size_t i = normSetSize; i < total_vectors; i++) { + TEST_DATA_T vec[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? static_cast(i + 1) : 1.0f; + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + + // Phase 3: Execute all jobs. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + ASSERT_EQ(tiered_index->indexSize(), total_vectors); + + // Phase 4: Query - use same direction as highest-label vector (should be nearest). + TEST_DATA_T query[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? static_cast(total_vectors) : 1.0f; + if constexpr (std::is_same_v) { + query[d] = val; + } else { + query[d] = vecsim_types::FP32_to_FP16(val); + } + } + + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 5, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(results), 5); + + // The nearest neighbor should be the vector with highest first-component + // (most similar direction to query). + auto it = VecSimQueryReply_GetIterator(results); + auto *entry = VecSimQueryReply_IteratorNext(it); + labelType top_label = VecSimQueryResult_GetId(entry); + double top_score = VecSimQueryResult_GetScore(entry); + + // Score for Cosine is 1 - cosine_similarity. Should be close to 0 for nearest. + ASSERT_LT(top_score, 0.01); + // The top result should be one of the vectors with the largest first component. + // With FP16+SQ8, vectors 48 and 49 are nearly indistinguishable, so allow some slack. + ASSERT_GE(top_label, total_vectors - 3) + << "Expected a high-label vector (near-identical direction to query)"; + + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8Single, CosineOverwriteDuringAccumulation) { + // Overwrite during accumulation with Cosine: subtractFromSum must use stored + // (normalized) data, matching what addToSum accumulated. + size_t dim = 8; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Add a vector with label 1 (magnitude = ~4). + TEST_DATA_T vec1[dim]; + for (size_t d = 0; d < dim; d++) { + float val = static_cast(d + 1) * 0.5f; + if constexpr (std::is_same_v) { + vec1[d] = val; + } else { + vec1[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec1, 1); + + // Record running sum after first insert. + std::vector sum_after_first(this->getRunningSumVec(tiered_index).begin(), + this->getRunningSumVec(tiered_index).end()); + + // Overwrite label 1 with a completely different vector (different direction). + TEST_DATA_T vec2[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? 10.0f : 0.01f; + if constexpr (std::is_same_v) { + vec2[d] = val; + } else { + vec2[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec2, 1); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + + // After overwrite: sum should reflect only vec2's stored (normalized) data. + // Since the running sum = 0 + stored(vec1) - stored(vec1) + stored(vec2) = stored(vec2), + // the sum should be the normalized form of vec2. + const auto &running_sum = this->getRunningSumVec(tiered_index); + float sum_norm_sq = 0.0f; + for (size_t d = 0; d < dim; d++) { + sum_norm_sq += running_sum[d] * running_sum[d]; + } + float sum_norm = std::sqrt(sum_norm_sq); + // With only one normalized vector in the sum, the norm should be ~1.0. + ASSERT_NEAR(sum_norm, 1.0f, 0.05f); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineDeleteDuringAccumulation) { + // Delete during accumulation with Cosine: verify running sum is correctly updated. + size_t dim = 8; + size_t normSetSize = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Add two vectors. + TEST_DATA_T vec1[dim], vec2[dim]; + for (size_t d = 0; d < dim; d++) { + float v1 = static_cast(d + 1); + float v2 = static_cast(dim - d); + if constexpr (std::is_same_v) { + vec1[d] = v1; + vec2[d] = v2; + } else { + vec1[d] = vecsim_types::FP32_to_FP16(v1); + vec2[d] = vecsim_types::FP32_to_FP16(v2); + } + } + VecSimIndex_AddVector(tiered_index, vec1, 1); + VecSimIndex_AddVector(tiered_index, vec2, 2); + + // Record sum with both vectors. + std::vector sum_with_both(this->getRunningSumVec(tiered_index).begin(), + this->getRunningSumVec(tiered_index).end()); + + // Delete label 1. + VecSimIndex_DeleteVector(tiered_index, 1); + + ASSERT_TRUE(this->getIsInAccumulationPhase(tiered_index)); + ASSERT_EQ(this->getFrontendIndex(tiered_index)->indexSize(), 1); + + // After deleting vec1, sum should equal just stored(vec2). + // stored(vec2) is the normalized version of vec2, so its norm ≈ 1. + const auto &running_sum = this->getRunningSumVec(tiered_index); + float sum_norm_sq = 0.0f; + for (size_t d = 0; d < dim; d++) { + sum_norm_sq += running_sum[d] * running_sum[d]; + } + float sum_norm = std::sqrt(sum_norm_sq); + ASSERT_NEAR(sum_norm, 1.0f, 0.05f); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineSearchAccuracyAfterTransition) { + // Verify that SQ8+Cosine produces reasonable search accuracy after transition. + // Compare results ordering against brute-force on the same index. + size_t dim = 32; + size_t normSetSize = 30; + size_t n = 100; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, + normSetSize, SIZE_MAX, 16, 200); + + // Insert vectors with different directions. + for (size_t i = 0; i < n; i++) { + TEST_DATA_T vec[dim]; + for (size_t d = 0; d < dim; d++) { + // Create vectors where the i-th vector has a strong d==i%dim component. + float val = (d == (i % dim)) ? 10.0f : 1.0f / (d + 1.0f); + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Execute all jobs to move vectors to backend. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + ASSERT_EQ(tiered_index->indexSize(), n); + + // Query for a vector similar to label 0 (strong component at dim 0). + TEST_DATA_T query[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? 10.0f : 0.5f / (d + 1.0f); + if constexpr (std::is_same_v) { + query[d] = val; + } else { + query[d] = vecsim_types::FP32_to_FP16(val); + } + } + + size_t k = 10; + auto *results = VecSimIndex_TopKQuery(tiered_index, query, k, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(results), k); + + // Verify results are sorted by increasing score (1 - cosine_sim). + auto it = VecSimQueryReply_GetIterator(results); + double prev_score = -1.0; + while (auto *entry = VecSimQueryReply_IteratorNext(it)) { + double score = VecSimQueryResult_GetScore(entry); + ASSERT_GE(score, 0.0); + ASSERT_LE(score, 2.0); // Cosine distance is in [0, 2] + ASSERT_GE(score, prev_score); + prev_score = score; + } + VecSimQueryReply_IteratorFree(it); + + // The top-1 result should be label 0 (same strong direction at dim 0). + it = VecSimQueryReply_GetIterator(results); + auto *first = VecSimQueryReply_IteratorNext(it); + // Labels with strong component at dim 0 are: 0, 32, 64, 96 + labelType top_label = VecSimQueryResult_GetId(first); + ASSERT_TRUE(top_label % dim == 0) + << "Top result label=" << top_label << " expected a vector with strong dim-0 component"; + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineRangeQueryAfterTransition) { + // Verify range query with Cosine metric after accumulation transition. + size_t dim = 8; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Insert parallel vectors (identical direction, different magnitudes) - should have distance 0. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + float scale = static_cast(i + 1); + for (size_t d = 0; d < dim; d++) { + float val = scale * (d + 1.0f); + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Execute all jobs. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + + // Query with same direction - all vectors should be at distance ~0. + TEST_DATA_T query[dim]; + for (size_t d = 0; d < dim; d++) { + float val = static_cast(d + 1); + if constexpr (std::is_same_v) { + query[d] = val; + } else { + query[d] = vecsim_types::FP32_to_FP16(val); + } + } + + // Range query with small radius should find all parallel vectors. + auto *results = VecSimIndex_RangeQuery(tiered_index, query, 0.1, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + // All vectors have the same direction, so cosine distance ≈ 0 for all. + ASSERT_EQ(VecSimQueryReply_Len(results), normSetSize); + VecSimQueryReply_Free(results); +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineMeanCorrectness) { + // Verify the mean used for SQ8 quantization is computed from normalized vectors. + // The mean of N unit vectors with the same direction should be that unit vector itself. + size_t dim = 4; + size_t normSetSize = 5; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Add parallel vectors (same direction [1,2,3,4], different magnitudes). + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + float scale = static_cast(i + 1); + for (size_t d = 0; d < dim; d++) { + float val = scale * (d + 1.0f); + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // The running sum / normSetSize should be the mean of the normalized vectors. + // Since all vectors have the same direction [1,2,3,4], after normalization they're all + // [1,2,3,4]/sqrt(1+4+9+16) = [1,2,3,4]/sqrt(30). The mean is the same unit vector. + const auto &running_sum = this->getRunningSumVec(tiered_index); + float norm_factor = std::sqrt(1.0f + 4.0f + 9.0f + 16.0f); + for (size_t d = 0; d < dim; d++) { + float expected_mean = (d + 1.0f) / norm_factor; + float actual_mean = running_sum[d] / normSetSize; + ASSERT_NEAR(actual_mean, expected_mean, 0.02f) << "Mean mismatch at dim " << d; + } +} + +TYPED_TEST(HNSWTieredIndexTestSQ8, CosineDeleteAndReinsertAfterTransition) { + // Delete from Cosine SQ8 backend and reinsert. + size_t dim = 16; + size_t normSetSize = 10; + auto mock_thread_pool = tieredIndexMock(); + auto *tiered_index = + this->CreateSQ8TieredIndex(mock_thread_pool, dim, VecSimMetric_Cosine, normSetSize); + + // Trigger transition. + for (size_t i = 0; i < normSetSize; i++) { + TEST_DATA_T vec[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == (i % dim)) ? 5.0f : 0.1f; + if constexpr (std::is_same_v) { + vec[d] = val; + } else { + vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, vec, i); + } + ASSERT_FALSE(this->getIsInAccumulationPhase(tiered_index)); + + // Move all to backend. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + ASSERT_EQ(tiered_index->indexSize(), normSetSize); + + // Delete label 0. + VecSimIndex_DeleteVector(tiered_index, 0); + // Process repair jobs. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + this->callExecuteReadySwapJobs(tiered_index); + ASSERT_EQ(tiered_index->indexSize(), normSetSize - 1); + + // Reinsert with same label, different vector. + TEST_DATA_T new_vec[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? 10.0f : 0.01f; + if constexpr (std::is_same_v) { + new_vec[d] = val; + } else { + new_vec[d] = vecsim_types::FP32_to_FP16(val); + } + } + VecSimIndex_AddVector(tiered_index, new_vec, 0); + + // Move to backend. + while (!mock_thread_pool.jobQ.empty()) { + mock_thread_pool.thread_iteration(); + } + ASSERT_EQ(tiered_index->indexSize(), normSetSize); + + // Query for the reinserted vector's direction. + TEST_DATA_T query[dim]; + for (size_t d = 0; d < dim; d++) { + float val = (d == 0) ? 1.0f : 0.0f; + if constexpr (std::is_same_v) { + query[d] = val; + } else { + query[d] = vecsim_types::FP32_to_FP16(val); + } + } + + auto *results = VecSimIndex_TopKQuery(tiered_index, query, 1, nullptr, BY_SCORE); + ASSERT_NE(results, nullptr); + ASSERT_EQ(VecSimQueryReply_Len(results), 1); + auto it = VecSimQueryReply_GetIterator(results); + auto *entry = VecSimQueryReply_IteratorNext(it); + // Label 0 has the strongest dim-0 component, should be top result. + ASSERT_EQ(VecSimQueryResult_GetId(entry), 0); + VecSimQueryReply_IteratorFree(it); + VecSimQueryReply_Free(results); +}