feat: external vector source support (#490)
This commit is contained in:
ZeFeng Yin
GitHub
parent
6df823f804
commit
e12bad5e37
|
|
@ -108,6 +108,10 @@ add_executable(core-example core/main.cc)
|
|||
target_link_libraries(core-example PRIVATE zvec-core-lib)
|
||||
list(APPEND ZVEC_EXAMPLE_TARGETS core-example)
|
||||
|
||||
add_executable(external-vector-example core/external_vector_example.cc)
|
||||
target_link_libraries(external-vector-example PRIVATE zvec-core-lib)
|
||||
list(APPEND ZVEC_EXAMPLE_TARGETS external-vector-example)
|
||||
|
||||
# Strip symbols to reduce executable size
|
||||
if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
|
||||
foreach(ZVEC_EXAMPLE_TARGET ${ZVEC_EXAMPLE_TARGETS})
|
||||
|
|
|
|||
|
|
@ -0,0 +1,219 @@
|
|||
// 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.
|
||||
|
||||
/// @file external_vector_example.cc
|
||||
/// @brief Demonstrates using HNSW index in external-vector mode.
|
||||
///
|
||||
/// In external-vector mode the index does NOT store raw vectors internally.
|
||||
/// Instead, a user-provided VectorSource is passed on every Add/Search call
|
||||
/// so the index can fetch vectors on demand. This is useful when vectors are
|
||||
/// already stored elsewhere (e.g. a columnar store, mmap file, remote storage)
|
||||
/// and you want to avoid duplicating them inside the index.
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <filesystem>
|
||||
#include <iostream>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
#include <zvec/core/interface/index.h>
|
||||
#include <zvec/core/interface/index_factory.h>
|
||||
#include <zvec/core/interface/index_param.h>
|
||||
#include <zvec/core/interface/index_param_builders.h>
|
||||
#include <zvec/core/interface/vector_source.h>
|
||||
|
||||
using namespace zvec::core_interface;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// A simple VectorSource backed by an in-memory float matrix.
|
||||
// In production this could be backed by mmap, a database, or remote storage.
|
||||
// ---------------------------------------------------------------------------
|
||||
class InMemoryVectorSource : public zvec::core::VectorSource {
|
||||
public:
|
||||
/// @param base Pointer to a contiguous float array of shape [n, dim].
|
||||
/// @param dim Dimensionality of each vector.
|
||||
InMemoryVectorSource(const float *base, uint32_t dim)
|
||||
: base_(base), dim_(dim) {}
|
||||
|
||||
const void *get_vector(uint32_t node_id) const override {
|
||||
return base_ + static_cast<size_t>(node_id) * dim_;
|
||||
}
|
||||
|
||||
private:
|
||||
const float *base_;
|
||||
uint32_t dim_;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper: generate random float vectors in [0, 1)
|
||||
// ---------------------------------------------------------------------------
|
||||
static std::vector<float> generate_random_vectors(uint32_t count,
|
||||
uint32_t dim) {
|
||||
std::vector<float> data(static_cast<size_t>(count) * dim);
|
||||
for (auto &v : data) {
|
||||
v = static_cast<float>(std::rand()) / static_cast<float>(RAND_MAX);
|
||||
}
|
||||
return data;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helper: brute-force kNN (L2) for recall verification
|
||||
// ---------------------------------------------------------------------------
|
||||
static std::vector<uint32_t> brute_force_knn(const float *base, uint32_t n,
|
||||
uint32_t dim, const float *query,
|
||||
uint32_t topk) {
|
||||
std::vector<std::pair<float, uint32_t>> dists(n);
|
||||
for (uint32_t i = 0; i < n; ++i) {
|
||||
float dist = 0.0f;
|
||||
for (uint32_t d = 0; d < dim; ++d) {
|
||||
float diff = base[static_cast<size_t>(i) * dim + d] - query[d];
|
||||
dist += diff * diff;
|
||||
}
|
||||
dists[i] = {dist, i};
|
||||
}
|
||||
std::partial_sort(dists.begin(), dists.begin() + topk, dists.end());
|
||||
std::vector<uint32_t> result(topk);
|
||||
for (uint32_t i = 0; i < topk; ++i) {
|
||||
result[i] = dists[i].second;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
int main() {
|
||||
constexpr uint32_t kDimension = 32;
|
||||
constexpr uint32_t kDocCount = 200;
|
||||
constexpr uint32_t kTopK = 5;
|
||||
const std::string index_path = "external_vector_example.index";
|
||||
|
||||
// Clean up any previous run
|
||||
std::filesystem::remove_all(index_path);
|
||||
|
||||
// ------ Step 1: Generate random vector data (simulating external storage)
|
||||
std::srand(42);
|
||||
auto vectors = generate_random_vectors(kDocCount, kDimension);
|
||||
|
||||
// Wrap data in our VectorSource
|
||||
InMemoryVectorSource source(vectors.data(), kDimension);
|
||||
|
||||
// ------ Step 2: Build HNSW index with external-vector mode enabled
|
||||
auto param = HNSWIndexParamBuilder()
|
||||
.WithMetricType(MetricType::kL2sq)
|
||||
.WithDataType(DataType::DT_FP32)
|
||||
.WithDimension(kDimension)
|
||||
.WithIsSparse(false)
|
||||
.WithUseExternalVector(true) // <-- key setting
|
||||
.Build();
|
||||
|
||||
auto index = IndexFactory::CreateAndInitIndex(*param);
|
||||
if (!index) {
|
||||
std::cerr << "Failed to create index." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ret = index->Open(
|
||||
index_path, StorageOptions{StorageOptions::StorageType::kMMAP, true});
|
||||
if (ret != 0) {
|
||||
std::cerr << "Failed to open index." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
// ------ Step 3: Add vectors using AddWithSource
|
||||
for (uint32_t i = 0; i < kDocCount; ++i) {
|
||||
VectorData vd;
|
||||
vd.vector =
|
||||
DenseVector{vectors.data() + static_cast<size_t>(i) * kDimension};
|
||||
ret = index->AddWithSource(vd, i, source);
|
||||
if (ret != 0) {
|
||||
std::cerr << "Failed to add doc " << i << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
std::cout << "[OK] Added " << kDocCount << " vectors in external mode."
|
||||
<< std::endl;
|
||||
|
||||
// ------ Step 4: Search using SearchWithSource
|
||||
auto query_param =
|
||||
HNSWQueryParamBuilder().with_topk(kTopK).with_ef_search(64).build();
|
||||
|
||||
// Use the first vector as query
|
||||
const float *query_vec = vectors.data();
|
||||
VectorData query;
|
||||
query.vector = DenseVector{query_vec};
|
||||
|
||||
SearchResult result;
|
||||
ret = index->SearchWithSource(query, query_param, source, &result);
|
||||
if (ret != 0) {
|
||||
std::cerr << "Search failed." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::cout << "[OK] Search returned " << result.doc_list_.size() << " results."
|
||||
<< std::endl;
|
||||
|
||||
// The closest vector to vectors[0] should be itself (doc_id=0)
|
||||
if (!result.doc_list_.empty() && result.doc_list_[0].key() == 0) {
|
||||
std::cout << "[OK] Nearest neighbor is doc_id=0 (self), score="
|
||||
<< result.doc_list_[0].score() << std::endl;
|
||||
}
|
||||
|
||||
// ------ Step 5: Verify recall against brute-force
|
||||
auto gt =
|
||||
brute_force_knn(vectors.data(), kDocCount, kDimension, query_vec, kTopK);
|
||||
uint32_t hits = 0;
|
||||
for (const auto &doc : result.doc_list_) {
|
||||
if (std::find(gt.begin(), gt.end(), static_cast<uint32_t>(doc.key())) !=
|
||||
gt.end()) {
|
||||
++hits;
|
||||
}
|
||||
}
|
||||
float recall = static_cast<float>(hits) / static_cast<float>(kTopK);
|
||||
std::cout << "[OK] Recall@" << kTopK << " = " << recall * 100.0f << "%"
|
||||
<< std::endl;
|
||||
|
||||
// ------ Step 6: Reopen index and search again (persistence verification)
|
||||
index->Close();
|
||||
std::cout << "[OK] Index closed." << std::endl;
|
||||
|
||||
// Must re-create index instance with same params before reopening
|
||||
index = IndexFactory::CreateAndInitIndex(*param);
|
||||
if (!index) {
|
||||
std::cerr << "Failed to re-create index for reopen." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
ret = index->Open(index_path,
|
||||
StorageOptions{StorageOptions::StorageType::kMMAP, false});
|
||||
if (ret != 0) {
|
||||
std::cerr << "Failed to reopen index." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
SearchResult result2;
|
||||
ret = index->SearchWithSource(query, query_param, source, &result2);
|
||||
if (ret != 0) {
|
||||
std::cerr << "Search after reopen failed." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::cout << "[OK] After reopen: search returned " << result2.doc_list_.size()
|
||||
<< " results, top1 doc_id=" << result2.doc_list_[0].key()
|
||||
<< std::endl;
|
||||
|
||||
// Cleanup
|
||||
index->Close();
|
||||
std::filesystem::remove_all(index_path);
|
||||
std::cout << "\n=== External Vector Example Complete ===" << std::endl;
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -693,6 +693,7 @@ void HnswAlgorithm<EntityType>::reverse_update_neighbors(
|
|||
template class HnswAlgorithm<HnswMmapStreamerEntity>;
|
||||
template class HnswAlgorithm<HnswBufferPoolStreamerEntity>;
|
||||
template class HnswAlgorithm<HnswContiguousStreamerEntity>;
|
||||
template class HnswAlgorithm<HnswExternalStreamerEntity>;
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
|
|||
|
|
@ -262,6 +262,9 @@ int HnswContext::update_context(ContextType type, const IndexMeta &meta,
|
|||
|
||||
entity_ = entity;
|
||||
dc_.update(entity_.get(), metric, meta.dimension());
|
||||
if (vector_source_) {
|
||||
entity_->set_vector_source(vector_source_);
|
||||
}
|
||||
magic_ = magic_num;
|
||||
level_topks_.clear();
|
||||
|
||||
|
|
|
|||
|
|
@ -121,6 +121,20 @@ class HnswContext : public IndexContext {
|
|||
return *entity_;
|
||||
}
|
||||
|
||||
//! Bind an external vector source to this context. It is stored so that it
|
||||
//! can be re-applied after the entity is re-cloned inside update_context,
|
||||
//! and immediately forwarded to the current entity clone.
|
||||
inline void set_vector_source(const VectorSource *src) {
|
||||
vector_source_ = src;
|
||||
if (entity_) {
|
||||
entity_->set_vector_source(src);
|
||||
}
|
||||
}
|
||||
|
||||
inline const VectorSource *vector_source() const {
|
||||
return vector_source_;
|
||||
}
|
||||
|
||||
inline void resize_results(size_t size) {
|
||||
if (group_by_search()) {
|
||||
group_results_.resize(size);
|
||||
|
|
@ -374,6 +388,7 @@ class HnswContext : public IndexContext {
|
|||
set_fetch_vector(false);
|
||||
set_group_params(0, 0);
|
||||
reset_group_by();
|
||||
set_vector_source(nullptr);
|
||||
}
|
||||
|
||||
inline std::map<std::string, TopkHeap> &group_topk_heaps() {
|
||||
|
|
@ -528,6 +543,7 @@ class HnswContext : public IndexContext {
|
|||
HnswEntity::Pointer entity_;
|
||||
HnswDistCalculator dc_;
|
||||
IndexMetric::Pointer metric_;
|
||||
const VectorSource *vector_source_{nullptr};
|
||||
|
||||
bool debug_mode_{false};
|
||||
bool force_padding_topk_{false};
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@
|
|||
#include <zvec/core/framework/index_dumper.h>
|
||||
#include <zvec/core/framework/index_error.h>
|
||||
#include <zvec/core/framework/index_storage.h>
|
||||
#include <zvec/core/interface/vector_source.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
|
@ -612,6 +613,11 @@ class HnswEntity {
|
|||
header_.hnsw.max_level = level;
|
||||
}
|
||||
|
||||
//! Bind an external vector source to this entity. The default
|
||||
//! implementation is a no-op; only entities that read vectors from an
|
||||
//! external source (e.g. HnswExternalStreamerEntity) override it.
|
||||
virtual void set_vector_source(const VectorSource * /*src*/) {}
|
||||
|
||||
virtual int load(const IndexStorage::Pointer & /*container*/,
|
||||
bool /*check_crc*/) {
|
||||
LOG_ERROR("Load not implemented");
|
||||
|
|
|
|||
|
|
@ -115,5 +115,8 @@ static const std::string PARAM_HNSW_REDUCER_EFCONSTRUCTION(
|
|||
static const std::string PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY(
|
||||
"proxima.hnsw.streamer.use_contiguous_memory");
|
||||
|
||||
static const std::string PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR(
|
||||
"proxima.hnsw.streamer.use_external_vector");
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) {
|
|||
params.get(PARAM_HNSW_STREAMER_USE_ID_MAP, &use_id_map_);
|
||||
params.get(PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY,
|
||||
&use_contiguous_memory_);
|
||||
params.get(PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR, &use_external_vector_);
|
||||
|
||||
params.get(PARAM_HNSW_STREAMER_DOCS_SOFT_LIMIT, &docs_soft_limit_);
|
||||
if (docs_soft_limit_ > 0 && docs_soft_limit_ > docs_hard_limit_) {
|
||||
|
|
@ -224,7 +225,11 @@ int HnswStreamer::setup_entity() {
|
|||
entity_->set_l0_neighbor_cnt(l0_max_neighbor_cnt_);
|
||||
entity_->set_scaling_factor(scaling_factor_);
|
||||
entity_->set_prune_cnt(prune_cnt_);
|
||||
entity_->set_vector_size(meta_.element_size());
|
||||
// For external-vector entities the per-node vector prefix is removed; set
|
||||
// vector_size to 0 so all inherited offset computations (key / neighbors /
|
||||
// node_size) are correct and add_vector writes no vector bytes. The distance
|
||||
// dimension is taken from meta.dimension(), not from vector_size().
|
||||
entity_->set_vector_size(use_external_vector_ ? 0 : meta_.element_size());
|
||||
entity_->set_chunk_size(chunk_size_);
|
||||
entity_->set_filter_same_key(filter_same_key_);
|
||||
entity_->set_get_vector(get_vector_enabled_);
|
||||
|
|
@ -252,7 +257,9 @@ int HnswStreamer::open(IndexStorage::Pointer stg) {
|
|||
break;
|
||||
}
|
||||
default: {
|
||||
if (use_contiguous_memory_) {
|
||||
if (use_external_vector_) {
|
||||
entity_ = std::make_unique<HnswExternalStreamerEntity>(stats_);
|
||||
} else if (use_contiguous_memory_) {
|
||||
entity_ = std::make_unique<HnswContiguousStreamerEntity>(stats_);
|
||||
} else {
|
||||
entity_ = std::make_unique<HnswMmapStreamerEntity>(stats_);
|
||||
|
|
@ -359,6 +366,11 @@ int HnswStreamer::open(IndexStorage::Pointer stg) {
|
|||
new HnswAlgorithm<HnswContiguousStreamerEntity>(contiguous_entity));
|
||||
break;
|
||||
}
|
||||
case HnswStorageMode::kExternal:
|
||||
alg_ = HnswAlgorithmBase::UPointer(
|
||||
new HnswAlgorithm<HnswExternalStreamerEntity>(
|
||||
static_cast<HnswExternalStreamerEntity &>(*entity_)));
|
||||
break;
|
||||
default:
|
||||
alg_ =
|
||||
HnswAlgorithmBase::UPointer(new HnswAlgorithm<HnswMmapStreamerEntity>(
|
||||
|
|
|
|||
|
|
@ -234,6 +234,7 @@ class HnswStreamer : public IndexStreamer {
|
|||
bool force_padding_topk_enabled_{false};
|
||||
bool use_id_map_{true};
|
||||
bool use_contiguous_memory_{false};
|
||||
bool use_external_vector_{false};
|
||||
|
||||
//! avoid add vector while dumping index
|
||||
ailego::SharedMutex shared_mutex_{};
|
||||
|
|
|
|||
|
|
@ -856,6 +856,41 @@ const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const {
|
|||
return HnswEntity::Pointer(entity);
|
||||
}
|
||||
|
||||
const HnswEntity::Pointer HnswExternalStreamerEntity::clone() const {
|
||||
std::vector<Chunk::Pointer> node_chunks;
|
||||
node_chunks.reserve(node_chunks_.size());
|
||||
for (size_t i = 0UL; i < node_chunks_.size(); ++i) {
|
||||
node_chunks.emplace_back(node_chunks_[i]->clone());
|
||||
if (ailego_unlikely(!node_chunks[i])) {
|
||||
LOG_ERROR("HnswExternalStreamerEntity get chunk failed in clone");
|
||||
return HnswEntity::Pointer();
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<Chunk::Pointer> upper_neighbor_chunks;
|
||||
upper_neighbor_chunks.reserve(upper_neighbor_chunks_.size());
|
||||
for (size_t i = 0UL; i < upper_neighbor_chunks_.size(); ++i) {
|
||||
upper_neighbor_chunks.emplace_back(upper_neighbor_chunks_[i]->clone());
|
||||
if (ailego_unlikely(!upper_neighbor_chunks[i])) {
|
||||
LOG_ERROR("HnswExternalStreamerEntity get chunk failed in clone");
|
||||
return HnswEntity::Pointer();
|
||||
}
|
||||
}
|
||||
|
||||
// Note: vec_src_ is intentionally NOT shared with the clone; it stays null
|
||||
// and is re-bound per add/search call via HnswContext::set_vector_source.
|
||||
auto *entity = new (std::nothrow) HnswExternalStreamerEntity(
|
||||
stats_, header(), chunk_size_, node_index_mask_bits_,
|
||||
upper_neighbor_mask_bits_, filter_same_key_, get_vector_enabled_,
|
||||
upper_neighbor_index_, upper_neighbor_rw_mutex_, keys_map_lock_,
|
||||
keys_map_, use_key_info_map_, std::move(node_chunks),
|
||||
std::move(upper_neighbor_chunks), broker_, nullptr, nullptr);
|
||||
if (ailego_unlikely(!entity)) {
|
||||
LOG_ERROR("HnswExternalStreamerEntity new failed");
|
||||
}
|
||||
return HnswEntity::Pointer(entity);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// HnswContiguousStreamerEntity implementation
|
||||
// ============================================================================
|
||||
|
|
|
|||
|
|
@ -36,7 +36,12 @@ namespace core {
|
|||
|
||||
|
||||
//! Storage mode for HnswStreamerEntity
|
||||
enum class HnswStorageMode { kMmap = 0, kBufferPool = 1, kContiguous = 2 };
|
||||
enum class HnswStorageMode {
|
||||
kMmap = 0,
|
||||
kBufferPool = 1,
|
||||
kContiguous = 2,
|
||||
kExternal = 3
|
||||
};
|
||||
|
||||
//! HnswStreamerEntity manage vector data, pkey, and node's neighbors
|
||||
class HnswStreamerEntity : public HnswEntity {
|
||||
|
|
@ -1098,5 +1103,93 @@ class HnswContiguousStreamerEntity : public HnswMmapStreamerEntity {
|
|||
static char *allocate_contiguous(size_t size);
|
||||
};
|
||||
|
||||
//! Typed entity subclass that reads vectors from an external vector source.
|
||||
//! The graph structure (key + neighbors) is stored in chunks just like
|
||||
//! HnswMmapStreamerEntity, but the per-node vector prefix is removed by
|
||||
//! setting vector_size() to 0 (see HnswStreamer setup). With vector_size()==0
|
||||
//! all inherited offset computations (key at node start, L0 neighbors right
|
||||
//! after the key, node_size == AlignSize(sizeof(key) + neighbor_size)) become
|
||||
//! automatically correct, and base add_vector writes a zero-byte vector (i.e.
|
||||
//! it skips vector storage). Vectors are instead read through the bound
|
||||
//! VectorSource, which is supplied per add/search call.
|
||||
class HnswExternalStreamerEntity : public HnswMmapStreamerEntity {
|
||||
public:
|
||||
using MemoryBlock = MmapMemoryBlock;
|
||||
using TypedNeighbors = NeighborsT<MmapMemoryBlock>;
|
||||
|
||||
using HnswMmapStreamerEntity::HnswMmapStreamerEntity;
|
||||
|
||||
HnswStorageMode storage_mode() const override {
|
||||
return HnswStorageMode::kExternal;
|
||||
}
|
||||
|
||||
//! Override clone to return the correct subclass type, so that
|
||||
//! static_cast<const HnswExternalStreamerEntity&> in the algorithm is safe.
|
||||
//! The external vector source is NOT shared with the clone; it is re-applied
|
||||
//! per add/search call (via HnswContext::set_vector_source).
|
||||
const HnswEntity::Pointer clone() const override;
|
||||
|
||||
//! Bind the external vector source for the current add/search call.
|
||||
void set_vector_source(const VectorSource *src) override {
|
||||
vec_src_ = src;
|
||||
}
|
||||
|
||||
//! Typed batch get_vector: zero-copy view into the external vector source.
|
||||
//! Hides HnswMmapStreamerEntity::get_vector_typed (non-virtual, used by the
|
||||
//! template algorithm via static_cast).
|
||||
inline int get_vector_typed(const node_id_t *ids, uint32_t count,
|
||||
std::vector<MmapMemoryBlock> &vec_blocks) const {
|
||||
if (ailego_unlikely(vec_src_ == nullptr)) {
|
||||
return IndexError_Runtime;
|
||||
}
|
||||
vec_blocks.resize(count);
|
||||
for (auto i = 0U; i < count; ++i) {
|
||||
vec_blocks[i].reset(const_cast<void *>(vec_src_->get_vector(ids[i])));
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Virtual get_vector overrides (distance-calculator / provider paths).
|
||||
const void *get_vector(node_id_t id) const override {
|
||||
return vec_src_ ? vec_src_->get_vector(id) : nullptr;
|
||||
}
|
||||
|
||||
int get_vector(const node_id_t *ids, uint32_t count,
|
||||
const void **vecs) const override {
|
||||
if (ailego_unlikely(vec_src_ == nullptr)) {
|
||||
return IndexError_Runtime;
|
||||
}
|
||||
vec_src_->get_vectors(ids, count, vecs);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int get_vector(const node_id_t id,
|
||||
IndexStorage::MemoryBlock &block) const override {
|
||||
if (ailego_unlikely(vec_src_ == nullptr)) {
|
||||
return IndexError_Runtime;
|
||||
}
|
||||
block.reset(const_cast<void *>(vec_src_->get_vector(id)));
|
||||
return 0;
|
||||
}
|
||||
|
||||
int get_vector(
|
||||
const node_id_t *ids, uint32_t count,
|
||||
std::vector<IndexStorage::MemoryBlock> &vec_blocks) const override {
|
||||
if (ailego_unlikely(vec_src_ == nullptr)) {
|
||||
return IndexError_Runtime;
|
||||
}
|
||||
vec_blocks.resize(count);
|
||||
for (auto i = 0U; i < count; ++i) {
|
||||
vec_blocks[i].reset(const_cast<void *>(vec_src_->get_vector(ids[i])));
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
//! Transient, per-call vector source. Never shared across clones; bound by
|
||||
//! HnswContext::set_vector_source before each add/search.
|
||||
const VectorSource *vec_src_{nullptr};
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
|
|||
|
|
@ -426,6 +426,19 @@ int Index::Add(const VectorData &vector_data, const uint32_t doc_id) {
|
|||
return ret;
|
||||
}
|
||||
|
||||
int Index::AddWithSource(const VectorData & /*vector*/, uint32_t /*doc_id*/,
|
||||
const core::VectorSource & /*src*/) {
|
||||
LOG_ERROR("AddWithSource is not supported by this index type");
|
||||
return core::IndexError_Unsupported;
|
||||
}
|
||||
|
||||
int Index::SearchWithSource(
|
||||
const VectorData & /*query*/,
|
||||
const BaseIndexQueryParam::Pointer & /*search_param*/,
|
||||
const core::VectorSource & /*src*/, SearchResult * /*result*/) {
|
||||
LOG_ERROR("SearchWithSource is not supported by this index type");
|
||||
return core::IndexError_Unsupported;
|
||||
}
|
||||
|
||||
int Index::Search(const VectorData &vector_data,
|
||||
const BaseIndexQueryParam::Pointer &search_param,
|
||||
|
|
|
|||
|
|
@ -49,6 +49,9 @@ ailego::JsonObject BaseIndexParam::SerializeToJsonObject(
|
|||
if (!omit_empty_value || is_huge_page) {
|
||||
json_obj.set("is_huge_page", ailego::JsonValue(is_huge_page));
|
||||
}
|
||||
if (!omit_empty_value || use_external_vector) {
|
||||
json_obj.set("use_external_vector", ailego::JsonValue(use_external_vector));
|
||||
}
|
||||
|
||||
// if (preprocess_param) {
|
||||
// json.set("preprocess_param", preprocess_param->SerializeToJson());
|
||||
|
|
@ -101,6 +104,7 @@ bool BaseIndexParam::DeserializeFromJsonObject(
|
|||
DESERIALIZE_VALUE_FIELD(json_obj, is_sparse);
|
||||
DESERIALIZE_VALUE_FIELD(json_obj, use_id_map);
|
||||
DESERIALIZE_VALUE_FIELD(json_obj, is_huge_page);
|
||||
DESERIALIZE_VALUE_FIELD(json_obj, use_external_vector);
|
||||
|
||||
ailego::JsonValue tmp_json_value;
|
||||
if (json_obj.has("quantizer_param")) {
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@
|
|||
#include <memory>
|
||||
#include <string>
|
||||
#include <zvec/core/interface/index.h>
|
||||
#include "algorithm/hnsw/hnsw_context.h"
|
||||
#include "algorithm/hnsw/hnsw_params.h"
|
||||
#include "algorithm/hnsw/hnsw_streamer.h"
|
||||
#include "algorithm/hnsw/hnsw_streamer_entity.h"
|
||||
|
|
@ -38,10 +39,40 @@ std::string HNSWIndex::storage_mode() const {
|
|||
return "buffer_pool";
|
||||
case core::HnswStorageMode::kContiguous:
|
||||
return "contiguous";
|
||||
case core::HnswStorageMode::kExternal:
|
||||
return "external";
|
||||
}
|
||||
return "";
|
||||
}
|
||||
|
||||
int HNSWIndex::AddWithSource(const VectorData &vector_data,
|
||||
const uint32_t doc_id,
|
||||
const core::VectorSource &src) {
|
||||
auto &context = acquire_context();
|
||||
if (!context) {
|
||||
LOG_ERROR("Failed to acquire context for AddWithSource");
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
if (auto *ctx = dynamic_cast<core::HnswContext *>(context.get())) {
|
||||
ctx->set_vector_source(&src);
|
||||
}
|
||||
return Index::Add(vector_data, doc_id);
|
||||
}
|
||||
|
||||
int HNSWIndex::SearchWithSource(
|
||||
const VectorData &query, const BaseIndexQueryParam::Pointer &search_param,
|
||||
const core::VectorSource &src, SearchResult *result) {
|
||||
auto &context = acquire_context();
|
||||
if (!context) {
|
||||
LOG_ERROR("Failed to acquire context for SearchWithSource");
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
if (auto *ctx = dynamic_cast<core::HnswContext *>(context.get())) {
|
||||
ctx->set_vector_source(&src);
|
||||
}
|
||||
return Index::Search(query, search_param, result);
|
||||
}
|
||||
|
||||
int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {
|
||||
param_ = dynamic_cast<const HNSWIndexParam &>(param);
|
||||
|
||||
|
|
@ -81,6 +112,8 @@ int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) {
|
|||
param_.use_id_map);
|
||||
proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_CONTIGUOUS_MEMORY,
|
||||
param_.use_contiguous_memory);
|
||||
proxima_index_params_.set(core::PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR,
|
||||
param_.use_external_vector);
|
||||
streamer_ = core::IndexFactory::CreateStreamer("HnswStreamer");
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@
|
|||
#include <zvec/core/framework/index_searcher.h>
|
||||
#include <zvec/core/framework/index_storage.h>
|
||||
#include <zvec/core/interface/index_param.h>
|
||||
#include <zvec/core/interface/vector_source.h>
|
||||
#include "zvec/core/framework/index_provider.h"
|
||||
|
||||
namespace zvec::core_interface {
|
||||
|
|
@ -126,12 +127,20 @@ class Index {
|
|||
// TODO: static reduce
|
||||
|
||||
virtual int Add(const VectorData &vector, uint32_t doc_id);
|
||||
|
||||
virtual int Fetch(const uint32_t doc_id,
|
||||
VectorDataBuffer *vector_data_buffer);
|
||||
virtual int Search(const VectorData &query,
|
||||
const BaseIndexQueryParam::Pointer &search_param,
|
||||
SearchResult *result);
|
||||
|
||||
virtual int AddWithSource(const VectorData &vector, uint32_t doc_id,
|
||||
const core::VectorSource &src);
|
||||
virtual int SearchWithSource(const VectorData &query,
|
||||
const BaseIndexQueryParam::Pointer &search_param,
|
||||
const core::VectorSource &src,
|
||||
SearchResult *result);
|
||||
|
||||
virtual BaseIndexParam::Pointer GetParam() const {
|
||||
return std::make_shared<BaseIndexParam>(param_);
|
||||
}
|
||||
|
|
@ -287,12 +296,19 @@ class HNSWIndex : public Index {
|
|||
HNSWIndex() = default;
|
||||
|
||||
//! Retrieve the storage mode of the underlying HNSW streamer entity.
|
||||
//! Returns a string among {"mmap", "buffer_pool", "contiguous"}.
|
||||
//! Returns a string among {"mmap", "buffer_pool", "contiguous", "external"}.
|
||||
//! Intended for introspection and debug/testing usage. Returns empty
|
||||
//! string when the streamer has not been initialized or is of an
|
||||
//! unexpected type (e.g. the sparse branch).
|
||||
std::string storage_mode() const;
|
||||
|
||||
int AddWithSource(const VectorData &vector, uint32_t doc_id,
|
||||
const core::VectorSource &src) override;
|
||||
int SearchWithSource(const VectorData &query,
|
||||
const BaseIndexQueryParam::Pointer &search_param,
|
||||
const core::VectorSource &src,
|
||||
SearchResult *result) override;
|
||||
|
||||
protected:
|
||||
int CreateAndInitStreamer(const BaseIndexParam ¶m) override;
|
||||
|
||||
|
|
|
|||
|
|
@ -251,6 +251,7 @@ class BaseIndexParam : public SerializableBase {
|
|||
bool is_huge_page = false;
|
||||
DataType data_type = DataType::DT_UNDEFINED;
|
||||
bool use_id_map = true;
|
||||
bool use_external_vector = false;
|
||||
|
||||
// IndexMeta meta;
|
||||
ailego::Params params;
|
||||
|
|
|
|||
|
|
@ -88,6 +88,11 @@ class BaseIndexParamBuilder { // : public
|
|||
return static_cast<ActualIndexParamBuilderType &>(*this);
|
||||
}
|
||||
|
||||
ActualIndexParamBuilderType &WithUseExternalVector(bool use_external_vector) {
|
||||
param->use_external_vector = use_external_vector;
|
||||
return static_cast<ActualIndexParamBuilderType &>(*this);
|
||||
}
|
||||
|
||||
virtual std::shared_ptr<ActualIndexParamType> Build() = 0;
|
||||
|
||||
protected:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,37 @@
|
|||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
class VectorSource {
|
||||
public:
|
||||
virtual ~VectorSource() = default;
|
||||
|
||||
virtual const void *get_vector(uint32_t node_id) const = 0;
|
||||
|
||||
virtual void get_vectors(const uint32_t *ids, uint32_t count,
|
||||
const void **out) const {
|
||||
for (uint32_t i = 0; i < count; ++i) {
|
||||
out[i] = get_vector(ids[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -1856,6 +1856,151 @@ TEST(IndexInterface, ContiguousMemoryEndToEnd) {
|
|||
.build());
|
||||
}
|
||||
|
||||
class TestVectorSource : public zvec::core::VectorSource {
|
||||
public:
|
||||
TestVectorSource(const float *base, uint32_t dim) : base_(base), dim_(dim) {}
|
||||
|
||||
const void *get_vector(uint32_t node_id) const override {
|
||||
return base_ + static_cast<size_t>(node_id) * dim_;
|
||||
}
|
||||
|
||||
private:
|
||||
const float *base_;
|
||||
uint32_t dim_;
|
||||
};
|
||||
|
||||
TEST(IndexInterface, ExternalVectorEndToEnd) {
|
||||
constexpr uint32_t kDimension = 64;
|
||||
constexpr uint32_t kNumVectors = 100;
|
||||
const std::string index_name{"test_external.index"};
|
||||
|
||||
std::vector<float> all_vectors(kDimension * kNumVectors);
|
||||
for (uint32_t i = 0; i < kNumVectors; ++i) {
|
||||
for (uint32_t d = 0; d < kDimension; ++d) {
|
||||
all_vectors[i * kDimension + d] =
|
||||
static_cast<float>(i * kDimension + d) * 0.01f;
|
||||
}
|
||||
}
|
||||
|
||||
TestVectorSource source(all_vectors.data(), kDimension);
|
||||
|
||||
zvec::test_util::RemoveTestFiles(index_name + "*");
|
||||
|
||||
auto param = HNSWIndexParamBuilder()
|
||||
.WithMetricType(MetricType::kL2sq)
|
||||
.WithDataType(DataType::DT_FP32)
|
||||
.WithDimension(kDimension)
|
||||
.WithIsSparse(false)
|
||||
.WithEFConstruction(100)
|
||||
.WithUseExternalVector(true)
|
||||
.Build();
|
||||
|
||||
auto index = IndexFactory::CreateAndInitIndex(*param);
|
||||
ASSERT_NE(nullptr, index);
|
||||
|
||||
index->Open(index_name, {StorageOptions::StorageType::kMMAP, true});
|
||||
|
||||
for (uint32_t i = 0; i < kNumVectors; ++i) {
|
||||
VectorData vector_data;
|
||||
vector_data.vector = DenseVector{all_vectors.data() + i * kDimension};
|
||||
int ret = index->AddWithSource(vector_data, i, source);
|
||||
ASSERT_EQ(0, ret) << "AddWithSource failed for doc_id=" << i;
|
||||
}
|
||||
|
||||
auto query_param = HNSWQueryParamBuilder()
|
||||
.with_topk(5)
|
||||
.with_fetch_vector(false)
|
||||
.with_ef_search(50)
|
||||
.build();
|
||||
|
||||
VectorData query;
|
||||
query.vector = DenseVector{all_vectors.data()};
|
||||
SearchResult result;
|
||||
int ret = index->SearchWithSource(query, query_param, source, &result);
|
||||
ASSERT_EQ(0, ret);
|
||||
ASSERT_GE(result.doc_list_.size(), 1u);
|
||||
ASSERT_EQ(0u, result.doc_list_[0].key());
|
||||
ASSERT_FLOAT_EQ(0.0f, result.doc_list_[0].score());
|
||||
|
||||
VectorData query2;
|
||||
query2.vector = DenseVector{all_vectors.data() + 50 * kDimension};
|
||||
SearchResult result2;
|
||||
ret = index->SearchWithSource(query2, query_param, source, &result2);
|
||||
ASSERT_EQ(0, ret);
|
||||
ASSERT_GE(result2.doc_list_.size(), 1u);
|
||||
ASSERT_EQ(50u, result2.doc_list_[0].key());
|
||||
ASSERT_FLOAT_EQ(0.0f, result2.doc_list_[0].score());
|
||||
|
||||
index->Close();
|
||||
|
||||
auto index2 = IndexFactory::CreateAndInitIndex(*param);
|
||||
ASSERT_NE(nullptr, index2);
|
||||
index2->Open(index_name, {StorageOptions::StorageType::kMMAP, false});
|
||||
|
||||
SearchResult result3;
|
||||
ret = index2->SearchWithSource(query, query_param, source, &result3);
|
||||
ASSERT_EQ(0, ret);
|
||||
ASSERT_GE(result3.doc_list_.size(), 1u);
|
||||
ASSERT_EQ(0u, result3.doc_list_[0].key());
|
||||
ASSERT_FLOAT_EQ(0.0f, result3.doc_list_[0].score());
|
||||
|
||||
index2->Close();
|
||||
zvec::test_util::RemoveTestFiles(index_name + "*");
|
||||
}
|
||||
|
||||
TEST(IndexInterface, ExternalVectorInnerProduct) {
|
||||
constexpr uint32_t kDimension = 16;
|
||||
constexpr uint32_t kNumVectors = 10;
|
||||
const std::string index_name{"test_external_ip.index"};
|
||||
|
||||
std::vector<float> all_vectors(kDimension * kNumVectors, 0.0f);
|
||||
for (uint32_t i = 0; i < kNumVectors; ++i) {
|
||||
all_vectors[i * kDimension + i % kDimension] = static_cast<float>(i + 1);
|
||||
}
|
||||
|
||||
TestVectorSource source(all_vectors.data(), kDimension);
|
||||
|
||||
zvec::test_util::RemoveTestFiles(index_name + "*");
|
||||
|
||||
auto param = HNSWIndexParamBuilder()
|
||||
.WithMetricType(MetricType::kInnerProduct)
|
||||
.WithDataType(DataType::DT_FP32)
|
||||
.WithDimension(kDimension)
|
||||
.WithIsSparse(false)
|
||||
.WithEFConstruction(100)
|
||||
.WithUseExternalVector(true)
|
||||
.Build();
|
||||
|
||||
auto index = IndexFactory::CreateAndInitIndex(*param);
|
||||
ASSERT_NE(nullptr, index);
|
||||
index->Open(index_name, {StorageOptions::StorageType::kMMAP, true});
|
||||
|
||||
for (uint32_t i = 0; i < kNumVectors; ++i) {
|
||||
VectorData vector_data;
|
||||
vector_data.vector = DenseVector{all_vectors.data() + i * kDimension};
|
||||
ASSERT_EQ(0, index->AddWithSource(vector_data, i, source));
|
||||
}
|
||||
|
||||
std::vector<float> query_vec(kDimension, 0.0f);
|
||||
query_vec[0] = 1.0f;
|
||||
VectorData query;
|
||||
query.vector = DenseVector{query_vec.data()};
|
||||
|
||||
auto query_param = HNSWQueryParamBuilder()
|
||||
.with_topk(1)
|
||||
.with_fetch_vector(false)
|
||||
.with_ef_search(50)
|
||||
.build();
|
||||
|
||||
SearchResult result;
|
||||
ASSERT_EQ(0, index->SearchWithSource(query, query_param, source, &result));
|
||||
ASSERT_EQ(1u, result.doc_list_.size());
|
||||
ASSERT_EQ(0u, result.doc_list_[0].key());
|
||||
ASSERT_FLOAT_EQ(1.0f, result.doc_list_[0].score());
|
||||
|
||||
index->Close();
|
||||
zvec::test_util::RemoveTestFiles(index_name + "*");
|
||||
}
|
||||
TEST(IndexInterface, IsDirty) {
|
||||
constexpr uint32_t kDimension = 16;
|
||||
const std::string index_name{"test_is_dirty.index"};
|
||||
|
|
|
|||
Loading…
Reference in New Issue