diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 1de9029..d58814d 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -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}) diff --git a/examples/c++/core/external_vector_example.cc b/examples/c++/core/external_vector_example.cc new file mode 100644 index 0000000..be37474 --- /dev/null +++ b/examples/c++/core/external_vector_example.cc @@ -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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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(node_id) * dim_; + } + + private: + const float *base_; + uint32_t dim_; +}; + +// --------------------------------------------------------------------------- +// Helper: generate random float vectors in [0, 1) +// --------------------------------------------------------------------------- +static std::vector generate_random_vectors(uint32_t count, + uint32_t dim) { + std::vector data(static_cast(count) * dim); + for (auto &v : data) { + v = static_cast(std::rand()) / static_cast(RAND_MAX); + } + return data; +} + +// --------------------------------------------------------------------------- +// Helper: brute-force kNN (L2) for recall verification +// --------------------------------------------------------------------------- +static std::vector brute_force_knn(const float *base, uint32_t n, + uint32_t dim, const float *query, + uint32_t topk) { + std::vector> 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(i) * dim + d] - query[d]; + dist += diff * diff; + } + dists[i] = {dist, i}; + } + std::partial_sort(dists.begin(), dists.begin() + topk, dists.end()); + std::vector 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(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(doc.key())) != + gt.end()) { + ++hits; + } + } + float recall = static_cast(hits) / static_cast(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; +} diff --git a/src/core/algorithm/hnsw/hnsw_algorithm.cc b/src/core/algorithm/hnsw/hnsw_algorithm.cc index 8c6fcfe..1fd6f65 100644 --- a/src/core/algorithm/hnsw/hnsw_algorithm.cc +++ b/src/core/algorithm/hnsw/hnsw_algorithm.cc @@ -693,6 +693,7 @@ void HnswAlgorithm::reverse_update_neighbors( template class HnswAlgorithm; template class HnswAlgorithm; template class HnswAlgorithm; +template class HnswAlgorithm; } // namespace core } // namespace zvec diff --git a/src/core/algorithm/hnsw/hnsw_context.cc b/src/core/algorithm/hnsw/hnsw_context.cc index 2310ffc..4710284 100644 --- a/src/core/algorithm/hnsw/hnsw_context.cc +++ b/src/core/algorithm/hnsw/hnsw_context.cc @@ -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(); diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index bc97bbf..cf18295 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -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 &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}; diff --git a/src/core/algorithm/hnsw/hnsw_entity.h b/src/core/algorithm/hnsw/hnsw_entity.h index 197d8e9..639fc4a 100644 --- a/src/core/algorithm/hnsw/hnsw_entity.h +++ b/src/core/algorithm/hnsw/hnsw_entity.h @@ -20,6 +20,7 @@ #include #include #include +#include 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"); diff --git a/src/core/algorithm/hnsw/hnsw_params.h b/src/core/algorithm/hnsw/hnsw_params.h index 6cb9e0c..2cd2ecf 100644 --- a/src/core/algorithm/hnsw/hnsw_params.h +++ b/src/core/algorithm/hnsw/hnsw_params.h @@ -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 diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index 64ff5a7..8bd03f2 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -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(stats_); + } else if (use_contiguous_memory_) { entity_ = std::make_unique(stats_); } else { entity_ = std::make_unique(stats_); @@ -359,6 +366,11 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { new HnswAlgorithm(contiguous_entity)); break; } + case HnswStorageMode::kExternal: + alg_ = HnswAlgorithmBase::UPointer( + new HnswAlgorithm( + static_cast(*entity_))); + break; default: alg_ = HnswAlgorithmBase::UPointer(new HnswAlgorithm( diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index 047cd6b..f06321b 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -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_{}; diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc index 50f15c3..1ad1ebd 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.cc @@ -856,6 +856,41 @@ const HnswEntity::Pointer HnswContiguousStreamerEntity::clone() const { return HnswEntity::Pointer(entity); } +const HnswEntity::Pointer HnswExternalStreamerEntity::clone() const { + std::vector 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 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 // ============================================================================ diff --git a/src/core/algorithm/hnsw/hnsw_streamer_entity.h b/src/core/algorithm/hnsw/hnsw_streamer_entity.h index 19f8ba1..9c53aaf 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer_entity.h +++ b/src/core/algorithm/hnsw/hnsw_streamer_entity.h @@ -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; + + using HnswMmapStreamerEntity::HnswMmapStreamerEntity; + + HnswStorageMode storage_mode() const override { + return HnswStorageMode::kExternal; + } + + //! Override clone to return the correct subclass type, so that + //! static_cast 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 &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(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(vec_src_->get_vector(id))); + return 0; + } + + int get_vector( + const node_id_t *ids, uint32_t count, + std::vector &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(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 diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index be6f5ff..00e4b4b 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -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, diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index f8371e6..9226eee 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -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")) { diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc index 0744da3..1c0dfca 100644 --- a/src/core/interface/indexes/hnsw_index.cc +++ b/src/core/interface/indexes/hnsw_index.cc @@ -15,6 +15,7 @@ #include #include #include +#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(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(context.get())) { + ctx->set_vector_source(&src); + } + return Index::Search(query, search_param, result); +} + int HNSWIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { param_ = dynamic_cast(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"); } diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 921221c..e9ed390 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -31,6 +31,7 @@ #include #include #include +#include #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(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; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 49051bc..2a91cd1 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -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; diff --git a/src/include/zvec/core/interface/index_param_builders.h b/src/include/zvec/core/interface/index_param_builders.h index 236e236..8b009c1 100644 --- a/src/include/zvec/core/interface/index_param_builders.h +++ b/src/include/zvec/core/interface/index_param_builders.h @@ -88,6 +88,11 @@ class BaseIndexParamBuilder { // : public return static_cast(*this); } + ActualIndexParamBuilderType &WithUseExternalVector(bool use_external_vector) { + param->use_external_vector = use_external_vector; + return static_cast(*this); + } + virtual std::shared_ptr Build() = 0; protected: diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h new file mode 100644 index 0000000..717ba51 --- /dev/null +++ b/src/include/zvec/core/interface/vector_source.h @@ -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 + +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 diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index f1e35af..85671d5 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -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(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 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(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 all_vectors(kDimension * kNumVectors, 0.0f); + for (uint32_t i = 0; i < kNumVectors; ++i) { + all_vectors[i * kDimension + i % kDimension] = static_cast(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 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"};