mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
ethanglaser
cfbf4ec5ff
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
1240 lines
40 KiB
C++
1240 lines
40 KiB
C++
/*
|
|
* Portions Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
/*
|
|
* Portions Copyright 2025 Intel Corporation
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include <faiss/Index.h>
|
|
#include <faiss/impl/AuxIndexStructures.h>
|
|
#include <faiss/impl/FaissException.h>
|
|
#include <faiss/impl/IDSelector.h>
|
|
#include <faiss/index_io.h>
|
|
#include <faiss/svs/IndexSVSFlat.h>
|
|
#include <faiss/svs/IndexSVSIVF.h>
|
|
#include <faiss/svs/IndexSVSIVFLVQ.h>
|
|
#include <faiss/svs/IndexSVSIVFLeanVec.h>
|
|
#include <faiss/svs/IndexSVSVamana.h>
|
|
#include <faiss/svs/IndexSVSVamanaLVQ.h>
|
|
#include <faiss/svs/IndexSVSVamanaLeanVec.h>
|
|
#include <gtest/gtest.h>
|
|
#include <random>
|
|
#include <type_traits>
|
|
|
|
#include "test_util.h"
|
|
|
|
namespace {
|
|
pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER;
|
|
|
|
// Test fixture class to manage shared test data
|
|
class SVS : public ::testing::Test {
|
|
protected:
|
|
static void SetUpTestSuite() {
|
|
// Generate test data once for all tests
|
|
constexpr size_t d = 64;
|
|
constexpr size_t n = 100;
|
|
test_data.resize(d * n);
|
|
|
|
std::mt19937 gen(123);
|
|
std::uniform_real_distribution<float> dis(0.0f, 1.0f);
|
|
for (size_t i = 0; i < test_data.size(); ++i) {
|
|
test_data[i] = dis(gen);
|
|
}
|
|
}
|
|
|
|
static std::vector<float> test_data;
|
|
static constexpr size_t d = 64;
|
|
static constexpr size_t n = 100;
|
|
};
|
|
|
|
// LVQ/LeanVec tests are only executed if LVQ/LeanVec support is available
|
|
class SVSLL : public SVS {
|
|
protected:
|
|
void SetUp() override {
|
|
if (!faiss::IndexSVSVamana::is_lvq_leanvec_enabled()) {
|
|
GTEST_SKIP() << "LVQ/LeanVec support not available on this "
|
|
"platform or build configuration";
|
|
}
|
|
}
|
|
};
|
|
|
|
// Consistency checks for behavior if LVQ/LeanVec are not available
|
|
// Only runs if LVQ/LeanVec is NOT available
|
|
class SVSNoLL : public SVS {
|
|
protected:
|
|
void SetUp() override {
|
|
if (faiss::IndexSVSVamana::is_lvq_leanvec_enabled()) {
|
|
GTEST_SKIP() << "LVQ/LeanVec support is available; skipping "
|
|
"NoLL tests";
|
|
}
|
|
}
|
|
};
|
|
|
|
// Define static members
|
|
std::vector<float> SVS::test_data;
|
|
} // namespace
|
|
|
|
template <typename T>
|
|
void write_and_read_index(T& index, const std::vector<float>& xb, size_t n) {
|
|
index.train(n, xb.data());
|
|
index.add(n, xb.data());
|
|
|
|
std::string temp_filename_template = "/tmp/faiss_svs_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
// Serialize
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
// Deserialize
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(faiss::read_index(filename.c_str()));
|
|
});
|
|
|
|
// Basic checks
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NE(loaded->impl, nullptr);
|
|
EXPECT_EQ(loaded->d, index.d);
|
|
EXPECT_EQ(loaded->metric_type, index.metric_type);
|
|
EXPECT_EQ(loaded->graph_max_degree, index.graph_max_degree);
|
|
EXPECT_EQ(loaded->alpha, index.alpha);
|
|
EXPECT_EQ(loaded->search_window_size, index.search_window_size);
|
|
EXPECT_EQ(loaded->search_buffer_capacity, index.search_buffer_capacity);
|
|
EXPECT_EQ(loaded->construction_window_size, index.construction_window_size);
|
|
EXPECT_EQ(loaded->max_candidate_pool_size, index.max_candidate_pool_size);
|
|
EXPECT_EQ(loaded->prune_to, index.prune_to);
|
|
EXPECT_EQ(loaded->use_full_search_history, index.use_full_search_history);
|
|
if constexpr (std::is_same_v<
|
|
std::decay_t<T>,
|
|
faiss::IndexSVSVamanaLeanVec>) {
|
|
auto* leanvec_loaded =
|
|
dynamic_cast<faiss::IndexSVSVamanaLeanVec*>(loaded);
|
|
ASSERT_NE(leanvec_loaded, nullptr);
|
|
EXPECT_EQ(leanvec_loaded->leanvec_d, index.leanvec_d);
|
|
|
|
EXPECT_NE(leanvec_loaded->training_data, nullptr);
|
|
}
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
template <typename T>
|
|
void train_save_load_and_add_index(
|
|
T& index,
|
|
const std::vector<float>& xb,
|
|
size_t n) {
|
|
index.train(n, xb.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()); });
|
|
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(faiss::read_index(filename.c_str()));
|
|
});
|
|
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NO_THROW({ loaded->add(n, xb.data()); });
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
template <typename T>
|
|
void save_and_load_index() {
|
|
T index;
|
|
|
|
std::string temp_filename_template = "/tmp/faiss_svs_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
faiss::write_index(&index, filename.c_str());
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(faiss::read_index(filename.c_str()));
|
|
});
|
|
delete loaded;
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVS) {
|
|
faiss::IndexSVSVamana index{d, 64ul};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSFP16) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSSQ8) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSLVQ4x0) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x0;
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSLVQ4x4) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x4;
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSLVQ4x8) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x8;
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSVamanaLeanVec4x4) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSVamanaLeanVec4x8) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x8};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSVamanaLeanVec8x8) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec8x8};
|
|
write_and_read_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LeanVecThrowsWithoutTraining) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
64,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
ASSERT_THROW(index.add(100, test_data.data()), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, VamanaTrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamana index{d, 64ul};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, VamanaFP16TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, VamanaSQ8TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LVQ4x0TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x0;
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LVQ4x4TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x4;
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LVQ4x8TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLVQ index{d, 64ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x8;
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LeanVec4x4TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LeanVec4x8TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x8};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, LeanVec8x8TrainSaveLoadAndAdd) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec8x8};
|
|
train_save_load_and_add_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, SaveAndLoadIndexSVSFlat) {
|
|
save_and_load_index<faiss::IndexSVSFlat>();
|
|
}
|
|
|
|
TEST_F(SVS, SaveAndLoadIndexSVSVamana) {
|
|
save_and_load_index<faiss::IndexSVSVamana>();
|
|
}
|
|
|
|
TEST_F(SVSLL, SaveAndLoadIndexSVSVamanaLVQ) {
|
|
save_and_load_index<faiss::IndexSVSVamanaLVQ>();
|
|
}
|
|
|
|
TEST_F(SVSLL, SaveAndLoadIndexSVSVamanaLeanVec) {
|
|
save_and_load_index<faiss::IndexSVSVamanaLeanVec>();
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSFlat) {
|
|
faiss::IndexSVSFlat index{d};
|
|
index.add(n, test_data.data());
|
|
|
|
std::string temp_filename_template = "/tmp/faiss_svs_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
// Serialize
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
// Deserialize
|
|
faiss::IndexSVSFlat* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<faiss::IndexSVSFlat*>(
|
|
faiss::read_index(filename.c_str()));
|
|
});
|
|
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NE(loaded->impl, nullptr);
|
|
EXPECT_EQ(loaded->d, index.d);
|
|
EXPECT_EQ(loaded->nlabels, index.nlabels);
|
|
EXPECT_EQ(loaded->metric_type, index.metric_type);
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
// Test search with IDSelector filtering
|
|
TEST_F(SVS, SearchWithIDSelector) {
|
|
faiss::IndexSVSVamana index{d, 64ul};
|
|
index.add(n, test_data.data());
|
|
|
|
const int nq = 8; // number of queries
|
|
const float* xq = test_data.data(); // reuse first nq vectors as queries
|
|
const int k = 10;
|
|
|
|
size_t min_id = n / 5; // inclusive
|
|
size_t max_id = n * 4 / 5; // exclusive
|
|
faiss::IDSelectorRange selector(min_id, max_id);
|
|
|
|
faiss::SearchParameters params; // generic search parameters with selector
|
|
params.sel = &selector;
|
|
|
|
std::vector<float> distances(nq * k);
|
|
std::vector<faiss::idx_t> labels(nq * k);
|
|
|
|
ASSERT_NO_THROW(
|
|
index.search(nq, xq, k, distances.data(), labels.data(), ¶ms));
|
|
|
|
// All returned labels must fall inside the selected range
|
|
for (int i = 0; i < nq * k; ++i) {
|
|
EXPECT_GE(labels[i], (faiss::idx_t)min_id);
|
|
EXPECT_LT(labels[i], (faiss::idx_t)max_id);
|
|
}
|
|
}
|
|
|
|
// Basic functional test for range_search and parameter override helper
|
|
TEST_F(SVS, RangeSearchFunctional) {
|
|
faiss::IndexSVSVamana index{d, 64ul};
|
|
index.add(n, test_data.data());
|
|
const int nq = 5;
|
|
const float* xq = test_data.data();
|
|
|
|
// Small radius
|
|
faiss::RangeSearchResult res_small(nq);
|
|
ASSERT_NO_THROW(index.range_search(nq, xq, 0.05f, &res_small));
|
|
|
|
// Larger radius to exercise loop continuation
|
|
faiss::RangeSearchResult res_big(nq);
|
|
ASSERT_NO_THROW(index.range_search(nq, xq, 5.0f, &res_big));
|
|
EXPECT_GE(res_big.lims[nq], res_small.lims[nq]);
|
|
|
|
// Provide custom params to ensure branch coverage (non-null params)
|
|
faiss::SearchParametersSVSVamana params;
|
|
params.search_window_size = 15;
|
|
params.search_buffer_capacity = 20;
|
|
|
|
// Provide IDSelector to ensure branch coverage (non-null params->sel)
|
|
size_t min_id = n / 5; // inclusive
|
|
size_t max_id = n * 4 / 5; // exclusive
|
|
faiss::IDSelectorRange selector(min_id, max_id);
|
|
params.sel = &selector;
|
|
|
|
faiss::RangeSearchResult res_params(nq);
|
|
ASSERT_NO_THROW(index.range_search(nq, xq, 1.0f, &res_params, ¶ms));
|
|
|
|
// All returned labels must fall inside the selected range
|
|
for (size_t i = 0; i < res_params.lims[nq]; ++i) {
|
|
EXPECT_GE(res_params.labels[i], (faiss::idx_t)min_id);
|
|
EXPECT_LT(res_params.labels[i], (faiss::idx_t)max_id);
|
|
}
|
|
}
|
|
|
|
TEST_F(SVSLL, LVQAndLeanVecDoNotThrowWhenEnabled) {
|
|
// explicit constructor with LVQ dataset
|
|
ASSERT_NO_THROW({
|
|
faiss::IndexSVSVamanaLVQ index(
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_LVQ4x4);
|
|
});
|
|
|
|
// default constructor, will initialize dataset on first add()
|
|
ASSERT_NO_THROW({
|
|
faiss::IndexSVSVamanaLVQ index;
|
|
index.add(n, test_data.data());
|
|
});
|
|
|
|
// explicit constructor with LeanVec dataset
|
|
ASSERT_NO_THROW({
|
|
faiss::IndexSVSVamanaLeanVec index(
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4);
|
|
});
|
|
|
|
// default constructor, will initialize dataset on first add()
|
|
ASSERT_NO_THROW({
|
|
faiss::IndexSVSVamanaLeanVec index;
|
|
index.d = 64;
|
|
index.leanvec_d = 32;
|
|
index.train(n, test_data.data());
|
|
index.add(n, test_data.data());
|
|
});
|
|
}
|
|
|
|
TEST_F(SVSNoLL, LVQAndLeanVecThrowWhenNotEnabled) {
|
|
// explicit constructor with LVQ dataset
|
|
ASSERT_THROW(
|
|
{
|
|
faiss::IndexSVSVamanaLVQ index(
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
faiss::SVSStorageKind::SVS_LVQ4x4);
|
|
},
|
|
faiss::FaissException);
|
|
|
|
// default constructor, will initialize dataset on first add()
|
|
ASSERT_THROW(
|
|
{
|
|
faiss::IndexSVSVamanaLVQ index;
|
|
index.add(n, test_data.data());
|
|
},
|
|
faiss::FaissException);
|
|
|
|
// explicit constructor with LeanVec dataset
|
|
ASSERT_THROW(
|
|
{
|
|
faiss::IndexSVSVamanaLeanVec index(
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4);
|
|
},
|
|
faiss::FaissException);
|
|
|
|
// default constructor, will initialize dataset on first add()
|
|
ASSERT_THROW(
|
|
{
|
|
faiss::IndexSVSVamanaLeanVec index;
|
|
index.d = 64;
|
|
index.leanvec_d = 32;
|
|
index.train(n, test_data.data());
|
|
},
|
|
faiss::FaissException);
|
|
}
|
|
|
|
/****************************************************************
|
|
* SVS IVF Tests
|
|
****************************************************************/
|
|
|
|
template <typename T>
|
|
void write_and_read_ivf_index(
|
|
T& index,
|
|
const std::vector<float>& xb,
|
|
size_t n) {
|
|
index.train(n, xb.data());
|
|
|
|
std::string temp_filename_template = "/tmp/faiss_svs_ivf_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
// Serialize
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
// Deserialize
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(faiss::read_index(filename.c_str()));
|
|
});
|
|
|
|
// Basic checks
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NE(loaded->impl, nullptr);
|
|
EXPECT_EQ(loaded->d, index.d);
|
|
EXPECT_EQ(loaded->metric_type, index.metric_type);
|
|
EXPECT_EQ(loaded->num_centroids, index.num_centroids);
|
|
EXPECT_EQ(loaded->n_probes, index.n_probes);
|
|
EXPECT_EQ(loaded->k_reorder, index.k_reorder);
|
|
EXPECT_EQ(loaded->num_iterations, index.num_iterations);
|
|
EXPECT_EQ(loaded->minibatch_size, index.minibatch_size);
|
|
EXPECT_EQ(loaded->storage_kind, index.storage_kind);
|
|
if constexpr (std::is_same_v<std::decay_t<T>, faiss::IndexSVSIVFLeanVec>) {
|
|
auto* leanvec_loaded = dynamic_cast<faiss::IndexSVSIVFLeanVec*>(loaded);
|
|
ASSERT_NE(leanvec_loaded, nullptr);
|
|
EXPECT_EQ(leanvec_loaded->leanvec_d, index.leanvec_d);
|
|
|
|
EXPECT_NE(leanvec_loaded->training_data, nullptr);
|
|
}
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
template <typename T>
|
|
void train_and_add_ivf_index(T& index, const std::vector<float>& xb, size_t n) {
|
|
index.train(n, xb.data());
|
|
|
|
// Additional data to add after training
|
|
ASSERT_NO_THROW({ index.add(n, xb.data()); });
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSIVF) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSIVFFP16) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadIndexSVSIVFSQ8) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLVQ4x0) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x0;
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLVQ4x4) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x4;
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLVQ4x8) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x8;
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLVQ8x0) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ8x0;
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLeanVec4x4) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLeanVec4x8) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec4x8};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadIndexSVSIVFLeanVec8x8) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec8x8};
|
|
write_and_read_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
// Each SVS family streams its payload to the third-party SVS reader through a
|
|
// ReaderStreambuf, and only a streambuf carrying the deserialization vector
|
|
// byte limit rejects an oversized read. Every family must therefore build one
|
|
// the same way.
|
|
template <typename T>
|
|
void expect_byte_limit_enforced(
|
|
T& index,
|
|
const std::vector<float>& xb,
|
|
size_t n) {
|
|
index.train(n, xb.data());
|
|
index.add(n, xb.data());
|
|
|
|
std::string temp_filename_template = "/tmp/faiss_svs_limit_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
const size_t old_limit = faiss::get_deserialization_vector_byte_limit();
|
|
faiss::set_deserialization_vector_byte_limit(64);
|
|
EXPECT_THROW(
|
|
{ delete faiss::read_index(filename.c_str()); },
|
|
faiss::FaissException);
|
|
faiss::set_deserialization_vector_byte_limit(old_limit);
|
|
|
|
faiss::Index* loaded = nullptr;
|
|
ASSERT_NO_THROW({ loaded = faiss::read_index(filename.c_str()); });
|
|
EXPECT_NE(loaded, nullptr);
|
|
delete loaded;
|
|
}
|
|
|
|
TEST_F(SVS, FlatHonorsDeserializationVectorByteLimit) {
|
|
faiss::IndexSVSFlat index{d};
|
|
expect_byte_limit_enforced(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, IVFHonorsDeserializationVectorByteLimit) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
expect_byte_limit_enforced(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, IVFTrainAndAdd) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, IVFFP16TrainAndAdd) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, IVFSQ8TrainAndAdd) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLVQ4x0TrainAndAdd) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x0;
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLVQ4x4TrainAndAdd) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x4;
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLVQ4x8TrainAndAdd) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ4x8;
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLVQ8x0TrainAndAdd) {
|
|
faiss::IndexSVSIVFLVQ index{d, 4ul};
|
|
index.storage_kind = faiss::SVSStorageKind::SVS_LVQ8x0;
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLeanVec4x4TrainAndAdd) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLeanVec4x8TrainAndAdd) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec4x8};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLeanVec8x8TrainAndAdd) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d, 4ul, faiss::METRIC_L2, 0, faiss::SVSStorageKind::SVS_LeanVec8x8};
|
|
train_and_add_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, IVFSearchWithIDSelector) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
index.train(n, test_data.data());
|
|
index.add(n, test_data.data());
|
|
|
|
const int nq = 8;
|
|
const float* xq = test_data.data();
|
|
const int k = 10;
|
|
|
|
size_t min_id = n / 5;
|
|
size_t max_id = n * 4 / 5;
|
|
faiss::IDSelectorRange selector(min_id, max_id);
|
|
|
|
faiss::SearchParameters params;
|
|
params.sel = &selector;
|
|
|
|
std::vector<float> distances(nq * k);
|
|
std::vector<faiss::idx_t> labels(nq * k);
|
|
|
|
// The current SVS runtime IVFIndex::search() API does not accept an
|
|
// IDFilter parameter (unlike VamanaIndex::search()), so the IDSelector
|
|
// passed via SearchParameters is silently ignored. This test verifies
|
|
// that search still works (without filtering). Once the SVS runtime
|
|
// exposes IDFilter support for IVF search, this can be wired up the
|
|
// same way Vamana uses make_faiss_id_filter().
|
|
ASSERT_NO_THROW(index.search(nq, xq, k, distances.data(), labels.data()));
|
|
}
|
|
|
|
TEST_F(SVSLL, IVFLeanVecThrowsWithoutTraining) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
64,
|
|
4ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4};
|
|
ASSERT_THROW(index.add(100, test_data.data()), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, IVFThrowsAddWithoutTraining) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
ASSERT_THROW(index.add(n, test_data.data()), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, SaveAndLoadIndexSVSIVF) {
|
|
save_and_load_index<faiss::IndexSVSIVF>();
|
|
}
|
|
|
|
TEST_F(SVSLL, SaveAndLoadIndexSVSIVFLVQ) {
|
|
save_and_load_index<faiss::IndexSVSIVFLVQ>();
|
|
}
|
|
|
|
TEST_F(SVSLL, SaveAndLoadIndexSVSIVFLeanVec) {
|
|
save_and_load_index<faiss::IndexSVSIVFLeanVec>();
|
|
}
|
|
|
|
/****************************************************************
|
|
* SVS Static IVF Tests
|
|
****************************************************************/
|
|
|
|
template <typename T>
|
|
void write_and_read_static_ivf_index(
|
|
T& index,
|
|
const std::vector<float>& xb,
|
|
size_t n) {
|
|
ASSERT_TRUE(index.is_static);
|
|
index.train(n, xb.data());
|
|
|
|
// Verify search works after training
|
|
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, xb.data(), k, distances.data(), labels.data()));
|
|
|
|
std::string temp_filename_template =
|
|
"/tmp/faiss_svs_static_ivf_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
// Serialize
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
// Deserialize
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(faiss::read_index(filename.c_str()));
|
|
});
|
|
|
|
// Basic checks
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NE(loaded->impl, nullptr);
|
|
EXPECT_TRUE(loaded->is_static);
|
|
EXPECT_EQ(loaded->d, index.d);
|
|
EXPECT_EQ(loaded->metric_type, index.metric_type);
|
|
EXPECT_EQ(loaded->num_centroids, index.num_centroids);
|
|
EXPECT_EQ(loaded->storage_kind, index.storage_kind);
|
|
if constexpr (std::is_same_v<std::decay_t<T>, faiss::IndexSVSIVFLeanVec>) {
|
|
auto* leanvec_loaded = dynamic_cast<faiss::IndexSVSIVFLeanVec*>(loaded);
|
|
ASSERT_NE(leanvec_loaded, nullptr);
|
|
EXPECT_EQ(leanvec_loaded->leanvec_d, index.leanvec_d);
|
|
EXPECT_NE(leanvec_loaded->training_data, nullptr);
|
|
}
|
|
|
|
// Verify search works on loaded index
|
|
ASSERT_NO_THROW(
|
|
loaded->search(nq, xb.data(), k, distances.data(), labels.data()));
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadStaticIVF) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
write_and_read_static_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadStaticIVFFP16) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16, true};
|
|
write_and_read_static_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadStaticIVFLVQ4x4) {
|
|
faiss::IndexSVSIVFLVQ index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_LVQ4x4, true};
|
|
write_and_read_static_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadStaticIVFLeanVec4x4) {
|
|
faiss::IndexSVSIVFLeanVec index{
|
|
d,
|
|
4ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4,
|
|
true};
|
|
write_and_read_static_ivf_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, StaticIVFAddThrows) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.train(n, test_data.data());
|
|
ASSERT_THROW(index.add(n, test_data.data()), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, StaticIVFRemoveThrows) {
|
|
faiss::IndexSVSIVF index{
|
|
d, 4ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.train(n, test_data.data());
|
|
faiss::IDSelectorRange selector(0, 10);
|
|
ASSERT_THROW(index.remove_ids(selector), faiss::FaissException);
|
|
}
|
|
|
|
// Verify that intra_query_threads must be set before train().
|
|
// The SVS runtime API does not currently expose a setter for
|
|
// intra_query_threads after index creation, so changes made between
|
|
// train() and search() are silently ignored. This test documents the
|
|
// expected usage pattern.
|
|
TEST_F(SVS, IVFIntraQueryThreadsSetBeforeTrain) {
|
|
faiss::IndexSVSIVF index{d, 4ul};
|
|
|
|
// Set intra_query_threads BEFORE train — this is the supported pattern.
|
|
index.intra_query_threads = 2;
|
|
index.train(n, test_data.data());
|
|
index.add(n, test_data.data());
|
|
|
|
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()));
|
|
|
|
// Changing intra_query_threads AFTER train() has no effect — this is a
|
|
// known limitation of the current SVS runtime API.
|
|
index.intra_query_threads = 4;
|
|
ASSERT_NO_THROW(index.search(
|
|
nq, test_data.data(), k, distances.data(), labels.data()));
|
|
}
|
|
|
|
/****************************************************************
|
|
* SVS Static Vamana Tests
|
|
****************************************************************/
|
|
|
|
template <typename T>
|
|
void write_and_read_static_vamana_index(
|
|
T& index,
|
|
const std::vector<float>& xb,
|
|
size_t n,
|
|
int read_io_flags = 0) {
|
|
ASSERT_TRUE(index.is_static);
|
|
|
|
// For LeanVec the training step seeds training_data; for plain Vamana
|
|
// it is a no-op. The static build itself happens inside add().
|
|
index.train(n, xb.data());
|
|
index.add(n, xb.data());
|
|
ASSERT_NE(index.impl, nullptr);
|
|
ASSERT_TRUE(index.is_trained);
|
|
|
|
// Verify search works after the static build
|
|
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, xb.data(), k, distances.data(), labels.data()));
|
|
|
|
std::string temp_filename_template =
|
|
"/tmp/faiss_svs_static_vamana_test_XXXXXX";
|
|
Tempfilename filename(&temp_file_mutex, temp_filename_template);
|
|
|
|
// Serialize
|
|
ASSERT_NO_THROW({ faiss::write_index(&index, filename.c_str()); });
|
|
|
|
// Deserialize
|
|
T* loaded = nullptr;
|
|
ASSERT_NO_THROW({
|
|
loaded = dynamic_cast<T*>(
|
|
faiss::read_index(filename.c_str(), read_io_flags));
|
|
});
|
|
|
|
// Basic checks
|
|
ASSERT_NE(loaded, nullptr);
|
|
ASSERT_NE(loaded->impl, nullptr);
|
|
EXPECT_TRUE(loaded->is_static);
|
|
EXPECT_EQ(loaded->d, index.d);
|
|
EXPECT_EQ(loaded->metric_type, index.metric_type);
|
|
EXPECT_EQ(loaded->graph_max_degree, index.graph_max_degree);
|
|
EXPECT_EQ(loaded->storage_kind, index.storage_kind);
|
|
if constexpr (std::is_same_v<
|
|
std::decay_t<T>,
|
|
faiss::IndexSVSVamanaLeanVec>) {
|
|
auto* leanvec_loaded =
|
|
dynamic_cast<faiss::IndexSVSVamanaLeanVec*>(loaded);
|
|
ASSERT_NE(leanvec_loaded, nullptr);
|
|
EXPECT_EQ(leanvec_loaded->leanvec_d, index.leanvec_d);
|
|
EXPECT_NE(leanvec_loaded->training_data, nullptr);
|
|
}
|
|
|
|
// Verify search works on the loaded index
|
|
ASSERT_NO_THROW(
|
|
loaded->search(nq, xb.data(), k, distances.data(), labels.data()));
|
|
|
|
delete loaded;
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadStaticVamana) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
write_and_read_static_vamana_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadStaticVamanaFP16) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16, true};
|
|
write_and_read_static_vamana_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndReadStaticVamanaSQ8) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8, true};
|
|
write_and_read_static_vamana_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadStaticVamanaLVQ4x4) {
|
|
faiss::IndexSVSVamanaLVQ index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_LVQ4x4, true};
|
|
write_and_read_static_vamana_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndReadStaticVamanaLeanVec4x4) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4,
|
|
true};
|
|
write_and_read_static_vamana_index(index, test_data, n);
|
|
}
|
|
|
|
TEST_F(SVS, StaticVamanaAddThrowsOnSecondCall) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.add(n, test_data.data());
|
|
ASSERT_THROW(index.add(n, test_data.data()), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, StaticVamanaRemoveThrows) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.add(n, test_data.data());
|
|
faiss::IDSelectorRange selector(0, 10);
|
|
ASSERT_THROW(index.remove_ids(selector), faiss::FaissException);
|
|
}
|
|
|
|
TEST_F(SVS, StaticVamanaResetAndRebuild) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.add(n, test_data.data());
|
|
EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
|
|
|
|
index.reset();
|
|
EXPECT_EQ(index.ntotal, 0);
|
|
EXPECT_FALSE(index.is_trained);
|
|
|
|
// After reset, a fresh add() must rebuild the static index.
|
|
ASSERT_NO_THROW(index.add(n, test_data.data()));
|
|
EXPECT_EQ(index.ntotal, static_cast<faiss::idx_t>(n));
|
|
EXPECT_TRUE(index.is_trained);
|
|
|
|
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()));
|
|
}
|
|
|
|
TEST_F(SVS, StaticVamanaReconstruct) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP32, true};
|
|
index.add(n, test_data.data());
|
|
|
|
std::vector<float> recons(d);
|
|
ASSERT_NO_THROW(index.reconstruct(0, recons.data()));
|
|
for (size_t i = 0; i < d; ++i) {
|
|
EXPECT_EQ(recons[i], test_data[i]);
|
|
}
|
|
}
|
|
|
|
// 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};
|
|
write_and_read_static_vamana_index(
|
|
index, test_data, n, faiss::IO_FLAG_MMAP_IFC);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndMapStaticVamanaFP16) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_FP16, true};
|
|
write_and_read_static_vamana_index(
|
|
index, test_data, n, faiss::IO_FLAG_MMAP_IFC);
|
|
}
|
|
|
|
TEST_F(SVS, WriteAndMapStaticVamanaSQ8) {
|
|
faiss::IndexSVSVamana index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_SQ8, true};
|
|
write_and_read_static_vamana_index(
|
|
index, test_data, n, faiss::IO_FLAG_MMAP_IFC);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndMapStaticVamanaLVQ4x4) {
|
|
faiss::IndexSVSVamanaLVQ index{
|
|
d, 64ul, faiss::METRIC_L2, faiss::SVSStorageKind::SVS_LVQ4x4, true};
|
|
write_and_read_static_vamana_index(
|
|
index, test_data, n, faiss::IO_FLAG_MMAP_IFC);
|
|
}
|
|
|
|
TEST_F(SVSLL, WriteAndMapStaticVamanaLeanVec4x4) {
|
|
faiss::IndexSVSVamanaLeanVec index{
|
|
d,
|
|
64ul,
|
|
faiss::METRIC_L2,
|
|
0,
|
|
faiss::SVSStorageKind::SVS_LeanVec4x4,
|
|
true};
|
|
write_and_read_static_vamana_index(
|
|
index, test_data, n, faiss::IO_FLAG_MMAP_IFC);
|
|
}
|