chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
|
||||
|
||||
cc_directories(cluster)
|
||||
cc_directories(flat)
|
||||
cc_directories(flat_sparse)
|
||||
cc_directories(ivf)
|
||||
cc_directories(hnsw)
|
||||
cc_directories(hnsw_sparse)
|
||||
cc_directories(vamana)
|
||||
|
||||
if(DISKANN_SUPPORTED)
|
||||
message(STATUS "build diskann tests")
|
||||
cc_directory(diskann)
|
||||
else()
|
||||
message(STATUS "skip diskann tests (unsupported platform)")
|
||||
endif()
|
||||
|
||||
if(RABITQ_SUPPORTED)
|
||||
cc_directories(hnsw_rabitq)
|
||||
endif()
|
||||
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_knn_cluster
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,106 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 <cmath>
|
||||
#include <random>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/params.h>
|
||||
#include "zvec/core/framework/index_framework.h"
|
||||
#include "zvec/core/framework/index_meta.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
TEST(KmeansCluster, General) {
|
||||
// Prepare index data
|
||||
const uint32_t count = 5000u;
|
||||
const uint32_t dimension = 33u;
|
||||
|
||||
IndexMeta index_meta;
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
index_meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(0.0, 5.0);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
// IndexCluster::Pointer cluster = std::make_shared<KmeansCluster>();
|
||||
IndexCluster::Pointer cluster = IndexFactory::CreateCluster("KmeansCluster");
|
||||
ASSERT_TRUE(!!cluster);
|
||||
|
||||
zvec::ailego::Params params;
|
||||
params.set("zvec.general.cluster.count", 1);
|
||||
params.set("zvec.kmeans.cluster.count", 56);
|
||||
|
||||
ASSERT_EQ(0, cluster->init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster->mount(features));
|
||||
cluster->suggest(64u);
|
||||
|
||||
auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
|
||||
std::cout << "---------- FIRST ----------\n";
|
||||
std::vector<IndexCluster::Centroid> centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
ASSERT_NE(0, cluster->classify(threads, centroids));
|
||||
ASSERT_NE(0, cluster->label(threads, centroids, &labels));
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- SECOND ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- THIRD ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster->classify(threads, centroids));
|
||||
ASSERT_EQ(0, cluster->label(threads, centroids, &labels));
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "cluster/multi_chunk_cluster.h"
|
||||
#include <cmath>
|
||||
#include <random>
|
||||
#include <ailego/algorithm/kmeans.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/params.h>
|
||||
#include "zvec/core/framework/index_framework.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace zvec::ailego;
|
||||
TEST(MultiChunkCluster, General) {
|
||||
const uint32_t count = 10000u;
|
||||
const uint32_t dimension = 960u;
|
||||
const uint32_t chunk_count = 480u;
|
||||
const uint32_t cluster_count = 256u;
|
||||
// const uint32_t thread_count = 4;
|
||||
const uint32_t thread_count = 16;
|
||||
|
||||
IndexMeta index_meta;
|
||||
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
index_meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(0.0, 5.0);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
MultiChunkCluster cluster = MultiChunkCluster();
|
||||
|
||||
Params params;
|
||||
params.set("zvec.cluster.multi_chunk_cluster.count", cluster_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.chunk_count", chunk_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.thread_count", thread_count);
|
||||
|
||||
ASSERT_EQ(0, cluster.init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster.mount(features));
|
||||
|
||||
IndexCluster::CentroidList centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
|
||||
ASSERT_EQ(0, cluster.cluster(nullptr, centroids));
|
||||
|
||||
for (size_t chunk = 0; chunk < chunk_count; ++chunk) {
|
||||
for (size_t cluster = 0; cluster < cluster_count; ++cluster) {
|
||||
size_t idx = chunk * cluster_count + cluster;
|
||||
const auto ¢ = centroids[idx];
|
||||
const auto &vec = cent.vector<float>();
|
||||
|
||||
std::cout << "chunk: " << chunk << ", cluster: " << cluster
|
||||
<< ", dim: " << vec.size() << ", count: " << cent.follows()
|
||||
<< " (" << cent.score() << ") { " << vec[0] << "," << vec[1]
|
||||
<< " }" << std::endl;
|
||||
ASSERT_EQ(0u, cent.similars().size());
|
||||
}
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster.label(nullptr, centroids, &labels));
|
||||
}
|
||||
|
||||
TEST(MultiChunkCluster, TestChunk) {
|
||||
const uint32_t count = 10000u;
|
||||
const uint32_t dimension = 95;
|
||||
const uint32_t chunk_count = 20u;
|
||||
const uint32_t cluster_count = 256u;
|
||||
// const uint32_t thread_count = 4;
|
||||
const uint32_t thread_count = 16;
|
||||
|
||||
IndexMeta index_meta;
|
||||
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
index_meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(0.0, 5.0);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
MultiChunkCluster cluster = MultiChunkCluster();
|
||||
|
||||
Params params;
|
||||
params.set("zvec.cluster.multi_chunk_cluster.count", cluster_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.chunk_count", chunk_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.thread_count", thread_count);
|
||||
|
||||
ASSERT_EQ(0, cluster.init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster.mount(features));
|
||||
|
||||
IndexCluster::CentroidList centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
|
||||
ASSERT_EQ(0, cluster.cluster(nullptr, centroids));
|
||||
|
||||
for (size_t chunk = 0; chunk < chunk_count; ++chunk) {
|
||||
for (size_t cluster = 0; cluster < cluster_count; ++cluster) {
|
||||
size_t idx = chunk * cluster_count + cluster;
|
||||
const auto ¢ = centroids[idx];
|
||||
const auto &vec = cent.vector<float>();
|
||||
|
||||
std::cout << "chunk: " << chunk << ", cluster: " << cluster
|
||||
<< ", dim: " << vec.size() << ", count: " << cent.follows()
|
||||
<< " (" << cent.score() << ") { " << vec[0] << "," << vec[1]
|
||||
<< " }" << std::endl;
|
||||
ASSERT_EQ(0u, cent.similars().size());
|
||||
}
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster.label(nullptr, centroids, &labels));
|
||||
}
|
||||
|
||||
TEST(MultiChunkCluster, General_InnerProduct) {
|
||||
const uint32_t count = 50000u;
|
||||
const uint32_t dimension = 96u;
|
||||
const uint32_t chunk_count = 12u;
|
||||
const uint32_t cluster_count = 16u;
|
||||
const uint32_t chain_length = 0;
|
||||
const uint32_t thread_count = 16;
|
||||
|
||||
IndexMeta index_meta;
|
||||
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
index_meta.set_metric("InnerProduct", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
// do normalize
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
|
||||
float norm = 0;
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
|
||||
norm += vec[j] * vec[j];
|
||||
}
|
||||
norm = sqrt(norm);
|
||||
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] /= norm;
|
||||
}
|
||||
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
MultiChunkCluster cluster = MultiChunkCluster();
|
||||
|
||||
Params params;
|
||||
params.set("zvec.cluster.multi_chunk_cluster.count", cluster_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.chunk_count", chunk_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.thread_count", thread_count);
|
||||
params.set("zvec.cluster.multi_chunk_cluster.markov_chain_length",
|
||||
chain_length);
|
||||
|
||||
ASSERT_EQ(0, cluster.init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster.mount(features));
|
||||
|
||||
IndexCluster::CentroidList centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
|
||||
ASSERT_EQ(0, cluster.cluster(nullptr, centroids));
|
||||
|
||||
for (size_t chunk = 0; chunk < chunk_count; ++chunk) {
|
||||
for (size_t cluster = 0; cluster < cluster_count; ++cluster) {
|
||||
size_t idx = chunk * cluster_count + cluster;
|
||||
const auto ¢ = centroids[idx];
|
||||
const auto &vec = cent.vector<float>();
|
||||
|
||||
std::cout << "chunk: " << chunk << ", cluster: " << cluster
|
||||
<< ", dim: " << vec.size() << ", count: " << cent.follows()
|
||||
<< " (" << cent.score() << ") { " << vec[0] << ", " << vec[1]
|
||||
<< ", " << vec[2] << ", ... , " << vec[vec.size() - 2] << ", "
|
||||
<< vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, cent.similars().size());
|
||||
}
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster.label(nullptr, centroids, &labels));
|
||||
}
|
||||
@@ -0,0 +1,502 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 <cmath>
|
||||
#include <random>
|
||||
#include <ailego/algorithm/kmeans.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/params.h>
|
||||
#include "zvec/core/framework/index_framework.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
TEST(OptKmeansCluster, General) {
|
||||
// Prepare index data
|
||||
const uint32_t count = 5000u;
|
||||
const uint32_t dimension = 33u;
|
||||
|
||||
IndexMeta index_meta;
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(0.0, 5.0);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
IndexCluster::Pointer cluster =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster);
|
||||
|
||||
Params params;
|
||||
params.set("zvec.general.cluster.count", 1);
|
||||
params.set("zvec.optkmeans.cluster.count", 56);
|
||||
|
||||
ASSERT_EQ(0, cluster->init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster->mount(features));
|
||||
cluster->suggest(64u);
|
||||
|
||||
auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
|
||||
std::cout << "---------- FIRST ----------\n";
|
||||
std::vector<IndexCluster::Centroid> centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
ASSERT_NE(0, cluster->classify(threads, centroids));
|
||||
ASSERT_NE(0, cluster->label(threads, centroids, &labels));
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- SECOND ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- THIRD ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster->classify(threads, centroids));
|
||||
ASSERT_EQ(0, cluster->label(threads, centroids, &labels));
|
||||
}
|
||||
|
||||
// TEST(OptKmeansCluster, NoEmptyCentroids) {
|
||||
// // Prepare index data
|
||||
// const uint32_t count = 500u;
|
||||
// const uint32_t dimension = 8u;
|
||||
|
||||
// IndexMeta index_meta;
|
||||
// index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
// index_meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
// std::shared_ptr<CompactIndexFeatures> features(
|
||||
// new CompactIndexFeatures(index_meta));
|
||||
|
||||
// std::random_device rd;
|
||||
// std::mt19937 gen(rd());
|
||||
// std::uniform_real_distribution<float> dist(0.0, 5.0);
|
||||
|
||||
// for (uint32_t i = 0; i < count; ++i) {
|
||||
// std::vector<float> vec(dimension);
|
||||
// for (size_t j = 0; j < dimension; ++j) {
|
||||
// vec[j] = dist(gen);
|
||||
// }
|
||||
// features->emplace(vec.data());
|
||||
// }
|
||||
|
||||
// // Create a Kmeans cluster
|
||||
// IndexCluster::Pointer cluster =
|
||||
// IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
// ASSERT_TRUE(!!cluster);
|
||||
|
||||
// Params params;
|
||||
// ASSERT_EQ(0, cluster->init(index_meta, params));
|
||||
// ASSERT_EQ(0, cluster->mount(features));
|
||||
// cluster->suggest(20u);
|
||||
|
||||
// auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
// std::vector<IndexCluster::Centroid> centroids;
|
||||
// for (uint32_t i = 0; i < 3; ++i) {
|
||||
// std::vector<float> vec(dimension);
|
||||
// for (size_t j = 0; j < dimension; ++j) {
|
||||
// vec[j] = NAN;
|
||||
// }
|
||||
// centroids.emplace_back(vec.data(), vec.size() * sizeof(float));
|
||||
// }
|
||||
// ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
// ASSERT_EQ(3u, centroids.size());
|
||||
|
||||
// for (uint32_t i = 0; i < 3; ++i) {
|
||||
// std::vector<float> vec(dimension);
|
||||
// for (size_t j = 0; j < dimension; ++j) {
|
||||
// vec[j] = dist(gen);
|
||||
// }
|
||||
// centroids.emplace_back(vec.data(), vec.size() * sizeof(float));
|
||||
// }
|
||||
// ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
// ASSERT_EQ(6u, centroids.size());
|
||||
|
||||
// for (uint32_t i = 0; i < 3; ++i) {
|
||||
// std::vector<float> vec(dimension);
|
||||
// for (size_t j = 0; j < dimension; ++j) {
|
||||
// vec[j] = NAN;
|
||||
// }
|
||||
// centroids.emplace_back(vec.data(), vec.size() * sizeof(float));
|
||||
// }
|
||||
// ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
// ASSERT_EQ(9u, centroids.size());
|
||||
|
||||
// for (uint32_t i = 0; i < 3; ++i) {
|
||||
// std::vector<float> vec(dimension);
|
||||
// for (size_t j = 0; j < dimension; ++j) {
|
||||
// vec[j] = dist(gen);
|
||||
// }
|
||||
// centroids.emplace_back(vec.data(), vec.size() * sizeof(float));
|
||||
// }
|
||||
// ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
// ASSERT_EQ(12u, centroids.size());
|
||||
|
||||
// for (const auto &it : centroids) {
|
||||
// const auto &vec = it.vector<float>();
|
||||
|
||||
// std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ",
|
||||
// "
|
||||
// << vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() -
|
||||
// 2]
|
||||
// << ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
// }
|
||||
|
||||
// params.set("zvec.optkmeans.cluster.purge_empty", true);
|
||||
// cluster->update(params);
|
||||
|
||||
// ASSERT_EQ(12u, centroids.size());
|
||||
// ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
// ASSERT_EQ(7u, centroids.size());
|
||||
// for (const auto &it : centroids) {
|
||||
// const auto &vec = it.vector<float>();
|
||||
|
||||
// std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ",
|
||||
// "
|
||||
// << vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() -
|
||||
// 2]
|
||||
// << ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
// }
|
||||
// }
|
||||
|
||||
TEST(OptKmeansCluster, IN4General) {
|
||||
// Prepare index data
|
||||
const uint32_t count = 5000u;
|
||||
const uint32_t dimension = 64u;
|
||||
const uint32_t dimension_wrong = 66u;
|
||||
|
||||
IndexMeta index_meta;
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_INT4, dimension);
|
||||
index_meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
IndexMeta index_meta_wrong;
|
||||
index_meta_wrong.set_meta(IndexMeta::DataType::DT_INT4, dimension_wrong);
|
||||
index_meta_wrong.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features_wrong(
|
||||
new CompactIndexFeatures(index_meta_wrong));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<unsigned short> dist(0, UINT8_MAX);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<uint8_t> vec(dimension / 2);
|
||||
std::vector<uint8_t> vec_wrong(dimension_wrong / 2);
|
||||
for (size_t j = 0; j < dimension / 2; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
for (size_t j = 0; j < dimension_wrong / 2; ++j) {
|
||||
vec_wrong[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
features_wrong->emplace(vec_wrong.data());
|
||||
}
|
||||
|
||||
// Create a OptKmeans cluster
|
||||
IndexCluster::Pointer cluster =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, cluster->init(index_meta_wrong, params));
|
||||
ASSERT_NE(0, cluster->mount(features_wrong));
|
||||
|
||||
params.set("zvec.general.cluster.count", 1);
|
||||
params.set("zvec.optkmeans.cluster.count", 56);
|
||||
|
||||
ASSERT_EQ(0, cluster->init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster->mount(features));
|
||||
cluster->suggest(64u);
|
||||
|
||||
auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
|
||||
std::cout << "---------- FIRST ----------\n";
|
||||
std::vector<IndexCluster::Centroid> centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
ASSERT_NE(0, cluster->classify(threads, centroids));
|
||||
ASSERT_NE(0, cluster->label(threads, centroids, &labels));
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- SECOND ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- THIRD ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster->classify(threads, centroids));
|
||||
ASSERT_EQ(0, cluster->label(threads, centroids, &labels));
|
||||
}
|
||||
|
||||
|
||||
TEST(OptKmeansCluster, IN4Correctness) {
|
||||
// Prepare index data
|
||||
const uint32_t count = 5000u;
|
||||
const uint32_t dimension = 64u;
|
||||
|
||||
IndexMeta index_meta1;
|
||||
index_meta1.set_meta(IndexMeta::DataType::DT_INT8, dimension);
|
||||
index_meta1.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
IndexMeta index_meta2;
|
||||
index_meta2.set_meta(IndexMeta::DataType::DT_INT4, dimension);
|
||||
index_meta2.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features1(
|
||||
new CompactIndexFeatures(index_meta1));
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features2(
|
||||
new CompactIndexFeatures(index_meta2));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<int> dist(-8, 7);
|
||||
|
||||
// Generate features
|
||||
for (size_t i = 0; i < count; ++i) {
|
||||
NumericalVector<int8_t> vec1(dimension);
|
||||
NibbleVector<int32_t> vec2(dimension);
|
||||
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
int8_t val = (int8_t)dist(gen);
|
||||
vec1[j] = val;
|
||||
vec2.set(j, val);
|
||||
}
|
||||
features1->emplace(vec1.data());
|
||||
features2->emplace(vec2.data());
|
||||
}
|
||||
|
||||
// Create a OptKmeans cluster of int8, and cluster only once
|
||||
IndexCluster::Pointer cluster_once =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster_once);
|
||||
|
||||
Params params_once;
|
||||
params_once.set("zvec.general.cluster.count", 65);
|
||||
params_once.set("zvec.optkmeans.cluster.count", 63);
|
||||
params_once.set("zvec.optkmeans.cluster.max_iterations", 1);
|
||||
// Use KMC2 to init centroids
|
||||
params_once.set("zvec.optkmeans.cluster.markov_chain_length", 20);
|
||||
|
||||
ASSERT_EQ(0, cluster_once->init(index_meta1, params_once));
|
||||
ASSERT_EQ(0, cluster_once->mount(features1));
|
||||
cluster_once->suggest(63);
|
||||
|
||||
auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
|
||||
// Cluster once and get centroids
|
||||
std::vector<IndexCluster::Centroid> centroids1;
|
||||
ASSERT_EQ(0, cluster_once->cluster(threads, centroids1));
|
||||
|
||||
// Use centroids_one as init centroids to both int8 and int4 cluster
|
||||
// Create a int8 cluster
|
||||
IndexCluster::Pointer cluster_int8 =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster_int8);
|
||||
|
||||
Params params_int8;
|
||||
params_int8.set("zvec.general.cluster.count", 65);
|
||||
params_int8.set("zvec.optkmeans.cluster.count", 63);
|
||||
|
||||
ASSERT_EQ(0, cluster_int8->init(index_meta1, params_int8));
|
||||
ASSERT_EQ(0, cluster_int8->mount(features1));
|
||||
cluster_int8->suggest(63u);
|
||||
|
||||
// Create a int4 cluster
|
||||
IndexCluster::Pointer cluster_int4 =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster_int4);
|
||||
|
||||
Params params_int4;
|
||||
params_int4.set("zvec.general.cluster.count", 65);
|
||||
params_int4.set("zvec.optkmeans.cluster.count", 63);
|
||||
|
||||
ASSERT_EQ(0, cluster_int4->init(index_meta2, params_int4));
|
||||
ASSERT_EQ(0, cluster_int4->mount(features2));
|
||||
cluster_int4->suggest(63u);
|
||||
|
||||
std::vector<IndexCluster::Centroid> centroids2;
|
||||
|
||||
// Use centroids of int8 to init centroids of int4
|
||||
for (size_t i = 0; i < centroids1.size(); ++i) {
|
||||
NibbleVector<int8_t> nvec;
|
||||
nvec.assign(reinterpret_cast<const int8_t *>(centroids1[i].feature()),
|
||||
dimension);
|
||||
IndexCluster::Centroid curr_centroid;
|
||||
curr_centroid.set_score(centroids1[i].score());
|
||||
curr_centroid.set_follows(centroids1[i].follows());
|
||||
curr_centroid.set_feature(nvec.data(), nvec.dimension() >> 1);
|
||||
centroids2.push_back(curr_centroid);
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster_int8->cluster(threads, centroids1));
|
||||
ASSERT_EQ(0, cluster_int4->cluster(threads, centroids2));
|
||||
|
||||
EXPECT_EQ(centroids1.size(), centroids2.size());
|
||||
for (size_t i = 0; i < centroids1.size(); ++i) {
|
||||
EXPECT_EQ(centroids1[i].follows(), centroids2[i].follows());
|
||||
EXPECT_DOUBLE_EQ(centroids1[i].score(), centroids2[i].score());
|
||||
}
|
||||
}
|
||||
|
||||
TEST(OptKmeansCluster, InnerProduct) {
|
||||
// Prepare index data
|
||||
const uint32_t count = 5000u;
|
||||
const uint32_t dimension = 33u;
|
||||
|
||||
IndexMeta index_meta;
|
||||
index_meta.set_meta(IndexMeta::DataType::DT_FP32, dimension);
|
||||
index_meta.set_metric("InnerProduct", 0, Params());
|
||||
|
||||
std::shared_ptr<CompactIndexFeatures> features(
|
||||
new CompactIndexFeatures(index_meta));
|
||||
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
std::vector<float> vec(dimension);
|
||||
for (size_t j = 0; j < dimension; ++j) {
|
||||
vec[j] = dist(gen);
|
||||
}
|
||||
features->emplace(vec.data());
|
||||
}
|
||||
|
||||
// Create a Kmeans cluster
|
||||
IndexCluster::Pointer cluster =
|
||||
IndexFactory::CreateCluster("OptKmeansCluster");
|
||||
ASSERT_TRUE(!!cluster);
|
||||
|
||||
Params params;
|
||||
params.set("zvec.general.cluster.count", 1);
|
||||
params.set("zvec.optkmeans.cluster.count", 56);
|
||||
|
||||
ASSERT_EQ(0, cluster->init(index_meta, params));
|
||||
ASSERT_EQ(0, cluster->mount(features));
|
||||
cluster->suggest(64u);
|
||||
|
||||
auto threads = std::make_shared<SingleQueueIndexThreads>();
|
||||
|
||||
std::cout << "---------- FIRST ----------\n";
|
||||
std::vector<IndexCluster::Centroid> centroids;
|
||||
std::vector<uint32_t> labels;
|
||||
ASSERT_NE(0, cluster->classify(threads, centroids));
|
||||
ASSERT_NE(0, cluster->label(threads, centroids, &labels));
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- SECOND ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
std::cout << "---------- THIRD ----------\n";
|
||||
ASSERT_EQ(0, cluster->cluster(threads, centroids));
|
||||
|
||||
for (const auto &it : centroids) {
|
||||
const auto &vec = it.vector<float>();
|
||||
|
||||
std::cout << it.follows() << " (" << it.score() << ") { " << vec[0] << ", "
|
||||
<< vec[1] << ", " << vec[2] << ", ... , " << vec[vec.size() - 2]
|
||||
<< ", " << vec[vec.size() - 1] << " }" << std::endl;
|
||||
ASSERT_EQ(0u, it.similars().size());
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, cluster->classify(threads, centroids));
|
||||
ASSERT_EQ(0, cluster->label(threads, centroids, &labels));
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_plugin core_knn_diskann
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/diskann
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,169 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "diskann_builder.h"
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <chrono>
|
||||
#include <future>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include "diskann_holder.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
constexpr size_t static dim = 64;
|
||||
|
||||
class DiskAnnBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string _dir;
|
||||
static shared_ptr<IndexMeta> _index_meta_ptr;
|
||||
};
|
||||
|
||||
std::string DiskAnnBuilderTest::_dir("DiskAnnBuilderTest");
|
||||
shared_ptr<IndexMeta> DiskAnnBuilderTest::_index_meta_ptr;
|
||||
|
||||
void DiskAnnBuilderTest::SetUp(void) {
|
||||
LoggerBroker::SetLevel(Logger::LEVEL_INFO);
|
||||
|
||||
_index_meta_ptr.reset(new (nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
_index_meta_ptr->set_metric("SquaredEuclidean", 0, Params());
|
||||
}
|
||||
|
||||
void DiskAnnBuilderTest::TearDown(void) {
|
||||
char cmdBuf[100];
|
||||
snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str());
|
||||
system(cmdBuf);
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnBuilderTest, TestGeneral) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 50);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestGeneral";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
}
|
||||
|
||||
// Regression test: building a small DiskAnn index must complete quickly.
|
||||
// A lost-wakeup bug in the condition-variable progress loops previously caused
|
||||
// 15–30 second stalls during train/build on small datasets because
|
||||
// notify_one() was either missing or racing against a wrong predicate.
|
||||
TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) {
|
||||
constexpr size_t kSmallDim = 4;
|
||||
constexpr size_t kSmallDocCnt = 12;
|
||||
|
||||
auto meta = make_shared<IndexMeta>(IndexMeta::DataType::DT_FP32, kSmallDim);
|
||||
meta->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder = make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(
|
||||
kSmallDim);
|
||||
for (size_t i = 0; i < kSmallDocCnt; ++i) {
|
||||
NumericalVector<float> vec(kSmallDim, static_cast<float>(i));
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 50);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 2);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*meta, params));
|
||||
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
auto t1 = std::chrono::steady_clock::now();
|
||||
|
||||
auto elapsed_ms =
|
||||
std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
|
||||
// Before the fix, this took 15–30 seconds. After the fix, it should
|
||||
// complete in well under 5 seconds even on slow CI machines.
|
||||
EXPECT_LT(elapsed_ms, 5000)
|
||||
<< "DiskAnn build with " << kSmallDocCnt << " vectors took " << elapsed_ms
|
||||
<< " ms — likely a lost-wakeup regression in progress loops.";
|
||||
}
|
||||
|
||||
// DiskAnn is now exposed implicitly: no caller ever invokes a
|
||||
// ``LoadDiskAnnPlugin`` / ``IsLibAioAvailable`` API (those were removed from
|
||||
// the public surface together with ``zvec.load_diskann_plugin()`` in Python).
|
||||
// The only contract this test validates is the UX guarantee: once the DiskAnn
|
||||
// module has been linked into the hosting binary (here, directly into the
|
||||
// test via the ``core_knn_diskann`` target), its factory entries are
|
||||
// registered automatically and the global ``IndexFactory`` can hand out a
|
||||
// ``DiskAnnBuilder`` without any explicit setup step. On hosts missing
|
||||
// libaio, DiskAnn would fail at the index-creation layer with a clear error
|
||||
// while other index types (HNSW/IVF/Flat/Vamana) remain unaffected; that
|
||||
// runtime branch lives in ``DiskAnnIndex::CreateAndInitStreamer`` and is
|
||||
// covered by the higher-level interface tests.
|
||||
TEST_F(DiskAnnBuilderTest, TestImplicitFactoryRegistration) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr)
|
||||
<< "DiskAnnBuilder factory entry missing: DiskAnn must be available "
|
||||
"without any manual plugin load step.";
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("DiskAnnStreamer");
|
||||
ASSERT_NE(streamer, nullptr)
|
||||
<< "DiskAnnStreamer factory entry missing: DiskAnn must be available "
|
||||
"without any manual plugin load step.";
|
||||
}
|
||||
@@ -0,0 +1,816 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "diskann_searcher.h"
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <ailego/math/distance.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include "diskann_holder.h"
|
||||
#include "diskann_params.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
constexpr size_t static dim = 64;
|
||||
|
||||
class DiskAnnSearcherTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string _dir;
|
||||
static shared_ptr<IndexMeta> _index_meta_ptr;
|
||||
};
|
||||
|
||||
std::string DiskAnnSearcherTest::_dir("DiskAnnSearcherTest/");
|
||||
shared_ptr<IndexMeta> DiskAnnSearcherTest::_index_meta_ptr;
|
||||
|
||||
void DiskAnnSearcherTest::SetUp(void) {
|
||||
LoggerBroker::SetLevel(Logger::LEVEL_INFO);
|
||||
|
||||
_index_meta_ptr.reset(new (nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
_index_meta_ptr->set_metric("SquaredEuclidean", 0, Params());
|
||||
}
|
||||
|
||||
void DiskAnnSearcherTest::TearDown(void) {
|
||||
char cmdBuf[100];
|
||||
snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str());
|
||||
system(cmdBuf);
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestGeneral) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestGeneral";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
auto ctx = searcher->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto linearByPKeysCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
ASSERT_TRUE(!!linearByPKeysCtx);
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
size_t topk = 200;
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
linearCtx->set_topk(topk);
|
||||
linearByPKeysCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
// do linear search test
|
||||
{
|
||||
float query[dim];
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
query[i] = 3.1f;
|
||||
}
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(query, qmeta, linearCtx));
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(3UL, linearResult[0].key());
|
||||
ASSERT_EQ(4UL, linearResult[1].key());
|
||||
ASSERT_EQ(2UL, linearResult[2].key());
|
||||
ASSERT_EQ(5UL, linearResult[3].key());
|
||||
ASSERT_EQ(1UL, linearResult[4].key());
|
||||
ASSERT_EQ(6UL, linearResult[5].key());
|
||||
ASSERT_EQ(0UL, linearResult[6].key());
|
||||
ASSERT_EQ(7UL, linearResult[7].key());
|
||||
for (size_t i = 8; i < topk; ++i) {
|
||||
ASSERT_EQ(i, linearResult[i].key());
|
||||
}
|
||||
}
|
||||
|
||||
// do linear search by p_keys test
|
||||
std::vector<std::vector<uint64_t>> p_keys;
|
||||
p_keys.resize(1);
|
||||
p_keys[0] = {8, 9, 10, 11, 3, 2, 1, 0};
|
||||
{
|
||||
float query[dim];
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
query[i] = 3.1f;
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(query, p_keys, qmeta,
|
||||
linearByPKeysCtx));
|
||||
auto &linearByPKeysResult = linearByPKeysCtx->result();
|
||||
ASSERT_EQ(8, linearByPKeysResult.size());
|
||||
ASSERT_EQ(3UL, linearByPKeysResult[0].key());
|
||||
ASSERT_EQ(2UL, linearByPKeysResult[1].key());
|
||||
ASSERT_EQ(1UL, linearByPKeysResult[2].key());
|
||||
ASSERT_EQ(0UL, linearByPKeysResult[3].key());
|
||||
ASSERT_EQ(8UL, linearByPKeysResult[4].key());
|
||||
ASSERT_EQ(9UL, linearByPKeysResult[5].key());
|
||||
ASSERT_EQ(10UL, linearByPKeysResult[6].key());
|
||||
ASSERT_EQ(11UL, linearByPKeysResult[7].key());
|
||||
}
|
||||
|
||||
size_t step = 500;
|
||||
for (size_t i = 0; i < doc_cnt; i += step) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
auto t1 = Realtime::MicroSeconds();
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
auto t2 = Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto t3 = Realtime::MicroSeconds();
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
// TODO: check
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float recall = totalHits * step * step * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * step * 1.0f / doc_cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
|
||||
EXPECT_GT(recall, 0.90f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestNodeCache) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestNodeCache";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.cache_node_num", 32);
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
auto ctx = searcher->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto linearByPKeysCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
ASSERT_TRUE(!!linearByPKeysCtx);
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
size_t topk = 200;
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
linearCtx->set_topk(topk);
|
||||
linearByPKeysCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
size_t step = 500;
|
||||
for (size_t i = 0; i < doc_cnt; i += step) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
auto t1 = Realtime::MicroSeconds();
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
auto t2 = Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto t3 = Realtime::MicroSeconds();
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
// TODO: check
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float recall = totalHits * step * step * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * step * 1.0f / doc_cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
|
||||
EXPECT_GT(recall, 0.90f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestFilter) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestFilter";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.cache_node_num", 32);
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
auto ctx = searcher->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto linearByPKeysCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
ASSERT_TRUE(!!linearByPKeysCtx);
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
linearByPKeysCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
size_t key = 50;
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = key + 0.1f;
|
||||
}
|
||||
|
||||
// no filter
|
||||
{
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
ASSERT_EQ(50UL, knnResult[0].key());
|
||||
ASSERT_EQ(51UL, knnResult[1].key());
|
||||
ASSERT_EQ(49UL, knnResult[2].key());
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(50UL, linearResult[0].key());
|
||||
ASSERT_EQ(51UL, linearResult[1].key());
|
||||
ASSERT_EQ(49UL, linearResult[2].key());
|
||||
}
|
||||
|
||||
// with filter
|
||||
{
|
||||
auto filterFunc = [](uint64_t key) {
|
||||
if (key == 50UL || key == 51UL || key == 49UL) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
|
||||
knnCtx->set_filter(filterFunc);
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
ASSERT_EQ(52UL, knnResult[0].key());
|
||||
ASSERT_EQ(48UL, knnResult[1].key());
|
||||
ASSERT_EQ(53UL, knnResult[2].key());
|
||||
|
||||
linearCtx->set_filter(filterFunc);
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(52UL, linearResult[0].key());
|
||||
ASSERT_EQ(48UL, linearResult[1].key());
|
||||
ASSERT_EQ(53UL, linearResult[2].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestGroup) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i / 10.0;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestGroup";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
auto ctx = searcher->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
size_t group_topk = 20;
|
||||
uint64_t total_time = 0;
|
||||
|
||||
auto groupbyFunc = [](uint64_t key) {
|
||||
uint32_t group_id = key / 10 % 10;
|
||||
|
||||
// std::cout << "key: " << key << ", group id: " << group_id << std::endl;
|
||||
|
||||
return std::string("g_") + std::to_string(group_id);
|
||||
};
|
||||
|
||||
size_t group_num = 5;
|
||||
|
||||
ctx->set_group_params(group_num, group_topk);
|
||||
ctx->set_group_by(groupbyFunc);
|
||||
|
||||
size_t query_value = doc_cnt / 2;
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = query_value / 10 + 0.1f;
|
||||
}
|
||||
|
||||
auto t1 = Realtime::MicroSeconds();
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx));
|
||||
auto t2 = Realtime::MicroSeconds();
|
||||
|
||||
total_time += t2 - t1;
|
||||
|
||||
auto &group_result = ctx->group_result();
|
||||
|
||||
for (uint32_t i = 0; i < group_result.size(); ++i) {
|
||||
const std::string &group_id = group_result[i].group_id();
|
||||
auto &result = group_result[i].docs();
|
||||
|
||||
ASSERT_GT(result.size(), 0);
|
||||
std::cout << "Group ID: " << group_id << std::endl;
|
||||
|
||||
for (uint32_t j = 0; j < result.size(); ++j) {
|
||||
std::cout << "\tKey: " << result[j].key() << std::fixed
|
||||
<< std::setprecision(3) << ", Score: " << result[j].score()
|
||||
<< std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
#if 0
|
||||
// do linear search by p_keys test
|
||||
auto groupbyFuncLinear = [](uint64_t key) {
|
||||
uint32_t group_id = key % 10;
|
||||
|
||||
return std::string("g_") + std::to_string(group_id);
|
||||
};
|
||||
|
||||
auto linear_pk_ctx = searcher->create_context();
|
||||
|
||||
linear_pk_ctx->set_group_params(group_num, group_topk);
|
||||
linear_pk_ctx->set_group_by(groupbyFuncLinear);
|
||||
|
||||
std::vector<std::vector<uint64_t>> p_keys;
|
||||
p_keys.resize(1);
|
||||
p_keys[0] = {4, 3, 2, 1, 5, 6, 7, 8, 9, 10};
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(vec.data(), p_keys, qmeta,
|
||||
linear_pk_ctx));
|
||||
auto &linear_by_pkeys_group_result = linear_pk_ctx->group_result();
|
||||
ASSERT_EQ(linear_by_pkeys_group_result.size(), group_num);
|
||||
|
||||
for (uint32_t i = 0; i < linear_by_pkeys_group_result.size(); ++i) {
|
||||
const std::string &group_id = linear_by_pkeys_group_result[i].group_id();
|
||||
auto &result = linear_by_pkeys_group_result[i].docs();
|
||||
|
||||
ASSERT_GT(result.size(), 0);
|
||||
std::cout << "Group ID: " << group_id << std::endl;
|
||||
|
||||
for (uint32_t j = 0; j < result.size(); ++j) {
|
||||
std::cout << "\tKey: " << result[j].key() << std::fixed
|
||||
<< std::setprecision(3) << ", Score: " << result[j].score()
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
ASSERT_EQ(10 - i, result[0].key());
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestFetchVector) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestFetchVector";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
|
||||
size_t query_cnt = 20U;
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
auto linearByPKeysCtx = searcher->create_context();
|
||||
knnCtx->set_fetch_vector(true);
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i += doc_cnt / 10) {
|
||||
std::string vec_value;
|
||||
ASSERT_EQ(0, searcher->get_vector(i, linearCtx, vec_value));
|
||||
|
||||
float vector_value = *(const float *)(vec_value.data());
|
||||
ASSERT_EQ(vector_value, i);
|
||||
}
|
||||
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t i = 0; i < query_cnt; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
|
||||
auto t1 = Realtime::MicroSeconds();
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
auto t2 = Realtime::MicroSeconds();
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto t3 = Realtime::MicroSeconds();
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
ASSERT_NE(knnResult[0].vector_string(), "");
|
||||
float vector_value = *((float *)(knnResult[0].vector_string().data()));
|
||||
ASSERT_EQ(vector_value, i);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(DiskAnnSearcherTest, TestRnnSearch) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 10000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
|
||||
params.set("zvec.diskann.builder.max_degree", 32);
|
||||
params.set("zvec.diskann.builder.list_size", 300);
|
||||
params.set("zvec.diskann.builder.max_pq_chunk_num", 32);
|
||||
params.set("zvec.diskann.builder.threads", 4);
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "/TestRnnSearch";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(doc_cnt, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_GT(stats.trained_costtime(), 0UL);
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
|
||||
// test searcher
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("DiskAnnSearcher");
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
|
||||
Params search_params;
|
||||
search_params.set("zvec.diskann.searcher.list_size", 500);
|
||||
|
||||
ASSERT_EQ(0, searcher->init(search_params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer()));
|
||||
|
||||
auto ctx = searcher->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 0.0;
|
||||
}
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
size_t topk = 50;
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &results = ctx->result();
|
||||
ASSERT_EQ(topk, results.size());
|
||||
|
||||
float radius = results[topk / 2].score();
|
||||
ctx->set_threshold(radius);
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx));
|
||||
ASSERT_GT(topk, results.size());
|
||||
for (size_t k = 0; k < results.size(); ++k) {
|
||||
ASSERT_GE(radius, results[k].score());
|
||||
}
|
||||
|
||||
// Test Reset Threshold
|
||||
ctx->reset_threshold();
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx));
|
||||
ASSERT_EQ(topk, results.size());
|
||||
ASSERT_LT(radius, results[topk - 1].score());
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_flat
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,334 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "flat/flat_builder.h"
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
static inline size_t RandomDimension(void) {
|
||||
std::mt19937 gen((std::random_device())());
|
||||
return (std::uniform_int_distribution<size_t>(1, 129))(gen);
|
||||
}
|
||||
|
||||
static size_t DIMENSION = RandomDimension();
|
||||
class FlatBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
public:
|
||||
static std::string dir_;
|
||||
static IndexMeta meta_;
|
||||
};
|
||||
|
||||
std::string FlatBuilderTest ::dir_("flat_builder_test/");
|
||||
IndexMeta FlatBuilderTest::meta_;
|
||||
|
||||
void FlatBuilderTest::SetUp(void) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
}
|
||||
|
||||
//! self-check column-major and row-major search.
|
||||
void FlatBuilderTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void build_process(IndexBuilder::Pointer &builder,
|
||||
IndexHolder::Pointer holder) {
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(FlatBuilderTest::meta_, params));
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
std::string path = FlatBuilderTest::dir_ + "TestGeneral";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInitSuccess) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInitFailedWithInvalidMeasure) {
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("invalid", 0, Params());
|
||||
Params params;
|
||||
int ret = builder->init(meta_, params);
|
||||
EXPECT_EQ(IndexError_InvalidArgument, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8InvalidColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim + 2);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
ASSERT_EQ(IndexMeta::MO_COLUMN, meta_.major_order());
|
||||
Params params;
|
||||
ASSERT_NE(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithRandomDimension) {
|
||||
size_t dim = DIMENSION;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_UNDEFINED);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestBuildWithRowMajor) {
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8BuildWithRowMajor) {
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_meta(IndexMeta::DT_INT8, DIMENSION);
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(DIMENSION);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestBuildWithColumnMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8BuildWithColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(meta_, params));
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder = std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(dim);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
int ret = builder->train(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder->build(holder);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestWithRowMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithRowMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_ROW);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(DIMENSION);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestWithColumnMajor) {
|
||||
meta_.set_meta(IndexMeta::DataType::DT_FP32, DIMENSION);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder =
|
||||
std::make_shared<OnePassIndexHolder<IndexMeta::DT_FP32>>(DIMENSION);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(DIMENSION);
|
||||
for (size_t j = 0; j < DIMENSION; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
TEST_F(FlatBuilderTest, TestInt8WithColumnMajor) {
|
||||
size_t dim = (DIMENSION + 3) / 4 * 4;
|
||||
meta_.set_meta(IndexMeta::DataType::DT_INT8, dim);
|
||||
meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
meta_.set_major_order(IndexMeta::MO_COLUMN);
|
||||
IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("FlatBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
Params params;
|
||||
std::string path = dir_ + "TestGeneral";
|
||||
|
||||
auto holder = std::make_shared<OnePassIndexHolder<IndexMeta::DT_INT8>>(dim);
|
||||
size_t doc_cnt = 128UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<int8_t> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = (int8_t)(i % 128);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
build_process(builder, holder);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,623 @@
|
||||
#include <future>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
constexpr size_t static dim = 16;
|
||||
|
||||
class FlatStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void hybrid_scale(std::vector<float> &dense_value,
|
||||
std::vector<float> &sparse_value, float alpha_scale);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatStreamerTest::dir_("flat_streamer_buffer_test_dir/");
|
||||
std::shared_ptr<IndexMeta> FlatStreamerTest::index_meta_ptr_;
|
||||
|
||||
void FlatStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (std::nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void FlatStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearch) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBuffer) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBufferMMap) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchWithLRU) {
|
||||
MemoryLimitPool::get_instance().init(100 * 1024UL * 1024UL);
|
||||
#ifdef __ANDROID__
|
||||
GTEST_SKIP()
|
||||
<< "Skipped on Android: requires ~6GB memory/disk (emulator limit)";
|
||||
#endif
|
||||
constexpr size_t static dim = 1600;
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
IndexMeta meta = IndexMeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, Params());
|
||||
ASSERT_EQ(0, write_streamer->init(meta, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/Test/LinearSearchWithLRU", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 50000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(meta, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "/Test/LinearSearchWithLRU", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < 10; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/LinearSearchMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/LinearSearchMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
const float *data = (float *)provider->get_vector(result1[0].key());
|
||||
EXPECT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
@@ -0,0 +1,231 @@
|
||||
#include <future>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/utility/file_helper.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
constexpr size_t static dim = 128;
|
||||
|
||||
class FlatStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void hybrid_scale(std::vector<float> &dense_value,
|
||||
std::vector<float> &sparse_value, float alpha_scale);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatStreamerTest::dir_("streamer_test/");
|
||||
std::shared_ptr<IndexMeta> FlatStreamerTest::index_meta_ptr_;
|
||||
|
||||
void FlatStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (std::nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
zvec::ailego::FileHelper::RemovePath(dir_.c_str());
|
||||
}
|
||||
|
||||
void FlatStreamerTest::TearDown(void) {
|
||||
zvec::ailego::FileHelper::RemovePath(dir_.c_str());
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/Test/LinearSearchMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t data_cnt = 300000UL, cnt = 500UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < data_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "/Test/LinearSearchMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 30;
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(FlatStreamerTest, TestLinearSearchBuffer) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/Test/LinearSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t data_cnt = 300000UL, cnt = 500UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < data_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "/Test/LinearSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 30;
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result1 = ctx->result();
|
||||
// ASSERT_EQ(topk, result1.size());
|
||||
// ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// vec[j] = i + 0.1f;
|
||||
// }
|
||||
// ctx->set_topk(topk);
|
||||
// ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
// auto &result2 = ctx->result();
|
||||
// ASSERT_EQ(topk, result2.size());
|
||||
// ASSERT_EQ(i, result2[0].key());
|
||||
// ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
// ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
cout << "Elapsed time: " << elapsed_time.micro_seconds() << " us" << endl;
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_flat_sparse
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,300 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "flat_sparse/flat_sparse_builder.h"
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
class FlatSparseBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string _dir;
|
||||
static shared_ptr<IndexMeta> _index_meta_ptr;
|
||||
};
|
||||
|
||||
std::string FlatSparseBuilderTest::_dir("FlatSparseBuilderTest/");
|
||||
shared_ptr<IndexMeta> FlatSparseBuilderTest::_index_meta_ptr;
|
||||
|
||||
void FlatSparseBuilderTest::SetUp(void) {
|
||||
_index_meta_ptr.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
_index_meta_ptr->set_metric("InnerProductSparse", 0, Params());
|
||||
}
|
||||
|
||||
void FlatSparseBuilderTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(_dir);
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseBuilderTest, TestGeneral) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("FlatSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder = make_shared<OnePassIndexSparseHolder<IndexMeta::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestGeneral";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_EQ(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 = make_shared<MultiPassIndexSparseHolder<IndexMeta::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, vec));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder->train(holder2));
|
||||
ASSERT_EQ(0, builder->build(holder2));
|
||||
auto dumper2 = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper2, nullptr);
|
||||
ASSERT_EQ(0, dumper2->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper2));
|
||||
ASSERT_EQ(0, dumper2->close());
|
||||
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_EQ(stats.built_costtime(), 0UL);
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseBuilderTest, TestIndexThreads) {
|
||||
IndexBuilder::Pointer builder1 =
|
||||
IndexFactory::CreateBuilder("FlatSparseBuilder");
|
||||
ASSERT_NE(builder1, nullptr);
|
||||
IndexBuilder::Pointer builder2 =
|
||||
IndexFactory::CreateBuilder("FlatSparseBuilder");
|
||||
ASSERT_NE(builder2, nullptr);
|
||||
|
||||
auto holder = make_shared<MultiPassIndexSparseHolder<IndexMeta::DT_FP32>>();
|
||||
|
||||
size_t doc_cnt = 1000UL;
|
||||
uint32_t sparse_count = 32;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
std::srand(Realtime::MilliSeconds());
|
||||
auto threads =
|
||||
std::make_shared<SingleQueueIndexThreads>(std::rand() % 4, false);
|
||||
ASSERT_EQ(0, builder1->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder2->init(*_index_meta_ptr, params));
|
||||
|
||||
auto build_index1 = [&]() {
|
||||
ASSERT_EQ(0, builder1->train(threads, holder));
|
||||
ASSERT_EQ(0, builder1->build(threads, holder));
|
||||
};
|
||||
auto build_index2 = [&]() {
|
||||
ASSERT_EQ(0, builder2->train(threads, holder));
|
||||
ASSERT_EQ(0, builder2->build(threads, holder));
|
||||
};
|
||||
|
||||
auto t1 = std::async(std::launch::async, build_index1);
|
||||
auto t2 = std::async(std::launch::async, build_index2);
|
||||
t1.wait();
|
||||
t2.wait();
|
||||
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndexThreads";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder1->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder2->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats1 = builder1->stats();
|
||||
ASSERT_EQ(doc_cnt, stats1.built_count());
|
||||
auto &stats2 = builder2->stats();
|
||||
ASSERT_EQ(doc_cnt, stats2.built_count());
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseBuilderTest, TestHalfFloatConverter) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("FlatSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder = make_shared<OnePassIndexSparseHolder<IndexMeta::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params converter_params;
|
||||
auto converter = IndexFactory::CreateConverter("HalfFloatSparseConverter");
|
||||
converter->init(*_index_meta_ptr, converter_params);
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
converter->transform(holder);
|
||||
|
||||
auto converted_holder = converter->sparse_result();
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, builder->init(index_meta, converter_params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(converted_holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(converted_holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestHalFloatConverter";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_EQ(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 = make_shared<MultiPassIndexSparseHolder<IndexMeta::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, vec));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder->train(holder2));
|
||||
ASSERT_EQ(0, builder->build(holder2));
|
||||
auto dumper2 = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper2, nullptr);
|
||||
ASSERT_EQ(0, dumper2->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper2));
|
||||
ASSERT_EQ(0, dumper2->close());
|
||||
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_EQ(stats.built_costtime(), 0UL);
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
@@ -0,0 +1,819 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 <future>
|
||||
#include <iostream>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
#include <ailego/math/norm2_matrix.h>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
#include "zvec/core/framework/index_factory.h"
|
||||
#include "zvec/core/framework/index_meta.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
constexpr size_t static sparse_dim_count = 16;
|
||||
|
||||
class FlatSparseSearcherTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatSparseSearcherTest::dir_("searcher_test/");
|
||||
std::shared_ptr<IndexMeta> FlatSparseSearcherTest::index_meta_ptr_;
|
||||
|
||||
void FlatSparseSearcherTest::generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm) {
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
// prepare sparse
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_vec(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_vec[j] = dist(gen);
|
||||
}
|
||||
|
||||
float norm;
|
||||
Norm2Matrix<float, 1>::Compute(sparse_vec.data(), sparse_dim_count, &norm);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_vec[j] = sparse_vec[j] / norm;
|
||||
}
|
||||
|
||||
sparse_indices_list.push_back(sparse_indices);
|
||||
sparse_vec_list.push_back(sparse_vec);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void FlatSparseSearcherTest::SetUp(void) {
|
||||
IndexLoggerBroker::SetLevel(2);
|
||||
|
||||
index_meta_ptr_.reset(new IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
index_meta_ptr_->set_metric("InnerProductSparse", 0, Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void FlatSparseSearcherTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestGeneral) {
|
||||
// init storage
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_TRUE(storage != nullptr);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestGeneral", true));
|
||||
|
||||
|
||||
// init streamer
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(streamer != nullptr);
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, Params());
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// generate sparse data
|
||||
size_t sparse_dim_count = 32;
|
||||
size_t cnt = 100U;
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
// test add data
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
|
||||
sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
// test get data
|
||||
uint32_t sparse_count;
|
||||
std::string sparse_indices_buffer;
|
||||
std::string sparse_values_buffer;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(
|
||||
0, streamer->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
|
||||
&sparse_values_buffer));
|
||||
ASSERT_EQ(sparse_dim_count, sparse_count);
|
||||
const uint32_t *sparse_indices_ptr =
|
||||
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
|
||||
const float *sparse_values_ptr =
|
||||
reinterpret_cast<const float *>(sparse_values_buffer.data());
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
|
||||
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
|
||||
// std::cout << "1: " << sparse_values_ptr[j]
|
||||
// << " 2: " << sparse_vec_list[i][j] << std::endl;
|
||||
}
|
||||
|
||||
// must clear ^_^
|
||||
sparse_indices_buffer.clear();
|
||||
sparse_values_buffer.clear();
|
||||
}
|
||||
|
||||
// test dump
|
||||
auto path = dir_ + "TestGeneral_dump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher get vector
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_TRUE(read_storage != nullptr);
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
|
||||
// test searcher get data
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(
|
||||
0, searcher->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
|
||||
&sparse_values_buffer));
|
||||
ASSERT_EQ(sparse_dim_count, sparse_count);
|
||||
const uint32_t *sparse_indices_ptr =
|
||||
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
|
||||
const float *sparse_values_ptr =
|
||||
reinterpret_cast<const float *>(sparse_values_buffer.data());
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
|
||||
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
|
||||
// std::cout << "1: " << sparse_values_ptr[j]
|
||||
// << " 2: " << sparse_vec_list[i][j] << std::endl;
|
||||
}
|
||||
|
||||
// must clear ^_^
|
||||
sparse_indices_buffer.clear();
|
||||
sparse_values_buffer.clear();
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestStreamerDump) {
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_NE(streamer, nullptr);
|
||||
|
||||
Params params;
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestStreamerDump.index", true));
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 10000U;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
|
||||
sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
auto path = dir_ + "TestStreamerDump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher knn
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 50) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, searcher->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0,
|
||||
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, linearCtx));
|
||||
|
||||
auto t3 = Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
auto &linearResult = linearCtx->result();
|
||||
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
ASSERT_EQ(linearResult[k].key(), knnResult[k].key());
|
||||
}
|
||||
}
|
||||
|
||||
printf("linear: %zu, knn: %zu\n", (size_t)linearTotalTime,
|
||||
(size_t)knnTotalTime);
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestLoadClose) {
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_TRUE(storage != nullptr);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestGeneral", true));
|
||||
|
||||
|
||||
// init streamer
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(streamer != nullptr);
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, Params());
|
||||
|
||||
Params params;
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// generate sparse data
|
||||
size_t sparse_dim_count = 32;
|
||||
size_t cnt = 100U;
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
// test add data
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count,
|
||||
sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
// test get data
|
||||
uint32_t sparse_count;
|
||||
std::string sparse_indices_buffer;
|
||||
std::string sparse_values_buffer;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(
|
||||
0, streamer->get_sparse_vector(i, &sparse_count, &sparse_indices_buffer,
|
||||
&sparse_values_buffer));
|
||||
ASSERT_EQ(sparse_dim_count, sparse_count);
|
||||
const uint32_t *sparse_indices_ptr =
|
||||
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
|
||||
const float *sparse_values_ptr =
|
||||
reinterpret_cast<const float *>(sparse_values_buffer.data());
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
|
||||
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
|
||||
// std::cout << "1: " << sparse_values_ptr[j]
|
||||
// << " 2: " << sparse_vec_list[i][j] << std::endl;
|
||||
}
|
||||
|
||||
// must clear ^_^
|
||||
sparse_indices_buffer.clear();
|
||||
sparse_values_buffer.clear();
|
||||
}
|
||||
|
||||
// test dump
|
||||
auto path = dir_ + "TestGeneral_dump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher get vector
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_TRUE(read_storage != nullptr);
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
|
||||
uint32_t loop = 5;
|
||||
while (loop--) {
|
||||
std::cout << "loop: " << loop << std::endl;
|
||||
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
|
||||
// test searcher get data
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, searcher->get_sparse_vector(i, &sparse_count,
|
||||
&sparse_indices_buffer,
|
||||
&sparse_values_buffer));
|
||||
ASSERT_EQ(sparse_dim_count, sparse_count);
|
||||
const uint32_t *sparse_indices_ptr =
|
||||
reinterpret_cast<const uint32_t *>(sparse_indices_buffer.data());
|
||||
const float *sparse_values_ptr =
|
||||
reinterpret_cast<const float *>(sparse_values_buffer.data());
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
ASSERT_EQ(sparse_indices_ptr[j], sparse_indices_list[i][j]);
|
||||
ASSERT_FLOAT_EQ(sparse_values_ptr[j], sparse_vec_list[i][j]);
|
||||
// std::cout << "1: " << sparse_values_ptr[j]
|
||||
// << " 2: " << sparse_vec_list[i][j] << std::endl;
|
||||
}
|
||||
|
||||
// must clear ^_^
|
||||
sparse_indices_buffer.clear();
|
||||
sparse_values_buffer.clear();
|
||||
}
|
||||
|
||||
ASSERT_EQ(searcher->unload(), 0);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestSearch) {
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearch.index", true));
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 5000UL;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_velues(sparse_dim_count);
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = -1.0 * i - 1.0f;
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
// test dump
|
||||
auto path = dir_ + "TestGeneral_dump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher get vector
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_TRUE(read_storage != nullptr);
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
|
||||
size_t step = 50;
|
||||
for (size_t i = 0; i < cnt; i += step) {
|
||||
// std::cout << "search " << i << std::endl;
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = i + 1.0f;
|
||||
}
|
||||
|
||||
ctx->set_topk(1U);
|
||||
ASSERT_EQ(0,
|
||||
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(1UL, result1.size());
|
||||
ASSERT_EQ(0, result1[0].key());
|
||||
// std::cout << result1[0].key() << " " << result1[0].score() << std::endl;
|
||||
|
||||
ctx->set_topk(3U);
|
||||
ASSERT_EQ(0,
|
||||
searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(3UL, result2.size());
|
||||
for (size_t i = 0; i < 3UL; ++i) {
|
||||
// std::cout << result2[i].key() << " " << result2[i].score() <<
|
||||
// std::endl;
|
||||
ASSERT_EQ(i, result2[i].key());
|
||||
}
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = 10.1f;
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
for (size_t i = 0; i < 100; ++i) {
|
||||
ASSERT_EQ(i, result[i].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestSearchPKeys) {
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearchByKeys.index", true));
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 5000UL;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
|
||||
std::vector<std::vector<uint64_t>> p_keys;
|
||||
p_keys.resize(1);
|
||||
p_keys[0].resize(cnt);
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_velues(sparse_dim_count);
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = -1.0 * i - 1.0f;
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, streamer->add_impl(i, sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
|
||||
p_keys[0][i] = i;
|
||||
}
|
||||
|
||||
// test dump
|
||||
auto path = dir_ + "TestGeneral_dump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher get vector
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("FileReadStorage");
|
||||
ASSERT_TRUE(read_storage != nullptr);
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
|
||||
size_t topk = 3;
|
||||
for (size_t i = 0; i < cnt; i += 50) {
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = i + 1.0f;
|
||||
}
|
||||
ctx->set_topk(1U);
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(1UL, result1.size());
|
||||
ASSERT_EQ(0, result1[0].key());
|
||||
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(0, result2[0].key());
|
||||
ASSERT_EQ(1, result2[1].key());
|
||||
ASSERT_EQ(2, result2[2].key());
|
||||
}
|
||||
|
||||
{
|
||||
ctx->set_topk(100U);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = 1.0f;
|
||||
}
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(0, result[0].key());
|
||||
ASSERT_EQ(1, result[1].key());
|
||||
ASSERT_EQ(10, result[10].key());
|
||||
ASSERT_EQ(20, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
}
|
||||
|
||||
{
|
||||
ctx->set_topk(100U);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = 10.0f;
|
||||
}
|
||||
|
||||
p_keys[0] = {{cnt + 1, 10, 1, 15, cnt + 2}};
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(3U, result.size());
|
||||
ASSERT_EQ(1, result[0].key());
|
||||
ASSERT_EQ(10, result[1].key());
|
||||
ASSERT_EQ(15, result[2].key());
|
||||
}
|
||||
|
||||
{
|
||||
ctx->set_topk(100U);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = 9.0f;
|
||||
}
|
||||
p_keys[0].clear();
|
||||
for (size_t j = 0; j < cnt; j += 10) {
|
||||
p_keys[0].push_back((uint64_t)j);
|
||||
}
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(0, result[0].key());
|
||||
ASSERT_EQ(10, result[1].key());
|
||||
ASSERT_EQ(100, result[10].key());
|
||||
ASSERT_EQ(200, result[20].key());
|
||||
ASSERT_EQ(300, result[30].key());
|
||||
ASSERT_EQ(350, result[35].key());
|
||||
ASSERT_EQ(990, result[99].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseSearcherTest, TestMultiThread) {
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
constexpr size_t static sparse_dim_count = 32;
|
||||
IndexMeta meta(IndexMeta::MetaType::MT_SPARSE, IndexMeta::DataType::DT_FP32);
|
||||
meta.set_metric("InnerProductSparse", 0, Params());
|
||||
ASSERT_EQ(0, streamer->init(meta, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TessKnnMultiThread", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto addVector = [&streamer](int baseKey, size_t addCnt) {
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
size_t succAdd = 0;
|
||||
auto ctx = streamer->create_context();
|
||||
for (size_t i = 0; i < addCnt; i++) {
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_velues(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = (float)i + baseKey;
|
||||
}
|
||||
|
||||
succAdd += !streamer->add_impl(baseKey + i, sparse_dim_count,
|
||||
sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx);
|
||||
}
|
||||
streamer->flush(0UL);
|
||||
return succAdd;
|
||||
};
|
||||
|
||||
auto t2 = std::async(std::launch::async, addVector, 1000, 1000);
|
||||
auto t3 = std::async(std::launch::async, addVector, 2000, 1000);
|
||||
auto t1 = std::async(std::launch::async, addVector, 0, 1000);
|
||||
ASSERT_EQ(1000U, t1.get());
|
||||
ASSERT_EQ(1000U, t2.get());
|
||||
ASSERT_EQ(1000U, t3.get());
|
||||
streamer->close();
|
||||
|
||||
// checking data
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
auto provider = streamer->create_sparse_provider();
|
||||
auto iter = provider->create_iterator();
|
||||
ASSERT_TRUE(!!iter);
|
||||
size_t total = 0;
|
||||
uint64_t min = 1000;
|
||||
uint64_t max = 0;
|
||||
|
||||
std::set<uint64_t> keys;
|
||||
|
||||
while (iter->is_valid()) {
|
||||
const uint32_t sparse_count = iter->sparse_count();
|
||||
ASSERT_EQ(sparse_count, sparse_dim_count);
|
||||
|
||||
const float *data = reinterpret_cast<const float *>(iter->sparse_data());
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
ASSERT_EQ((float)iter->key(), data[j]);
|
||||
}
|
||||
total++;
|
||||
min = std::min(min, iter->key());
|
||||
max = std::max(max, iter->key());
|
||||
keys.insert(iter->key());
|
||||
iter->next();
|
||||
}
|
||||
|
||||
ASSERT_EQ(3000, keys.size());
|
||||
|
||||
ASSERT_EQ(3000, total);
|
||||
ASSERT_EQ(0, min);
|
||||
ASSERT_EQ(2999, max);
|
||||
|
||||
// test dump
|
||||
auto path = dir_ + "TestGeneral_dump";
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, streamer->dump(dumper));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
// do searcher get vector
|
||||
IndexSearcher::Pointer searcher =
|
||||
IndexFactory::CreateSearcher("FlatSparseSearcher");
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileReadStorage");
|
||||
ASSERT_TRUE(read_storage != nullptr);
|
||||
ASSERT_TRUE(searcher != nullptr);
|
||||
ASSERT_EQ(0, read_storage->open(path, false));
|
||||
ASSERT_EQ(0, searcher->init(Params()));
|
||||
ASSERT_EQ(0, searcher->load(read_storage, IndexMetric::Pointer()));
|
||||
|
||||
// ====== multi thread search
|
||||
size_t topk = 10;
|
||||
size_t cnt = 3000;
|
||||
auto knnSearch = [&]() {
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto linearByPkeysCtx = searcher->create_context();
|
||||
auto ctx = searcher->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32);
|
||||
linearCtx->set_topk(topk);
|
||||
linearByPkeysCtx->set_topk(topk);
|
||||
ctx->set_topk(topk);
|
||||
size_t totalCnts = 0;
|
||||
size_t totalHits = 0;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_velues(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_velues[j] = ((float)i + 1.1f);
|
||||
}
|
||||
|
||||
ASSERT_EQ(0,
|
||||
searcher->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, ctx));
|
||||
ASSERT_EQ(
|
||||
0, searcher->search_bf_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), qmeta, linearCtx));
|
||||
std::vector<std::vector<uint64_t>> p_keys = {{cnt - 1, cnt - 2, cnt - 3}};
|
||||
ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl(
|
||||
sparse_dim_count, sparse_indices.data(),
|
||||
sparse_velues.data(), p_keys, qmeta, linearByPkeysCtx));
|
||||
auto &r1 = ctx->result();
|
||||
ASSERT_EQ(topk, r1.size());
|
||||
// std::cout << "r1 top1: " << r1[0].key() << ", score: " << r1[0].score()
|
||||
// << std::endl;
|
||||
ASSERT_EQ(cnt - 1, r1[0].key());
|
||||
auto &r2 = linearCtx->result();
|
||||
ASSERT_EQ(topk, r2.size());
|
||||
// std::cout << "r2 top1: " << r2[0].key() << ", score: " << r2[0].score()
|
||||
// << std::endl;
|
||||
ASSERT_EQ(cnt - 1, r2[0].key());
|
||||
auto &r3 = linearByPkeysCtx->result();
|
||||
ASSERT_EQ(std::min(topk, p_keys[0].size()), r3.size());
|
||||
#if 0
|
||||
printf("linear: %zd => %zd %zd %zd %zd %zd\n", i, r2[0].key,
|
||||
r2[1].key, r2[2].key, r2[3].key, r2[4].key);
|
||||
printf("knn: %zd => %zd %zd %zd %zd %zd\n", i, r1[0].key, r1[1].key,
|
||||
r1[2].key, r1[3].key, r1[4].key);
|
||||
#endif
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (r2[j].key() == r1[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
printf("%f\n", totalHits * 1.0f / totalCnts);
|
||||
ASSERT_FLOAT_EQ(1.0f, totalHits * 1.0f / totalCnts);
|
||||
};
|
||||
|
||||
auto s1 = std::async(std::launch::async, knnSearch);
|
||||
auto s2 = std::async(std::launch::async, knnSearch);
|
||||
auto s3 = std::async(std::launch::async, knnSearch);
|
||||
s1.wait();
|
||||
s2.wait();
|
||||
s3.wait();
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
@@ -0,0 +1,230 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 <string>
|
||||
#include <vector>
|
||||
#include <ailego/math/norm2_matrix.h>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <algorithm/flat_sparse/flat_sparse_utility.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class FlatSparseStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm);
|
||||
|
||||
static std::string dir_;
|
||||
static shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string FlatSparseStreamerTest::dir_("FlatSparseStreamerTest/");
|
||||
shared_ptr<IndexMeta> FlatSparseStreamerTest::index_meta_ptr_;
|
||||
|
||||
void FlatSparseStreamerTest::generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm) {
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
// prepare sparse
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_vec(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_vec[j] = dist(gen);
|
||||
}
|
||||
|
||||
float norm;
|
||||
ailego::Norm2Matrix<float, 1>::Compute(sparse_vec.data(), sparse_dim_count,
|
||||
&norm);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_vec[j] = sparse_vec[j] / norm;
|
||||
}
|
||||
|
||||
sparse_indices_list.push_back(sparse_indices);
|
||||
sparse_vec_list.push_back(sparse_vec);
|
||||
}
|
||||
}
|
||||
|
||||
void FlatSparseStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
index_meta_ptr_->set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void FlatSparseStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(FlatSparseStreamerTest, TestGeneral) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
size_t sparse_dim_count = 32;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto write_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, write_storage->init(stg_params));
|
||||
ASSERT_EQ(0, write_storage->open(dir_ + "Test/FlatSparseSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, write_streamer->open(write_storage));
|
||||
|
||||
size_t cnt = 20000U;
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, write_streamer->add_impl(
|
||||
i, sparse_dim_count, sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
write_storage->close();
|
||||
write_storage.reset();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("FlatSparseStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/FlatSparseSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
|
||||
auto linearCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
|
||||
auto knnCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
// streamer->print_debug_info();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 100) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(
|
||||
0, read_streamer->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(
|
||||
sparse_dim_count, sparse_indices.data(), sparse_vec.data(),
|
||||
qmeta, linearCtx));
|
||||
|
||||
auto t3 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
// std::cout << "i: " << i << std::endl;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 100.0f / cnt;
|
||||
// float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
|
||||
std::cout << "knnTotalTime=" << knnTotalTime
|
||||
<< " linearTotalTime=" << linearTotalTime << std::endl;
|
||||
|
||||
#if 0
|
||||
printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d "
|
||||
"R@%zd=%f R@1=%f cost=%f\n",
|
||||
knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall,
|
||||
topk1Recall, cost);
|
||||
#endif
|
||||
EXPECT_GT(recall, 0.80f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
// EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_hnsw core_knn_flat
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/hnsw
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,517 @@
|
||||
#include <future>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <ailego/utility/math_helper.h>
|
||||
#include <ailego/utility/memory_helper.h>
|
||||
#include <algorithm/hnsw/hnsw_params.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/core/framework/index_framework.h>
|
||||
#include <zvec/core/framework/index_streamer.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
constexpr size_t static dim = 16;
|
||||
|
||||
class HnswStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
void hybrid_scale(std::vector<float> &dense_value,
|
||||
std::vector<float> &sparse_value, float alpha_scale);
|
||||
|
||||
static std::string dir_;
|
||||
static std::shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string HnswStreamerTest::dir_("hnsw_streamer_buffer_test_dir/");
|
||||
std::shared_ptr<IndexMeta> HnswStreamerTest::index_meta_ptr_;
|
||||
|
||||
void HnswStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (std::nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void HnswStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(HnswStreamerTest, TestHnswSearch) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/HnswSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
TEST_F(HnswStreamerTest, TestHnswSearchBuffer) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/TestHnswSearchBuffer", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBuffer", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
TEST_F(HnswStreamerTest, TestHnswSearchBufferMMap) {
|
||||
MemoryLimitPool::get_instance().init(2 * 1024UL * 1024UL * 1024UL);
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/TestHnswSearchBufferMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/TestHnswSearchBufferMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
TEST_F(HnswStreamerTest, TestHnswSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
ASSERT_EQ(0, write_streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "Test/HnswSearchMMap", true));
|
||||
ASSERT_EQ(0, write_streamer->open(storage));
|
||||
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
size_t cnt = 10000UL;
|
||||
IndexQueryMeta qmeta(IndexMeta::DT_FP32, dim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
write_streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
storage->close();
|
||||
|
||||
ElapsedTime elapsed_time;
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSearchMMap", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
size_t topk = 3;
|
||||
auto provider = read_streamer->create_provider();
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(topk, result1.size());
|
||||
IndexStorage::MemoryBlock block;
|
||||
ASSERT_EQ(0, provider->get_vector(result1[0].key(), block));
|
||||
const float *data = (float *)block.data();
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
ASSERT_FLOAT_EQ(data[j], i);
|
||||
}
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, read_streamer->search_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
|
||||
ctx->set_topk(100U);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 10.1f;
|
||||
}
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result = ctx->result();
|
||||
ASSERT_EQ(100U, result.size());
|
||||
ASSERT_EQ(10, result[0].key());
|
||||
ASSERT_EQ(11, result[1].key());
|
||||
ASSERT_EQ(5, result[10].key());
|
||||
ASSERT_EQ(0, result[20].key());
|
||||
ASSERT_EQ(30, result[30].key());
|
||||
ASSERT_EQ(35, result[35].key());
|
||||
ASSERT_EQ(99, result[99].key());
|
||||
|
||||
read_streamer->close();
|
||||
read_streamer.reset();
|
||||
cout << "Elapsed time: " << elapsed_time.milli_seconds() << " ms" << endl;
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,34 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
include(${PROJECT_ROOT_DIR}/cmake/option.cmake)
|
||||
|
||||
if(APPLE)
|
||||
set(APPLE_FRAMEWORK_LIBS
|
||||
-framework CoreFoundation
|
||||
-framework CoreGraphics
|
||||
-framework CoreData
|
||||
-framework CoreText
|
||||
-framework Security
|
||||
-framework Foundation
|
||||
-Wl,-U,_MallocExtension_ReleaseFreeMemory
|
||||
-Wl,-U,_ProfilerStart
|
||||
-Wl,-U,_ProfilerStop
|
||||
-Wl,-U,_RegisterThriftProtocol
|
||||
)
|
||||
endif()
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_hnsw_rabitq core_knn_flat core_knn_cluster
|
||||
${CMAKE_THREAD_LIBS_INIT}
|
||||
${CMAKE_DL_LIBS}
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/hnsw_rabitq
|
||||
LDFLAGS ${APPLE_FRAMEWORK_LIBS}
|
||||
)
|
||||
cc_test_suite(hnsw_rabitq ${CC_TARGET})
|
||||
endforeach()
|
||||
@@ -0,0 +1,660 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "hnsw_rabitq_streamer.h"
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <gtest/gtest.h>
|
||||
#include "zvec/ailego/container/params.h"
|
||||
#include "zvec/ailego/utility/file_helper.h"
|
||||
#include "zvec/core/framework/index_holder.h"
|
||||
#include "zvec/core/framework/index_streamer.h"
|
||||
#include "hnsw_rabitq_streamer.h"
|
||||
#include "rabitq_converter.h"
|
||||
#include "rabitq_reformer.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
constexpr size_t static dim = 128;
|
||||
|
||||
class HnswRabitqStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string dir_;
|
||||
static shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string HnswRabitqStreamerTest::dir_("hnswRabitqStreamerTest");
|
||||
shared_ptr<IndexMeta> HnswRabitqStreamerTest::index_meta_ptr_;
|
||||
|
||||
void HnswRabitqStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, dim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
}
|
||||
|
||||
void HnswRabitqStreamerTest::TearDown(void) {
|
||||
ailego::FileHelper::RemovePath(dir_.c_str());
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestBuildAndSearch) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i * dim + j) / 1000.0f;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/Test/AddVector", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
for (auto it = holder->create_iterator(); it->is_valid(); it->next()) {
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
ASSERT_EQ(0,
|
||||
streamer->add_impl(it->key(), it->data(), query_meta, context));
|
||||
}
|
||||
streamer->flush(0UL);
|
||||
|
||||
// Perform search verification
|
||||
NumericalVector<float> query_vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = static_cast<float>(j) / 1000.0f;
|
||||
}
|
||||
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
|
||||
context->set_topk(10);
|
||||
ASSERT_EQ(0, streamer->search_impl(query_vec.data(), query_meta, 1, context));
|
||||
|
||||
const auto &result = context->result(0);
|
||||
ASSERT_GT(result.size(), 0UL);
|
||||
ASSERT_LE(result.size(), 10UL);
|
||||
|
||||
// reopen and load reformer from storage
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
IndexStreamer::Pointer new_streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder);
|
||||
ASSERT_EQ(0, new_streamer->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, new_streamer->open(storage));
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestLinearSearch) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestLinearSearch", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
for (auto it = holder->create_iterator(); it->is_valid(); it->next()) {
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
ASSERT_EQ(0,
|
||||
streamer->add_impl(it->key(), it->data(), query_meta, context));
|
||||
}
|
||||
|
||||
// Test linear search with exact match
|
||||
size_t topk = 3;
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> query_vec(dim);
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i += 100) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = static_cast<float>(i);
|
||||
}
|
||||
context->set_topk(1U);
|
||||
ASSERT_EQ(0,
|
||||
streamer->search_bf_impl(query_vec.data(), query_meta, context));
|
||||
auto &result1 = context->result();
|
||||
ASSERT_EQ(1UL, result1.size());
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
// Test with slight offset
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
context->set_topk(topk);
|
||||
ASSERT_EQ(0,
|
||||
streamer->search_bf_impl(query_vec.data(), query_meta, context));
|
||||
auto &result2 = context->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestKnnSearch) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 10U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.efconstruction", 100U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.ef", 50U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestKnnSearch", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
for (auto it = holder->create_iterator(); it->is_valid(); it->next()) {
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
ASSERT_EQ(0,
|
||||
streamer->add_impl(it->key(), it->data(), query_meta, context));
|
||||
}
|
||||
|
||||
// Compare KNN search with brute force search
|
||||
auto linear_ctx = streamer->create_context();
|
||||
auto knn_ctx = streamer->create_context();
|
||||
size_t topk = 50;
|
||||
linear_ctx->set_topk(topk);
|
||||
knn_ctx->set_topk(topk);
|
||||
|
||||
int total_hits = 0;
|
||||
int total_cnts = 0;
|
||||
int topk1_hits = 0;
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> query_vec(dim);
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i += 100) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
|
||||
ASSERT_EQ(0,
|
||||
streamer->search_impl(query_vec.data(), query_meta, 1, knn_ctx));
|
||||
ASSERT_EQ(
|
||||
0, streamer->search_bf_impl(query_vec.data(), query_meta, linear_ctx));
|
||||
|
||||
auto &knn_result = knn_ctx->result(0);
|
||||
ASSERT_EQ(topk, knn_result.size());
|
||||
topk1_hits += (i == knn_result[0].key());
|
||||
|
||||
auto &linear_result = linear_ctx->result();
|
||||
ASSERT_EQ(topk, linear_result.size());
|
||||
ASSERT_EQ(i, linear_result[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
total_cnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linear_result[j].key() == knn_result[k].key()) {
|
||||
total_hits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float recall = total_hits * 1.0f / total_cnts;
|
||||
float topk1_recall = topk1_hits * 100.0f / static_cast<float>(doc_cnt);
|
||||
EXPECT_GT(recall, 0.60f);
|
||||
// actual: no guarantee
|
||||
// TODO(jiliang.ljl): check if ok?
|
||||
EXPECT_GT(topk1_recall, 0.00f);
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestRandomData) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 1500UL;
|
||||
|
||||
// Add random vectors
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 32U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 20U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.efconstruction", 200U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.ef", 100U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestRandomData", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
for (auto it = holder->create_iterator(); it->is_valid(); it->next()) {
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
ASSERT_EQ(0,
|
||||
streamer->add_impl(it->key(), it->data(), query_meta, context));
|
||||
}
|
||||
|
||||
// Test with random queries
|
||||
auto linear_ctx = streamer->create_context();
|
||||
auto knn_ctx = streamer->create_context();
|
||||
size_t topk = 50;
|
||||
linear_ctx->set_topk(topk);
|
||||
knn_ctx->set_topk(topk);
|
||||
|
||||
int total_hits = 0;
|
||||
int total_cnts = 0;
|
||||
int topk1_hits = 0;
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> query_vec(dim);
|
||||
|
||||
size_t query_cnt = 200;
|
||||
for (size_t i = 0; i < query_cnt; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
|
||||
}
|
||||
|
||||
ASSERT_EQ(
|
||||
0, streamer->search_bf_impl(query_vec.data(), query_meta, linear_ctx));
|
||||
ASSERT_EQ(0,
|
||||
streamer->search_impl(query_vec.data(), query_meta, 1, knn_ctx));
|
||||
|
||||
auto &knn_result = knn_ctx->result(0);
|
||||
ASSERT_EQ(topk, knn_result.size());
|
||||
|
||||
auto &linear_result = linear_ctx->result();
|
||||
ASSERT_EQ(topk, linear_result.size());
|
||||
|
||||
topk1_hits += (linear_result[0].key() == knn_result[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
total_cnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linear_result[j].key() == knn_result[k].key()) {
|
||||
total_hits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float recall = total_hits * 1.0f / total_cnts;
|
||||
float topk1_recall = topk1_hits * 1.0f / query_cnt;
|
||||
EXPECT_GT(recall, 0.50f);
|
||||
EXPECT_GT(topk1_recall, 0.70f);
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestOpenClose) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 500UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestOpenClose", true));
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
|
||||
// Add first half of vectors
|
||||
for (size_t i = 0; i < doc_cnt / 2; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), query_meta, context));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, streamer->flush(0UL));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
|
||||
// Reopen and add second half
|
||||
IndexStreamer::Pointer streamer2 =
|
||||
std::make_shared<HnswRabitqStreamer>(holder);
|
||||
ASSERT_EQ(0, streamer2->init(*index_meta_ptr_, params));
|
||||
ASSERT_EQ(0, streamer2->open(storage));
|
||||
|
||||
auto context2 = streamer2->create_context();
|
||||
for (size_t i = doc_cnt / 2; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer2->add_impl(i, vec.data(), query_meta, context2));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, streamer2->flush(0UL));
|
||||
|
||||
// Verify search works after reopen
|
||||
NumericalVector<float> query_vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
query_vec[j] = 10.0f;
|
||||
}
|
||||
|
||||
context2->set_topk(5);
|
||||
ASSERT_EQ(0,
|
||||
streamer2->search_impl(query_vec.data(), query_meta, 1, context2));
|
||||
const auto &result = context2->result(0);
|
||||
ASSERT_EQ(5UL, result.size());
|
||||
ASSERT_EQ(10UL, result[0].key());
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestCreateIterator) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 300UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestCreateIterator", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto context = streamer->create_context();
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), query_meta, context));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
|
||||
// Test iterator
|
||||
auto provider = streamer->create_provider();
|
||||
auto iter = provider->create_iterator();
|
||||
ASSERT_TRUE(!!iter);
|
||||
|
||||
size_t count = 0;
|
||||
while (iter->is_valid()) {
|
||||
ASSERT_EQ(count, iter->key());
|
||||
// const float *data = (const float *)iter->data();
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// ASSERT_EQ(static_cast<float>(count), data[j]);
|
||||
// }
|
||||
iter->next();
|
||||
count++;
|
||||
}
|
||||
ASSERT_EQ(doc_cnt, count);
|
||||
|
||||
// Test get_vector
|
||||
// for (size_t i = 0; i < doc_cnt; i++) {
|
||||
// const float *data = (const float *)provider->get_vector(i);
|
||||
// ASSERT_NE(data, nullptr);
|
||||
// for (size_t j = 0; j < dim; ++j) {
|
||||
// ASSERT_EQ(static_cast<float>(i), data[j]);
|
||||
// }
|
||||
// }
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestDimensions) {
|
||||
std::vector<size_t> dimensions = {1, 2, 4, 8, 16, 32, 33,
|
||||
63, 64, 128, 256, 512, 1024, 2047,
|
||||
2048, 2049, 4095, 4096, 4097, 8192, 16384};
|
||||
size_t doc_cnt = 100;
|
||||
|
||||
for (size_t test_dim : dimensions) {
|
||||
std::cout << "Testing dimension: " << test_dim << std::endl;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::DataType::DT_FP32, test_dim);
|
||||
index_meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(
|
||||
test_dim);
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", test_dim);
|
||||
|
||||
int ret = streamer->init(index_meta, params);
|
||||
|
||||
// dimension <= 63 or >= 4096: init() should return -31
|
||||
if (test_dim <= 63 || test_dim >= 4096) {
|
||||
ASSERT_EQ(-31, ret) << "expected init to fail with -31, dim=" << test_dim;
|
||||
std::cout << "Dimension " << test_dim
|
||||
<< " correctly rejected with ret=" << ret << std::endl;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Valid dimensions: verify full streaming build succeeds
|
||||
ASSERT_EQ(0, ret) << "init failed, dim=" << test_dim;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(test_dim);
|
||||
for (size_t j = 0; j < test_dim; ++j) {
|
||||
vec[j] = static_cast<float>(i * test_dim + j) / 1000.0f;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, std::move(vec))) << "dim=" << test_dim;
|
||||
}
|
||||
|
||||
RabitqConverter converter;
|
||||
converter.init(index_meta, ailego::Params());
|
||||
ASSERT_EQ(0, converter.train(holder))
|
||||
<< "converter train failed, dim=" << test_dim;
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(0, converter.to_reformer(&index_reformer)) << "dim=" << test_dim;
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
|
||||
// Recreate streamer with reformer
|
||||
streamer = std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params))
|
||||
<< "init with reformer failed, dim=" << test_dim;
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
std::string storage_path =
|
||||
dir_ + "/TestDimensions_" + std::to_string(test_dim);
|
||||
ASSERT_EQ(0, storage->open(storage_path, true))
|
||||
<< "storage open failed, dim=" << test_dim;
|
||||
ASSERT_EQ(0, streamer->open(storage))
|
||||
<< "streamer open failed, dim=" << test_dim;
|
||||
|
||||
auto context = streamer->create_context();
|
||||
IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, test_dim);
|
||||
for (auto it = holder->create_iterator(); it->is_valid(); it->next()) {
|
||||
ASSERT_EQ(0,
|
||||
streamer->add_impl(it->key(), it->data(), query_meta, context))
|
||||
<< "add failed, dim=" << test_dim << ", key=" << it->key();
|
||||
}
|
||||
ASSERT_EQ(0, streamer->flush(0UL)) << "flush failed, dim=" << test_dim;
|
||||
|
||||
std::cout << "Dimension " << test_dim << " passed" << std::endl;
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(HnswRabitqStreamerTest, TestExBitsMismatch) {
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexProvider<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 500UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i * dim + j) / 1000.0f;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
// Train converter with default total_bits (7, ex_bits=6)
|
||||
RabitqConverter converter;
|
||||
converter.init(*index_meta_ptr_, ailego::Params());
|
||||
ASSERT_EQ(converter.train(holder), 0);
|
||||
std::shared_ptr<IndexReformer> index_reformer;
|
||||
ASSERT_EQ(converter.to_reformer(&index_reformer), 0);
|
||||
auto reformer = std::dynamic_pointer_cast<RabitqReformer>(index_reformer);
|
||||
|
||||
// Create streamer with total_bits=2 (ex_bits=1), mismatching the reformer
|
||||
IndexStreamer::Pointer streamer =
|
||||
std::make_shared<HnswRabitqStreamer>(holder, reformer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw_rabitq.streamer.max_neighbor_count", 16U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.upper_neighbor_count", 8U);
|
||||
params.set("proxima.hnsw_rabitq.streamer.scaling_factor", 5U);
|
||||
params.set("proxima.hnsw_rabitq.general.dimension", dim);
|
||||
params.set(PARAM_RABITQ_TOTAL_BITS, 2U);
|
||||
ASSERT_EQ(0, streamer->init(*index_meta_ptr_, params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "/TestExBitsMismatch", true));
|
||||
|
||||
// open() should detect ex_bits mismatch and return IndexError_Mismatch
|
||||
ASSERT_EQ(IndexError_Mismatch, streamer->open(storage));
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_hnsw_sparse
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/hnsw_sparse
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,469 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "hnsw_sparse_builder.h"
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <future>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
#include "zvec/core/framework/index_framework.h"
|
||||
#include "hnsw_sparse_params.h"
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class HnswSparseBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
static std::string _dir;
|
||||
static shared_ptr<IndexMeta> _index_meta_ptr;
|
||||
};
|
||||
|
||||
std::string HnswSparseBuilderTest::_dir("HnswSparseBuilderTest/");
|
||||
shared_ptr<IndexMeta> HnswSparseBuilderTest::_index_meta_ptr;
|
||||
|
||||
void HnswSparseBuilderTest::SetUp(void) {
|
||||
_index_meta_ptr.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
_index_meta_ptr->set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
}
|
||||
|
||||
void HnswSparseBuilderTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(_dir);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestGeneral) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestGeneral";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, vec));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder->train(holder2));
|
||||
ASSERT_EQ(0, builder->build(holder2));
|
||||
auto dumper2 = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper2, nullptr);
|
||||
ASSERT_EQ(0, dumper2->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper2));
|
||||
ASSERT_EQ(0, dumper2->close());
|
||||
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestMemquota) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt = 1000UL;
|
||||
uint32_t sparse_count = 32;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set("proxima.hnsw.sparse_builder.memory_quota", 100000UL);
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder->train(holder));
|
||||
ASSERT_EQ(IndexError_NoMemory, builder->build(holder));
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestIndexThreads) {
|
||||
IndexBuilder::Pointer builder1 =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder1, nullptr);
|
||||
IndexBuilder::Pointer builder2 =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder2, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
|
||||
size_t doc_cnt = 1000UL;
|
||||
uint32_t sparse_count = 32;
|
||||
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
std::srand(ailego::Realtime::MilliSeconds());
|
||||
auto threads =
|
||||
std::make_shared<SingleQueueIndexThreads>(std::rand() % 4, false);
|
||||
ASSERT_EQ(0, builder1->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder2->init(*_index_meta_ptr, params));
|
||||
|
||||
auto build_index1 = [&]() {
|
||||
ASSERT_EQ(0, builder1->train(threads, holder));
|
||||
ASSERT_EQ(0, builder1->build(threads, holder));
|
||||
};
|
||||
auto build_index2 = [&]() {
|
||||
ASSERT_EQ(0, builder2->train(threads, holder));
|
||||
ASSERT_EQ(0, builder2->build(threads, holder));
|
||||
};
|
||||
|
||||
auto t1 = std::async(std::launch::async, build_index1);
|
||||
auto t2 = std::async(std::launch::async, build_index2);
|
||||
t1.wait();
|
||||
t2.wait();
|
||||
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndexThreads";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder1->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder2->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats1 = builder1->stats();
|
||||
ASSERT_EQ(doc_cnt, stats1.built_count());
|
||||
auto &stats2 = builder2->stats();
|
||||
ASSERT_EQ(doc_cnt, stats2.built_count());
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestHalfFloatConverter) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
auto holder =
|
||||
make_shared<OnePassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
ailego::Params converter_params;
|
||||
auto converter = IndexFactory::CreateConverter("HalfFloatSparseConverter");
|
||||
converter->init(*_index_meta_ptr, converter_params);
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
converter->transform(holder);
|
||||
|
||||
auto converted_holder = converter->sparse_result();
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(index_meta, converter_params));
|
||||
|
||||
ASSERT_EQ(0, builder->train(converted_holder));
|
||||
|
||||
ASSERT_EQ(0, builder->build(converted_holder));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestHalFloatConverter";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
|
||||
// cleanup and rebuild
|
||||
ASSERT_EQ(0, builder->cleanup());
|
||||
|
||||
auto holder2 =
|
||||
make_shared<MultiPassIndexSparseHolder<IndexMeta::DataType::DT_FP32>>();
|
||||
size_t doc_cnt2 = 2000UL;
|
||||
for (size_t i = 0; i < doc_cnt2; i++) {
|
||||
SparseVector<float> vec;
|
||||
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_count);
|
||||
NumericalVector<float> sparse_values(sparse_count);
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices[j] = 20 * j;
|
||||
sparse_values[j] = i;
|
||||
}
|
||||
|
||||
vec.add_sparses(sparse_indices, sparse_values);
|
||||
|
||||
ASSERT_TRUE(holder2->emplace(i, vec));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
ASSERT_EQ(0, builder->train(holder2));
|
||||
ASSERT_EQ(0, builder->build(holder2));
|
||||
auto dumper2 = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper2, nullptr);
|
||||
ASSERT_EQ(0, dumper2->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper2));
|
||||
ASSERT_EQ(0, dumper2->close());
|
||||
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt2, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestIndptr) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
|
||||
std::vector<uint64_t> keys;
|
||||
keys.reserve(doc_cnt);
|
||||
|
||||
std::vector<uint64_t> sparse_indptr;
|
||||
sparse_indptr.reserve(doc_cnt + 1);
|
||||
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
sparse_indices.reserve(doc_cnt * sparse_count);
|
||||
|
||||
std::vector<float> sparse_values;
|
||||
sparse_values.reserve(doc_cnt * sparse_count);
|
||||
|
||||
size_t sparse_count_total = 0;
|
||||
sparse_indptr.push_back(0);
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices.push_back(20 * j);
|
||||
sparse_values.push_back(i);
|
||||
}
|
||||
|
||||
keys.push_back(i);
|
||||
|
||||
sparse_count_total += sparse_count;
|
||||
sparse_indptr.push_back(sparse_count_total);
|
||||
}
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(*_index_meta_ptr, params));
|
||||
|
||||
ASSERT_EQ(0, builder->build(doc_cnt, keys.data(), sparse_indptr.data(),
|
||||
sparse_indices.data(), sparse_values.data()));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndptr";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseBuilderTest, TestIndptrFp16) {
|
||||
IndexBuilder::Pointer builder =
|
||||
IndexFactory::CreateBuilder("HnswSparseBuilder");
|
||||
ASSERT_NE(builder, nullptr);
|
||||
|
||||
uint32_t sparse_count = 4;
|
||||
size_t doc_cnt = 1000UL;
|
||||
|
||||
std::vector<uint64_t> keys;
|
||||
keys.reserve(doc_cnt);
|
||||
|
||||
std::vector<uint64_t> sparse_indptr;
|
||||
sparse_indptr.reserve(doc_cnt + 1);
|
||||
|
||||
std::vector<uint32_t> sparse_indices;
|
||||
sparse_indices.reserve(doc_cnt * sparse_count);
|
||||
|
||||
std::vector<float> sparse_values;
|
||||
sparse_values.reserve(doc_cnt * sparse_count);
|
||||
|
||||
size_t sparse_count_total = 0;
|
||||
sparse_indptr.push_back(0);
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
for (size_t j = 0; j < sparse_count; ++j) {
|
||||
sparse_indices.push_back(20 * j);
|
||||
sparse_values.push_back(i);
|
||||
}
|
||||
|
||||
keys.push_back(i);
|
||||
|
||||
sparse_count_total += sparse_count;
|
||||
sparse_indptr.push_back(sparse_count_total);
|
||||
}
|
||||
|
||||
IndexMeta meta(IndexMeta::MetaType::MT_SPARSE, IndexMeta::DataType::DT_FP16);
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_BUILDER_THREAD_COUNT, 1);
|
||||
ASSERT_EQ(0, builder->init(meta, params));
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
ASSERT_EQ(0, builder->build(qmeta, doc_cnt, keys.data(), sparse_indptr.data(),
|
||||
sparse_indices.data(), sparse_values.data()));
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
string path = _dir + "TestIndptrFp16";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats = builder->stats();
|
||||
ASSERT_EQ(0UL, stats.trained_count());
|
||||
ASSERT_EQ(doc_cnt, stats.built_count());
|
||||
ASSERT_EQ(doc_cnt, stats.dumped_count());
|
||||
ASSERT_EQ(0UL, stats.discarded_count());
|
||||
ASSERT_EQ(0UL, stats.trained_costtime());
|
||||
ASSERT_GT(stats.built_costtime(), 0UL);
|
||||
// ASSERT_GT(stats.dumped_costtime(), 0UL);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,360 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#include <fcntl.h>
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <ailego/math/norm_matrix.h>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
#include "hnsw_sparse_streamer.h"
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class HnswSparseStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void);
|
||||
void TearDown(void);
|
||||
void generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm);
|
||||
|
||||
static std::string dir_;
|
||||
static shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string HnswSparseStreamerTest::dir_(
|
||||
"hnsw_sparse_streamer_buffer_test_dir/");
|
||||
shared_ptr<IndexMeta> HnswSparseStreamerTest::index_meta_ptr_;
|
||||
|
||||
void HnswSparseStreamerTest::generate_sparse_data(
|
||||
size_t cnt, uint32_t sparse_dim_count,
|
||||
std::vector<NumericalVector<uint32_t>> &sparse_indices_list,
|
||||
std::vector<NumericalVector<float>> &sparse_vec_list, bool norm) {
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_real_distribution<float> dist(-1.0, 1.0);
|
||||
|
||||
for (size_t i = 0; i < cnt; ++i) {
|
||||
// prepare sparse
|
||||
NumericalVector<uint32_t> sparse_indices(sparse_dim_count);
|
||||
NumericalVector<float> sparse_vec(sparse_dim_count);
|
||||
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_indices[j] = j * 20;
|
||||
sparse_vec[j] = dist(gen);
|
||||
}
|
||||
|
||||
float norm;
|
||||
ailego::Norm2Matrix<float, 1>::Compute(sparse_vec.data(), sparse_dim_count,
|
||||
&norm);
|
||||
for (size_t j = 0; j < sparse_dim_count; ++j) {
|
||||
sparse_vec[j] = sparse_vec[j] / norm;
|
||||
}
|
||||
|
||||
sparse_indices_list.push_back(sparse_indices);
|
||||
sparse_vec_list.push_back(sparse_vec);
|
||||
}
|
||||
}
|
||||
|
||||
void HnswSparseStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (nothrow) IndexMeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32));
|
||||
index_meta_ptr_->set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void HnswSparseStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseStreamerTest, TestGeneral) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
size_t sparse_dim_count = 32;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, 20);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, 16);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, 10);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EF, 5);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto write_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, write_storage->init(stg_params));
|
||||
ASSERT_EQ(0, write_storage->open(dir_ + "Test/HnswSparseSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, write_streamer->open(write_storage));
|
||||
|
||||
size_t cnt = 20000U;
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, write_streamer->add_impl(
|
||||
i, sparse_dim_count, sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
write_storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("BufferStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSparseSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
|
||||
auto linearCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
|
||||
auto knnCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
// streamer->print_debug_info();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 100) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(
|
||||
0, read_streamer->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(
|
||||
sparse_dim_count, sparse_indices.data(), sparse_vec.data(),
|
||||
qmeta, linearCtx));
|
||||
|
||||
auto t3 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
// std::cout << "i: " << i << std::endl;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 100.0f / cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
#if 0
|
||||
printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d "
|
||||
"R@%zd=%f R@1=%f cost=%f\n",
|
||||
knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall,
|
||||
topk1Recall, cost);
|
||||
#endif
|
||||
EXPECT_GT(recall, 0.80f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
// EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
TEST_F(HnswSparseStreamerTest, TestHnswSearchMMap) {
|
||||
IndexStreamer::Pointer write_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_TRUE(write_streamer != nullptr);
|
||||
|
||||
size_t sparse_dim_count = 32;
|
||||
|
||||
IndexMeta index_meta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
index_meta.set_metric("InnerProductSparse", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_MAX_NEIGHBOR_COUNT, 20);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_SCALING_FACTOR, 16);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EFCONSTRUCTION, 10);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_EF, 5);
|
||||
params.set(PARAM_HNSW_SPARSE_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto write_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, write_storage->init(stg_params));
|
||||
ASSERT_EQ(0, write_storage->open(dir_ + "Test/HnswSparseSearch", true));
|
||||
ASSERT_EQ(0, write_streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, write_streamer->open(write_storage));
|
||||
|
||||
size_t cnt = 20000U;
|
||||
auto ctx = write_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
std::vector<NumericalVector<uint32_t>> sparse_indices_list;
|
||||
std::vector<NumericalVector<float>> sparse_vec_list;
|
||||
|
||||
generate_sparse_data(cnt, sparse_dim_count, sparse_indices_list,
|
||||
sparse_vec_list, true);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::MetaType::MT_SPARSE,
|
||||
IndexMeta::DataType::DT_FP32);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
ASSERT_EQ(0, write_streamer->add_impl(
|
||||
i, sparse_dim_count, sparse_indices_list[i].data(),
|
||||
sparse_vec_list[i].data(), qmeta, ctx));
|
||||
}
|
||||
write_streamer->flush(0UL);
|
||||
write_streamer->close();
|
||||
write_streamer.reset();
|
||||
write_storage->close();
|
||||
|
||||
IndexStreamer::Pointer read_streamer =
|
||||
IndexFactory::CreateStreamer("HnswSparseStreamer");
|
||||
ASSERT_EQ(0, read_streamer->init(*index_meta_ptr_, params));
|
||||
auto read_storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, read_storage);
|
||||
ASSERT_EQ(0, read_storage->init(stg_params));
|
||||
ASSERT_EQ(0, read_storage->open(dir_ + "Test/HnswSparseSearch", false));
|
||||
ASSERT_EQ(0, read_streamer->open(read_storage));
|
||||
|
||||
auto linearCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!linearCtx);
|
||||
|
||||
auto knnCtx = read_streamer->create_context();
|
||||
ASSERT_TRUE(!!knnCtx);
|
||||
|
||||
// streamer->print_debug_info();
|
||||
size_t topk = 200;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
|
||||
uint64_t knnTotalTime = 0;
|
||||
uint64_t linearTotalTime = 0;
|
||||
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
|
||||
for (size_t i = 0; i < cnt; i += 100) {
|
||||
const auto &sparse_indices = sparse_indices_list[i];
|
||||
const auto &sparse_vec = sparse_vec_list[i];
|
||||
|
||||
auto t1 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(
|
||||
0, read_streamer->search_impl(sparse_dim_count, sparse_indices.data(),
|
||||
sparse_vec.data(), qmeta, knnCtx));
|
||||
|
||||
auto t2 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
ASSERT_EQ(0, read_streamer->search_bf_impl(
|
||||
sparse_dim_count, sparse_indices.data(), sparse_vec.data(),
|
||||
qmeta, linearCtx));
|
||||
|
||||
auto t3 = ailego::Realtime::MicroSeconds();
|
||||
|
||||
knnTotalTime += t2 - t1;
|
||||
linearTotalTime += t3 - t2;
|
||||
|
||||
// std::cout << "i: " << i << std::endl;
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 100.0f / cnt;
|
||||
float cost = linearTotalTime * 1.0f / knnTotalTime;
|
||||
#if 0
|
||||
printf("knnTotalTime=%zd linearTotalTime=%zd totalHits=%d totalCnts=%d "
|
||||
"R@%zd=%f R@1=%f cost=%f\n",
|
||||
knnTotalTime, linearTotalTime, totalHits, totalCnts, topk, recall,
|
||||
topk1Recall, cost);
|
||||
#endif
|
||||
EXPECT_GT(recall, 0.80f);
|
||||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
// EXPECT_GT(cost, 2.0f);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_cluster core_knn_flat core_knn_ivf
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/ivf
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,528 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "ivf_builder.h"
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
|
||||
using namespace zvec::core;
|
||||
using namespace zvec::ailego;
|
||||
using namespace std;
|
||||
|
||||
class IVFBuilderTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp() override;
|
||||
void TearDown() override;
|
||||
|
||||
void prepare_index_holder(uint32_t base_key, uint32_t num);
|
||||
|
||||
IndexMeta index_meta_;
|
||||
Params params_;
|
||||
uint32_t dimension_;
|
||||
IndexHolder::Pointer holder_;
|
||||
IndexThreads::Pointer threads_{};
|
||||
};
|
||||
|
||||
void IVFBuilderTest::SetUp() {
|
||||
dimension_ = 8U;
|
||||
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
index_meta_.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
params_.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "8");
|
||||
params_.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
|
||||
std::mt19937 gen((std::random_device())());
|
||||
bool v = std::uniform_int_distribution<size_t>(0, 1)(gen);
|
||||
if (v) {
|
||||
threads_ = std::make_shared<SingleQueueIndexThreads>();
|
||||
}
|
||||
}
|
||||
|
||||
void IVFBuilderTest::TearDown() {}
|
||||
|
||||
void IVFBuilderTest::prepare_index_holder(uint32_t base_key, uint32_t num) {
|
||||
MultiPassIndexHolder<IndexMeta::DataType::DT_FP32> *holder =
|
||||
new MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>(dimension_);
|
||||
uint32_t key = base_key;
|
||||
for (size_t i = 0; i < num; ++i) {
|
||||
NumericalVector<float> vec(dimension_);
|
||||
for (size_t j = 0; j < dimension_; ++j) {
|
||||
vec[j] = 1.0f * i;
|
||||
}
|
||||
holder->emplace(key + i, vec);
|
||||
}
|
||||
|
||||
holder_.reset(holder);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestInitSuccess) {
|
||||
IVFBuilder builder;
|
||||
int ret = builder.init(index_meta_, params_);
|
||||
EXPECT_EQ(0, ret);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestInitFailedWithInvalidMetric) {
|
||||
IVFBuilder builder;
|
||||
index_meta_.set_metric("invalid", 0, Params());
|
||||
int ret = builder.init(index_meta_, params_);
|
||||
EXPECT_EQ(IndexError_NoExist, ret);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestInitFailedWithInvalidCentroidsNum) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(IndexError_InvalidArgument, ret);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestTrainWithHolder1Level) {
|
||||
IVFBuilder builder;
|
||||
int ret = builder.init(index_meta_, params_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_GT(centroid_index->centroids_count(), 0u);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestTrainWithHolder2Level) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_EQ(centroid_index->centroids_count(), 8);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestTrainWithTrainer2Level) {
|
||||
IndexTrainer::Pointer trainer =
|
||||
IndexFactory::CreateTrainer("StratifiedClusterTrainer");
|
||||
ASSERT_TRUE(!!trainer);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
Params params;
|
||||
params.set("zvec.stratified.trainer.cluster_count", "4*2");
|
||||
ASSERT_EQ(0, trainer->init(index_meta_, params));
|
||||
ASSERT_EQ(0, trainer->train(threads_, holder_));
|
||||
|
||||
IVFBuilder builder;
|
||||
int ret = builder.init(index_meta_, params_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
|
||||
ret = builder.train(trainer);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_EQ(centroid_index->centroids_count(), 8);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestTrainWithTrainer1Level) {
|
||||
IVFBuilder builder;
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
IndexTrainer::Pointer trainer =
|
||||
IndexFactory::CreateTrainer("StratifiedClusterTrainer");
|
||||
ASSERT_TRUE(!!trainer);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
Params params1;
|
||||
params1.set("zvec.stratified.trainer.cluster_count", "4");
|
||||
ASSERT_EQ(0, trainer->init(index_meta_, params1));
|
||||
ASSERT_EQ(0, trainer->train(threads_, holder_));
|
||||
|
||||
ret = builder.train(trainer);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_EQ(centroid_index->centroids_count(), 4);
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWith2Level) {
|
||||
IVFBuilder builder;
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWith1Level) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestDump) {
|
||||
IVFBuilder builder;
|
||||
int ret = builder.init(index_meta_, params_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
ret = dumper->create("path");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
}
|
||||
|
||||
#if 0
|
||||
TEST_F(IVFBuilderTest, TestBuildWithNoEnoughMemory)
|
||||
{
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
dimension_ = 256;
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(IndexError_IndexFull, ret);
|
||||
}
|
||||
#endif
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWithEnoughMemory) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
dimension_ = 256;
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
ret = dumper->create("path");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
}
|
||||
|
||||
#if 0
|
||||
TEST_F(IVFBuilderTest, TestBuildWithRowMajorAndNoEnoughMemory)
|
||||
{
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
dimension_ = 256;
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
index_meta_.set_major_order(IndexMeta::MajorOrder::MO_ROW);
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(IndexError_IndexFull, ret);
|
||||
}
|
||||
#endif
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWithRowMajorAndMemory) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
dimension_ = 256;
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
index_meta_.set_major_order(IndexMeta::MajorOrder::MO_ROW);
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
ret = dumper->create("path");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWithEmptyCentroid) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "2*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
|
||||
dimension_ = 256;
|
||||
index_meta_.set_meta(IndexMeta::DataType::DT_FP32, dimension_);
|
||||
index_meta_.set_major_order(IndexMeta::MajorOrder::MO_ROW);
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
size_t doc_cnt = 10;
|
||||
|
||||
MultiPassIndexHolder<IndexMeta::DataType::DT_FP32> *holder =
|
||||
new MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>(dimension_);
|
||||
for (size_t i = 0; i < doc_cnt; ++i) {
|
||||
NumericalVector<float> vec(dimension_);
|
||||
for (size_t j = 0; j < dimension_; ++j) {
|
||||
vec[j] = 1.0f;
|
||||
}
|
||||
holder->emplace(i, vec);
|
||||
}
|
||||
holder_.reset(holder);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
ret = dumper->create("path");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)10, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)10, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestTrainClusterParams) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "2*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
|
||||
prepare_index_holder(0, 1000);
|
||||
EXPECT_EQ(0, builder.init(index_meta_, params));
|
||||
EXPECT_EQ(0, builder.train(threads_, holder_));
|
||||
EXPECT_EQ(0, builder.build(threads_, holder_));
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
EXPECT_EQ(0, dumper->create("test.index"));
|
||||
EXPECT_EQ(0, builder.dump(dumper));
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWithConverterClass) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
|
||||
params.set(PARAM_IVF_BUILDER_CONVERTER_CLASS, "HalfFloatConverter");
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_GT(centroid_index->centroids_count(), 0u);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
|
||||
ret = dumper->create("path");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestBuildWithConverterClassMultiLevel) {
|
||||
IVFBuilder builder;
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
|
||||
params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
|
||||
params.set(PARAM_IVF_BUILDER_CONVERTER_CLASS, "HalfFloatConverter");
|
||||
|
||||
int ret = builder.init(index_meta_, params);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
prepare_index_holder(0, 1000);
|
||||
|
||||
ret = builder.train(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.build(threads_, holder_);
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
auto centroid_index = builder.centroid_index();
|
||||
EXPECT_EQ(centroid_index->centroids_count(), 8);
|
||||
|
||||
IndexDumper::Pointer dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ret = dumper->create("./ivf_converter_test.index");
|
||||
EXPECT_EQ(0, ret);
|
||||
|
||||
ret = builder.dump(dumper);
|
||||
EXPECT_EQ((size_t)1000, builder.stats().built_count());
|
||||
EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
|
||||
EXPECT_EQ((size_t)0, builder.stats().discarded_count());
|
||||
EXPECT_EQ(0, dumper->close());
|
||||
File::RemovePath("./ivf_converter_test.index");
|
||||
}
|
||||
|
||||
TEST_F(IVFBuilderTest, TestIndexThreads) {
|
||||
IndexBuilder::Pointer builder1 = IndexFactory::CreateBuilder("IVFBuilder");
|
||||
ASSERT_NE(builder1, nullptr);
|
||||
IndexBuilder::Pointer builder2 = IndexFactory::CreateBuilder("IVFBuilder");
|
||||
ASSERT_NE(builder2, nullptr);
|
||||
|
||||
size_t dim = 128UL;
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
std::srand(Realtime::MilliSeconds());
|
||||
auto holder =
|
||||
std::make_shared<MultiPassIndexHolder<IndexMeta::DataType::DT_FP32>>(dim);
|
||||
size_t doc_cnt = 1000;
|
||||
for (size_t i = 0; i < doc_cnt; i++) {
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = i;
|
||||
}
|
||||
ASSERT_TRUE(holder->emplace(i, vec));
|
||||
}
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "2*2");
|
||||
ASSERT_EQ(0, builder1->init(meta, params));
|
||||
ASSERT_EQ(0, builder2->init(meta, params));
|
||||
|
||||
auto threads =
|
||||
std::make_shared<SingleQueueIndexThreads>(std::rand() % 4, false);
|
||||
auto build_index1 = [&]() {
|
||||
ASSERT_EQ(0, builder1->train(threads, holder));
|
||||
ASSERT_EQ(0, builder1->build(threads, holder));
|
||||
};
|
||||
auto build_index2 = [&]() {
|
||||
ASSERT_EQ(0, builder2->train(threads, holder));
|
||||
ASSERT_EQ(0, builder2->build(threads, holder));
|
||||
};
|
||||
|
||||
auto t1 = std::async(std::launch::async, build_index1);
|
||||
auto t2 = std::async(std::launch::async, build_index2);
|
||||
t1.wait();
|
||||
t2.wait();
|
||||
|
||||
|
||||
auto dumper = IndexFactory::CreateDumper("FileDumper");
|
||||
ASSERT_NE(dumper, nullptr);
|
||||
|
||||
std::string path = "./hc_index";
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder1->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
ASSERT_EQ(0, dumper->create(path));
|
||||
ASSERT_EQ(0, builder2->dump(dumper));
|
||||
ASSERT_EQ(0, dumper->close());
|
||||
|
||||
auto &stats1 = builder1->stats();
|
||||
ASSERT_EQ(doc_cnt, stats1.built_count());
|
||||
auto &stats2 = builder2->stats();
|
||||
ASSERT_EQ(doc_cnt, stats2.built_count());
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake)
|
||||
|
||||
file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc)
|
||||
|
||||
foreach(CC_SRCS ${ALL_TEST_SRCS})
|
||||
get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE)
|
||||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_utility core_metric core_quantizer core_knn_vamana core_knn_hnsw core_knn_flat
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/vamana
|
||||
)
|
||||
endforeach()
|
||||
@@ -0,0 +1,892 @@
|
||||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// 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 "vamana_streamer.h"
|
||||
#include <sys/stat.h>
|
||||
#include <sys/types.h>
|
||||
#ifndef _MSC_VER
|
||||
#include <fcntl.h>
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
#include <future>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "tests/test_util.h"
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wunused-result"
|
||||
#endif
|
||||
|
||||
using namespace std;
|
||||
using namespace testing;
|
||||
using namespace zvec::ailego;
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
constexpr size_t kDim = 16;
|
||||
|
||||
class VamanaStreamerTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp(void) override;
|
||||
void TearDown(void) override;
|
||||
|
||||
IndexStreamer::Pointer CreateVamanaStreamer(
|
||||
const ailego::Params &extra_params = ailego::Params());
|
||||
|
||||
static std::string dir_;
|
||||
static shared_ptr<IndexMeta> index_meta_ptr_;
|
||||
};
|
||||
|
||||
std::string VamanaStreamerTest::dir_("vamana_streamer_test_dir/");
|
||||
shared_ptr<IndexMeta> VamanaStreamerTest::index_meta_ptr_;
|
||||
|
||||
void VamanaStreamerTest::SetUp(void) {
|
||||
index_meta_ptr_.reset(new (nothrow)
|
||||
IndexMeta(IndexMeta::DataType::DT_FP32, kDim));
|
||||
index_meta_ptr_->set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
void VamanaStreamerTest::TearDown(void) {
|
||||
zvec::test_util::RemoveTestPath(dir_);
|
||||
}
|
||||
|
||||
IndexStreamer::Pointer VamanaStreamerTest::CreateVamanaStreamer(
|
||||
const ailego::Params &extra_params) {
|
||||
auto streamer = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
if (!streamer) return nullptr;
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
|
||||
params.merge(extra_params);
|
||||
|
||||
if (streamer->init(*index_meta_ptr_, params) != 0) {
|
||||
return nullptr;
|
||||
}
|
||||
return streamer;
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestAddVector) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestAddVector", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
for (size_t i = 0; i < 1000UL; i++) {
|
||||
NumericalVector<float> vec(kDim);
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
streamer.reset();
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestLinearSearch) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestLinearSearch.index", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 5000UL;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
NumericalVector<float> vec(kDim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
size_t topk = 3;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ctx->set_topk(1U);
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result1 = ctx->result();
|
||||
ASSERT_EQ(1UL, result1.size());
|
||||
ASSERT_EQ(i, result1[0].key());
|
||||
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
ctx->set_topk(topk);
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, ctx));
|
||||
auto &result2 = ctx->result();
|
||||
ASSERT_EQ(topk, result2.size());
|
||||
ASSERT_EQ(i, result2[0].key());
|
||||
ASSERT_EQ(i == cnt - 1 ? i - 1 : i + 1, result2[1].key());
|
||||
ASSERT_EQ(i == 0 ? 2 : (i == cnt - 1 ? i - 2 : i - 1), result2[2].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestKnnSearch) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestKnnSearch.index", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
NumericalVector<float> vec(kDim);
|
||||
size_t cnt = 5000U;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
auto linearCtx = streamer->create_context();
|
||||
auto knnCtx = streamer->create_context();
|
||||
size_t topk = 100;
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
int topk1Hits = 0;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx));
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
topk1Hits += i == knnResult[0].key();
|
||||
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
float topk1Recall = topk1Hits * 1.0f / cnt;
|
||||
EXPECT_GT(recall, 0.90f);
|
||||
EXPECT_GT(topk1Recall, 0.95f);
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestOpenClose) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
constexpr size_t dim_large = 128;
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim_large);
|
||||
meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
|
||||
streamer = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
ASSERT_EQ(0, streamer->init(meta, params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestOpenClose.index", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t testCnt = 200;
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim_large);
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
for (size_t i = 0; i < testCnt; i++) {
|
||||
std::vector<float> vec(dim_large);
|
||||
for (size_t d = 0; d < dim_large; ++d) {
|
||||
vec[d] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
ASSERT_EQ(0, streamer->flush(0UL));
|
||||
ASSERT_EQ(0, streamer->close());
|
||||
|
||||
// Re-open and verify data
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
auto provider = streamer->create_provider();
|
||||
auto iter = provider->create_iterator();
|
||||
ASSERT_TRUE(!!iter);
|
||||
size_t total = 0;
|
||||
while (iter->is_valid()) {
|
||||
float *data = (float *)iter->data();
|
||||
for (size_t d = 0; d < dim_large; ++d) {
|
||||
ASSERT_FLOAT_EQ(static_cast<float>(iter->key()), data[d]);
|
||||
}
|
||||
total++;
|
||||
iter->next();
|
||||
}
|
||||
ASSERT_EQ(testCnt, total);
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestKnnMultiThread) {
|
||||
constexpr size_t dim = 32;
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 500U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 200U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U);
|
||||
params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
auto streamer = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
ASSERT_EQ(0, streamer->init(meta, params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestKnnMultiThread", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
auto addVector = [&streamer, dim](int baseKey, size_t addCnt) {
|
||||
NumericalVector<float> vec(dim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
size_t succAdd = 0;
|
||||
auto ctx = streamer->create_context();
|
||||
for (size_t i = 0; i < addCnt; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i + baseKey);
|
||||
}
|
||||
succAdd += !streamer->add_impl(baseKey + i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
streamer->flush(0UL);
|
||||
return succAdd;
|
||||
};
|
||||
auto t1 = std::async(std::launch::async, addVector, 0, 1000);
|
||||
auto t2 = std::async(std::launch::async, addVector, 1000, 1000);
|
||||
auto t3 = std::async(std::launch::async, addVector, 2000, 1000);
|
||||
ASSERT_EQ(1000U, t1.get());
|
||||
ASSERT_EQ(1000U, t2.get());
|
||||
ASSERT_EQ(1000U, t3.get());
|
||||
streamer->close();
|
||||
|
||||
// Verify data
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
auto provider = streamer->create_provider();
|
||||
auto iter = provider->create_iterator();
|
||||
ASSERT_TRUE(!!iter);
|
||||
size_t total = 0;
|
||||
uint64_t minKey = 10000;
|
||||
uint64_t maxKey = 0;
|
||||
while (iter->is_valid()) {
|
||||
float *data = (float *)iter->data();
|
||||
for (size_t d = 0; d < dim; ++d) {
|
||||
ASSERT_FLOAT_EQ(static_cast<float>(iter->key()), data[d]);
|
||||
}
|
||||
total++;
|
||||
minKey = std::min(minKey, iter->key());
|
||||
maxKey = std::max(maxKey, iter->key());
|
||||
iter->next();
|
||||
}
|
||||
ASSERT_EQ(3000, total);
|
||||
ASSERT_EQ(0, minKey);
|
||||
ASSERT_EQ(2999, maxKey);
|
||||
|
||||
// Multi-thread search
|
||||
size_t topk = 100;
|
||||
size_t cnt = 3000;
|
||||
auto knnSearch = [&]() {
|
||||
NumericalVector<float> vec(dim);
|
||||
auto linearCtx = streamer->create_context();
|
||||
auto knnCtx = streamer->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
size_t totalCnts = 0;
|
||||
size_t totalHits = 0;
|
||||
for (size_t i = 0; i < cnt; i += 1) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx));
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ASSERT_TRUE((totalHits * 1.0f / totalCnts) > 0.80f);
|
||||
};
|
||||
auto s1 = std::async(std::launch::async, knnSearch);
|
||||
auto s2 = std::async(std::launch::async, knnSearch);
|
||||
auto s3 = std::async(std::launch::async, knnSearch);
|
||||
s1.wait();
|
||||
s2.wait();
|
||||
s3.wait();
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestContiguousMemory) {
|
||||
ailego::Params extra;
|
||||
extra.set(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, true);
|
||||
extra.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 2000U);
|
||||
auto streamer = CreateVamanaStreamer(extra);
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestContiguous.index", true));
|
||||
|
||||
// First build with default mmap mode
|
||||
{
|
||||
auto builder_streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, builder_streamer);
|
||||
ASSERT_EQ(0, builder_streamer->open(storage));
|
||||
auto ctx = builder_streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
NumericalVector<float> vec(kDim);
|
||||
size_t cnt = 3000UL;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, builder_streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
ASSERT_EQ(0, builder_streamer->flush(0UL));
|
||||
ASSERT_EQ(0, builder_streamer->close());
|
||||
}
|
||||
|
||||
// Re-open with contiguous memory mode for search
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 3000UL;
|
||||
size_t topk = 50;
|
||||
NumericalVector<float> vec(kDim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
auto linearCtx = streamer->create_context();
|
||||
auto knnCtx = streamer->create_context();
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
int totalHits = 0;
|
||||
int totalCnts = 0;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, knnCtx));
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
float recall = totalHits * 1.0f / totalCnts;
|
||||
EXPECT_GT(recall, 0.90f);
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestContiguousMultiThreadSearch) {
|
||||
constexpr size_t dim = 32;
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
// Build with mmap mode
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestContiguousMT", true));
|
||||
|
||||
{
|
||||
ailego::Params build_params;
|
||||
build_params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U);
|
||||
build_params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U);
|
||||
build_params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
build_params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
build_params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U);
|
||||
build_params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
auto builder = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, builder);
|
||||
ASSERT_EQ(0, builder->init(meta, build_params));
|
||||
ASSERT_EQ(0, builder->open(storage));
|
||||
|
||||
auto ctx = builder->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t i = 0; i < 3000; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, builder->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
ASSERT_EQ(0, builder->flush(0UL));
|
||||
ASSERT_EQ(0, builder->close());
|
||||
}
|
||||
|
||||
// Re-open with contiguous memory
|
||||
ailego::Params search_params;
|
||||
search_params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
search_params.set(PARAM_VAMANA_STREAMER_USE_CONTIGUOUS_MEMORY, true);
|
||||
|
||||
auto searcher = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, searcher);
|
||||
ASSERT_EQ(0, searcher->init(meta, search_params));
|
||||
ASSERT_EQ(0, searcher->open(storage));
|
||||
|
||||
size_t topk = 50;
|
||||
size_t cnt = 3000;
|
||||
auto knnSearch = [&]() {
|
||||
NumericalVector<float> vec(dim);
|
||||
auto linearCtx = searcher->create_context();
|
||||
auto knnCtx = searcher->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
linearCtx->set_topk(topk);
|
||||
knnCtx->set_topk(topk);
|
||||
size_t totalCnts = 0;
|
||||
size_t totalHits = 0;
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + 0.1f;
|
||||
}
|
||||
ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, knnCtx));
|
||||
ASSERT_EQ(0, searcher->search_bf_impl(vec.data(), qmeta, linearCtx));
|
||||
auto &knnResult = knnCtx->result();
|
||||
ASSERT_EQ(topk, knnResult.size());
|
||||
auto &linearResult = linearCtx->result();
|
||||
ASSERT_EQ(topk, linearResult.size());
|
||||
ASSERT_EQ(i, linearResult[0].key());
|
||||
for (size_t k = 0; k < topk; ++k) {
|
||||
totalCnts++;
|
||||
for (size_t j = 0; j < topk; ++j) {
|
||||
if (linearResult[j].key() == knnResult[k].key()) {
|
||||
totalHits++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ASSERT_TRUE((totalHits * 1.0f / totalCnts) > 0.80f);
|
||||
};
|
||||
auto s1 = std::async(std::launch::async, knnSearch);
|
||||
auto s2 = std::async(std::launch::async, knnSearch);
|
||||
auto s3 = std::async(std::launch::async, knnSearch);
|
||||
s1.wait();
|
||||
s2.wait();
|
||||
s3.wait();
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestProvider) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestProvider", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
size_t cnt = 500;
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
NumericalVector<float> vec(kDim);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
ASSERT_EQ(0, streamer->flush(0UL));
|
||||
|
||||
auto provider = streamer->create_provider();
|
||||
ASSERT_NE(nullptr, provider);
|
||||
auto iter = provider->create_iterator();
|
||||
ASSERT_TRUE(!!iter);
|
||||
size_t total = 0;
|
||||
while (iter->is_valid()) {
|
||||
ASSERT_NE(nullptr, iter->data());
|
||||
float *data = (float *)iter->data();
|
||||
for (size_t d = 0; d < kDim; ++d) {
|
||||
ASSERT_FLOAT_EQ(static_cast<float>(iter->key()), data[d]);
|
||||
}
|
||||
total++;
|
||||
iter->next();
|
||||
}
|
||||
ASSERT_EQ(cnt, total);
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestAddAndSearch) {
|
||||
auto streamer = CreateVamanaStreamer();
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestAddAndSearch.index", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
NumericalVector<float> vec(kDim);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kDim);
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
|
||||
// Add and search interleaved
|
||||
for (size_t batch = 0; batch < 5; batch++) {
|
||||
size_t base = batch * 200;
|
||||
for (size_t i = 0; i < 200; i++) {
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(base + i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(base + i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
|
||||
// Search for recently added vectors
|
||||
size_t current_cnt = (batch + 1) * 200;
|
||||
size_t topk = std::min(current_cnt, (size_t)10);
|
||||
auto searchCtx = streamer->create_context();
|
||||
searchCtx->set_topk(topk);
|
||||
for (size_t j = 0; j < kDim; ++j) {
|
||||
vec[j] = static_cast<float>(base);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->search_bf_impl(vec.data(), qmeta, searchCtx));
|
||||
auto &result = searchCtx->result();
|
||||
ASSERT_EQ(topk, result.size());
|
||||
ASSERT_EQ(base, result[0].key());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VamanaStreamerTest, TestKnnConcurrentAddAndSearch) {
|
||||
constexpr size_t dim = 32;
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 128U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 30U * 1024U * 1024U);
|
||||
params.set(PARAM_VAMANA_STREAMER_GET_VECTOR_ENABLE, true);
|
||||
|
||||
auto streamer = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
ASSERT_EQ(0, streamer->init(meta, params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestConcurrentAddSearch", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// First add some base data
|
||||
{
|
||||
auto ctx = streamer->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t i = 0; i < 2000; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
ASSERT_EQ(0, streamer->add_impl(i, vec.data(), qmeta, ctx));
|
||||
}
|
||||
}
|
||||
|
||||
std::atomic<bool> stop_search{false};
|
||||
|
||||
// Concurrent add
|
||||
auto addFuture = std::async(std::launch::async, [&]() {
|
||||
auto ctx = streamer->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t i = 2000; i < 3000; i++) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i);
|
||||
}
|
||||
streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
}
|
||||
stop_search.store(true);
|
||||
});
|
||||
|
||||
// Concurrent search
|
||||
auto searchFuture = std::async(std::launch::async, [&]() {
|
||||
auto ctx = streamer->create_context();
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
ctx->set_topk(10);
|
||||
while (!stop_search.load()) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 100.1f;
|
||||
}
|
||||
int ret = streamer->search_impl(vec.data(), qmeta, ctx);
|
||||
ASSERT_EQ(0, ret);
|
||||
auto &result = ctx->result();
|
||||
ASSERT_GT(result.size(), 0UL);
|
||||
}
|
||||
});
|
||||
|
||||
addFuture.wait();
|
||||
searchFuture.wait();
|
||||
}
|
||||
|
||||
// Test concurrent build (parallel add_impl) which was crashing due to
|
||||
// unprotected node_chunks_ / node_chunk_bases_ access during chunk allocation.
|
||||
TEST_F(VamanaStreamerTest, TestConcurrentBuild) {
|
||||
constexpr size_t dim = kDim;
|
||||
constexpr size_t total_vectors = 5000;
|
||||
constexpr size_t thread_count = 4;
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_INDEX_SIZE, 50U * 1024U * 1024U);
|
||||
|
||||
IndexMeta meta(IndexMeta::DataType::DT_FP32, dim);
|
||||
meta.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
auto streamer = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
ASSERT_EQ(0, streamer->init(meta, params));
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestConcurrentBuild", true));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// Parallel insertion from multiple threads (mimics local_builder behavior)
|
||||
std::atomic<int> error_count{0};
|
||||
std::vector<std::future<void>> futures;
|
||||
|
||||
for (size_t t = 0; t < thread_count; ++t) {
|
||||
futures.push_back(std::async(std::launch::async, [&, t]() {
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
|
||||
for (size_t i = t; i < total_vectors; i += thread_count) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = static_cast<float>(i) + static_cast<float>(j) * 0.01f;
|
||||
}
|
||||
int ret = streamer->add_impl(i, vec.data(), qmeta, ctx);
|
||||
if (ret != 0) {
|
||||
error_count.fetch_add(1);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
for (auto &f : futures) {
|
||||
f.wait();
|
||||
}
|
||||
ASSERT_EQ(0, error_count.load());
|
||||
|
||||
// Verify search still works correctly after concurrent build
|
||||
auto search_ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!search_ctx);
|
||||
search_ctx->set_topk(1);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim);
|
||||
NumericalVector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
vec[j] = 0.0f;
|
||||
}
|
||||
ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, search_ctx));
|
||||
auto &result = search_ctx->result();
|
||||
ASSERT_GT(result.size(), 0UL);
|
||||
}
|
||||
|
||||
// Test Vamana + INT8 quantization + rotation end-to-end
|
||||
TEST_F(VamanaStreamerTest, TestInt8WithRotate) {
|
||||
constexpr size_t kTestDim = 128;
|
||||
constexpr size_t kCnt = 2000U;
|
||||
constexpr size_t kTopk = 10;
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
|
||||
|
||||
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
index_meta_raw.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
// Create INT8 converter with rotation enabled
|
||||
Params converter_params;
|
||||
converter_params.set("integer_streaming.converter.enable_rotate", true);
|
||||
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
|
||||
ASSERT_NE(nullptr, converter);
|
||||
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer);
|
||||
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
|
||||
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// Add 2000 vectors
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
|
||||
for (size_t i = 0; i < kCnt; i++) {
|
||||
NumericalVector<float> vec(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
|
||||
|
||||
std::string new_vec;
|
||||
IndexQueryMeta new_meta;
|
||||
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
|
||||
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
streamer.reset();
|
||||
storage.reset();
|
||||
|
||||
// Reopen: reformer should auto-detect rotator from storage
|
||||
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage2);
|
||||
ASSERT_EQ(0, storage2->init(stg_params));
|
||||
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
|
||||
|
||||
auto streamer2 = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer2);
|
||||
ASSERT_EQ(0, streamer2->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer2->open(storage2));
|
||||
|
||||
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer2);
|
||||
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
|
||||
ASSERT_EQ(0, reformer2->load(storage2));
|
||||
|
||||
// Search: verify knn results are non-empty
|
||||
auto knnCtx = streamer2->create_context();
|
||||
knnCtx->set_topk(kTopk);
|
||||
auto linearCtx = streamer2->create_context();
|
||||
linearCtx->set_topk(kTopk);
|
||||
|
||||
NumericalVector<float> query(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
|
||||
|
||||
std::string new_query;
|
||||
IndexQueryMeta new_qmeta;
|
||||
ASSERT_EQ(0,
|
||||
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
|
||||
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
|
||||
ASSERT_EQ(0,
|
||||
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
|
||||
|
||||
EXPECT_EQ(kTopk, knnCtx->result().size());
|
||||
EXPECT_EQ(kTopk, linearCtx->result().size());
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
Reference in New Issue
Block a user