diff --git a/src/VecSim/index_factories/hnsw_factory.cpp b/src/VecSim/index_factories/hnsw_factory.cpp index d577f57a1..54b4c333f 100644 --- a/src/VecSim/index_factories/hnsw_factory.cpp +++ b/src/VecSim/index_factories/hnsw_factory.cpp @@ -17,6 +17,7 @@ using bfloat16 = vecsim_types::bfloat16; using float16 = vecsim_types::float16; +using sq8 = vecsim_types::sq8; namespace HNSWFactory { @@ -34,11 +35,127 @@ NewIndex_ChooseMultiOrSingle(const HNSWParams *params, HNSWIndex_Single(params, abstractInitParams, components); } +template +[[nodiscard]] constexpr size_t GetSQ8StoredDataSize(size_t dim, bool with_norm) { + static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP); + + // WithNorm is a template parameter, so dispatch the runtime flag to the two instantiations. + return with_norm ? sq8::storage_bytes_count(dim) + : sq8::storage_bytes_count(dim); +} + +// Alignment required by a query blob of type DataType. Per the asymmetric-types contract in +// spaces.h, the hint returned alongside an asymmetric distance function describes its first +// (storage) operand, so the query side must be obtained from the symmetric dispatcher for the +// query's own type. Only that hint is wanted here, never the function it returns, so the call is +// contained in this adapter instead of leaving a discarded value at the call site. +template +[[nodiscard]] unsigned char GetQueryAlignment(VecSimMetric metric, size_t dim) { + unsigned char alignment = 0; + spaces::GetDistFunc(metric, dim, &alignment); + return alignment; +} + +// Helper to build an SQ8-quantized HNSW index given compile-time DataType and Metric. +template +VecSimIndex *NewIndex_SQ8(const HNSWParams *hnswParams, AbstractIndexInitParams abstractInitParams, + const float *mean_ptr) { + auto &allocator = abstractInitParams.allocator; + const size_t dim = abstractInitParams.dim; + const bool with_norm = mean_ptr != nullptr; + unsigned char storage_alignment = 0, asym_storage_alignment = 0; + + // Override blob size for the SQ8 storage layout. + abstractInitParams.storedDataSize = GetSQ8StoredDataSize(dim, with_norm); + + // Symmetric: both stored vectors are SQ8 blobs. + auto sym_func = spaces::GetDistFunc(Metric, dim, &storage_alignment); + // Asymmetric: stored vector is SQ8 blob, query is DataType. + auto asym_func = + spaces::GetDistFunc(Metric, dim, &asym_storage_alignment); + // Both hints describe the same stored blob, so they must be combined rather than overwritten. + storage_alignment = spaces::combineAlignments(storage_alignment, asym_storage_alignment); + // Queries stay in DataType and are compared against stored blobs by asym_func. + const unsigned char query_alignment = GetQueryAlignment(Metric, dim); + + PreprocessorInterface *pp = nullptr; + IndexCalculatorInterface *calc = nullptr; + + if (with_norm) { + // Mean-centered SQ8 quantization with norm correction. + vecsim_stl::vector mean_vec(allocator); + mean_vec.assign(mean_ptr, mean_ptr + dim); + + float mean_sum_squares = 0.0f; + for (float v : mean_vec) { + mean_sum_squares += v * v; + } + + pp = new (allocator) QuantPreprocessor(allocator, dim, mean_vec); + calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_squares); + } else { + // Plain SQ8 quantization without mean centering. + pp = new (allocator) QuantPreprocessor(allocator, dim); + // sym_func for storage-storage; asym_func for query-storage. + calc = new (allocator) DistanceCalculatorCommon(allocator, sym_func, asym_func); + } + + auto *container = new (allocator) + MultiPreprocessorsContainer(allocator, query_alignment, storage_alignment); + [[maybe_unused]] const int ret = container->addPreprocessor(pp); + assert(ret != -1 && "SQ8 preprocessor was not added correctly"); + + IndexComponents components{calc, container}; + return NewIndex_ChooseMultiOrSingle(hnswParams, abstractInitParams, + components); +} + VecSimIndex *NewIndex(const VecSimParams *params, bool is_normalized) { const HNSWParams *hnswParams = ¶ms->algoParams.hnswParams; + AbstractIndexInitParams abstractInitParams = VecSimFactory::NewAbstractInitParams(hnswParams, params->logCtx, is_normalized); + if (hnswParams->quantType == VecSimQuant_SQ8) { + if (hnswParams->type != VecSimType_FLOAT32 && hnswParams->type != VecSimType_FLOAT16) { + return NULL; // SQ8 supports FP32 and FP16 only. + } + + VecSimMetric metric = hnswParams->metric; + if (is_normalized && metric == VecSimMetric_Cosine) { + metric = VecSimMetric_IP; + } + + if (metric == VecSimMetric_Cosine) { + return NULL; // SQ8 does not support cosine metric. + } + + const float *mean_ptr = static_cast(hnswParams->quantParams); + + if (hnswParams->type == VecSimType_FLOAT32) { + if (metric == VecSimMetric_L2) { + return NewIndex_SQ8(hnswParams, abstractInitParams, + mean_ptr); + } else if (metric == VecSimMetric_IP) { + return NewIndex_SQ8(hnswParams, abstractInitParams, + mean_ptr); + } + } else if (hnswParams->type == VecSimType_FLOAT16) { + if (metric == VecSimMetric_L2) { + return NewIndex_SQ8(hnswParams, abstractInitParams, + mean_ptr); + } else if (metric == VecSimMetric_IP) { + return NewIndex_SQ8(hnswParams, abstractInitParams, + mean_ptr); + } + } + + // Unreachable today: the checks above leave only FP32/FP16 x L2/IP. Kept so that adding a + // type or metric cannot silently fall through and build an unquantized index instead. + return NULL; + } + if (hnswParams->type == VecSimType_FLOAT32) { IndexComponents indexComponents = CreateIndexComponents( abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); @@ -94,7 +211,27 @@ size_t EstimateInitialSize(const HNSWParams *params, bool is_normalized) { size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize(); size_t est = sizeof(VecSimAllocator) + allocations_overhead; - if (params->type == VecSimType_FLOAT32) { + + if (params->quantType == VecSimQuant_SQ8) { + if (params->type != VecSimType_FLOAT32 && params->type != VecSimType_FLOAT16) { + throw std::invalid_argument("Invalid params->type for VecSimQuant_SQ8"); + } + // Calculator + preprocessor container + preprocessor. + // Use representative types; sizeof is independent of the template parameters. + if (params->quantParams) { // mean provided, WithNorm = true + est += allocations_overhead + + sizeof(DistanceCalculatorWithNorm); + est += allocations_overhead + sizeof(MultiPreprocessorsContainer); + est += allocations_overhead + sizeof(QuantPreprocessor); + est += allocations_overhead + + params->dim * sizeof(float); // mean vector in QuantPreprocessor + } else { + est += allocations_overhead + sizeof(DistanceCalculatorCommon); + est += allocations_overhead + sizeof(MultiPreprocessorsContainer); + est += allocations_overhead + sizeof(QuantPreprocessor); + } + est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); + } else if (params->type == VecSimType_FLOAT32) { est += EstimateComponentsMemory(params->metric, is_normalized); est += EstimateInitialSize_ChooseMultiOrSingle(params->multi); } else if (params->type == VecSimType_FLOAT64) { @@ -125,9 +262,20 @@ size_t EstimateElementSize(const HNSWParams *params) { size_t M = (params->M) ? params->M : HNSW_DEFAULT_M; size_t elementGraphDataSize = sizeof(ElementGraphData) + sizeof(idType) * M * 2; - size_t size_total_data_per_element = - elementGraphDataSize + - VecSimParams_GetStoredDataSize(params->type, params->dim, params->metric); + size_t stored_data_size; + if (params->quantType == VecSimQuant_SQ8) { + bool with_norm = params->quantParams != nullptr; + if (params->metric == VecSimMetric_L2) { + stored_data_size = GetSQ8StoredDataSize(params->dim, with_norm); + } else { + stored_data_size = GetSQ8StoredDataSize(params->dim, with_norm); + } + } else { + stored_data_size = + VecSimParams_GetStoredDataSize(params->type, params->dim, params->metric); + } + + size_t size_total_data_per_element = elementGraphDataSize + stored_data_size; // when reserving space for new labels in the lookup hash table, each entry is a pointer to a // label node (bucket). diff --git a/src/VecSim/index_factories/tiered_factory.cpp b/src/VecSim/index_factories/tiered_factory.cpp index 337db6cc3..c9b129faa 100644 --- a/src/VecSim/index_factories/tiered_factory.cpp +++ b/src/VecSim/index_factories/tiered_factory.cpp @@ -95,6 +95,14 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) { } VecSimIndex *NewIndex(const TieredIndexParams *params) { + // Quantization is not wired into the tiered index yet (MOD-14957). Reject it here rather than + // let it through: the primary index would be built from these params and quantize its storage, + // while NewBFParams does not carry quantType, so the frontend would stay unquantized and the + // two would disagree on the stored blob layout. + if (params->primaryIndexParams->algoParams.hnswParams.quantType != VecSimQuant_NONE) { + return nullptr; + } + // Tiered index that contains HNSW index as primary index VecSimType type = params->primaryIndexParams->algoParams.hnswParams.type; if (type == VecSimType_FLOAT32) { diff --git a/src/VecSim/spaces/computer/preprocessors.h b/src/VecSim/spaces/computer/preprocessors.h index 195be1418..11de96998 100644 --- a/src/VecSim/spaces/computer/preprocessors.h +++ b/src/VecSim/spaces/computer/preprocessors.h @@ -426,8 +426,7 @@ class QuantPreprocessor : public PreprocessorInterface { QuantPreprocessor(std::shared_ptr allocator, size_t dim) requires(!WithNorm) : PreprocessorInterface(allocator), dim(dim), - storage_bytes_count(dim * sizeof(OUTPUT_TYPE) + - sq8::storage_metadata_count() * sizeof(MetadataType)), + storage_bytes_count(sq8::storage_bytes_count(dim)), query_bytes_count(dim * sizeof(DataType) + sq8::query_metadata_count() * sizeof(MetadataType)) {} @@ -436,9 +435,7 @@ class QuantPreprocessor : public PreprocessorInterface { const vecsim_stl::vector &mean_vec) requires(WithNorm) : PreprocessorInterface(allocator), mean(mean_vec), dim(dim), - storage_bytes_count(dim * sizeof(OUTPUT_TYPE) + - sq8::storage_metadata_count() * - sizeof(MetadataType)), + storage_bytes_count(sq8::storage_bytes_count(dim)), query_bytes_count(dim * sizeof(DataType) + sq8::query_metadata_count() * sizeof(MetadataType)) { assert(this->mean.size() == dim && "mean vector size must equal dim"); diff --git a/src/VecSim/types/sq8.h b/src/VecSim/types/sq8.h index c1e9c40b8..9f9e04508 100644 --- a/src/VecSim/types/sq8.h +++ b/src/VecSim/types/sq8.h @@ -47,6 +47,15 @@ struct sq8 { ((WithNorm && Metric == VecSimMetric_IP) ? 1 : 0); } + // Size of a stored SQ8 blob: one byte per dimension, followed by FP32 metadata. Single source + // of truth for the storage layout: every caller that sizes or allocates a stored blob must use + // this, so the layout cannot drift between the preprocessor and the index factories. + template + static constexpr size_t storage_bytes_count(size_t dim) { + return dim * sizeof(value_type) + + storage_metadata_count() * sizeof(float); + } + // Index of x_mean_ip / y_mean_ip in the last slot in metadata array template static constexpr size_t mean_ip_index() { diff --git a/src/VecSim/vec_sim_common.h b/src/VecSim/vec_sim_common.h index fe10a5a0c..26dc2841d 100644 --- a/src/VecSim/vec_sim_common.h +++ b/src/VecSim/vec_sim_common.h @@ -68,6 +68,13 @@ typedef enum { VecSimType_INT64 } VecSimType; +// Quantization type for HNSW indices. +typedef enum { + VecSimQuant_NONE = 0, // No quantization (default). + // 8-bit scalar quantization. Mean normalization is optional, selected by quantParams below. + VecSimQuant_SQ8 = 1, +} VecSimQuantType; + // Algorithm type/library. typedef enum { VecSimAlgo_BF, VecSimAlgo_HNSWLIB, VecSimAlgo_TIERED, VecSimAlgo_SVS } VecSimAlgo; @@ -156,6 +163,10 @@ typedef struct { size_t efConstruction; size_t efRuntime; double epsilon; + VecSimQuantType quantType; // Quantization type. Default: VecSimQuant_NONE. + // For VecSimQuant_SQ8: pointer to float mean[dim], or NULL for zero mean. Read only, never + // retained: the index copies the mean vector during construction. + const void *quantParams; } HNSWParams; typedef struct { diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index c3e1cc987..4eeeae443 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -36,6 +36,7 @@ endif() add_executable(test_hnsw ../utils/test_main_with_timeout.cpp ../utils/mock_thread_pool.cpp test_hnsw.cpp test_hnsw_multi.cpp test_hnsw_tiered.cpp unit_test_utils.cpp) add_executable(test_hnsw_parallel ../utils/test_main_with_timeout.cpp test_hnsw_parallel.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) +add_executable(test_hnsw_sq8 ../utils/test_main_with_timeout.cpp ../utils/mock_thread_pool.cpp test_hnsw_sq8.cpp unit_test_utils.cpp) add_executable(test_bruteforce ../utils/test_main_with_timeout.cpp test_bruteforce.cpp test_bruteforce_multi.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) add_executable(test_allocator ../utils/test_main_with_timeout.cpp test_allocator.cpp ../utils/mock_thread_pool.cpp unit_test_utils.cpp) add_executable(test_spaces ../utils/test_main_with_timeout.cpp test_spaces.cpp) @@ -51,6 +52,7 @@ add_executable(test_svs ../utils/test_main_with_timeout.cpp ../utils/mock_thread target_link_libraries(test_hnsw PUBLIC gtest VectorSimilarity) target_link_libraries(test_hnsw_parallel PUBLIC gtest VectorSimilarity) +target_link_libraries(test_hnsw_sq8 PUBLIC gtest VectorSimilarity) target_link_libraries(test_bruteforce PUBLIC gtest VectorSimilarity) target_link_libraries(test_allocator PUBLIC gtest VectorSimilarity) target_link_libraries(test_spaces PUBLIC gtest VectorSimilarity) @@ -68,6 +70,7 @@ include(GoogleTest) gtest_discover_tests(test_hnsw) gtest_discover_tests(test_hnsw_parallel) +gtest_discover_tests(test_hnsw_sq8) gtest_discover_tests(test_bruteforce) gtest_discover_tests(test_allocator) gtest_discover_tests(test_spaces) diff --git a/tests/unit/test_hnsw_sq8.cpp b/tests/unit/test_hnsw_sq8.cpp new file mode 100644 index 000000000..089e5d5a0 --- /dev/null +++ b/tests/unit/test_hnsw_sq8.cpp @@ -0,0 +1,413 @@ +/* + * 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 + * GNU Affero General Public License v3 (AGPLv3). + */ + +#include "gtest/gtest.h" +#include "VecSim/algorithms/hnsw/hnsw_single.h" +#include "VecSim/types/float16.h" +#include "VecSim/types/sq8.h" +#include "VecSim/vec_sim.h" +#include "unit_test_utils.h" + +#include +#include +#include +#include + +template +struct HNSWSQ8IndexType : IndexType { + static constexpr bool with_quant_params = WithQuantParams; +}; + +using HNSWSQ8DataTypeSet = + ::testing::Types, + HNSWSQ8IndexType, + HNSWSQ8IndexType, + HNSWSQ8IndexType>; + +template +class HNSWSQ8Test : public ::testing::Test { +public: + using data_t = typename index_type_t::data_t; + +protected: + static constexpr float quantization_mean_value = 1.0f; + + static data_t ToDataType(float value) { + if constexpr (std::is_same_v) { + return vecsim_types::FP32_to_FP16(value); + } else { + return value; + } + } + + void SetUp(HNSWParams ¶ms) { + params.type = index_type_t::get_index_type(); + params.quantType = VecSimQuant_SQ8; + if constexpr (index_type_t::with_quant_params) { + quantization_mean.assign(params.dim, quantization_mean_value); + params.quantParams = quantization_mean.data(); + } + VecSimParams vecsim_params = CreateParams(params); + index = VecSimIndex_New(&vecsim_params); + ASSERT_NE(index, nullptr); + dim = params.dim; + } + + void TearDown() override { + if (index) { + VecSimIndex_Free(index); + } + } + + HNSWIndex *CastToHNSW() { + return dynamic_cast *>(index); + } + + void GenerateVector(data_t *out_vec, float initial_value = 0.25f, float step = 0.0f) { + for (size_t i = 0; i < dim; i++) { + out_vec[i] = ToDataType(initial_value + step * static_cast(i)); + } + } + + int GenerateAndAddVector(size_t label, float initial_value = 0.25f, float step = 0.0f) { + std::vector vector(dim); + GenerateVector(vector.data(), initial_value, step); + return VecSimIndex_AddVector(index, vector.data(), label); + } + + void create_index_test(); + void search_by_id_test(); + void search_by_score_test(); + void search_empty_index_test(); + void test_override(); + void test_range_query(); + void test_get_distance(VecSimMetric metric); + void test_batch_iterator_basic(); + + VecSimIndex *index = nullptr; + size_t dim = 0; + std::vector quantization_mean; +}; + +TYPED_TEST_SUITE(HNSWSQ8Test, HNSWSQ8DataTypeSet); + +/* ---------------------------- Create index tests ---------------------------- */ + +template +void HNSWSQ8Test::create_index_test() { + HNSWParams params = {.dim = 40, .M = 16, .efConstruction = 200}; + SetUp(params); + + constexpr float initial_value = 0.5f; + constexpr float step = 1.0f; + ASSERT_EQ(VecSimIndex_IndexSize(index), 0u); + ASSERT_EQ(GenerateAndAddVector(0, initial_value, step), 1); + ASSERT_EQ(VecSimIndex_IndexSize(index), 1u); + + auto *hnsw_index = CastToHNSW(); + ASSERT_NE(hnsw_index, nullptr); + const auto *stored = reinterpret_cast(hnsw_index->getDataByInternalId(0)); + EXPECT_EQ(stored[0], 0); + EXPECT_EQ(stored[dim - 1], 255); + + // The quantized vector is followed by the minimum value and quantization delta. + float stored_min; + float stored_delta; + std::memcpy(&stored_min, stored + dim + sq8::MIN_VAL * sizeof(float), sizeof(float)); + std::memcpy(&stored_delta, stored + dim + sq8::DELTA * sizeof(float), sizeof(float)); + const float expected_min = + initial_value - (index_type_t::with_quant_params ? quantization_mean_value : 0.0f); + EXPECT_FLOAT_EQ(stored_min, expected_min); + EXPECT_FLOAT_EQ(stored_delta, step * static_cast(dim - 1) / 255.0f); + + EXPECT_EQ(index->basicInfo().type, index_type_t::get_index_type()); + EXPECT_EQ(index->basicInfo().algo, VecSimAlgo_HNSWLIB); +} + +TYPED_TEST(HNSWSQ8Test, CreateIndex) { this->create_index_test(); } + +TYPED_TEST(HNSWSQ8Test, RejectStandaloneCosine) { + HNSWParams params = {.type = TypeParam::get_index_type(), + .dim = 4, + .metric = VecSimMetric_Cosine, + .quantType = VecSimQuant_SQ8}; + if constexpr (TypeParam::with_quant_params) { + this->quantization_mean.assign(params.dim, this->quantization_mean_value); + params.quantParams = this->quantization_mean.data(); + } + + VecSimParams vecsim_params = CreateParams(params); + this->index = VecSimIndex_New(&vecsim_params); + EXPECT_EQ(this->index, nullptr); +} + +/* ---------------------------- Size Estimation tests ---------------------------- */ + +TYPED_TEST(HNSWSQ8Test, SizeEstimation) { + constexpr size_t block_size = 256; + HNSWParams params = {.dim = 128, .blockSize = block_size, .M = 64}; + this->SetUp(params); + + // EstimateInitialSize is called after creating the index because index creation normalizes + // the parameters. + EXPECT_EQ(EstimateInitialSize(params), this->index->getAllocationSize()); + + size_t label = 0; + while (this->index->indexSize() < 200 || this->index->indexSize() % block_size != 0) { + ASSERT_EQ(this->GenerateAndAddVector(label, static_cast(label)), 1); + label++; + } + + // Estimate the memory delta of adding a vector that requires a full new block. + const size_t estimation = EstimateElementSize(params) * block_size; + const size_t before = this->index->getAllocationSize(); + ASSERT_EQ(this->GenerateAndAddVector(label, static_cast(label)), 1); + const size_t actual = this->index->getAllocationSize() - before; + + // Check that the actual size is within 1% of the estimation. + EXPECT_GE(estimation, actual * 0.99); + EXPECT_LE(estimation, actual * 1.01); +} + +/* ---------------------------- Functionality tests ---------------------------- */ + +template +void HNSWSQ8Test::search_by_id_test() { + HNSWParams params = { + .dim = 4, .initialCapacity = 200, .M = 16, .efConstruction = 200, .efRuntime = 100}; + SetUp(params); + + for (size_t i = 0; i < 100; i++) { + ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 1); + } + + data_t query[4]; + GenerateVector(query, 50.0f); + // Vector values are equal to their labels, so the closest vectors have labels 45 through 55. + static constexpr size_t expected[] = {45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55}; + auto verify = [&](size_t id, double score, size_t result_index) { + // Results are sorted by ID. + EXPECT_EQ(id, expected[result_index]); + EXPECT_FLOAT_EQ(score, 4.0f * (50.0f - id) * (50.0f - id)); // L2 distance. + }; + runTopKSearchTest(index, query, std::size(expected), verify, nullptr, BY_ID); +} + +TYPED_TEST(HNSWSQ8Test, SearchByID) { this->search_by_id_test(); } + +template +void HNSWSQ8Test::search_by_score_test() { + HNSWParams params = { + .dim = 4, .initialCapacity = 200, .M = 16, .efConstruction = 200, .efRuntime = 100}; + SetUp(params); + + for (size_t i = 0; i < 100; i++) { + ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 1); + } + + data_t query[4]; + GenerateVector(query, 50.0f); + // Vector values are equal to their labels, so results are ordered by distance from label 50. + static constexpr size_t expected[] = {50, 49, 51, 48, 52, 47, 53, 46, 54, 45, 55}; + auto verify = [&](size_t id, double score, size_t result_index) { + EXPECT_EQ(id, expected[result_index]); + EXPECT_FLOAT_EQ(score, 4.0f * (50.0f - id) * (50.0f - id)); + }; + runTopKSearchTest(index, query, std::size(expected), verify); +} + +TYPED_TEST(HNSWSQ8Test, SearchByScore) { this->search_by_score_test(); } + +template +void HNSWSQ8Test::search_empty_index_test() { + HNSWParams params = {.dim = 4, .initialCapacity = 0}; + SetUp(params); + + data_t query[4]; + GenerateVector(query, 50.0f); + + // We do not expect any results. + VecSimQueryReply *reply = VecSimIndex_TopKQuery(index, query, 11, nullptr, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(reply), 0u); + VecSimQueryReply_Free(reply); + + reply = VecSimIndex_RangeQuery(index, query, 1.0, nullptr, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(reply), 0u); + VecSimQueryReply_Free(reply); + + // Add some vectors and remove them all from the index, so it will be empty again. + for (size_t i = 0; i < 100; i++) { + GenerateAndAddVector(i, static_cast(i)); + } + 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); + ASSERT_EQ(VecSimQueryReply_Len(reply), 0u); + VecSimQueryReply_Free(reply); + + reply = VecSimIndex_RangeQuery(index, query, 1.0, nullptr, BY_SCORE); + ASSERT_EQ(VecSimQueryReply_Len(reply), 0u); + VecSimQueryReply_Free(reply); +} + +TYPED_TEST(HNSWSQ8Test, SearchEmptyIndex) { this->search_empty_index_test(); } + +template +void HNSWSQ8Test::test_override() { + constexpr size_t count = 250; + HNSWParams params = { + .dim = 4, .initialCapacity = 100, .M = 8, .efConstruction = 20, .efRuntime = count}; + 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); + } + // Add vectors up to count. + for (size_t i = 100; i < count; i++) { + ASSERT_EQ(GenerateAndAddVector(i, static_cast(i)), 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. + 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)); + }; + runTopKSearchTest(index, query, count, verify); +} + +TYPED_TEST(HNSWSQ8Test, Override) { this->test_override(); } + +template +void HNSWSQ8Test::test_range_query() { + constexpr size_t count = 100; + constexpr size_t close_count = 20; + HNSWParams params = {.dim = 4, .initialCapacity = count, .efRuntime = count}; + SetUp(params); + + constexpr float pivot = 1.0f; + constexpr float value_radius = 1.5f; + std::mt19937 generator(42); + std::uniform_real_distribution distribution(pivot - value_radius, pivot + value_radius); + // Insert close_count vectors near the pivot vector. + for (size_t i = 0; i < close_count; i++) { + GenerateAndAddVector(i, distribution(generator)); + } + // Add the remaining vectors far from the pivot vector. + for (size_t i = close_count; i < count; i++) { + GenerateAndAddVector(i, 5.0f + distribution(generator)); + } + + data_t query[4]; + GenerateVector(query, pivot); + constexpr double max_distance = 4.0 * value_radius * value_radius; + auto verify = [&](size_t id, double score, size_t) { + EXPECT_LT(id, close_count); + EXPECT_LE(score, max_distance); + }; + runRangeQueryTest(index, query, max_distance, verify, close_count, BY_SCORE); +} + +TYPED_TEST(HNSWSQ8Test, RangeQuery) { this->test_range_query(); } + +template +void HNSWSQ8Test::test_get_distance(VecSimMetric metric) { + HNSWParams params = {.dim = 4, .metric = metric, .initialCapacity = 1}; + SetUp(params); + + ASSERT_EQ(GenerateAndAddVector(0, 0.25f, 0.25f), 1); + data_t query[4]; + GenerateVector(query, 0.5f, 0.25f); + auto processed_query = CastToHNSW()->preprocessQuery(query); + const double expected = metric == VecSimMetric_L2 ? 0.25 : -1.5; + // Values were chosen so the expected distances can be calculated exactly. + EXPECT_NEAR(VecSimIndex_GetDistanceFrom_Unsafe(index, 0, processed_query.get()), expected, + 1e-5); +} + +TYPED_TEST(HNSWSQ8Test, GetDistanceL2) { this->test_get_distance(VecSimMetric_L2); } +TYPED_TEST(HNSWSQ8Test, GetDistanceIP) { this->test_get_distance(VecSimMetric_IP); } + +/* ---------------------------- Batch iterator tests ---------------------------- */ + +template +void HNSWSQ8Test::test_batch_iterator_basic() { + constexpr size_t count = 250; + constexpr size_t batch_size = 5; + 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); + } + + data_t query[4]; + GenerateVector(query, static_cast(count)); + VecSimBatchIterator *iterator = VecSimBatchIterator_New(index, query, nullptr); + ASSERT_NE(iterator, nullptr); + + // Get the five largest remaining labels in each iteration. Since vector values equal their + // labels, this is also their order by distance from the query vector. + size_t iteration = 0; + while (VecSimBatchIterator_HasNext(iterator)) { + auto verify = [&](size_t id, double, size_t result_index) { + EXPECT_EQ(id, count - iteration * batch_size - result_index - 1); + }; + runBatchIteratorSearchTest(iterator, batch_size, verify); + iteration++; + } + EXPECT_EQ(iteration, count / batch_size); + VecSimBatchIterator_Free(iterator); +} + +TYPED_TEST(HNSWSQ8Test, BatchIteratorBasic) { this->test_batch_iterator_basic(); } + +// SQ8 quantizes to uint8 with FP32 metadata and only has kernels for FP32 and FP16 sources, so +// every other data type must be rejected outright rather than produce an index. Note that +// EstimateElementSize deliberately does not re-check this: like VecSimParams_GetStoredDataSize on +// the unquantized path, it answers for whatever params it is handed, so index creation is the +// boundary that enforces the supported set. +TEST(HNSWSQ8ParamsTest, RejectsUnsupportedDataType) { + for (auto type : {VecSimType_FLOAT64, VecSimType_BFLOAT16, VecSimType_INT8, VecSimType_UINT8}) { + HNSWParams hnsw_params = { + .type = type, .dim = 4, .metric = VecSimMetric_L2, .quantType = VecSimQuant_SQ8}; + VecSimParams params = CreateParams(hnsw_params); + + EXPECT_EQ(VecSimIndex_New(¶ms), nullptr) << "data type " << type; + } +} + +// SQ8 is not wired into the tiered index yet (MOD-14957), so the tiered factory must reject it +// instead of building a quantized primary index against an unquantized frontend. Without the +// guard this aborts on a debug build and silently mismatches the two blob layouts on a release +// one. MOD-14957 should replace this expectation rather than delete it. +TEST(HNSWSQ8TieredTest, RejectsQuantizedTieredIndex) { + HNSWParams hnsw_params = {.type = VecSimType_FLOAT32, + .dim = 4, + .metric = VecSimMetric_L2, + .quantType = VecSimQuant_SQ8}; + VecSimParams primary_params = CreateParams(hnsw_params); + // No job queue or thread pool is needed: the factory rejects these params before it reaches + // anything that would use them. + TieredIndexParams tiered_params = {.primaryIndexParams = &primary_params}; + VecSimParams params = CreateParams(tiered_params); + + EXPECT_EQ(VecSimIndex_New(¶ms), nullptr); +}