SVS: allow option to disable vector storing (#5525)

Summary:
Allow avoiding storing original vectors, which may not necessarily be used, in which case they would be wasting memory

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5525

Reviewed By: alibeklfc

Differential Revision: D118722620

Pulled By: mnorris11

fbshipit-source-id: 9757fa3d272187308c750fb6c6d0e3a6e8ab2bb3

faiss/python/__init__.pyi

@@ -3906,6 +3906,7 @@ class IndexSVSVamana(Index):
     use_full_search_history: bool
     is_static: bool
     storage_kind: SVSStorageKind
+    store_vectors: bool
 
     def __init__(
         self,
@@ -3914,6 +3915,7 @@ class IndexSVSVamana(Index):
         metric: MetricType = METRIC_L2,
         storage: SVSStorageKind = SVS_FP32,
         is_static: bool = False,
+        store_vectors: bool = True,
     ) -> None: ...
     @staticmethod
     def is_lvq_leanvec_enabled() -> bool: ...
@@ -3926,6 +3928,7 @@ class IndexSVSVamanaLVQ(IndexSVSVamana):
         metric: MetricType = METRIC_L2,
         storage: SVSStorageKind = SVS_LVQ4x0,
         is_static: bool = False,
+        store_vectors: bool = True,
     ) -> None: ...
 
 class IndexSVSVamanaLeanVec(IndexSVSVamana):
@@ -3939,6 +3942,7 @@ class IndexSVSVamanaLeanVec(IndexSVSVamana):
         leanvec_dims: int = 0,
         storage: SVSStorageKind = SVS_LeanVec4x4,
         is_static: bool = False,
+        store_vectors: bool = True,
     ) -> None: ...
 
 class IndexSVSIVF(Index):

faiss/svs/IndexSVSVamana.cpp

@@ -67,11 +67,13 @@ IndexSVSVamana::IndexSVSVamana(
         size_t degree,
         MetricType metric,
         SVSStorageKind storage,
-        bool is_static)
+        bool is_static,
+        bool store_vectors)
         : Index(d, metric),
           graph_max_degree{degree},
           is_static{is_static},
-          storage_kind{storage} {
+          storage_kind{storage},
+          store_vectors{store_vectors} {
     prune_to = graph_max_degree < 4 ? graph_max_degree : graph_max_degree - 4;
     alpha = metric == METRIC_L2 ? 1.2f : 0.95f;
 
@@ -124,6 +126,14 @@ IndexSVSVamana::~IndexSVSVamana() {
 }
 
 void IndexSVSVamana::add(idx_t n, const float* x) {
+    // Opting out after data has been added would leave stored_vectors
+    // misaligned with the ids in the index, so release it instead of growing.
+    if (!store_vectors && stored_vectors_valid) {
+        stored_vectors.clear();
+        stored_vectors.shrink_to_fit();
+        stored_vectors_valid = false;
+    }
+
     if (is_static) {
         FAISS_THROW_IF_MSG(
                 impl,
@@ -167,7 +177,8 @@ void IndexSVSVamana::reconstruct(idx_t key, float* recons) const {
     FAISS_THROW_IF_NOT_MSG(
             stored_vectors_valid && !stored_vectors.empty(),
             "IndexSVSVamana::reconstruct: stored_vectors unavailable "
-            "(invalidated by remove_ids or not restored after deserialization)");
+            "(store_vectors disabled, invalidated by remove_ids, or not "
+            "restored after deserialization)");
     std::memcpy(recons, stored_vectors.data() + key * d, sizeof(float) * d);
 }
 
@@ -187,7 +198,8 @@ void IndexSVSVamana::reset() {
         }
     }
     stored_vectors.clear();
-    stored_vectors_valid = true;
+    stored_vectors.shrink_to_fit();
+    stored_vectors_valid = store_vectors;
     mmap_owner.reset(); // Release the memory mapping
     is_trained = false;
     ntotal = 0;

faiss/svs/IndexSVSVamana.h

@@ -111,7 +111,8 @@ struct IndexSVSVamana : Index {
             size_t degree,
             MetricType metric = METRIC_L2,
             SVSStorageKind storage = SVSStorageKind::SVS_FP32,
-            bool is_static = false);
+            bool is_static = false,
+            bool store_vectors = true);
 
     ~IndexSVSVamana() override;
 
@@ -164,6 +165,12 @@ struct IndexSVSVamana : Index {
     std::vector<float> stored_vectors;
     bool stored_vectors_valid{true};
 
+    // Set to false before the first add() to skip the stored_vectors copy,
+    // saving ntotal * d * 4 bytes at the cost of reconstruct() support and
+    // hence of use as an IVF coarse quantizer. Clearing it after vectors have
+    // been added drops the copy, which can no longer be aligned with the ids.
+    bool store_vectors{true};
+
    protected:
     /* Initializes the implementation. For static indexes the data is consumed
        at build time; for dynamic indexes n/x are ignored and add() populates

faiss/svs/IndexSVSVamanaLVQ.cpp

@@ -34,7 +34,9 @@ IndexSVSVamanaLVQ::IndexSVSVamanaLVQ(
         size_t degree,
         MetricType metric,
         SVSStorageKind storage,
-        bool is_static)
-        : IndexSVSVamana(d, degree, metric, storage, is_static) {}
+        bool is_static,
+        bool store_vectors)
+        : IndexSVSVamana(d, degree, metric, storage, is_static, store_vectors) {
+}
 
 } // namespace faiss

faiss/svs/IndexSVSVamanaLVQ.h

@@ -35,7 +35,8 @@ struct IndexSVSVamanaLVQ : IndexSVSVamana {
             size_t degree,
             MetricType metric = METRIC_L2,
             SVSStorageKind storage = SVSStorageKind::SVS_LVQ4x0,
-            bool is_static = false);
+            bool is_static = false,
+            bool store_vectors = true);
 
     ~IndexSVSVamanaLVQ() override = default;
 };

faiss/svs/IndexSVSVamanaLeanVec.cpp

@@ -44,8 +44,15 @@ IndexSVSVamanaLeanVec::IndexSVSVamanaLeanVec(
         MetricType metric,
         size_t leanvec_dims,
         SVSStorageKind storage_kind,
-        bool is_static)
-        : IndexSVSVamana(d, degree, metric, storage_kind, is_static) {
+        bool is_static,
+        bool store_vectors)
+        : IndexSVSVamana(
+                  d,
+                  degree,
+                  metric,
+                  storage_kind,
+                  is_static,
+                  store_vectors) {
     is_trained = false;
     leanvec_d = leanvec_dims == 0 ? d / 2 : leanvec_dims;
 }

faiss/svs/IndexSVSVamanaLeanVec.h

@@ -36,7 +36,8 @@ struct IndexSVSVamanaLeanVec : IndexSVSVamana {
             MetricType metric = METRIC_L2,
             size_t leanvec_dims = 0,
             SVSStorageKind storage = SVSStorageKind::SVS_LeanVec4x4,
-            bool is_static = false);
+            bool is_static = false,
+            bool store_vectors = true);
 
     ~IndexSVSVamanaLeanVec() override;
 

tests/test_svs.cpp

@@ -1036,6 +1036,168 @@ TEST_F(SVS, StaticVamanaReconstruct) {
     }
 }
 
+// store_vectors defaults to true, so the fp32 copy backing reconstruct() is
+// still kept unless the caller opts out.
+TEST_F(SVS, VamanaStoresVectorsByDefault) {
+    faiss::IndexSVSVamana index{d, 64ul};
+    EXPECT_TRUE(index.store_vectors);
+    index.add(n, test_data.data());
+    EXPECT_EQ(index.stored_vectors.size(), n * d);
+
+    std::vector<float> recons(d);
+    ASSERT_NO_THROW(index.reconstruct(0, recons.data()));
+}
+
+TEST_F(SVS, VamanaStoreVectorsDisabled) {
+    faiss::IndexSVSVamana index{
+            d,
+            64ul,
+            faiss::METRIC_L2,
+            faiss::SVSStorageKind::SVS_FP32,
+            false,
+            false};
+    EXPECT_FALSE(index.store_vectors);
+    index.add(n, test_data.data());
+    EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(index.stored_vectors.empty());
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(index.reconstruct(0, recons.data()), faiss::FaissException);
+
+    // Dropping the copy must not affect search.
+    const int nq = 4;
+    const int k = 5;
+    std::vector<float> distances(nq * k);
+    std::vector<faiss::idx_t> labels(nq * k);
+    ASSERT_NO_THROW(index.search(
+            nq, test_data.data(), k, distances.data(), labels.data()));
+
+    // The opt-out survives reset().
+    index.reset();
+    index.add(n, test_data.data());
+    EXPECT_TRUE(index.stored_vectors.empty());
+}
+
+TEST_F(SVS, StaticVamanaStoreVectorsDisabled) {
+    faiss::IndexSVSVamana index{
+            d,
+            64ul,
+            faiss::METRIC_L2,
+            faiss::SVSStorageKind::SVS_FP32,
+            true,
+            false};
+    index.add(n, test_data.data());
+    EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(index.stored_vectors.empty());
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(index.reconstruct(0, recons.data()), faiss::FaissException);
+}
+
+// The flag is a plain public field, so setting it after construction (as the
+// Python bindings do) must take effect too.
+TEST_F(SVS, VamanaStoreVectorsDisabledAfterConstruction) {
+    faiss::IndexSVSVamana index{d, 64ul};
+    index.store_vectors = false;
+    index.add(n, test_data.data());
+    EXPECT_TRUE(index.stored_vectors.empty());
+}
+
+// Opting out mid-stream cannot keep the remaining copy aligned with the ids,
+// so it is dropped rather than left partially populated.
+TEST_F(SVS, VamanaStoreVectorsDroppedOnMidStreamOptOut) {
+    faiss::IndexSVSVamana index{d, 64ul};
+    index.add(n / 2, test_data.data());
+    EXPECT_EQ(index.stored_vectors.size(), (n / 2) * d);
+
+    index.store_vectors = false;
+    index.add(n / 2, test_data.data() + (n / 2) * d);
+    EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(index.stored_vectors.empty());
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(index.reconstruct(0, recons.data()), faiss::FaissException);
+}
+
+// An index built without the copy is written as ISVD, loads back and searches
+// identically; only reconstruct() is unavailable.
+TEST_F(SVS, VamanaStoreVectorsDisabledWriteAndRead) {
+    faiss::IndexSVSVamana index{
+            d,
+            64ul,
+            faiss::METRIC_L2,
+            faiss::SVSStorageKind::SVS_FP32,
+            false,
+            false};
+    index.add(n, test_data.data());
+
+    const int nq = 4;
+    const int k = 5;
+    std::vector<float> distances(nq * k), distances_loaded(nq * k);
+    std::vector<faiss::idx_t> labels(nq * k), labels_loaded(nq * k);
+    index.search(nq, test_data.data(), k, distances.data(), labels.data());
+
+    std::string temp_filename_template = "/tmp/faiss_svs_test_XXXXXX";
+    Tempfilename filename(&temp_file_mutex, temp_filename_template);
+    ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
+
+    faiss::IndexSVSVamana* loaded = nullptr;
+    ASSERT_NO_THROW({
+        loaded = dynamic_cast<faiss::IndexSVSVamana*>(
+                faiss::read_index(filename.c_str()));
+    });
+    ASSERT_NE(loaded, nullptr);
+    EXPECT_EQ(loaded->ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(loaded->stored_vectors.empty());
+
+    loaded->search(
+            nq,
+            test_data.data(),
+            k,
+            distances_loaded.data(),
+            labels_loaded.data());
+    EXPECT_EQ(labels, labels_loaded);
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(loaded->reconstruct(0, recons.data()), faiss::FaissException);
+
+    delete loaded;
+}
+
+TEST_F(SVSLL, LVQStoreVectorsDisabled) {
+    faiss::IndexSVSVamanaLVQ index{
+            d,
+            64ul,
+            faiss::METRIC_L2,
+            faiss::SVSStorageKind::SVS_LVQ4x8,
+            false,
+            false};
+    index.add(n, test_data.data());
+    EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(index.stored_vectors.empty());
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(index.reconstruct(0, recons.data()), faiss::FaissException);
+}
+
+TEST_F(SVSLL, LeanVecStoreVectorsDisabled) {
+    faiss::IndexSVSVamanaLeanVec index{
+            d,
+            64ul,
+            faiss::METRIC_L2,
+            0,
+            faiss::SVSStorageKind::SVS_LeanVec4x8,
+            false,
+            false};
+    index.train(n, test_data.data());
+    index.add(n, test_data.data());
+    EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
+    EXPECT_TRUE(index.stored_vectors.empty());
+
+    std::vector<float> recons(d);
+    EXPECT_THROW(index.reconstruct(0, recons.data()), faiss::FaissException);
+}
+
 TEST_F(SVS, WriteAndMapStaticVamana) {
     faiss::IndexSVSVamana index{
             d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};

tests/test_svs_py.py

@@ -264,6 +264,63 @@ class TestSVSAdapter(unittest.TestCase):
         self.assertEqual(D.shape, (self.nq, 4))
 
 
+@unittest.skipIf(_SKIP_SVS, _SKIP_REASON)
+class TestSVSStoreVectors(unittest.TestCase):
+    """Test the store_vectors opt-out for the fp32 reconstruct() copy"""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.d = 32
+        cls.nb = 500
+        cls.nq = 10
+        np.random.seed(1234)
+        cls.xb = np.random.random((cls.nb, cls.d)).astype("float32")
+        cls.xq = np.random.random((cls.nq, cls.d)).astype("float32")
+
+    def test_stores_vectors_by_default(self):
+        index = faiss.IndexSVSVamana(self.d, 64)
+        self.assertTrue(index.store_vectors)
+        index.add(self.xb)
+        self.assertEqual(index.stored_vectors.size(), self.nb * self.d)
+        np.testing.assert_array_equal(index.reconstruct(0), self.xb[0])
+
+    def test_store_vectors_disabled_in_constructor(self):
+        index = faiss.IndexSVSVamana(
+            self.d, 64, faiss.METRIC_L2, faiss.SVS_FP32, False, False
+        )
+        self.assertFalse(index.store_vectors)
+        index.add(self.xb)
+        self.assertEqual(index.ntotal, self.nb)
+        self.assertEqual(index.stored_vectors.size(), 0)
+        with self.assertRaises(RuntimeError):
+            index.reconstruct(0)
+
+    def test_store_vectors_disabled_by_attribute(self):
+        """Search must be unaffected by dropping the fp32 copy"""
+        reference = faiss.IndexSVSVamana(self.d, 64)
+        reference.add(self.xb)
+        _, I_ref = reference.search(self.xq, 4)
+
+        index = faiss.IndexSVSVamana(self.d, 64)
+        index.store_vectors = False
+        index.add(self.xb)
+        self.assertEqual(index.stored_vectors.size(), 0)
+
+        _, I = index.search(self.xq, 4)
+        np.testing.assert_array_equal(I_ref, I)
+
+    def test_store_vectors_disabled_survives_serialization(self):
+        index = faiss.IndexSVSVamana(self.d, 64)
+        index.store_vectors = False
+        index.add(self.xb)
+
+        loaded = faiss.deserialize_index(faiss.serialize_index(index))
+        self.assertEqual(loaded.ntotal, self.nb)
+        self.assertEqual(loaded.stored_vectors.size(), 0)
+        with self.assertRaises(RuntimeError):
+            loaded.reconstruct(0)
+
+
 @unittest.skipIf(_SKIP_SVS, _SKIP_REASON)
 class TestSVSFactory(unittest.TestCase):
     """Test that SVS factory works correctly"""