diff --git a/src/VecSim/algorithms/hnsw/hnsw_multi.h b/src/VecSim/algorithms/hnsw/hnsw_multi.h index db978d28c..a1e81e83f 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_multi.h +++ b/src/VecSim/algorithms/hnsw/hnsw_multi.h @@ -118,6 +118,11 @@ class HNSWIndex_Multi : public HNSWIndex { int addVector(const void *vector_data, labelType label) override; vecsim_stl::vector markDelete(labelType label) override; double getDistanceFrom_Unsafe(labelType label, const void *vector_data) const override { + // See the note in hnsw_single.h: a quantized index cannot answer a raw-blob distance query + // without reading past the caller's vector. + if (this->isQuantized) { + return INVALID_SCORE; + } return getDistanceFromInternal(label, vector_data); } int removeLabel(labelType label) override { return labelLookup.erase(label); } diff --git a/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h b/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h index 1666b7b37..b895b580c 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h +++ b/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h @@ -41,6 +41,17 @@ HNSWIndex::HNSWIndex(std::ifstream &input, const HNSWParams template void HNSWIndex::saveIndexIMP(std::ofstream &output) { + // The V4 format records type, dim and metric, but neither quantType nor the mean vector, and + // the loading path always builds unquantized components. A saved quantized index would + // therefore reload with the wrong stride and consume graph bytes as vector data. Refuse instead + // of emitting a file that cannot be decoded. Note the caller has already written the encoding + // version by this point, so a rejected save leaves a stub file behind; that still fails closed + // on load, unlike a full file with a layout the loader misreads. MOD-14957 adds SQ8 + // serialization. + if (this->isQuantized) { + throw std::runtime_error( + "Cannot save index: serialization of quantized indexes is not supported"); + } this->saveIndexFields(output); this->saveGraph(output); } diff --git a/src/VecSim/algorithms/hnsw/hnsw_single.h b/src/VecSim/algorithms/hnsw/hnsw_single.h index e6892dcab..9f1a3a372 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_single.h +++ b/src/VecSim/algorithms/hnsw/hnsw_single.h @@ -88,6 +88,13 @@ class HNSWIndex_Single : public HNSWIndex { int addVector(const void *vector_data, labelType label) override; vecsim_stl::vector markDelete(labelType label) override; double getDistanceFrom_Unsafe(labelType label, const void *vector_data) const override { + // The public API documents vector_data as a raw dim-by-type vector, but a quantized index's + // kernels read query metadata appended past that, so honouring the documented contract here + // would read out of bounds. There is no public API for producing a quantized query blob; + // MOD-14958 owns that decision. Report "no answer" rather than read past the caller's blob. + if (this->isQuantized) { + return INVALID_SCORE; + } return getDistanceFromInternal(label, vector_data); } int removeLabel(labelType label) override { return labelLookup.erase(label); } diff --git a/src/VecSim/index_factories/hnsw_factory.cpp b/src/VecSim/index_factories/hnsw_factory.cpp index d577f57a1..f68421c5e 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,141 @@ 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); + abstractInitParams.isQuantized = true; + + // 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); + + // Mean-centred FP16 L2 is not supported: QuantPreprocessor centres the query and narrows + // the result back into the FP16 query body, while storage keeps its centred min/delta in + // FP32. The two then disagree, so an identical vector and query pair yields a non-zero + // distance (mean 10000 gives a per-component error of 1.0), and a large enough mean + // overflows FP16 to infinity. Enabling this needs an asymmetric kernel that takes an FP32 + // centred query. + if (hnswParams->type == VecSimType_FLOAT16 && mean_ptr != nullptr && + metric == VecSimMetric_L2) { + return NULL; + } + + 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. The assert makes a + // debug build shout if a new type or metric ever reaches here, and the return keeps a + // release build failing closed rather than falling through and silently building an + // unquantized index instead. + assert(false && "unhandled SQ8 data type and metric combination"); + return NULL; + } + if (hnswParams->type == VecSimType_FLOAT32) { IndexComponents indexComponents = CreateIndexComponents( abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized); @@ -94,7 +225,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 +276,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/src/VecSim/vec_sim_index.h b/src/VecSim/vec_sim_index.h index dcf2ac30d..f57993bd8 100644 --- a/src/VecSim/vec_sim_index.h +++ b/src/VecSim/vec_sim_index.h @@ -48,6 +48,10 @@ struct AbstractIndexInitParams { size_t blockSize; bool multi; bool isDisk; // Whether the index stores vectors on disk + // Whether stored vectors are quantized. A quantized index's blobs, both stored and query, carry + // metadata the distance kernels read, so a caller's raw dim-by-type vector is not a usable + // query blob for it. + bool isQuantized; void *logCtx; size_t inputBlobSize; }; @@ -82,6 +86,7 @@ struct VecSimIndexAbstract : public VecSimIndexInterface { mutable VecSearchMode lastMode; // The last search mode in RediSearch (used for debug/testing). bool isMulti; // Determines if the index should multi-index or not. bool isDisk; // Whether the index stores vectors on disk. + bool isQuantized; // Whether stored vectors are quantized. void *logCallbackCtx; // Context for the log callback. RawDataContainer *vectors; // The raw vectors data container. private: @@ -125,8 +130,8 @@ struct VecSimIndexAbstract : public VecSimIndexInterface { : VecSimIndexInterface(params.allocator), dim(params.dim), vecType(params.vecType), metric(params.metric), blockSize(params.blockSize ? params.blockSize : DEFAULT_BLOCK_SIZE), lastMode(EMPTY_MODE), - isMulti(params.multi), isDisk(params.isDisk), logCallbackCtx(params.logCtx), - indexCalculator(components.indexCalculator), + isMulti(params.multi), isDisk(params.isDisk), isQuantized(params.isQuantized), + logCallbackCtx(params.logCtx), indexCalculator(components.indexCalculator), storedDistanceDispatch( components.indexCalculator ? components.indexCalculator->getDistanceDispatch(DistanceMode::StoredToStored) 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..079dca377 --- /dev/null +++ b/tests/unit/test_hnsw_sq8.cpp @@ -0,0 +1,493 @@ +/* + * 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; +}; + +// FLOAT16 with a mean vector is absent on purpose: the functional tests below all use L2, and +// mean-centred FP16 L2 is rejected at construction (see HNSWFactory::NewIndex). That combination is +// covered explicitly by HNSWSQ8ParamsTest.RejectsMeanCenteredFP16L2 instead. +using HNSWSQ8DataTypeSet = + ::testing::Types, + 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); + + // Values were chosen so the expected distances can be calculated exactly. Comparing against a + // stored SQ8 blob needs a preprocessed query, which only the index itself can produce. + auto *hnsw_index = CastToHNSW(); + auto processed_query = hnsw_index->preprocessQuery(query); + const double expected = metric == VecSimMetric_L2 ? 0.25 : -1.5; + EXPECT_NEAR( + hnsw_index->calcDistanceForQuery(hnsw_index->getDataByInternalId(0), processed_query.get()), + expected, 1e-5); + + // The public API documents blob as a raw dim-by-type vector, which is not a usable query blob + // for a quantized index: the kernels read query metadata appended past it. It must report no + // answer rather than read past the caller's buffer. + EXPECT_TRUE(std::isnan(VecSimIndex_GetDistanceFrom_Unsafe(index, 0, query))); +} + +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; + } +} + +// Mean-centred FP16 with L2 must be rejected: QuantPreprocessor narrows the centred query back into +// the FP16 query body while storage keeps its centred min/delta in FP32, so identical vector and +// query pairs diverge and a large mean overflows FP16 to infinity. The same combination with IP is +// supported, because that path does not centre the query. +TEST(HNSWSQ8ParamsTest, RejectsMeanCenteredFP16L2) { + std::vector mean(4, 1.0f); + + HNSWParams l2 = {.type = VecSimType_FLOAT16, + .dim = 4, + .metric = VecSimMetric_L2, + .quantType = VecSimQuant_SQ8, + .quantParams = mean.data()}; + VecSimParams l2_params = CreateParams(l2); + EXPECT_EQ(VecSimIndex_New(&l2_params), nullptr); + + HNSWParams ip = {.type = VecSimType_FLOAT16, + .dim = 4, + .metric = VecSimMetric_IP, + .quantType = VecSimQuant_SQ8, + .quantParams = mean.data()}; + VecSimParams ip_params = CreateParams(ip); + VecSimIndex *ip_index = VecSimIndex_New(&ip_params); + ASSERT_NE(ip_index, nullptr); + VecSimIndex_Free(ip_index); +} + +// Serialization does not record quantType or the mean vector, and the loading path always builds +// unquantized components, so saving a quantized index would produce a file the loader misreads. +// saveIndex must refuse rather than emit one. +TYPED_TEST(HNSWSQ8Test, RejectsSerialization) { + HNSWParams params = {.dim = 4, .initialCapacity = 1}; + this->SetUp(params); + ASSERT_EQ(this->GenerateAndAddVector(0, 0.25f, 0.25f), 1); + + const auto file_name = std::string(getenv("ROOT")) + "/tests/unit/sq8_should_not_be_written"; + EXPECT_THROW(this->CastToHNSW()->saveIndex(file_name), std::runtime_error); + std::remove(file_name.c_str()); +} + +// Every other functional test uses L2, so without this the symmetric SQ8-to-SQ8 IP kernel that +// graph construction selects for an IP index would never run. Vectors vary per component as well as +// per label, so quantization does not collapse into the degenerate min == max branch. +TYPED_TEST(HNSWSQ8Test, GraphConstructionIP) { + constexpr size_t n = 100; + constexpr size_t dim = 16; + HNSWParams params = { + .dim = dim, .metric = VecSimMetric_IP, .initialCapacity = n, .M = 16, .efRuntime = n}; + this->SetUp(params); + + // Each label i gets a vector whose components ramp from i upward, so no two vectors share a + // quantization range and every vector has a non-zero delta. + for (size_t i = 0; i < n; i++) { + ASSERT_EQ(this->GenerateAndAddVector(i, static_cast(i) * 0.5f, 0.25f), 1); + } + ASSERT_EQ(VecSimIndex_IndexSize(this->index), n); + + // This is plain inner product, not cosine: the distance is 1 - IP, so the closest vector is the + // one with the largest projection onto the query rather than the query's own twin. Every vector + // and the query are positive and magnitude grows with the label, so IP is strictly increasing + // in the label and results must come back from the highest label downward. + std::vector query(dim); + this->GenerateVector(query.data(), 1.0f, 0.25f); + + auto verify = [&](size_t id, double, size_t result_index) { + EXPECT_EQ(id, n - 1 - result_index); + }; + runTopKSearchTest(this->index, query.data(), 10, verify); +} + +// 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); +}