[MOD-14956] Add SQ8 quantization support for HNSW index

src/VecSim/algorithms/hnsw/hnsw_multi.h

@@ -118,6 +118,11 @@ class HNSWIndex_Multi : public HNSWIndex<DataType, DistType> {
     int addVector(const void *vector_data, labelType label) override;
     vecsim_stl::vector<idType> 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); }

src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h

@@ -41,6 +41,17 @@ HNSWIndex<DataType, DistType>::HNSWIndex(std::ifstream &input, const HNSWParams
 
 template <typename DataType, typename DistType>
 void HNSWIndex<DataType, DistType>::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, and worse, it has
+    // already truncated whatever was at that path. Both are tracked separately; whoever adds
+    // quantized serialization should move this check ahead of the file being opened.
+    if (this->isQuantized) {
+        throw std::runtime_error(
+            "Cannot save index: serialization of quantized indexes is not supported");
+    }
     this->saveIndexFields(output);
     this->saveGraph(output);
 }

src/VecSim/algorithms/hnsw/hnsw_single.h

@@ -88,6 +88,14 @@ class HNSWIndex_Single : public HNSWIndex<DataType, DistType> {
     int addVector(const void *vector_data, labelType label) override;
     vecsim_stl::vector<idType> 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,
+        // and adding one is a separate decision. Report "no answer" rather than read past the
+        // caller's blob. See the note on VecSimIndex_GetDistanceFrom_Unsafe in vec_sim.h.
+        if (this->isQuantized) {
+            return INVALID_SCORE;
+        }
         return getDistanceFromInternal(label, vector_data);
     }
     int removeLabel(labelType label) override { return labelLookup.erase(label); }

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,166 @@ NewIndex_ChooseMultiOrSingle(const HNSWParams *params,
             HNSWIndex_Single<DataType, DistType>(params, abstractInitParams, components);
 }
 
+template <VecSimMetric Metric>
+[[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<Metric, true>(dim)
+                     : sq8::storage_bytes_count<Metric, false>(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 <typename DataType>
+[[nodiscard]] unsigned char GetQueryAlignment(VecSimMetric metric, size_t dim) {
+    unsigned char alignment = 0;
+    spaces::GetDistFunc<DataType, float>(metric, dim, &alignment);
+    return alignment;
+}
+
+// The metric the SQ8 components are actually built for. A Cosine index whose blobs are already
+// normalized by someone else (the frontend of a tiered index) is built as IP.
+[[nodiscard]] constexpr VecSimMetric ResolveSQ8Metric(VecSimMetric metric, bool is_normalized) {
+    return (is_normalized && metric == VecSimMetric_Cosine) ? VecSimMetric_IP : metric;
+}
+
+// Single definition of the SQ8 parameter combinations this factory can build, taking the metric
+// already resolved by ResolveSQ8Metric. NewIndex fails closed when this returns false, and
+// EstimateInitialSize rejects the same set, so it cannot report a size for parameters that cannot
+// produce an index.
+[[nodiscard]] constexpr bool SQ8ParamsSupported(VecSimType type, VecSimMetric resolved_metric,
+                                                bool with_mean) {
+    // Kernels exist for FP32 and FP16 sources only.
+    if (type != VecSimType_FLOAT32 && type != VecSimType_FLOAT16) {
+        return false;
+    }
+    // Only L2 and IP have SQ8 kernels. Cosine needs a normalization step the SQ8 preprocessor does
+    // not perform, and anything outside the enum has to be rejected here as well: the dispatch in
+    // NewIndex would otherwise fall through to its unreachable-branch assert and abort an
+    // assertions-enabled host, where the unquantized path throws and is caught by VecSimIndex_New.
+    if (resolved_metric != VecSimMetric_L2 && resolved_metric != VecSimMetric_IP) {
+        return false;
+    }
+    // Mean-centred FP16 L2 loses correctness: 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. Only the
+    // WithNorm && L2 branch centres the query, so FP16 with a mean and IP is unaffected.
+    if (type == VecSimType_FLOAT16 && with_mean && resolved_metric == VecSimMetric_L2) {
+        return false;
+    }
+    return true;
+}
+
+// Helper to build an SQ8-quantized HNSW index given compile-time DataType and Metric.
+template <typename DataType, VecSimMetric Metric>
+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<Metric>(dim, with_norm);
+    abstractInitParams.isQuantized = true;
+
+    // Symmetric: both stored vectors are SQ8 blobs.
+    auto sym_func = spaces::GetDistFunc<sq8, float>(Metric, dim, &storage_alignment);
+    // Asymmetric: stored vector is SQ8 blob, query is DataType.
+    auto asym_func =
+        spaces::GetDistFunc<sq8, float, DataType>(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<DataType>(Metric, dim);
+
+    PreprocessorInterface *pp = nullptr;
+    IndexCalculatorInterface<float> *calc = nullptr;
+
+    if (with_norm) {
+        // Mean-centered SQ8 quantization with norm correction.
+        vecsim_stl::vector<float> 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<DataType, Metric, true>(allocator, dim, mean_vec);
+        calc = new (allocator) DistanceCalculatorWithNorm<DataType, float, Metric>(
+            allocator, asym_func, sym_func, mean_sum_squares);
+    } else {
+        // Plain SQ8 quantization without mean centering.
+        pp = new (allocator) QuantPreprocessor<DataType, Metric>(allocator, dim);
+        // sym_func for storage-storage; asym_func for query-storage.
+        calc = new (allocator) DistanceCalculatorCommon<float>(allocator, sym_func, asym_func);
+    }
+
+    auto *container = new (allocator)
+        MultiPreprocessorsContainer<DataType, 1>(allocator, query_alignment, storage_alignment);
+    [[maybe_unused]] const int ret = container->addPreprocessor(pp);
+    assert(ret != -1 && "SQ8 preprocessor was not added correctly");
+
+    IndexComponents<DataType, float> components{calc, container};
+    return NewIndex_ChooseMultiOrSingle<DataType, float>(hnswParams, abstractInitParams,
+                                                         components);
+}
+
 VecSimIndex *NewIndex(const VecSimParams *params, bool is_normalized) {
     const HNSWParams *hnswParams = &params->algoParams.hnswParams;
     AbstractIndexInitParams abstractInitParams =
         VecSimFactory::NewAbstractInitParams(hnswParams, params->logCtx, is_normalized);
 
+    if (hnswParams->quantType != VecSimQuant_NONE) {
+        // Any quantization type this factory does not implement must fail closed. Falling through
+        // to the plain path below would silently build a full-precision index for a caller that
+        // asked for a quantized one. Unreachable while VecSimQuantType holds only NONE and SQ8, and
+        // not covered by a test because forming an out-of-range enumerator is undefined behaviour;
+        // the guard is what keeps adding SQ4 to the enum from becoming a silent fallthrough.
+        if (hnswParams->quantType != VecSimQuant_SQ8) {
+            return NULL;
+        }
+
+        const VecSimMetric metric = ResolveSQ8Metric(hnswParams->metric, is_normalized);
+        const float *mean_ptr = static_cast<const float *>(hnswParams->quantParams);
+
+        if (!SQ8ParamsSupported(hnswParams->type, metric, mean_ptr != nullptr)) {
+            return NULL;
+        }
+
+        if (hnswParams->type == VecSimType_FLOAT32) {
+            if (metric == VecSimMetric_L2) {
+                return NewIndex_SQ8<float, VecSimMetric_L2>(hnswParams, abstractInitParams,
+                                                            mean_ptr);
+            } else if (metric == VecSimMetric_IP) {
+                return NewIndex_SQ8<float, VecSimMetric_IP>(hnswParams, abstractInitParams,
+                                                            mean_ptr);
+            }
+        } else if (hnswParams->type == VecSimType_FLOAT16) {
+            if (metric == VecSimMetric_L2) {
+                return NewIndex_SQ8<float16, VecSimMetric_L2>(hnswParams, abstractInitParams,
+                                                              mean_ptr);
+            } else if (metric == VecSimMetric_IP) {
+                return NewIndex_SQ8<float16, VecSimMetric_IP>(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<float, float> indexComponents = CreateIndexComponents<float, float>(
             abstractInitParams.allocator, hnswParams->metric, hnswParams->dim, is_normalized);
@@ -94,7 +250,31 @@ 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_NONE) {
+        // Reject exactly what NewIndex rejects. Reporting a size for a combination that cannot be
+        // built lets a caller size its capacity from an index it will then fail to create.
+        if (params->quantType != VecSimQuant_SQ8 ||
+            !SQ8ParamsSupported(params->type, ResolveSQ8Metric(params->metric, is_normalized),
+                                params->quantParams != nullptr)) {
+            throw std::invalid_argument("Unsupported quantization params for HNSW index");
+        }
+        // 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<float, float, VecSimMetric_L2>);
+            est += allocations_overhead + sizeof(MultiPreprocessorsContainer<float, 1>);
+            est += allocations_overhead + sizeof(QuantPreprocessor<float, VecSimMetric_L2, true>);
+            est += allocations_overhead +
+                   params->dim * sizeof(float); // mean vector in QuantPreprocessor
+        } else {
+            est += allocations_overhead + sizeof(DistanceCalculatorCommon<float>);
+            est += allocations_overhead + sizeof(MultiPreprocessorsContainer<float, 1>);
+            est += allocations_overhead + sizeof(QuantPreprocessor<float, VecSimMetric_L2>);
+        }
+        est += EstimateInitialSize_ChooseMultiOrSingle<float>(params->multi);
+    } else if (params->type == VecSimType_FLOAT32) {
         est += EstimateComponentsMemory<float, float>(params->metric, is_normalized);
         est += EstimateInitialSize_ChooseMultiOrSingle<float>(params->multi);
     } else if (params->type == VecSimType_FLOAT64) {
@@ -125,9 +305,25 @@ 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;
+    // Unlike EstimateInitialSize, this does not reject unsupported combinations: the return type is
+    // size_t with no sentinel, and VecSimIndex_EstimateElementSize is extern "C", so throwing here
+    // would carry an exception into the C host. It therefore answers for whatever params it is
+    // handed, exactly as VecSimParams_GetStoredDataSize does on the unquantized path. Index
+    // creation is the boundary that enforces the supported set (see SQ8ParamsSupported).
+    if (params->quantType == VecSimQuant_SQ8) {
+        bool with_norm = params->quantParams != nullptr;
+        if (params->metric == VecSimMetric_L2) {
+            stored_data_size = GetSQ8StoredDataSize<VecSimMetric_L2>(params->dim, with_norm);
+        } else {
+            stored_data_size = GetSQ8StoredDataSize<VecSimMetric_IP>(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).

src/VecSim/index_factories/tiered_factory.cpp

@@ -66,6 +66,13 @@ inline VecSimIndex *NewIndex(const TieredIndexParams *params) {
 inline size_t EstimateInitialSize(const TieredIndexParams *params) {
     HNSWParams hnsw_params = params->primaryIndexParams->algoParams.hnswParams;
 
+    // NewIndex below rejects quantization until the tiered index handles it, so refuse to size it
+    // too. The primary index on its own would accept these params, so without this check the
+    // estimate reports a number for an index that cannot be created.
+    if (hnsw_params.quantType != VecSimQuant_NONE) {
+        throw std::invalid_argument("Quantization is not supported for tiered HNSW indexes");
+    }
+
     // Add size estimation of VecSimTieredIndex sub indexes.
     // Normalization is done by the frontend index.
     size_t est = HNSWFactory::EstimateInitialSize(&hnsw_params, true);
@@ -95,6 +102,14 @@ inline size_t EstimateInitialSize(const TieredIndexParams *params) {
 }
 
 VecSimIndex *NewIndex(const TieredIndexParams *params) {
+    // Quantization is not wired into the tiered index yet. 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) {
@@ -233,6 +248,8 @@ size_t EstimateInitialSize(const TieredIndexParams *params) {
 }
 
 size_t EstimateElementSize(const TieredIndexParams *params) {
+    // Deliberately not validated here, unlike EstimateInitialSize above: see the note in
+    // HNSWFactory::EstimateElementSize on why this function has no error channel.
     size_t est = 0;
     if (params->primaryIndexParams->algo == VecSimAlgo_HNSWLIB) {
         est = HNSWFactory::EstimateElementSize(&params->primaryIndexParams->algoParams.hnswParams);

src/VecSim/spaces/computer/preprocessors.h

@@ -426,8 +426,7 @@ class QuantPreprocessor : public PreprocessorInterface {
     QuantPreprocessor(std::shared_ptr<VecSimAllocator> allocator, size_t dim)
         requires(!WithNorm)
         : PreprocessorInterface(allocator), dim(dim),
-          storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
-                              sq8::storage_metadata_count<Metric>() * sizeof(MetadataType)),
+          storage_bytes_count(sq8::storage_bytes_count<Metric>(dim)),
           query_bytes_count(dim * sizeof(DataType) +
                             sq8::query_metadata_count<Metric>() * sizeof(MetadataType)) {}
 
@@ -436,9 +435,7 @@ class QuantPreprocessor : public PreprocessorInterface {
                       const vecsim_stl::vector<float> &mean_vec)
         requires(WithNorm)
         : PreprocessorInterface(allocator), mean(mean_vec), dim(dim),
-          storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
-                              sq8::storage_metadata_count<Metric, WithNorm>() *
-                                  sizeof(MetadataType)),
+          storage_bytes_count(sq8::storage_bytes_count<Metric, WithNorm>(dim)),
           query_bytes_count(dim * sizeof(DataType) +
                             sq8::query_metadata_count<Metric, WithNorm>() * sizeof(MetadataType)) {
         assert(this->mean.size() == dim && "mean vector size must equal dim");