[MOD-17189] Share SQ8 training state for RAM and disk indexes

src/VecSim/algorithms/hnsw/hnsw_tiered.h

@@ -18,7 +18,7 @@
 #include <utility>
 
 #include "VecSim/algorithms/brute_force/brute_force_single.h"
-#include "VecSim/spaces/computer/preprocessors.h"
+#include "VecSim/quantization/sq8_training_state.h"
 #include "VecSim/vec_sim_tiered_index.h"
 #include "hnsw.h"
 
@@ -70,7 +70,6 @@ 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
@@ -150,17 +149,7 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
 
     // Writer-owned state exists only while SQ8 vectors are accumulating in FLAT.
     // Reads and maintenance always use the real, initially empty backend.
-    struct SQAccumulationState {
-        size_t normalizationSetSize;
-        vecsim_stl::vector<double> runningSumVec;
-
-        SQAccumulationState(const std::shared_ptr<VecSimAllocator> &allocator, size_t dim,
-                            size_t normalizationSetSize)
-            : normalizationSetSize(
-                  std::min(normalizationSetSize, MAX_QUANT_NORMALIZATION_SET_SIZE)),
-              runningSumVec(dim, 0.0, allocator) {}
-    };
-    std::optional<SQAccumulationState> sqAccumulationState;
+    std::optional<SQ8TrainingState> sqAccumulationState;
 
 #ifdef BUILD_TESTS
     std::function<void()> beforeQuantizationFinalization;
@@ -171,25 +160,16 @@ class TieredHNSWIndex : public VecSimTieredIndex<DataType, DistType> {
 
     void addToSum(std::span<const DataType> vector) {
         if constexpr (QuantInput<DataType>) {
-            auto &sum = sqAccumulationState->runningSumVec;
-            assert(vector.size() == sum.size());
-            for (size_t i = 0; i < vector.size(); ++i) {
-                sum[i] += to_fp32(vector[i]);
-            }
+            sqAccumulationState->add(vector);
         }
     }
 
     void subtractFromSum(std::span<const DataType> vector) {
         if constexpr (QuantInput<DataType>) {
-            auto &sum = sqAccumulationState->runningSumVec;
-            assert(vector.size() == sum.size());
-            for (size_t i = 0; i < vector.size(); ++i) {
-                sum[i] -= to_fp32(vector[i]);
-            }
+            sqAccumulationState->subtract(vector);
         }
     }
 
-    [[nodiscard]] vecsim_stl::vector<float> calculateQuantizationMean(size_t count) const;
     int addVectorDuringAccumulation(const void *blob, labelType label);
     int deleteVectorDuringAccumulation(labelType label);
     void finalizeQuantizationAndSubmitJobs(size_t count);
@@ -1130,7 +1110,8 @@ int TieredHNSWIndex<DataType, DistType>::addVectorDuringAccumulation(const void
             }
         }
         const size_t count = this->frontendIndex->indexSize();
-        const bool should_finalize = count >= sqAccumulationState->normalizationSetSize;
+        assert(sqAccumulationState->getCount() == count);
+        const bool should_finalize = sqAccumulationState->shouldFinalize();
         return std::tuple{result, count, should_finalize};
     }();
 
@@ -1158,25 +1139,17 @@ int TieredHNSWIndex<DataType, DistType>::deleteVectorDuringAccumulation(labelTyp
     for (const auto &[new_id, previous] : updated_ids) {
         updateInsertJobInternalId(previous.first, new_id, previous.second);
     }
+    assert(sqAccumulationState->getCount() == this->frontendIndex->indexSize());
     return removed;
 }
 
-template <typename DataType, typename DistType>
-vecsim_stl::vector<float>
-TieredHNSWIndex<DataType, DistType>::calculateQuantizationMean(size_t count) const {
-    assert(sqAccumulationState);
-    assert(count >= sqAccumulationState->normalizationSetSize);
-    const auto &sum = sqAccumulationState->runningSumVec;
-    vecsim_stl::vector<float> mean(sum.size(), this->allocator);
-    for (size_t i = 0; i < sum.size(); ++i) {
-        mean[i] = static_cast<float>(sum[i] / static_cast<double>(count));
-    }
-    return mean;
-}
-
 template <typename DataType, typename DistType>
 void TieredHNSWIndex<DataType, DistType>::finalizeQuantizationAndSubmitJobs(size_t count) {
-    const auto mean = calculateQuantizationMean(count);
+    assert(sqAccumulationState);
+    assert(sqAccumulationState->getCount() == count);
+    assert(sqAccumulationState->getCount() == this->frontendIndex->indexSize());
+    assert(sqAccumulationState->shouldFinalize());
+    const auto mean = sqAccumulationState->calculateMean();
     // Allocate the snapshot before finishing accumulation. Synchronous execution removes jobs
     // from labelToInsertJobs, so iterating the map while dispatching would invalidate iterators.
     vecsim_stl::vector<AsyncJob *> jobs(this->allocator);

src/VecSim/quantization/sq8_training_state.h

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2006-Present, Redis Ltd.
+ * All rights reserved.
+ *
+ * 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).
+ */
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <span>
+
+#include "VecSim/spaces/computer/preprocessors.h"
+#include "VecSim/utils/vecsim_stl.h"
+
+static const size_t MAX_QUANT_NORMALIZATION_SET_SIZE = 100 * DEFAULT_BLOCK_SIZE;
+
+/** Writer-owned SQ8 training state. Inputs must be the stored, preprocessed vectors.
+ * Callers synchronize mutations and subtract a vector before overwriting or deleting it.
+ */
+class SQ8TrainingState {
+    size_t normalizationSetSize;
+    vecsim_stl::vector<double> runningSumVec;
+    size_t count{0};
+
+public:
+    SQ8TrainingState(const std::shared_ptr<VecSimAllocator> &allocator, size_t dim,
+                     size_t normalization_set_size)
+        : normalizationSetSize(std::min(normalization_set_size, MAX_QUANT_NORMALIZATION_SET_SIZE)),
+          runningSumVec(dim, 0.0, allocator) {}
+
+    template <QuantInput T>
+    void add(std::span<const T> vector) {
+        assert(vector.size() == runningSumVec.size());
+        for (size_t i = 0; i < vector.size(); ++i) {
+            runningSumVec[i] += to_fp32(vector[i]);
+        }
+        ++count;
+    }
+
+    template <QuantInput T>
+    void subtract(std::span<const T> vector) {
+        assert(count > 0);
+        assert(vector.size() == runningSumVec.size());
+        for (size_t i = 0; i < vector.size(); ++i) {
+            runningSumVec[i] -= to_fp32(vector[i]);
+        }
+        --count;
+    }
+
+    bool shouldFinalize() const { return count > 0 && count >= normalizationSetSize; }
+
+    /// Requires at least one live training vector.
+    [[nodiscard]] vecsim_stl::vector<float> calculateMean() const {
+        assert(count > 0);
+        vecsim_stl::vector<float> mean(runningSumVec.size(), runningSumVec.getAllocator());
+        for (size_t i = 0; i < runningSumVec.size(); ++i) {
+            mean[i] = static_cast<float>(runningSumVec[i] / static_cast<double>(count));
+        }
+        return mean;
+    }
+
+    size_t getCount() const { return count; }
+    size_t getNormalizationSetSize() const { return normalizationSetSize; }
+    const vecsim_stl::vector<double> &getRunningSum() const { return runningSumVec; }
+};

src/VecSim/vec_sim_common.h

@@ -254,7 +254,8 @@ typedef struct {
 // A struct that contains HNSW Disk tiered index specific params.
 // Consider removing and use TieredHNSWParams instead if they both share swapJobThreshold
 typedef struct {
-    char _placeholder; // Reserved for future fields and avoid compiler errors
+    size_t QuantNormalizationSetSize; // Training threshold, with the semantics of
+                                      // TieredHNSWParams::QuantNormalizationSetSize.
 } TieredHNSWDiskParams;
 
 // A struct that contains SVS tiered index specific params.

tests/unit/test_components.cpp

@@ -12,6 +12,7 @@
 #include <cmath>
 #include "gtest/gtest.h"
 #include "VecSim/vec_sim.h"
+#include "VecSim/quantization/sq8_training_state.h"
 #include "VecSim/spaces/computer/preprocessor_container.h"
 #include "VecSim/spaces/computer/preprocessors.h"
 #include "VecSim/spaces/computer/calculator.h"
@@ -2644,3 +2645,85 @@ TEST(DistanceCalculatorWithNormTest, OddDimension_MatchesBruteForce) {
         delete calc;
     }
 }
+
+template <typename T>
+class SQ8TrainingStateTest : public ::testing::Test {};
+
+using SQ8TrainingInputTypes = ::testing::Types<float, vecsim_types::float16>;
+TYPED_TEST_SUITE(SQ8TrainingStateTest, SQ8TrainingInputTypes);
+
+TYPED_TEST(SQ8TrainingStateTest, SumMeanAndRemoval) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 2, 2);
+    const TypeParam first[] = {from_fp32<TypeParam>(1.5f), from_fp32<TypeParam>(-2.0f)};
+    const TypeParam second[] = {from_fp32<TypeParam>(-0.5f), from_fp32<TypeParam>(4.0f)};
+    EXPECT_EQ(state.getCount(), 0);
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const TypeParam>(first));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const TypeParam>(second));
+    EXPECT_EQ(state.getCount(), 2);
+    EXPECT_TRUE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 1.0);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[1], 2.0);
+    const auto mean = state.calculateMean();
+    EXPECT_EQ(mean.getAllocator(), allocator);
+    EXPECT_FLOAT_EQ(mean[0], 0.5f);
+    EXPECT_FLOAT_EQ(mean[1], 1.0f);
+
+    state.subtract(std::span<const TypeParam>(first));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    const auto remaining_mean = state.calculateMean();
+    EXPECT_FLOAT_EQ(remaining_mean[0], -0.5f);
+    EXPECT_FLOAT_EQ(remaining_mean[1], 4.0f);
+    state.subtract(std::span<const TypeParam>(second));
+    EXPECT_EQ(state.getCount(), 0);
+    EXPECT_FALSE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 0.0);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[1], 0.0);
+}
+
+TYPED_TEST(SQ8TrainingStateTest, ReplacementKeepsLiveCount) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, 2);
+    const TypeParam original[] = {from_fp32<TypeParam>(2.0f)};
+    const TypeParam replacement[] = {from_fp32<TypeParam>(8.0f)};
+    state.add(std::span<const TypeParam>(original));
+    state.subtract(std::span<const TypeParam>(original));
+    state.add(std::span<const TypeParam>(replacement));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_FALSE(state.shouldFinalize());
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 8.0);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 8.0f);
+}
+
+TEST(SQ8TrainingStateTest, ThresholdIsCapped) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, std::numeric_limits<size_t>::max());
+    EXPECT_EQ(state.getNormalizationSetSize(), MAX_QUANT_NORMALIZATION_SET_SIZE);
+    const float vector[] = {1.0f};
+    for (size_t i = 0; i < MAX_QUANT_NORMALIZATION_SET_SIZE - 1; ++i) {
+        state.add(std::span<const float>(vector));
+    }
+    EXPECT_FALSE(state.shouldFinalize());
+    state.add(std::span<const float>(vector));
+    EXPECT_TRUE(state.shouldFinalize());
+    EXPECT_EQ(state.getCount(), MAX_QUANT_NORMALIZATION_SET_SIZE);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 1.0f);
+}
+
+TEST(SQ8TrainingStateTest, AccumulatesInFP64) {
+    auto allocator = VecSimAllocator::newVecsimAllocator();
+    SQ8TrainingState state(allocator, 1, 3);
+    const float large[] = {16777216.0f};
+    const float small[] = {1.0f};
+    state.add(std::span<const float>(large));
+    state.add(std::span<const float>(small));
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 16777217.0);
+    state.subtract(std::span<const float>(large));
+    EXPECT_EQ(state.getCount(), 1);
+    EXPECT_DOUBLE_EQ(state.getRunningSum()[0], 1.0);
+    EXPECT_FLOAT_EQ(state.calculateMean()[0], 1.0f);
+}

tests/unit/test_hnsw_tiered.cpp

@@ -6784,11 +6784,11 @@ class HNSWTieredIndexTestSQ8 : public ::testing::Test {
     }
 
     const vecsim_stl::vector<double> &getRunningSumVec(TieredHNSWIndex<data_t, dist_t> *idx) {
-        return idx->sqAccumulationState->runningSumVec;
+        return idx->sqAccumulationState->getRunningSum();
     }
 
     size_t getNormalizationSetSize(TieredHNSWIndex<data_t, dist_t> *idx) {
-        return idx->sqAccumulationState->normalizationSetSize;
+        return idx->sqAccumulationState->getNormalizationSetSize();
     }
 
     // Compare the actual cached dispatches and preprocessing with a backend constructed with