From 2b9987d9b8a40f1d227f468b94152518e2c703b7 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Thu, 8 Oct 2026 17:49:48 +0800 Subject: [PATCH 01/16] feat(flat): support external vector sources with batched iteration --- examples/c++/CMakeLists.txt | 4 + examples/c++/core/external_flat_example.cc | 120 ++++ src/core/algorithm/flat/flat_streamer.cc | 220 +++++++ src/core/algorithm/flat/flat_streamer.h | 11 + .../algorithm/flat/flat_streamer_context.h | 18 +- .../algorithm/flat/flat_streamer_entity.cc | 21 +- .../algorithm/flat/flat_streamer_entity.h | 32 +- src/core/algorithm/flat/flat_utility.h | 2 + src/core/interface/index.cc | 20 + src/core/interface/indexes/flat_index.cc | 83 +++ src/core/interface/vector_source.cc | 22 + src/include/zvec/core/interface/index.h | 26 + src/include/zvec/core/interface/index_param.h | 3 + .../zvec/core/interface/vector_source.h | 38 ++ .../interface/external_flat_index_test.cc | 571 ++++++++++++++++++ 15 files changed, 1184 insertions(+), 7 deletions(-) create mode 100644 examples/c++/core/external_flat_example.cc create mode 100644 tests/core/interface/external_flat_index_test.cc diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index ecb138c173..771502bac0 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -135,6 +135,10 @@ 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) +add_executable(external-flat-example core/external_flat_example.cc) +target_link_libraries(external-flat-example PRIVATE zvec-core-lib) +list(APPEND ZVEC_EXAMPLE_TARGETS external-flat-example) + # This is also a link-integration check: DiskANN must work through the # all-in-one core library without linking core_knn_diskann separately. add_executable(diskann-core-example core/diskann_main.cc) diff --git a/examples/c++/core/external_flat_example.cc b/examples/c++/core/external_flat_example.cc new file mode 100644 index 0000000000..683eb23b6d --- /dev/null +++ b/examples/c++/core/external_flat_example.cc @@ -0,0 +1,120 @@ +// 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. + +// External Flat scans logical batches supplied by the source. A disk-backed +// source may read physical blocks and attach page pins to Batch::lease. +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace zvec::core_interface; + +class ExternalRows : public zvec::core::VectorSource { + public: + static constexpr uint32_t kDimension = 4; + struct Row { + uint32_t id; + std::array vector; + }; + std::vector rows; + + const void *get_vector(uint32_t id) const override { + for (const auto &row : rows) { + if (row.id == id) return row.vector.data(); + } + return nullptr; + } + + class Cursor : public Iterator { + public: + explicit Cursor(const ExternalRows &source) : source_(source) {} + int next_batch(uint32_t max_count, Batch *out) override { + if (!out || !max_count) return zvec::core::IndexError_InvalidArgument; + out->clear(); + const size_t end = std::min(source_.rows.size(), position_ + max_count); + for (; position_ < end; ++position_) { + const auto &row = source_.rows[position_]; + out->ids.push_back(row.id); + out->vectors.push_back(row.vector.data()); + } + // No lease is needed: this example keeps rows immutable throughout + // every request. Empty output signals end of scan. + return 0; + } + + private: + const ExternalRows &source_; + size_t position_{0}; + }; + + Iterator::Pointer create_iterator() const override { + return std::make_unique(*this); + } +}; + +int main() { + const std::string path = "external_flat_example.index"; + if (std::filesystem::exists(path)) { + std::cerr << "Remove " << path << " before running this example.\n"; + return 1; + } + ExternalRows source; + source.rows = {{42, {1, 0, 0, 0}}, {7, {0, 1, 0, 0}}, {105, {0, 0, 1, 0}}}; + auto param = FlatIndexParamBuilder() + .with_data_type(DataType::DT_FP32) + .with_dimension(ExternalRows::kDimension) + .with_metric_type(MetricType::kL2sq) + .with_use_external_vector(true) + .build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + if (!index || index->open(path, {StorageOptions::StorageType::kMMAP, true})) + return 1; + // Register IDs only; source owns vector data and subsequent updates. + for (const auto &row : source.rows) { + if (index->add_with_source(VectorData{DenseVector{row.vector.data()}}, + row.id, source)) + return 1; + } + VectorData query{DenseVector{source.rows[0].vector.data()}}; + auto qp = + FlatQueryParamBuilder().with_topk(1).with_fetch_vector(true).build(); + SearchResult result; + if (index->search_with_source(query, qp, source, &result) || + result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42 || + result.doc_list_[0].score() != 0) + return 1; + std::cout << "Full scan: id=42, distance=0\n"; + if (index->close()) return 1; + + // Reopen the persisted ID registry and bind the source on the next request. + index = IndexFactory::CreateAndInitIndex(*param); + if (!index || index->open(path, {StorageOptions::StorageType::kMMAP, false})) + return 1; + qp->bf_pks = std::make_shared>( + std::initializer_list{105, 42}); + if (index->search_with_source(query, qp, source, &result) || + result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42) + return 1; + std::cout << "Candidate search after reopen: id=42\n"; + VectorDataBuffer fetched; + if (index->fetch_with_source(42, source, &fetched) || index->close()) + return 1; + std::filesystem::remove(path); + return 0; +} diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index d233f425ed..1075a3d85c 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -109,6 +109,21 @@ int FlatStreamer::init(const IndexMeta &imeta, params.get(PARAM_FLAT_USE_ID_MAP, &use_key_info_map_); params.get(PARAM_FLAT_USE_CONTIGUOUS_MEMORY, &use_contiguous_memory_); + use_external_vector_ = false; + params.get(PARAM_FLAT_USE_EXTERNAL_VECTOR, &use_external_vector_); + if (use_external_vector_) { + if (use_contiguous_memory_ || meta_.major_order() != IndexMeta::MO_ROW || + meta_.data_type() != IndexMeta::DT_FP32 || + meta_.meta_type() != IndexMeta::MT_DENSE || meta_.dimension() == 0 || + meta_.element_size() != meta_.dimension() * sizeof(float) || + (meta_.metric_name() != "SquaredEuclidean" && + meta_.metric_name() != "InnerProduct") || + (quantizer_ && meta_.quantizer_name() != "Fp32Quantizer")) { + return IndexError_Unsupported; + } + use_key_info_map_ = true; + } + state_ = STATE_INITED; return 0; @@ -182,6 +197,7 @@ int FlatStreamer::open(IndexStorage::Pointer stg) { entity_->enable_filter_same_key(true); entity_->set_linear_list_count(1); entity_->set_use_key_info_map(use_key_info_map_); + entity_->set_use_external_vector(use_external_vector_); entity_->set_quantizer(quantizer_); *entity_->mutable_meta() = meta_; @@ -233,6 +249,7 @@ int FlatStreamer::flush(uint64_t checkpoint) { template int FlatStreamer::dump(const IndexDumper::Pointer &dumper) { + if (use_external_vector_) return IndexError_Unsupported; std::string searcher_name = "FlatSearcher"; if constexpr (BATCH_SIZE == 16) { searcher_name = "FlatSearcher16"; @@ -285,6 +302,10 @@ int FlatStreamer::add_impl(uint64_t pkey, const void *query, (*stats_.mutable_discarded_count())++; return IndexError_Cast; } + if (use_external_vector_) { + if (pkey >= kInvalidNodeId) return IndexError_InvalidArgument; + return add_with_id_impl(static_cast(pkey), query, qmeta, context); + } READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); if (!dump_lock.try_lock()) { @@ -334,6 +355,18 @@ int FlatStreamer::add_with_id_impl(uint32_t id, const void *query, (*stats_.mutable_discarded_count())++; return IndexError_Cast; } + if (use_external_vector_) { + if (state_ != STATE_OPENED || !ctx->vector_source()) + return IndexError_NoReady; + if (id == kInvalidNodeId) return IndexError_InvalidArgument; + VectorSource::Batch batch; + int ret = ctx->vector_source()->get_vector_batch(&id, 1, &batch); + if (ret != 0) return ret; + if (batch.ids.size() != 1 || batch.ids[0] != id || + batch.vectors.size() != 1 || !batch.vectors[0]) { + return IndexError_ReadData; + } + } READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); if (!dump_lock.try_lock()) { @@ -357,6 +390,10 @@ int FlatStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, nullptr, context); + } if (state_ != STATE_OPENED) return IndexError_NoReady; ailego_assert(query && count && !!context); ailego_assert(quantizer_ || metric_->is_matched(meta_, qmeta)); @@ -405,6 +442,10 @@ int FlatStreamer::search_bf_by_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, &p_keys, context); + } if (state_ != STATE_OPENED) return IndexError_NoReady; if (count == 0 || count > p_keys.size()) return IndexError_InvalidArgument; ailego_assert(query && count && !!context); @@ -449,10 +490,185 @@ int FlatStreamer::search_bf_by_p_keys_impl( return 0; } +template +int FlatStreamer::search_external( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + const std::vector> *keys, + Context::Pointer &context) const { + auto *ctx = dynamic_cast *>(context.get()); + if (!ctx || !query || !count || (keys && keys->size() < count) || + qmeta.meta_type() != meta_.meta_type() || + qmeta.data_type() != meta_.data_type() || + qmeta.dimension() != meta_.dimension() || + qmeta.element_size() != meta_.element_size()) { + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) ctx->reset(this); + const auto *source = ctx->vector_source(); + if (!source || !ctx->threshold_is_valid()) return IndexError_NoReady; + const bool grouped = ctx->group_by_search(); + if (grouped && !ctx->group_by().is_valid()) return IndexError_InvalidArgument; + ctx->reset_results(count); + ctx->resize_group_results(count); + // Bound both pointer scratch space and per-request pinned vector bytes. + const uint32_t batch_limit = static_cast(std::min( + 1024, + std::max(BATCH_SIZE, read_block_size_ / meta_.element_size()))); + auto *scratch = ctx->search_scratch(); + + for (uint32_t q = 0; q < count; ++q) { + const void *current_query = + static_cast(query) + q * qmeta.element_size(); + auto *heap = ctx->result_heap(); + heap->clear(); + ctx->group_topk_heaps().clear(); + // Destroy the outstanding lease before its cursor, including on errors. + VectorSource::Iterator::Pointer cursor; + VectorSource::Batch batch; + auto consume = [&]() -> int { + if (batch.ids.size() != batch.vectors.size() || + batch.ids.size() > batch_limit) + return IndexError_ReadData; + scratch->vector_ptrs.clear(); + scratch->vector_keys.clear(); + if (!keys) entity_->key_presence(batch.ids, &scratch->key_presence); + for (size_t i = 0; i < batch.ids.size(); ++i) { + const auto id = batch.ids[i]; + if (id == kInvalidNodeId || !batch.vectors[i]) + return IndexError_ReadData; + if (!keys && (!scratch->key_presence[i] || + (ctx->filter().is_valid() && ctx->filter()(id)))) { + ++*ctx->mutable_stats(q)->mutable_filtered_count(); + continue; + } + scratch->vector_ptrs.push_back(batch.vectors[i]); + scratch->vector_keys.push_back(id); + } + const size_t size = scratch->vector_ptrs.size(); + if (!size) return 0; + scratch->distances.resize(size); + if (quantizer_) { + quantizer_->calc_distance_dp_query_batch( + scratch->vector_ptrs.data(), static_cast(size), current_query, + scratch->distances.data()); + } else if (auto distance = metric_->batch_distance()) { + distance(scratch->vector_ptrs.data(), current_query, size, + meta_.dimension(), scratch->distances.data(), nullptr); + } else { + auto scalar_distance = metric_->distance(); + for (size_t i = 0; i < size; ++i) { + scalar_distance(current_query, scratch->vector_ptrs[i], + meta_.dimension(), &scratch->distances[i]); + } + } + *ctx->mutable_stats(q)->mutable_dist_calced_count() += size; + for (size_t i = 0; i < size; ++i) { + const auto id = scratch->vector_keys[i]; + const auto distance = scratch->distances[i]; + if (grouped) { + auto &group_heap = ctx->group_topk_heaps()[ctx->group_by()(id)]; + if (group_heap.empty()) group_heap.limit(ctx->group_topk()); + group_heap.emplace(id, distance); + } else { + heap->emplace(id, distance); + } + } + return 0; + }; + + if (keys) { + std::vector ids; + ids.reserve(batch_limit); + auto read_candidates = [&]() -> int { + if (ids.empty()) return 0; + batch.clear(); + int ret = source->get_vector_batch(ids.data(), ids.size(), &batch); + if (ret != 0) return ret; + if (batch.ids != ids) return IndexError_ReadData; + ret = consume(); + ids.clear(); + return ret; + }; + for (auto key : (*keys)[q]) { + if (key >= kInvalidNodeId || !entity_->contains_key(key) || + (ctx->filter().is_valid() && ctx->filter()(key))) + continue; + ids.push_back(static_cast(key)); + if (ids.size() == batch_limit) { + int ret = read_candidates(); + if (ret != 0) return ret; + } + } + int ret = read_candidates(); + if (ret != 0) return ret; + } else { + cursor = source->create_iterator(); + if (!cursor) return IndexError_Unsupported; + while (true) { + batch.clear(); + int ret = cursor->next_batch(batch_limit, &batch); + if (ret != 0) return ret; + ret = consume(); + if (ret != 0) return ret; + if (batch.ids.empty()) break; + } + } + batch.clear(); + if (grouped) + ctx->topk_to_group_result(q); + else + ctx->take_topk_result(q); + + // Results outlive the source snapshot, cursor and index. Own the selected + // raw vectors instead of retaining borrowed pointers or disk page pins. + auto fetch_results = [&](IndexDocumentList *docs) -> int { + for (size_t begin = 0; begin < docs->size(); begin += batch_limit) { + const size_t end = std::min(docs->size(), begin + batch_limit); + std::vector ids; + ids.reserve(end - begin); + for (size_t i = begin; i < end; ++i) ids.push_back((*docs)[i].key()); + batch.clear(); + int ret = source->get_vector_batch(ids.data(), ids.size(), &batch); + if (ret != 0) return ret; + if (batch.ids != ids || batch.vectors.size() != ids.size()) { + return IndexError_ReadData; + } + for (size_t i = 0; i < ids.size(); ++i) { + if (!batch.vectors[i]) return IndexError_ReadData; + auto owned = std::make_shared( + static_cast(batch.vectors[i]), + meta_.element_size()); + auto block = + IndexStorage::MemoryBlock::MakeSharedView(owned->data(), owned); + auto &doc = (*docs)[begin + i]; + doc = IndexDocument(doc.key(), doc.score(), doc.index(), block); + } + } + return 0; + }; + if (ctx->fetch_vector()) { + if (grouped) { + for (auto &group : *ctx->mutable_group_result(q)) { + int ret = fetch_results(group.mutable_docs()); + if (ret != 0) return ret; + } + } else { + int ret = fetch_results(ctx->mutable_result(q)); + if (ret != 0) return ret; + } + } + } + return 0; +} + template int FlatStreamer::group_by_search_impl( const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, nullptr, context); + } FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); if (!bf_context) { @@ -501,6 +717,10 @@ int FlatStreamer::group_by_search_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, &p_keys, context); + } if (count == 0 || count > p_keys.size()) return IndexError_InvalidArgument; FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); diff --git a/src/core/algorithm/flat/flat_streamer.h b/src/core/algorithm/flat/flat_streamer.h index 26823172d7..38c5c480e1 100644 --- a/src/core/algorithm/flat/flat_streamer.h +++ b/src/core/algorithm/flat/flat_streamer.h @@ -17,6 +17,7 @@ #include #include #include +#include #include "flat_streamer_entity.h" #include "flat_utility.h" @@ -131,6 +132,10 @@ class FlatStreamer : public IndexStreamer { return *entity_; } + bool use_external_vector() const { + return use_external_vector_; + } + //! Retrieve the turbo quantizer const std::shared_ptr &quantizer() const { return quantizer_; @@ -176,6 +181,11 @@ class FlatStreamer : public IndexStreamer { } private: + int search_external(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + const std::vector> *keys, + Context::Pointer &context) const; + //! Constants static constexpr uint32_t kDefaultBlockVecCount = 32u; static constexpr uint32_t kDefaultSegmentSize = 4 * 1024 * 1024u; @@ -198,6 +208,7 @@ class FlatStreamer : public IndexStreamer { bool use_key_info_map_{true}; uint32_t read_block_size_{0}; bool use_contiguous_memory_{false}; + bool use_external_vector_{false}; std::shared_ptr quantizer_{}; std::unique_ptr entity_{}; }; diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h index 3080448718..3eb4691900 100644 --- a/src/core/algorithm/flat/flat_streamer_context.h +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -128,7 +128,7 @@ class FlatStreamerContext : public IndexStreamer::Context { } key_t key = result_heap_[i].key(); - if (fetch_vector_) { + if (fetch_vector_ && !owner_->use_external_vector()) { IndexStorage::MemoryBlock block; owner_->entity().get_vector_by_key(key, block); results_[idx].emplace_back(key, score, key, block); @@ -200,8 +200,8 @@ class FlatStreamerContext : public IndexStreamer::Context { break; } node_id_t id = group_topk_list[i].second[j].first; - auto provider = owner_->create_provider(); - if (fetch_vector_) { + if (fetch_vector_ && !owner_->use_external_vector()) { + auto provider = owner_->create_provider(); IndexStorage::MemoryBlock block; provider->get_vector(id, block); group_results_[idx][i].mutable_docs()->emplace_back(id, score, id, @@ -224,7 +224,16 @@ class FlatStreamerContext : public IndexStreamer::Context { group_topk_heaps_.clear(); } + void set_vector_source(const VectorSource *source) { + vector_source_ = source; + } + + const VectorSource *vector_source() const { + return vector_source_; + } + void reset() override { + vector_source_ = nullptr; for (auto &it : results_) { it.clear(); } @@ -235,7 +244,9 @@ class FlatStreamerContext : public IndexStreamer::Context { //! Reset the context void reset(const FlatStreamer *owner) { + const auto *source = vector_source_; this->reset(); + vector_source_ = source; magic_ = owner->magic(); feature_size_ = owner->meta().element_size(); @@ -270,6 +281,7 @@ class FlatStreamerContext : public IndexStreamer::Context { private: const FlatStreamer *owner_{nullptr}; + const VectorSource *vector_source_{nullptr}; FlatSearchScratch search_scratch_{}; std::vector stats_vec_{}; uint32_t magic_{0}; diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc index 91f1e0ef6e..ac9139a5fc 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -155,7 +155,7 @@ int FlatStreamerEntity::open(IndexStorage::Pointer storage, vec_cols_ = index_meta_.element_size() / vec_unit_size_; meta_.header.block_size = ailego_align(sizeof(BlockHeader) + sizeof(DeletionMap) + - (index_meta_.element_size() + sizeof(uint64_t)) * + (stored_vector_size() + sizeof(uint64_t)) * meta_.header.block_vector_count, 32); @@ -348,6 +348,7 @@ int FlatStreamerEntity::search(const void *query, const IndexFilter &filter, IndexContext::Stats *context_stats, FlatSearchScratch * /*scratch*/, size_t /*batch_size*/) const { + if (use_external_vector_) return IndexError_Unsupported; IndexStorage::MemoryBlock head_block; this->get_head_block(head_block); const BlockLocation *bl = @@ -394,6 +395,7 @@ int FlatStreamerEntity::search_by_p_keys(const void *query, IndexDocumentHeap *heap, FlatSearchScratch * /*scratch*/, size_t /*batch_size*/) const { + if (use_external_vector_) return IndexError_Unsupported; for (uint64_t key : p_keys) { if (filter.is_valid() && filter(key)) { continue; @@ -498,6 +500,7 @@ FlatStreamerEntity::Pointer FlatStreamerEntity::clone() const { LOG_ERROR("Failed to New FlatStreamerEntity object"); return nullptr; } + entity->use_external_vector_ = this->use_external_vector_; entity->index_meta_ = this->index_meta_; entity->storage_ = this->storage_; // entity->reformer_ = this->reformer_; @@ -716,6 +719,7 @@ const void *FlatContiguousStreamerEntity::get_vector_ptr( } const void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const { + if (use_external_vector_) return nullptr; VectorLocation loc{}; key_info_map_lock_->lock_shared(); if (use_key_info_map_) { @@ -747,6 +751,7 @@ const void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const { int FlatStreamerEntity::get_vector_by_key( const uint64_t key, IndexStorage::MemoryBlock &block) const { + if (use_external_vector_) return IndexError_Unsupported; VectorLocation loc{}; key_info_map_lock_->lock_shared(); if (use_key_info_map_) { @@ -778,6 +783,7 @@ int FlatStreamerEntity::get_vector_by_key( int FlatStreamerEntity::get_vectors_by_key( const uint64_t *keys, uint32_t count, std::vector &blocks) const { + if (use_external_vector_) return IndexError_Unsupported; std::vector locations(count); key_info_map_lock_->lock_shared(); if (use_key_info_map_) { @@ -815,6 +821,8 @@ int FlatStreamerEntity::get_vectors_by_key( } IndexProvider::Iterator::Pointer FlatStreamerEntity::creater_iterator() const { + // An external scan requires a request-bound source. + if (use_external_vector_) return nullptr; auto entity = this->clone(); if (!entity) { LOG_ERROR("Failed to clone entity"); @@ -961,6 +969,10 @@ int FlatStreamerEntity::load_linear_meta(IndexStorage::Pointer storage) { return IndexError_InvalidFormat; } auto *mt = reinterpret_cast(data_block.data()); + if (mt->header.reserved_[0] != meta_.header.reserved_[0]) { + LOG_ERROR("Unmatched Flat external-vector storage mode"); + return IndexError_Mismatch; + } if (mt->header.block_vector_count != meta_.header.block_vector_count) { LOG_ERROR("Unmatched BlockVecCount Setting, Index %u vs Setting %u", mt->header.block_vector_count, meta_.header.block_vector_count); @@ -1293,7 +1305,7 @@ int FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key, size_t vector_off = get_block_vector_offset(block.block_index, header->vector_count); - if (segment->write(vector_off, data, size) != size) { + if (!use_external_vector_ && segment->write(vector_off, data, size) != size) { LOG_ERROR("Failed to write vector, off=%zu size=%zu", vector_off, size); return IndexError_WriteData; } @@ -1327,6 +1339,11 @@ int FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key, int FlatStreamerEntity::add_vector_with_id(const uint32_t id, const void *query, const uint32_t size) { + if (use_external_vector_) { + // Keep sparse external IDs compact. Updates belong to the source. + int ret = add(id, query, size); + return ret == IndexError_Duplicate ? 0 : ret; + } std::lock_guard lock(mutex_); // if (filter_same_key_) { // key_info_map_lock_->lock_shared(); diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index fb4c8c666e..d0bed8e8a8 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -39,6 +39,7 @@ struct FlatSearchScratch { std::vector vector_keys{}; std::vector distances{}; std::vector query_buffer{}; + std::vector key_presence{}; }; /*! Flat Streamer Entity @@ -157,6 +158,28 @@ class FlatStreamerEntity { LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); } + void set_use_external_vector(bool enabled) { + use_external_vector_ = enabled; + // Persist the physical layout independently of the logical vector meta. + meta_.header.reserved_[0] = enabled ? 1 : 0; + } + + bool contains_key(uint64_t key) const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + return key_info_map_.find(key) != key_info_map_.end(); + } + + void key_presence(const std::vector &keys, + std::vector *present) const { + present->resize(keys.size()); + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + for (size_t i = 0; i < keys.size(); ++i) { + (*present)[i] = key_info_map_.find(keys[i]) != key_info_map_.end(); + } + } + //! Set params void set_segment_size(uint32_t size) { meta_.segment_size = size; @@ -370,6 +393,10 @@ class FlatStreamerEntity { size_t size); private: + size_t stored_vector_size() const { + return use_external_vector_ ? 0 : index_meta_.element_size(); + } + size_t get_block_offset(uint32_t block_index) const { return block_index * linear_block_size(); } @@ -386,14 +413,14 @@ class FlatStreamerEntity { size_t get_block_key_offset(uint32_t block_index, uint32_t vector_index) const { return get_block_offset(block_index) + - block_vector_count() * index_meta_.element_size() + + block_vector_count() * stored_vector_size() + sizeof(uint64_t) * vector_index; } size_t get_block_vector_offset(uint32_t block_index, uint32_t vector_index) const { return this->get_block_offset(block_index) + - vector_index * index_meta_.element_size(); + vector_index * stored_vector_size(); } //! Get header block of an linear list @@ -477,6 +504,7 @@ class FlatStreamerEntity { // Open/close and initial loading require external lifecycle exclusion. mutable std::mutex segments_mutex_{}; IndexMeta index_meta_{}; + bool use_external_vector_{false}; IndexStorage::Pointer storage_{}; IndexMetric::MatrixDistance row_distance_{}, column_distance_{}; IndexMetric::MatrixBatchDistance batch_distance_{}; diff --git a/src/core/algorithm/flat/flat_utility.h b/src/core/algorithm/flat/flat_utility.h index 4020c8c9a5..31b8c0869d 100644 --- a/src/core/algorithm/flat/flat_utility.h +++ b/src/core/algorithm/flat/flat_utility.h @@ -34,6 +34,8 @@ static const std::string FLAT_SEGMENT_FEATURES_SEG_ID("flat.features"); static const std::string FLAT_SEGMENT_MAPPING_SEG_ID("flat.mapping"); // index params +static const std::string PARAM_FLAT_USE_EXTERNAL_VECTOR( + "proxima.flat.use_external_vector"); static const std::string PARAM_FLAT_COLUMN_MAJOR_ORDER( "proxima.flat.column_major_order"); static const std::string PARAM_FLAT_BATCH_SIZE("proxima.flat.batch_size"); diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 1932d40135..d78b955235 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -685,6 +685,11 @@ bool Index::is_dirty() const { return storage_->is_dirty(); } +int Index::fetch_with_source(uint32_t, const core::VectorSource &, + VectorDataBuffer *) { + return core::IndexError_Unsupported; +} + int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { if (!is_open_) { LOG_ERROR("Index is not open"); @@ -1095,6 +1100,14 @@ int Index::_refine_dense_candidates( auto &context = acquire_context(); if (!context) return core::IndexError_Runtime; + if (param_.use_external_vector) { + int ret = bind_vector_source( + search_param->refiner_param->reference_vector_source, context); + if (ret != 0) { + context->reset(); + return ret; + } + } context->set_topk(search_param->topk); context->set_fetch_vector(search_param->fetch_vector); if (search_param->filter && search_param->filter->is_valid()) { @@ -1415,6 +1428,13 @@ int Index::_sparse_search(const VectorData &vector_data, int Index::merge(const std::vector &indexes, const IndexFilter &filter, const MergeOptions &options) { + // Merge has no request-bound source from which to materialize vectors. + if (param_.use_external_vector) return core::IndexError_Unsupported; + for (const auto &index : indexes) { + if (index && index->param_.use_external_vector) { + return core::IndexError_Unsupported; + } + } if (indexes.empty()) { return core::IndexError_Success; } diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index 1fed9d4333..dee682c8ee 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -18,12 +18,20 @@ #include #include #include "algorithm/flat/flat_index_format.h" +#include "algorithm/flat/flat_streamer_context.h" #include "algorithm/flat/flat_utility.h" namespace zvec::core_interface { namespace { +struct FlatSourceScope { + core::IndexContext::Pointer &context; + ~FlatSourceScope() { + if (context) context->reset(); + } +}; + //! Read the IndexMeta persisted in the flat linear meta segment without //! initializing a streamer. Returns non-zero when the storage or segment //! cannot be read (e.g. the index file does not exist yet). @@ -134,6 +142,64 @@ ailego::Params MakeTurboQuantizerParams(const QuantizerParam &quantizer_param, } // namespace +int FlatIndex::bind_vector_source(const core::VectorSource *source, + core::IndexContext::Pointer &context) { + if (!param_.use_external_vector) return core::IndexError_Unsupported; + if (!source) return core::IndexError_InvalidArgument; + auto *ctx = dynamic_cast *>(context.get()); + auto *streamer = dynamic_cast *>(streamer_.get()); + if (!ctx || !streamer) return core::IndexError_Cast; + if (ctx->magic() != streamer->magic()) ctx->reset(streamer); + ctx->set_vector_source(source); + return 0; +} + +int FlatIndex::add_with_source(const VectorData &vector, uint32_t doc_id, + const core::VectorSource &source) { + if (!is_open_) return core::IndexError_NoReady; + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + FlatSourceScope scope{context}; + int ret = bind_vector_source(&source, context); + return ret == 0 ? Index::add(vector, doc_id) : ret; +} + +int FlatIndex::search_with_source( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + const core::VectorSource &source, SearchResult *result) { + if (!is_open_) return core::IndexError_NoReady; + if (!search_param || !result) return core::IndexError_InvalidArgument; + *result = SearchResult{}; + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + FlatSourceScope scope{context}; + int ret = bind_vector_source(&source, context); + return ret == 0 ? Index::search(query, search_param, result) : ret; +} + +int FlatIndex::fetch_with_source(uint32_t doc_id, + const core::VectorSource &source, + VectorDataBuffer *result) { + if (!is_open_) return core::IndexError_NoReady; + if (!param_.use_external_vector) return core::IndexError_Unsupported; + if (!result) return core::IndexError_InvalidArgument; + auto *streamer = dynamic_cast *>(streamer_.get()); + if (!streamer) return core::IndexError_Cast; + if (!streamer->entity().contains_key(doc_id)) return core::IndexError_NoExist; + core::VectorSource::Batch batch; + int ret = source.get_vector_batch(&doc_id, 1, &batch); + if (ret != 0) return ret; + if (batch.ids.size() != 1 || batch.ids[0] != doc_id || + batch.vectors.size() != 1 || !batch.vectors[0]) { + return core::IndexError_ReadData; + } + DenseVectorBuffer buffer; + buffer.data.assign(static_cast(batch.vectors[0]), + input_vector_meta_.element_size()); + result->vector_buffer = std::move(buffer); + return 0; +} + int FlatIndex::open(const std::string &file_path, StorageOptions storage_options) { // Restore the persisted encoding while keeping the configured structural @@ -220,6 +286,21 @@ int FlatIndex::fallback_to_legacy_pipeline() { int FlatIndex::create_and_init_converter_reformer( const QuantizerParam &quantizer_param, const BaseIndexParam &index_param) { const auto &flat_param = dynamic_cast(index_param); + if (flat_param.use_external_vector && + (flat_param.is_sparse || flat_param.data_type != DataType::DT_FP32 || + flat_param.dimension <= 0 || flat_param.dimension > MAX_DIMENSION || + (flat_param.metric_type != MetricType::kL2sq && + flat_param.metric_type != MetricType::kInnerProduct) || + flat_param.major_order == IndexMeta::MO_COLUMN || + flat_param.use_contiguous_memory || quantizer_param.enable_rotate || + quantizer_param.type != QuantizerType::kNone || + flat_param.preprocess_param.type != PreprocessorType::kNone || + (flat_param.storage_data_type != DataType::DT_UNDEFINED && + flat_param.storage_data_type != DataType::DT_FP32))) { + LOG_ERROR("External Flat requires raw row-major FP32 L2/IP vectors"); + return core::IndexError_Unsupported; + } + // Prefer the turbo quantizer path (quantized records + SIMD batch distance // kernels funneled through the streamer entity) whenever a turbo quantizer // can express the configuration; only the remaining combinations fall @@ -299,6 +380,8 @@ int FlatIndex::create_and_init_streamer(const BaseIndexParam ¶m) { proxima_index_params_.set(core::PARAM_FLAT_COLUMN_MAJOR_ORDER, param_.major_order == IndexMeta::MO_COLUMN); proxima_index_params_.set(core::PARAM_FLAT_USE_ID_MAP, param_.use_id_map); + proxima_index_params_.set(core::PARAM_FLAT_USE_EXTERNAL_VECTOR, + param_.use_external_vector); proxima_index_params_.set(core::PARAM_FLAT_USE_CONTIGUOUS_MEMORY, param_.use_contiguous_memory); if (is_sparse_) { diff --git a/src/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index ea19016e11..fa218c7874 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include namespace zvec { @@ -21,6 +22,27 @@ VectorSource::VectorSource() = default; VectorSource::~VectorSource() = default; +VectorSource::Iterator::Pointer VectorSource::create_iterator() const { + return nullptr; +} + +int VectorSource::get_vector_batch(const uint32_t *ids, uint32_t count, + Batch *out) const { + if (!out || (count && !ids)) return IndexError_InvalidArgument; + out->clear(); + if (!count) return 0; + out->ids.assign(ids, ids + count); + out->vectors.resize(count); + get_vectors(ids, count, out->vectors.data()); + for (auto *vector : out->vectors) { + if (!vector) { + out->clear(); + return IndexError_NoExist; + } + } + return 0; +} + void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, const void **out) const { for (uint32_t i = 0; i < count; ++i) { diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index db014b0b97..71ed3cd784 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -148,6 +148,9 @@ class ZVEC_CORE_API Index { const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, const core::VectorSource &src, SearchResult *result); + virtual int fetch_with_source(uint32_t doc_id, const core::VectorSource &src, + VectorDataBuffer *result); + virtual BaseIndexParam::Pointer get_param() const; virtual bool is_trained() const; @@ -168,6 +171,11 @@ class ZVEC_CORE_API Index { } protected: + virtual int bind_vector_source(const core::VectorSource *, + core::IndexContext::Pointer &) { + return core::IndexError_Unsupported; + } + int _sparse_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer); virtual int _dense_fetch(const uint32_t doc_id, @@ -269,6 +277,21 @@ class ZVEC_CORE_API Index { class ZVEC_CORE_API FlatIndex : public Index { public: FlatIndex() = default; + + // External Flat supports raw row-major FP32 L2/IP. Register IDs after + // writing the source; vector bytes and subsequent updates stay in the source. + // doc_id is the source ID with either use_id_map setting. + int add_with_source(const VectorData &vector, uint32_t doc_id, + const core::VectorSource &src) override; + // Full scans enumerate every registered ID once (extra IDs are ignored). + // Candidate searches/refinement use random batches. Hold one stable source + // snapshot across enumeration and random reads for the entire request. + int search_with_source(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + const core::VectorSource &src, + SearchResult *result) override; + int fetch_with_source(uint32_t doc_id, const core::VectorSource &src, + VectorDataBuffer *result) override; // FlatIndex(const FlatIndexParam ¶m) : param_(param) {} // FlatIndex(FlatIndexParam &¶m) : param(std::move(param)) {} @@ -304,6 +327,9 @@ class ZVEC_CORE_API FlatIndex : public Index { int create_and_init_legacy_converter_reformer( const QuantizerParam ¶m, const BaseIndexParam &index_param); + int bind_vector_source(const core::VectorSource *source, + core::IndexContext::Pointer &context) override; + FlatIndexParam param_{}; }; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 2f6146689a..4e7937b00a 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -24,6 +24,7 @@ #include #include #include +#include #include #include "zvec/core/framework/index_framework.h" @@ -204,6 +205,8 @@ struct RefinerParam { float scale_factor_{0}; std::shared_ptr reference_index = nullptr; + // Borrowed for this request; independent of the coarse index source. + const core::VectorSource *reference_vector_source = nullptr; }; // --- GroupBy Parameters --- diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 66b4da5dd2..2ac175a3a5 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -15,6 +15,8 @@ #pragma once #include +#include +#include #include namespace zvec { @@ -22,6 +24,31 @@ namespace core { class ZVEC_CORE_API VectorSource { public: + // Logical row-major vectors in the same ID namespace as get_vector(). + // Pointers need not be contiguous. External Flat currently expects raw FP32. + struct Batch { + std::vector ids; + std::vector vectors; + // Optional backing allocation or page pin. Data must remain valid until + // the batch is released, or (for scans) the next next_batch() call. + std::shared_ptr lease; + + void clear() { + vectors.clear(); + ids.clear(); + lease.reset(); + } + }; + + class Iterator { + public: + using Pointer = std::unique_ptr; + virtual ~Iterator() = default; + // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an + // IndexError on failure. Enumerate every visible ID exactly once. + virtual int next_batch(uint32_t max_count, Batch *out) = 0; + }; + VectorSource(); virtual ~VectorSource(); @@ -29,6 +56,17 @@ class ZVEC_CORE_API VectorSource { virtual void get_vectors(const uint32_t *ids, uint32_t count, const void **out) const; + + // Optional full scan. Each call creates an independent cursor. The caller + // holds a stable source snapshot for the whole request, including subsequent + // random reads. A disk source can scan physical blocks behind this interface. + virtual Iterator::Pointer create_iterator() const; + + // Return IDs in input order. The default borrows get_vectors() pointers, + // which must all remain valid until the caller consumes the batch. Sources + // with transient buffers should override this method and attach a lease. + virtual int get_vector_batch(const uint32_t *ids, uint32_t count, + Batch *out) const; }; } // namespace core diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc new file mode 100644 index 0000000000..46a891336c --- /dev/null +++ b/tests/core/interface/external_flat_index_test.cc @@ -0,0 +1,571 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/flat/flat_streamer_context.h" +#include "tests/test_util.h" + +using namespace zvec::core_interface; +namespace core = zvec::core; + +namespace { +constexpr uint32_t kDimension = 8; + +// Each batch owns a temporary page. Release must precede advancing the cursor. +// Reverse scan order and noncontiguous IDs catch assumptions about row offsets. +class BatchSource : public core::VectorSource { + public: + std::map> rows; + mutable std::atomic scans{0}, reads{0}; + size_t page_rows{7}; + bool supports_scan{true}, fail_scan{false}, malformed_scan{false}; + bool null_scan{false}, fail_read{false}, wrong_read_ids{false}; + bool legacy_reads{false}; + + explicit BatchSource(uint32_t count = 89, float offset = 0) { + for (uint32_t i = 0; i < count; ++i) { + auto &row = rows[i * 13 + 3]; + for (uint32_t d = 0; d < kDimension; ++d) { + row.push_back(((i * 19 + d * 7) % 101) / 101.0f + offset); + } + } + } + + const void *get_vector(uint32_t id) const override { + auto it = rows.find(id); + return it == rows.end() ? nullptr : it->second.data(); + } + + int get_vector_batch(const uint32_t *ids, uint32_t count, + Batch *out) const override { + ++reads; + if (legacy_reads) return VectorSource::get_vector_batch(ids, count, out); + if (fail_read) return core::IndexError_ReadData; + out->clear(); + auto page = std::make_shared>(count * kDimension); + for (uint32_t i = 0; i < count; ++i) { + const void *data = get_vector(ids[i]); + if (!data) return core::IndexError_NoExist; + std::memcpy(page->data() + i * kDimension, data, + kDimension * sizeof(float)); + out->ids.push_back(ids[i] + (wrong_read_ids ? 1 : 0)); + out->vectors.push_back(page->data() + i * kDimension); + } + out->lease = page; + return 0; + } + + class Cursor : public Iterator { + public: + explicit Cursor(const BatchSource &source) : source_(source) { + for (auto it = source.rows.rbegin(); it != source.rows.rend(); ++it) { + ids_.push_back(it->first); + } + } + int next_batch(uint32_t max_count, Batch *out) override { + EXPECT_TRUE(previous_.expired()); + out->clear(); + if (source_.fail_scan && position_) return core::IndexError_ReadData; + const size_t count = std::min( + {size_t{max_count}, source_.page_rows, ids_.size() - position_}); + if (!count) return 0; + auto page = std::make_shared>(count * kDimension); + for (size_t i = 0; i < count; ++i) { + auto id = ids_[position_++]; + std::memcpy(page->data() + i * kDimension, source_.get_vector(id), + kDimension * sizeof(float)); + out->ids.push_back(id); + out->vectors.push_back( + source_.null_scan ? nullptr : page->data() + i * kDimension); + } + out->lease = page; + previous_ = page; + if (source_.malformed_scan) out->vectors.pop_back(); + return 0; + } + + private: + const BatchSource &source_; + std::vector ids_; + size_t position_{0}; + std::weak_ptr previous_; + }; + + Iterator::Pointer create_iterator() const override { + ++scans; + return supports_scan ? std::make_unique(*this) : nullptr; + } +}; + +VectorData Vector(const std::vector &row) { + return VectorData{DenseVector{row.data()}}; +} + +FlatIndexParam::Pointer Param(bool external, + MetricType metric = MetricType::kL2sq, + bool id_map = true) { + return FlatIndexParamBuilder() + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_metric_type(metric) + .with_use_external_vector(external) + .with_use_id_map(id_map) + .build(); +} + +void SameResults(const SearchResult &want, const SearchResult &got) { + ASSERT_EQ(want.doc_list_.size(), got.doc_list_.size()); + for (size_t i = 0; i < want.doc_list_.size(); ++i) { + EXPECT_EQ(want.doc_list_[i].key(), got.doc_list_[i].key()); + EXPECT_NEAR(want.doc_list_[i].score(), got.doc_list_[i].score(), 1e-5); + } +} + +class ExternalFlatTest : public testing::Test { + protected: + std::vector opened_; + std::vector paths_; + + std::string Path(const std::string &suffix) { + std::string path = + "test_external_flat_" + + std::string( + testing::UnitTest::GetInstance()->current_test_info()->name()) + + suffix; + paths_.push_back(path); + zvec::test_util::RemoveTestFiles(path); + return path; + } + + Index::Pointer Open(const BaseIndexParam::Pointer ¶m, + const std::string &path, bool create = true) { + auto index = IndexFactory::CreateAndInitIndex(*param); + EXPECT_NE(nullptr, index); + if (!index) return index; + int ret = index->open(path, {StorageOptions::StorageType::kMMAP, create}); + EXPECT_EQ(0, ret); + if (!ret) opened_.push_back(index); + return index; + } + + void Add(Index *index, const BatchSource &source, bool external = true) { + for (auto &[id, row] : source.rows) { + ASSERT_EQ(0, external ? index->add_with_source(Vector(row), id, source) + : index->add(Vector(row), id)); + } + } + + void Close(const Index::Pointer &index) { + EXPECT_EQ(0, index->close()); + opened_.erase(std::remove(opened_.begin(), opened_.end(), index), + opened_.end()); + } + + void TearDown() override { + for (auto &index : opened_) index->close(); + for (auto &path : paths_) zvec::test_util::RemoveTestFiles(path); + } +}; + +TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { + int case_id = 0; + for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { + for (bool id_map : {false, true}) { + BatchSource source; + auto suffix = std::to_string(case_id++); + auto external = Open(Param(true, metric, id_map), Path(suffix + "ext")); + auto embedded = Open(Param(false, metric, id_map), Path(suffix + "base")); + Add(external.get(), source); + Add(embedded.get(), source, false); + ASSERT_EQ(source.rows.size(), external->get_doc_count()); + source.rows[99999] = std::vector(kDimension, 0.3f); + std::vector query(kDimension, 0.3f); + auto qp = + FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(1, source.scans); + auto expected = source.rows; + source.rows.clear(); + Close(external); + ASSERT_EQ(11, got.doc_list_.size()); + // External vectors are returned directly in the input layout, without + // creating a redundant dequantized copy in reverted_vector_list_. + EXPECT_TRUE(got.reverted_vector_list_.empty()); + for (size_t i = 0; i < got.doc_list_.size(); ++i) { + const auto &doc = got.doc_list_[i]; + const auto &row = expected.at(doc.key()); + EXPECT_EQ(0, std::memcmp(row.data(), doc.vector(), + kDimension * sizeof(float))); + } + } + } +} + +TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { + BatchSource source; + source.supports_scan = false; + auto index = Open(Param(true), Path("index")); + Add(index.get(), source); + auto qp = + FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); + qp->radius = 0.5f; + qp->filter = std::make_shared(); + qp->filter->set([](uint64_t id) { return id == 3; }); + qp->bf_pks = std::make_shared>(); + for (auto &[id, row] : source.rows) qp->bf_pks->push_back(id); + qp->bf_pks->push_back(999999); + qp->bf_pks->push_back(1ull << 40); + std::vector query(kDimension, 0.3f); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_EQ(0, source.scans); + ASSERT_FALSE(result.doc_list_.empty()); + for (auto &doc : result.doc_list_) { + EXPECT_NE(3, doc.key()); + EXPECT_LE(doc.score(), qp->radius); + VectorDataBuffer buffer; + ASSERT_EQ(0, index->fetch_with_source(doc.key(), source, &buffer)); + EXPECT_EQ(0, + std::memcmp( + source.get_vector(doc.key()), + std::get(buffer.vector_buffer).data.data(), + kDimension * sizeof(float))); + } + VectorDataBuffer buffer; + EXPECT_NE(0, index->fetch_with_source(999999, source, &buffer)); + EXPECT_NE(0, index->fetch(3, &buffer)); + EXPECT_NE(0, index->add(Vector(query), 3)); + EXPECT_NE(0, index->search(Vector(query), qp, &result)); +} + +TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { + BatchSource source; + auto path = Path("index"); + auto index = Open(Param(true, MetricType::kL2sq, false), path); + Add(index.get(), source); + auto *streamer = + dynamic_cast *>(index->index_searcher().get()); + ASSERT_NE(nullptr, streamer); + EXPECT_EQ(ailego_align(sizeof(core::BlockHeader) + sizeof(core::DeletionMap) + + 32 * sizeof(uint64_t), + 32), + streamer->entity().linear_block_size()); + source.rows[3] = std::vector(kDimension, 12.0f); + ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(3)), 3, source)); + EXPECT_EQ(89, index->get_doc_count()); + ASSERT_EQ(0, index->flush()); + Close(index); + index = Open(Param(true, MetricType::kL2sq, false), path, false); + ASSERT_EQ(89, index->get_doc_count()); + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(source.rows.at(3)), qp, source, + &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); + source.rows[1000000] = std::vector(kDimension, 24.0f); + ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(1000000)), 1000000, + source)); + EXPECT_EQ(90, index->get_doc_count()); + Close(index); + auto wrong = IndexFactory::CreateAndInitIndex(*Param(false)); + ASSERT_NE(nullptr, wrong); + EXPECT_NE(0, wrong->open(path, {StorageOptions::StorageType::kMMAP, false})); + wrong.reset(); + auto embedded_path = Path("embedded"); + auto embedded = Open(Param(false), embedded_path); + Close(embedded); + wrong = IndexFactory::CreateAndInitIndex(*Param(true)); + EXPECT_NE(0, wrong->open(embedded_path, + {StorageOptions::StorageType::kMMAP, false})); +} + +TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { + BatchSource source; + auto index = Open(Param(true), Path("index")); + Add(index.get(), source); + auto qp = FlatQueryParamBuilder().with_topk(3).build(); + std::vector query(kDimension, 0.3f); + SearchResult result; + for (int failure = 0; failure < 4; ++failure) { + source.fail_scan = failure == 0; + source.malformed_scan = failure == 1; + source.null_scan = failure == 2; + source.supports_scan = failure != 3; + result.doc_list_.emplace_back(999, 0); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + EXPECT_NE(0, index->search(Vector(query), qp, &result)); + } + source.fail_scan = source.malformed_scan = source.null_scan = false; + source.supports_scan = true; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + qp->bf_pks = std::make_shared>( + std::initializer_list{3}); + source.fail_read = true; + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + source.fail_read = false; + source.wrong_read_ids = true; + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + source.wrong_read_ids = false; + source.rows.erase(3); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_NE(0, index->add_with_source(Vector(query), 3, source)); +} + +TEST_F(ExternalFlatTest, GroupByUsesExternalVectors) { + BatchSource source; + auto index = Open(Param(true), Path("ext")); + auto embedded = Open(Param(false), Path("base")); + Add(index.get(), source); + Add(embedded.get(), source, false); + auto qp = + FlatQueryParamBuilder().with_topk(6).with_fetch_vector(true).build(); + qp->group_by_param = std::make_shared(); + qp->group_by_param->group_count = 3; + qp->group_by_param->group_topk = 2; + qp->group_by_param->group_by = [](uint64_t id) { + return std::to_string(id % 3); + }; + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &got)); + ASSERT_EQ(3, got.group_doc_list_.size()); + for (size_t i = 0; i < want.group_doc_list_.size(); ++i) { + const auto &a = want.group_doc_list_[i]; + const auto &b = got.group_doc_list_[i]; + EXPECT_EQ(a.group_id(), b.group_id()); + ASSERT_EQ(2, b.docs().size()); + for (size_t j = 0; j < b.docs().size(); ++j) { + EXPECT_EQ(a.docs()[j].key(), b.docs()[j].key()); + EXPECT_NEAR(a.docs()[j].score(), b.docs()[j].score(), 1e-5); + EXPECT_EQ(0, + std::memcmp(source.get_vector(b.docs()[j].key()), + b.docs()[j].vector(), kDimension * sizeof(float))); + } + } +} + +TEST_F(ExternalFlatTest, RefineBindsItsOwnSource) { + BatchSource coarse_source(30), reference_source(30, 0.8f); + auto coarse = Open(Param(true), Path("coarse")); + auto reference = Open(Param(true), Path("reference")); + Add(coarse.get(), coarse_source); + Add(reference.get(), reference_source); + auto qp = + FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, reference->search_with_source(Vector(query), qp, + reference_source, &want)); + reference_source.supports_scan = false; + qp->refiner_param = std::make_shared(); + qp->refiner_param->scale_factor_ = 6; + qp->refiner_param->reference_index = reference; + qp->refiner_param->reference_vector_source = &reference_source; + ASSERT_EQ(0, + coarse->search_with_source(Vector(query), qp, coarse_source, &got)); + SameResults(want, got); + for (auto &doc : got.doc_list_) { + EXPECT_EQ(0, std::memcmp(reference_source.get_vector(doc.key()), + doc.vector(), kDimension * sizeof(float))); + } + qp->refiner_param->reference_vector_source = nullptr; + EXPECT_NE(0, + coarse->search_with_source(Vector(query), qp, coarse_source, &got)); + qp->refiner_param.reset(); + EXPECT_NE(0, reference->search(Vector(query), qp, &got)); +} + +TEST_F(ExternalFlatTest, HnswExternalRefinesWithExternalFlat) { + BatchSource source(30); + auto original = std::move(source.rows); + uint32_t id = 0; + for (auto &entry : original) source.rows[id++] = entry.second; + auto coarse_param = HNSWIndexParamBuilder() + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_metric_type(MetricType::kL2sq) + .with_use_external_vector(true) + .build(); + auto coarse = Open(coarse_param, Path("coarse")); + auto reference = Open(Param(true), Path("reference")); + Add(coarse.get(), source); + Add(reference.get(), source); + std::vector query(kDimension, 0.3f); + auto flat_qp = + FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); + SearchResult want, got; + ASSERT_EQ( + 0, reference->search_with_source(Vector(query), flat_qp, source, &want)); + auto qp = HNSWQueryParamBuilder() + .with_topk(5) + .with_ef_search(64) + .with_fetch_vector(true) + .build(); + qp->refiner_param = std::make_shared(); + qp->refiner_param->scale_factor_ = 6; + qp->refiner_param->reference_index = reference; + qp->refiner_param->reference_vector_source = &source; + source.supports_scan = false; + ASSERT_EQ(0, coarse->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); +} + +TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { + BatchSource first(50), second(50, 2.0f); + auto index = Open(Param(true), Path("index")); + Add(index.get(), first); + std::vector query(kDimension, 0.3f); + SearchResult expected[2]; + BatchSource *sources[] = {&first, &second}; + auto qp = FlatQueryParamBuilder().with_topk(5).build(); + for (int i = 0; i < 2; ++i) { + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, *sources[i], + &expected[i])); + } + std::vector threads; + for (int i = 0; i < 4; ++i) { + threads.emplace_back([&, i] { + auto local_qp = FlatQueryParamBuilder().with_topk(5).build(); + for (int j = 0; j < 30; ++j) { + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), local_qp, + *sources[i % 2], &result)); + SameResults(expected[i % 2], result); + } + }); + } + for (auto &thread : threads) thread.join(); + EXPECT_EQ(61, first.scans); + EXPECT_EQ(61, second.scans); +} + +TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { + for (int i = 0; i < 6; ++i) { + auto param = Param(true); + if (i == 0) param->metric_type = MetricType::kCosine; + if (i == 1) + param->quantizer_param = + std::make_shared(QuantizerType::kInt8); + if (i == 2) param->storage_data_type = DataType::DT_FP16; + if (i == 3) param->major_order = core::IndexMeta::MO_COLUMN; + if (i == 4) param->use_contiguous_memory = true; + if (i == 5) param->is_sparse = true; + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)); + } + BatchSource source(1); + auto index = Open(Param(true), Path("index")); + Add(index.get(), source); + EXPECT_NE(0, index->index_searcher()->dump(nullptr)); + auto provider = index->create_index_provider(); + ASSERT_NE(nullptr, provider); + EXPECT_EQ(1, provider->count()); + EXPECT_EQ(nullptr, provider->create_iterator()); + EXPECT_EQ(nullptr, provider->get_vector(3)); + EXPECT_NE(0, index->merge({index}, IndexFilter{})); +} + +TEST_F(ExternalFlatTest, LegacyRandomAccessAndBufferPoolReopen) { + BatchSource source(40); + source.legacy_reads = true; + auto path = Path("index"); + auto index = Open(Param(true), path); + Add(index.get(), source); + Close(index); + index = IndexFactory::CreateAndInitIndex(*Param(true)); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kBufferPool, + false, true})); + opened_.push_back(index); + auto qp = + FlatQueryParamBuilder().with_topk(1).with_fetch_vector(true).build(); + SearchResult result; + auto vector = Vector(source.rows.at(3)); + EXPECT_NE(0, index->add_with_source(vector, 3, source)); + EXPECT_NE(0, index->search(vector, qp, &result)); + ASSERT_EQ(0, index->search_with_source(vector, qp, source, &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); +} + +TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { + BatchSource source(20); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(zvec::ailego::Params{})); + ASSERT_EQ(0, storage->open(Path("index"), true)); + core::IndexMeta meta; + meta.set_meta(core::IndexMeta::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params{}); + zvec::ailego::Params params; + params.set(core::PARAM_FLAT_USE_EXTERNAL_VECTOR, true); + auto streamer = core::IndexFactory::CreateStreamer("FlatStreamer16"); + ASSERT_EQ(0, streamer->init(meta, params)); + ASSERT_EQ(0, streamer->open(storage)); + auto context = streamer->create_context(); + auto *ctx = dynamic_cast *>(context.get()); + ASSERT_NE(nullptr, ctx); + ctx->set_vector_source(&source); + ctx->set_topk(1); + core::IndexQueryMeta qmeta(core::IndexMeta::DT_FP32, kDimension); + for (auto &[id, row] : source.rows) { + ASSERT_EQ(0, streamer->add_impl(id, row.data(), qmeta, context)); + } + auto query = source.rows.at(3); + query.insert(query.end(), source.rows.at(16).begin(), + source.rows.at(16).end()); + ASSERT_EQ(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + ASSERT_EQ(1, context->result(0).size()); + ASSERT_EQ(1, context->result(1).size()); + EXPECT_EQ(3, context->result(0)[0].key()); + EXPECT_EQ(16, context->result(1)[0].key()); + EXPECT_EQ(2, source.scans); + core::IndexQueryMeta wrong(core::IndexMeta::DT_FP32, kDimension - 1); + EXPECT_NE(0, streamer->search_bf_impl(query.data(), wrong, 2, context)); + context->reset(); + EXPECT_NE(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + EXPECT_EQ(0, streamer->close()); + context.reset(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(ExternalFlatTest, EmptyIndexAndEmptyCandidates) { + BatchSource source(0); + auto index = Open(Param(true), Path("index")); + auto qp = FlatQueryParamBuilder().with_topk(3).build(); + std::vector query(kDimension, 0.3f); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + source.supports_scan = false; + qp->bf_pks = std::make_shared>(); + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); +} +} // namespace From 93443bfd6cdb819b8d45dd13fd2ab7553e9892b3 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Thu, 8 Oct 2026 20:14:18 +0800 Subject: [PATCH 02/16] test(flat): fix helper naming for clang-tidy --- .../interface/external_flat_index_test.cc | 93 ++++++++++--------- 1 file changed, 48 insertions(+), 45 deletions(-) diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 46a891336c..5cf4740af0 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -145,7 +145,7 @@ class ExternalFlatTest : public testing::Test { std::vector opened_; std::vector paths_; - std::string Path(const std::string &suffix) { + std::string test_path(const std::string &suffix) { std::string path = "test_external_flat_" + std::string( @@ -156,8 +156,8 @@ class ExternalFlatTest : public testing::Test { return path; } - Index::Pointer Open(const BaseIndexParam::Pointer ¶m, - const std::string &path, bool create = true) { + Index::Pointer open_index(const BaseIndexParam::Pointer ¶m, + const std::string &path, bool create = true) { auto index = IndexFactory::CreateAndInitIndex(*param); EXPECT_NE(nullptr, index); if (!index) return index; @@ -167,14 +167,15 @@ class ExternalFlatTest : public testing::Test { return index; } - void Add(Index *index, const BatchSource &source, bool external = true) { + void add_vectors(Index *index, const BatchSource &source, + bool external = true) { for (auto &[id, row] : source.rows) { ASSERT_EQ(0, external ? index->add_with_source(Vector(row), id, source) : index->add(Vector(row), id)); } } - void Close(const Index::Pointer &index) { + void close_index(const Index::Pointer &index) { EXPECT_EQ(0, index->close()); opened_.erase(std::remove(opened_.begin(), opened_.end(), index), opened_.end()); @@ -192,10 +193,12 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { for (bool id_map : {false, true}) { BatchSource source; auto suffix = std::to_string(case_id++); - auto external = Open(Param(true, metric, id_map), Path(suffix + "ext")); - auto embedded = Open(Param(false, metric, id_map), Path(suffix + "base")); - Add(external.get(), source); - Add(embedded.get(), source, false); + auto external = + open_index(Param(true, metric, id_map), test_path(suffix + "ext")); + auto embedded = + open_index(Param(false, metric, id_map), test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); ASSERT_EQ(source.rows.size(), external->get_doc_count()); source.rows[99999] = std::vector(kDimension, 0.3f); std::vector query(kDimension, 0.3f); @@ -209,7 +212,7 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { EXPECT_EQ(1, source.scans); auto expected = source.rows; source.rows.clear(); - Close(external); + close_index(external); ASSERT_EQ(11, got.doc_list_.size()); // External vectors are returned directly in the input layout, without // creating a redundant dequantized copy in reverted_vector_list_. @@ -227,8 +230,8 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { BatchSource source; source.supports_scan = false; - auto index = Open(Param(true), Path("index")); - Add(index.get(), source); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); auto qp = FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); qp->radius = 0.5f; @@ -263,9 +266,9 @@ TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { BatchSource source; - auto path = Path("index"); - auto index = Open(Param(true, MetricType::kL2sq, false), path); - Add(index.get(), source); + auto path = test_path("index"); + auto index = open_index(Param(true, MetricType::kL2sq, false), path); + add_vectors(index.get(), source); auto *streamer = dynamic_cast *>(index->index_searcher().get()); ASSERT_NE(nullptr, streamer); @@ -277,8 +280,8 @@ TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(3)), 3, source)); EXPECT_EQ(89, index->get_doc_count()); ASSERT_EQ(0, index->flush()); - Close(index); - index = Open(Param(true, MetricType::kL2sq, false), path, false); + close_index(index); + index = open_index(Param(true, MetricType::kL2sq, false), path, false); ASSERT_EQ(89, index->get_doc_count()); auto qp = FlatQueryParamBuilder().with_topk(1).build(); SearchResult result; @@ -291,14 +294,14 @@ TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(1000000)), 1000000, source)); EXPECT_EQ(90, index->get_doc_count()); - Close(index); + close_index(index); auto wrong = IndexFactory::CreateAndInitIndex(*Param(false)); ASSERT_NE(nullptr, wrong); EXPECT_NE(0, wrong->open(path, {StorageOptions::StorageType::kMMAP, false})); wrong.reset(); - auto embedded_path = Path("embedded"); - auto embedded = Open(Param(false), embedded_path); - Close(embedded); + auto embedded_path = test_path("embedded"); + auto embedded = open_index(Param(false), embedded_path); + close_index(embedded); wrong = IndexFactory::CreateAndInitIndex(*Param(true)); EXPECT_NE(0, wrong->open(embedded_path, {StorageOptions::StorageType::kMMAP, false})); @@ -306,8 +309,8 @@ TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { BatchSource source; - auto index = Open(Param(true), Path("index")); - Add(index.get(), source); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); auto qp = FlatQueryParamBuilder().with_topk(3).build(); std::vector query(kDimension, 0.3f); SearchResult result; @@ -339,10 +342,10 @@ TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { TEST_F(ExternalFlatTest, GroupByUsesExternalVectors) { BatchSource source; - auto index = Open(Param(true), Path("ext")); - auto embedded = Open(Param(false), Path("base")); - Add(index.get(), source); - Add(embedded.get(), source, false); + auto index = open_index(Param(true), test_path("ext")); + auto embedded = open_index(Param(false), test_path("base")); + add_vectors(index.get(), source); + add_vectors(embedded.get(), source, false); auto qp = FlatQueryParamBuilder().with_topk(6).with_fetch_vector(true).build(); qp->group_by_param = std::make_shared(); @@ -373,10 +376,10 @@ TEST_F(ExternalFlatTest, GroupByUsesExternalVectors) { TEST_F(ExternalFlatTest, RefineBindsItsOwnSource) { BatchSource coarse_source(30), reference_source(30, 0.8f); - auto coarse = Open(Param(true), Path("coarse")); - auto reference = Open(Param(true), Path("reference")); - Add(coarse.get(), coarse_source); - Add(reference.get(), reference_source); + auto coarse = open_index(Param(true), test_path("coarse")); + auto reference = open_index(Param(true), test_path("reference")); + add_vectors(coarse.get(), coarse_source); + add_vectors(reference.get(), reference_source); auto qp = FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); std::vector query(kDimension, 0.3f); @@ -413,10 +416,10 @@ TEST_F(ExternalFlatTest, HnswExternalRefinesWithExternalFlat) { .with_metric_type(MetricType::kL2sq) .with_use_external_vector(true) .build(); - auto coarse = Open(coarse_param, Path("coarse")); - auto reference = Open(Param(true), Path("reference")); - Add(coarse.get(), source); - Add(reference.get(), source); + auto coarse = open_index(coarse_param, test_path("coarse")); + auto reference = open_index(Param(true), test_path("reference")); + add_vectors(coarse.get(), source); + add_vectors(reference.get(), source); std::vector query(kDimension, 0.3f); auto flat_qp = FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); @@ -439,8 +442,8 @@ TEST_F(ExternalFlatTest, HnswExternalRefinesWithExternalFlat) { TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { BatchSource first(50), second(50, 2.0f); - auto index = Open(Param(true), Path("index")); - Add(index.get(), first); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), first); std::vector query(kDimension, 0.3f); SearchResult expected[2]; BatchSource *sources[] = {&first, &second}; @@ -480,8 +483,8 @@ TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)); } BatchSource source(1); - auto index = Open(Param(true), Path("index")); - Add(index.get(), source); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); EXPECT_NE(0, index->index_searcher()->dump(nullptr)); auto provider = index->create_index_provider(); ASSERT_NE(nullptr, provider); @@ -494,10 +497,10 @@ TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { TEST_F(ExternalFlatTest, LegacyRandomAccessAndBufferPoolReopen) { BatchSource source(40); source.legacy_reads = true; - auto path = Path("index"); - auto index = Open(Param(true), path); - Add(index.get(), source); - Close(index); + auto path = test_path("index"); + auto index = open_index(Param(true), path); + add_vectors(index.get(), source); + close_index(index); index = IndexFactory::CreateAndInitIndex(*Param(true)); ASSERT_NE(nullptr, index); ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kBufferPool, @@ -519,7 +522,7 @@ TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { BatchSource source(20); auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); ASSERT_EQ(0, storage->init(zvec::ailego::Params{})); - ASSERT_EQ(0, storage->open(Path("index"), true)); + ASSERT_EQ(0, storage->open(test_path("index"), true)); core::IndexMeta meta; meta.set_meta(core::IndexMeta::DT_FP32, kDimension); meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params{}); @@ -557,7 +560,7 @@ TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { TEST_F(ExternalFlatTest, EmptyIndexAndEmptyCandidates) { BatchSource source(0); - auto index = Open(Param(true), Path("index")); + auto index = open_index(Param(true), test_path("index")); auto qp = FlatQueryParamBuilder().with_topk(3).build(); std::vector query(kDimension, 0.3f); SearchResult result; From b732cce14fb53b5b591f6470d89e1c834989e150 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Thu, 8 Oct 2026 20:51:15 +0800 Subject: [PATCH 03/16] refactor(flat): trim external vector integration to core search --- examples/c++/core/external_flat_example.cc | 5 +- src/core/algorithm/flat/flat_streamer.cc | 88 ++------- .../algorithm/flat/flat_streamer_context.h | 4 +- src/core/interface/index.cc | 28 +-- src/core/interface/indexes/flat_index.cc | 36 +--- src/core/interface/vector_source.cc | 18 -- src/include/zvec/core/interface/index.h | 20 +- src/include/zvec/core/interface/index_param.h | 3 - .../zvec/core/interface/vector_source.h | 8 +- .../interface/external_flat_index_test.cc | 186 ++++-------------- 10 files changed, 90 insertions(+), 306 deletions(-) diff --git a/examples/c++/core/external_flat_example.cc b/examples/c++/core/external_flat_example.cc index 683eb23b6d..b877e532b5 100644 --- a/examples/c++/core/external_flat_example.cc +++ b/examples/c++/core/external_flat_example.cc @@ -94,6 +94,7 @@ int main() { VectorData query{DenseVector{source.rows[0].vector.data()}}; auto qp = FlatQueryParamBuilder().with_topk(1).with_fetch_vector(true).build(); + // fetch_vector returns borrowed pointers; source outlives result. SearchResult result; if (index->search_with_source(query, qp, source, &result) || result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42 || @@ -112,9 +113,7 @@ int main() { result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42) return 1; std::cout << "Candidate search after reopen: id=42\n"; - VectorDataBuffer fetched; - if (index->fetch_with_source(42, source, &fetched) || index->close()) - return 1; + if (index->close()) return 1; std::filesystem::remove(path); return 0; } diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 1075a3d85c..5aca246dfd 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -359,13 +359,7 @@ int FlatStreamer::add_with_id_impl(uint32_t id, const void *query, if (state_ != STATE_OPENED || !ctx->vector_source()) return IndexError_NoReady; if (id == kInvalidNodeId) return IndexError_InvalidArgument; - VectorSource::Batch batch; - int ret = ctx->vector_source()->get_vector_batch(&id, 1, &batch); - if (ret != 0) return ret; - if (batch.ids.size() != 1 || batch.ids[0] != id || - batch.vectors.size() != 1 || !batch.vectors[0]) { - return IndexError_ReadData; - } + if (!ctx->vector_source()->get_vector(id)) return IndexError_NoExist; } READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); @@ -506,10 +500,8 @@ int FlatStreamer::search_external( if (ctx->magic() != magic_) ctx->reset(this); const auto *source = ctx->vector_source(); if (!source || !ctx->threshold_is_valid()) return IndexError_NoReady; - const bool grouped = ctx->group_by_search(); - if (grouped && !ctx->group_by().is_valid()) return IndexError_InvalidArgument; + if (ctx->group_by_search()) return IndexError_Unsupported; ctx->reset_results(count); - ctx->resize_group_results(count); // Bound both pointer scratch space and per-request pinned vector bytes. const uint32_t batch_limit = static_cast(std::min( 1024, @@ -521,7 +513,6 @@ int FlatStreamer::search_external( static_cast(query) + q * qmeta.element_size(); auto *heap = ctx->result_heap(); heap->clear(); - ctx->group_topk_heaps().clear(); // Destroy the outstanding lease before its cursor, including on errors. VectorSource::Iterator::Pointer cursor; VectorSource::Batch batch; @@ -565,13 +556,7 @@ int FlatStreamer::search_external( for (size_t i = 0; i < size; ++i) { const auto id = scratch->vector_keys[i]; const auto distance = scratch->distances[i]; - if (grouped) { - auto &group_heap = ctx->group_topk_heaps()[ctx->group_by()(id)]; - if (group_heap.empty()) group_heap.limit(ctx->group_topk()); - group_heap.emplace(id, distance); - } else { - heap->emplace(id, distance); - } + heap->emplace(id, distance); } return 0; }; @@ -582,10 +567,10 @@ int FlatStreamer::search_external( auto read_candidates = [&]() -> int { if (ids.empty()) return 0; batch.clear(); - int ret = source->get_vector_batch(ids.data(), ids.size(), &batch); - if (ret != 0) return ret; - if (batch.ids != ids) return IndexError_ReadData; - ret = consume(); + batch.ids = ids; + batch.vectors.resize(ids.size()); + source->get_vectors(ids.data(), ids.size(), batch.vectors.data()); + int ret = consume(); ids.clear(); return ret; }; @@ -614,47 +599,16 @@ int FlatStreamer::search_external( } } batch.clear(); - if (grouped) - ctx->topk_to_group_result(q); - else - ctx->take_topk_result(q); - - // Results outlive the source snapshot, cursor and index. Own the selected - // raw vectors instead of retaining borrowed pointers or disk page pins. - auto fetch_results = [&](IndexDocumentList *docs) -> int { - for (size_t begin = 0; begin < docs->size(); begin += batch_limit) { - const size_t end = std::min(docs->size(), begin + batch_limit); - std::vector ids; - ids.reserve(end - begin); - for (size_t i = begin; i < end; ++i) ids.push_back((*docs)[i].key()); - batch.clear(); - int ret = source->get_vector_batch(ids.data(), ids.size(), &batch); - if (ret != 0) return ret; - if (batch.ids != ids || batch.vectors.size() != ids.size()) { - return IndexError_ReadData; - } - for (size_t i = 0; i < ids.size(); ++i) { - if (!batch.vectors[i]) return IndexError_ReadData; - auto owned = std::make_shared( - static_cast(batch.vectors[i]), - meta_.element_size()); - auto block = - IndexStorage::MemoryBlock::MakeSharedView(owned->data(), owned); - auto &doc = (*docs)[begin + i]; - doc = IndexDocument(doc.key(), doc.score(), doc.index(), block); - } - } - return 0; - }; + ctx->take_topk_result(q); if (ctx->fetch_vector()) { - if (grouped) { - for (auto &group : *ctx->mutable_group_result(q)) { - int ret = fetch_results(group.mutable_docs()); - if (ret != 0) return ret; - } - } else { - int ret = fetch_results(ctx->mutable_result(q)); - if (ret != 0) return ret; + // Match HNSW's borrowed random-access result contract. Scan pages have + // already been released, so result pointers must come from the source. + for (auto &doc : *ctx->mutable_result(q)) { + const void *vector = source->get_vector(doc.key()); + if (!vector) return IndexError_NoExist; + IndexStorage::MemoryBlock block; + block.reset(const_cast(vector)); + doc = IndexDocument(doc.key(), doc.score(), doc.index(), block); } } } @@ -665,10 +619,7 @@ template int FlatStreamer::group_by_search_impl( const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { - if (use_external_vector_) { - if (state_ != STATE_OPENED) return IndexError_NoReady; - return search_external(query, qmeta, count, nullptr, context); - } + if (use_external_vector_) return IndexError_Unsupported; FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); if (!bf_context) { @@ -717,10 +668,7 @@ int FlatStreamer::group_by_search_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { - if (use_external_vector_) { - if (state_ != STATE_OPENED) return IndexError_NoReady; - return search_external(query, qmeta, count, &p_keys, context); - } + if (use_external_vector_) return IndexError_Unsupported; if (count == 0 || count > p_keys.size()) return IndexError_InvalidArgument; FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h index 3eb4691900..bafd90a15d 100644 --- a/src/core/algorithm/flat/flat_streamer_context.h +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -200,8 +200,8 @@ class FlatStreamerContext : public IndexStreamer::Context { break; } node_id_t id = group_topk_list[i].second[j].first; - if (fetch_vector_ && !owner_->use_external_vector()) { - auto provider = owner_->create_provider(); + auto provider = owner_->create_provider(); + if (fetch_vector_) { IndexStorage::MemoryBlock block; provider->get_vector(id, block); group_results_[idx][i].mutable_docs()->emplace_back(id, score, id, diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index d78b955235..dc95a8b938 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -685,11 +685,6 @@ bool Index::is_dirty() const { return storage_->is_dirty(); } -int Index::fetch_with_source(uint32_t, const core::VectorSource &, - VectorDataBuffer *) { - return core::IndexError_Unsupported; -} - int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { if (!is_open_) { LOG_ERROR("Index is not open"); @@ -1100,14 +1095,6 @@ int Index::_refine_dense_candidates( auto &context = acquire_context(); if (!context) return core::IndexError_Runtime; - if (param_.use_external_vector) { - int ret = bind_vector_source( - search_param->refiner_param->reference_vector_source, context); - if (ret != 0) { - context->reset(); - return ret; - } - } context->set_topk(search_param->topk); context->set_fetch_vector(search_param->fetch_vector); if (search_param->filter && search_param->filter->is_valid()) { @@ -1428,16 +1415,19 @@ int Index::_sparse_search(const VectorData &vector_data, int Index::merge(const std::vector &indexes, const IndexFilter &filter, const MergeOptions &options) { - // Merge has no request-bound source from which to materialize vectors. - if (param_.use_external_vector) return core::IndexError_Unsupported; + if (indexes.empty()) { + return core::IndexError_Success; + } + // External Flat providers contain IDs only; merge has no bound source. + if (param_.index_type == IndexType::kFlat && param_.use_external_vector) { + return core::IndexError_Unsupported; + } for (const auto &index : indexes) { - if (index && index->param_.use_external_vector) { + if (index && index->param_.index_type == IndexType::kFlat && + index->param_.use_external_vector) { return core::IndexError_Unsupported; } } - if (indexes.empty()) { - return core::IndexError_Success; - } // ivf need builder auto reducer = core::IndexFactory::CreateStreamerReducer("MixedStreamerReducer"); diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index dee682c8ee..e9251a77e6 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -142,15 +142,14 @@ ailego::Params MakeTurboQuantizerParams(const QuantizerParam &quantizer_param, } // namespace -int FlatIndex::bind_vector_source(const core::VectorSource *source, +int FlatIndex::bind_vector_source(const core::VectorSource &source, core::IndexContext::Pointer &context) { if (!param_.use_external_vector) return core::IndexError_Unsupported; - if (!source) return core::IndexError_InvalidArgument; auto *ctx = dynamic_cast *>(context.get()); auto *streamer = dynamic_cast *>(streamer_.get()); if (!ctx || !streamer) return core::IndexError_Cast; if (ctx->magic() != streamer->magic()) ctx->reset(streamer); - ctx->set_vector_source(source); + ctx->set_vector_source(&source); return 0; } @@ -160,7 +159,7 @@ int FlatIndex::add_with_source(const VectorData &vector, uint32_t doc_id, auto &context = acquire_context(); if (!context) return core::IndexError_Runtime; FlatSourceScope scope{context}; - int ret = bind_vector_source(&source, context); + int ret = bind_vector_source(source, context); return ret == 0 ? Index::add(vector, doc_id) : ret; } @@ -173,33 +172,10 @@ int FlatIndex::search_with_source( auto &context = acquire_context(); if (!context) return core::IndexError_Runtime; FlatSourceScope scope{context}; - int ret = bind_vector_source(&source, context); + int ret = bind_vector_source(source, context); return ret == 0 ? Index::search(query, search_param, result) : ret; } -int FlatIndex::fetch_with_source(uint32_t doc_id, - const core::VectorSource &source, - VectorDataBuffer *result) { - if (!is_open_) return core::IndexError_NoReady; - if (!param_.use_external_vector) return core::IndexError_Unsupported; - if (!result) return core::IndexError_InvalidArgument; - auto *streamer = dynamic_cast *>(streamer_.get()); - if (!streamer) return core::IndexError_Cast; - if (!streamer->entity().contains_key(doc_id)) return core::IndexError_NoExist; - core::VectorSource::Batch batch; - int ret = source.get_vector_batch(&doc_id, 1, &batch); - if (ret != 0) return ret; - if (batch.ids.size() != 1 || batch.ids[0] != doc_id || - batch.vectors.size() != 1 || !batch.vectors[0]) { - return core::IndexError_ReadData; - } - DenseVectorBuffer buffer; - buffer.data.assign(static_cast(batch.vectors[0]), - input_vector_meta_.element_size()); - result->vector_buffer = std::move(buffer); - return 0; -} - int FlatIndex::open(const std::string &file_path, StorageOptions storage_options) { // Restore the persisted encoding while keeping the configured structural @@ -423,6 +399,10 @@ int FlatIndex::_prepare_for_search( return core::IndexError_Runtime; } + if (param_.use_external_vector && flat_search_param->group_by_param) { + return core::IndexError_Unsupported; + } + context->set_topk(flat_search_param->topk); context->set_fetch_vector(flat_search_param->fetch_vector); if (flat_search_param->filter && flat_search_param->filter->is_valid()) { diff --git a/src/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index fa218c7874..98307cb76b 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include #include namespace zvec { @@ -26,23 +25,6 @@ VectorSource::Iterator::Pointer VectorSource::create_iterator() const { return nullptr; } -int VectorSource::get_vector_batch(const uint32_t *ids, uint32_t count, - Batch *out) const { - if (!out || (count && !ids)) return IndexError_InvalidArgument; - out->clear(); - if (!count) return 0; - out->ids.assign(ids, ids + count); - out->vectors.resize(count); - get_vectors(ids, count, out->vectors.data()); - for (auto *vector : out->vectors) { - if (!vector) { - out->clear(); - return IndexError_NoExist; - } - } - return 0; -} - void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, const void **out) const { for (uint32_t i = 0; i < count; ++i) { diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 71ed3cd784..ed27ab143c 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -148,9 +148,6 @@ class ZVEC_CORE_API Index { const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, const core::VectorSource &src, SearchResult *result); - virtual int fetch_with_source(uint32_t doc_id, const core::VectorSource &src, - VectorDataBuffer *result); - virtual BaseIndexParam::Pointer get_param() const; virtual bool is_trained() const; @@ -171,11 +168,6 @@ class ZVEC_CORE_API Index { } protected: - virtual int bind_vector_source(const core::VectorSource *, - core::IndexContext::Pointer &) { - return core::IndexError_Unsupported; - } - int _sparse_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer); virtual int _dense_fetch(const uint32_t doc_id, @@ -284,14 +276,14 @@ class ZVEC_CORE_API FlatIndex : public Index { int add_with_source(const VectorData &vector, uint32_t doc_id, const core::VectorSource &src) override; // Full scans enumerate every registered ID once (extra IDs are ignored). - // Candidate searches/refinement use random batches. Hold one stable source - // snapshot across enumeration and random reads for the entire request. + // Candidate searches use get_vectors(). Hold a stable source snapshot for + // the request. With fetch_vector, returned pointers borrow source memory; + // keep that memory valid while using the results, as with external HNSW. + // External Flat does not support group-by or binding a source for refinement. int search_with_source(const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, const core::VectorSource &src, SearchResult *result) override; - int fetch_with_source(uint32_t doc_id, const core::VectorSource &src, - VectorDataBuffer *result) override; // FlatIndex(const FlatIndexParam ¶m) : param_(param) {} // FlatIndex(FlatIndexParam &¶m) : param(std::move(param)) {} @@ -327,8 +319,8 @@ class ZVEC_CORE_API FlatIndex : public Index { int create_and_init_legacy_converter_reformer( const QuantizerParam ¶m, const BaseIndexParam &index_param); - int bind_vector_source(const core::VectorSource *source, - core::IndexContext::Pointer &context) override; + int bind_vector_source(const core::VectorSource &source, + core::IndexContext::Pointer &context); FlatIndexParam param_{}; }; diff --git a/src/include/zvec/core/interface/index_param.h b/src/include/zvec/core/interface/index_param.h index 4e7937b00a..2f6146689a 100644 --- a/src/include/zvec/core/interface/index_param.h +++ b/src/include/zvec/core/interface/index_param.h @@ -24,7 +24,6 @@ #include #include #include -#include #include #include "zvec/core/framework/index_framework.h" @@ -205,8 +204,6 @@ struct RefinerParam { float scale_factor_{0}; std::shared_ptr reference_index = nullptr; - // Borrowed for this request; independent of the coarse index source. - const core::VectorSource *reference_vector_source = nullptr; }; // --- GroupBy Parameters --- diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 2ac175a3a5..455d66c365 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -52,6 +52,8 @@ class ZVEC_CORE_API VectorSource { VectorSource(); virtual ~VectorSource(); + // Borrowed random-access pointers. Keep all get_vectors() results valid + // through consumption; fetch_vector results also borrow this memory. virtual const void *get_vector(uint32_t node_id) const = 0; virtual void get_vectors(const uint32_t *ids, uint32_t count, @@ -61,12 +63,6 @@ class ZVEC_CORE_API VectorSource { // holds a stable source snapshot for the whole request, including subsequent // random reads. A disk source can scan physical blocks behind this interface. virtual Iterator::Pointer create_iterator() const; - - // Return IDs in input order. The default borrows get_vectors() pointers, - // which must all remain valid until the caller consumes the batch. Sources - // with transient buffers should override this method and attach a lease. - virtual int get_vector_batch(const uint32_t *ids, uint32_t count, - Batch *out) const; }; } // namespace core diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 5cf4740af0..185c0353d2 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -35,11 +35,10 @@ constexpr uint32_t kDimension = 8; class BatchSource : public core::VectorSource { public: std::map> rows; - mutable std::atomic scans{0}, reads{0}; + mutable std::atomic scans{0}; size_t page_rows{7}; bool supports_scan{true}, fail_scan{false}, malformed_scan{false}; - bool null_scan{false}, fail_read{false}, wrong_read_ids{false}; - bool legacy_reads{false}; + bool null_scan{false}, fail_read{false}; explicit BatchSource(uint32_t count = 89, float offset = 0) { for (uint32_t i = 0; i < count; ++i) { @@ -51,29 +50,11 @@ class BatchSource : public core::VectorSource { } const void *get_vector(uint32_t id) const override { + if (fail_read) return nullptr; auto it = rows.find(id); return it == rows.end() ? nullptr : it->second.data(); } - int get_vector_batch(const uint32_t *ids, uint32_t count, - Batch *out) const override { - ++reads; - if (legacy_reads) return VectorSource::get_vector_batch(ids, count, out); - if (fail_read) return core::IndexError_ReadData; - out->clear(); - auto page = std::make_shared>(count * kDimension); - for (uint32_t i = 0; i < count; ++i) { - const void *data = get_vector(ids[i]); - if (!data) return core::IndexError_NoExist; - std::memcpy(page->data() + i * kDimension, data, - kDimension * sizeof(float)); - out->ids.push_back(ids[i] + (wrong_read_ids ? 1 : 0)); - out->vectors.push_back(page->data() + i * kDimension); - } - out->lease = page; - return 0; - } - class Cursor : public Iterator { public: explicit Cursor(const BatchSource &source) : source_(source) { @@ -91,7 +72,7 @@ class BatchSource : public core::VectorSource { auto page = std::make_shared>(count * kDimension); for (size_t i = 0; i < count; ++i) { auto id = ids_[position_++]; - std::memcpy(page->data() + i * kDimension, source_.get_vector(id), + std::memcpy(page->data() + i * kDimension, source_.rows.at(id).data(), kDimension * sizeof(float)); out->ids.push_back(id); out->vectors.push_back( @@ -187,7 +168,7 @@ class ExternalFlatTest : public testing::Test { } }; -TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { +TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { for (bool id_map : {false, true}) { @@ -210,8 +191,6 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { external->search_with_source(Vector(query), qp, source, &got)); SameResults(want, got); EXPECT_EQ(1, source.scans); - auto expected = source.rows; - source.rows.clear(); close_index(external); ASSERT_EQ(11, got.doc_list_.size()); // External vectors are returned directly in the input layout, without @@ -219,9 +198,9 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndOwnsReturnedVectors) { EXPECT_TRUE(got.reverted_vector_list_.empty()); for (size_t i = 0; i < got.doc_list_.size(); ++i) { const auto &doc = got.doc_list_[i]; - const auto &row = expected.at(doc.key()); - EXPECT_EQ(0, std::memcmp(row.data(), doc.vector(), - kDimension * sizeof(float))); + // Returned vectors borrow the stable random-access source, not the + // transient scan pages. The source remains alive after index close. + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); } } } @@ -249,16 +228,9 @@ TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { for (auto &doc : result.doc_list_) { EXPECT_NE(3, doc.key()); EXPECT_LE(doc.score(), qp->radius); - VectorDataBuffer buffer; - ASSERT_EQ(0, index->fetch_with_source(doc.key(), source, &buffer)); - EXPECT_EQ(0, - std::memcmp( - source.get_vector(doc.key()), - std::get(buffer.vector_buffer).data.data(), - kDimension * sizeof(float))); + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); } VectorDataBuffer buffer; - EXPECT_NE(0, index->fetch_with_source(999999, source, &buffer)); EXPECT_NE(0, index->fetch(3, &buffer)); EXPECT_NE(0, index->add(Vector(query), 3)); EXPECT_NE(0, index->search(Vector(query), qp, &result)); @@ -327,119 +299,20 @@ TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { source.fail_scan = source.malformed_scan = source.null_scan = false; source.supports_scan = true; ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + source.fail_read = true; + qp->fetch_vector = true; + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + qp->fetch_vector = false; qp->bf_pks = std::make_shared>( std::initializer_list{3}); - source.fail_read = true; EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); source.fail_read = false; - source.wrong_read_ids = true; - EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); - source.wrong_read_ids = false; source.rows.erase(3); EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_NE(0, index->add_with_source(Vector(query), 3, source)); } -TEST_F(ExternalFlatTest, GroupByUsesExternalVectors) { - BatchSource source; - auto index = open_index(Param(true), test_path("ext")); - auto embedded = open_index(Param(false), test_path("base")); - add_vectors(index.get(), source); - add_vectors(embedded.get(), source, false); - auto qp = - FlatQueryParamBuilder().with_topk(6).with_fetch_vector(true).build(); - qp->group_by_param = std::make_shared(); - qp->group_by_param->group_count = 3; - qp->group_by_param->group_topk = 2; - qp->group_by_param->group_by = [](uint64_t id) { - return std::to_string(id % 3); - }; - std::vector query(kDimension, 0.3f); - SearchResult want, got; - ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); - ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &got)); - ASSERT_EQ(3, got.group_doc_list_.size()); - for (size_t i = 0; i < want.group_doc_list_.size(); ++i) { - const auto &a = want.group_doc_list_[i]; - const auto &b = got.group_doc_list_[i]; - EXPECT_EQ(a.group_id(), b.group_id()); - ASSERT_EQ(2, b.docs().size()); - for (size_t j = 0; j < b.docs().size(); ++j) { - EXPECT_EQ(a.docs()[j].key(), b.docs()[j].key()); - EXPECT_NEAR(a.docs()[j].score(), b.docs()[j].score(), 1e-5); - EXPECT_EQ(0, - std::memcmp(source.get_vector(b.docs()[j].key()), - b.docs()[j].vector(), kDimension * sizeof(float))); - } - } -} - -TEST_F(ExternalFlatTest, RefineBindsItsOwnSource) { - BatchSource coarse_source(30), reference_source(30, 0.8f); - auto coarse = open_index(Param(true), test_path("coarse")); - auto reference = open_index(Param(true), test_path("reference")); - add_vectors(coarse.get(), coarse_source); - add_vectors(reference.get(), reference_source); - auto qp = - FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); - std::vector query(kDimension, 0.3f); - SearchResult want, got; - ASSERT_EQ(0, reference->search_with_source(Vector(query), qp, - reference_source, &want)); - reference_source.supports_scan = false; - qp->refiner_param = std::make_shared(); - qp->refiner_param->scale_factor_ = 6; - qp->refiner_param->reference_index = reference; - qp->refiner_param->reference_vector_source = &reference_source; - ASSERT_EQ(0, - coarse->search_with_source(Vector(query), qp, coarse_source, &got)); - SameResults(want, got); - for (auto &doc : got.doc_list_) { - EXPECT_EQ(0, std::memcmp(reference_source.get_vector(doc.key()), - doc.vector(), kDimension * sizeof(float))); - } - qp->refiner_param->reference_vector_source = nullptr; - EXPECT_NE(0, - coarse->search_with_source(Vector(query), qp, coarse_source, &got)); - qp->refiner_param.reset(); - EXPECT_NE(0, reference->search(Vector(query), qp, &got)); -} - -TEST_F(ExternalFlatTest, HnswExternalRefinesWithExternalFlat) { - BatchSource source(30); - auto original = std::move(source.rows); - uint32_t id = 0; - for (auto &entry : original) source.rows[id++] = entry.second; - auto coarse_param = HNSWIndexParamBuilder() - .with_data_type(DataType::DT_FP32) - .with_dimension(kDimension) - .with_metric_type(MetricType::kL2sq) - .with_use_external_vector(true) - .build(); - auto coarse = open_index(coarse_param, test_path("coarse")); - auto reference = open_index(Param(true), test_path("reference")); - add_vectors(coarse.get(), source); - add_vectors(reference.get(), source); - std::vector query(kDimension, 0.3f); - auto flat_qp = - FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); - SearchResult want, got; - ASSERT_EQ( - 0, reference->search_with_source(Vector(query), flat_qp, source, &want)); - auto qp = HNSWQueryParamBuilder() - .with_topk(5) - .with_ef_search(64) - .with_fetch_vector(true) - .build(); - qp->refiner_param = std::make_shared(); - qp->refiner_param->scale_factor_ = 6; - qp->refiner_param->reference_index = reference; - qp->refiner_param->reference_vector_source = &source; - source.supports_scan = false; - ASSERT_EQ(0, coarse->search_with_source(Vector(query), qp, source, &got)); - SameResults(want, got); -} - TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { BatchSource first(50), second(50, 2.0f); auto index = open_index(Param(true), test_path("index")); @@ -492,11 +365,31 @@ TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { EXPECT_EQ(nullptr, provider->create_iterator()); EXPECT_EQ(nullptr, provider->get_vector(3)); EXPECT_NE(0, index->merge({index}, IndexFilter{})); + EXPECT_EQ(0, index->merge({}, IndexFilter{})); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(embedded.get(), source, false); + EXPECT_NE(0, embedded->merge({index}, IndexFilter{})); + + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + qp->group_by_param = std::make_shared(); + qp->group_by_param->group_count = 1; + qp->group_by_param->group_topk = 1; + qp->group_by_param->group_by = [](uint64_t) { return std::string("group"); }; + SearchResult result; + const auto query = Vector(source.rows.at(3)); + EXPECT_EQ(core::IndexError_Unsupported, + index->search_with_source(query, qp, source, &result)); + EXPECT_TRUE(result.group_doc_list_.empty()); + qp->group_by_param.reset(); + // The existing refine entry point has no source binding for external Flat. + qp->refiner_param = std::make_shared(); + qp->refiner_param->scale_factor_ = 1; + qp->refiner_param->reference_index = index; + EXPECT_NE(0, embedded->search(query, qp, &result)); } -TEST_F(ExternalFlatTest, LegacyRandomAccessAndBufferPoolReopen) { +TEST_F(ExternalFlatTest, BufferPoolReopen) { BatchSource source(40); - source.legacy_reads = true; auto path = test_path("index"); auto index = open_index(Param(true), path); add_vectors(index.get(), source); @@ -551,6 +444,13 @@ TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { EXPECT_EQ(2, source.scans); core::IndexQueryMeta wrong(core::IndexMeta::DT_FP32, kDimension - 1); EXPECT_NE(0, streamer->search_bf_impl(query.data(), wrong, 2, context)); + ctx->set_group_params(1, 1); + EXPECT_EQ(core::IndexError_Unsupported, + streamer->search_bf_impl(query.data(), qmeta, 2, context)); + EXPECT_EQ(core::IndexError_Unsupported, + streamer->search_bf_by_p_keys_impl(query.data(), {{3}, {16}}, qmeta, + 2, context)); + ctx->set_group_params(0, 0); context->reset(); EXPECT_NE(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); EXPECT_EQ(0, streamer->close()); From b33cee558fdbc5f41435d24365e8898b1b9348cf Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Fri, 9 Oct 2026 11:57:35 +0800 Subject: [PATCH 04/16] test(metric): validate INT4 L2 against decoded vectors --- .../metric/quantized_integer_metric_test.cc | 130 +++++++++++------- 1 file changed, 77 insertions(+), 53 deletions(-) diff --git a/tests/core/metric/quantized_integer_metric_test.cc b/tests/core/metric/quantized_integer_metric_test.cc index 57541fcd41..09ec8437ed 100644 --- a/tests/core/metric/quantized_integer_metric_test.cc +++ b/tests/core/metric/quantized_integer_metric_test.cc @@ -344,62 +344,86 @@ TEST(QuantizedIntegerMetric, TestInt8SquaredEuclideanMetric) { } TEST(QuantizedIntegerMetric, TestInt4SquaredEuclidean) { - std::random_device rd; - std::mt19937 gen(rd()); + std::mt19937 gen(42); std::uniform_real_distribution dist(-1.0, 2.0); - - const size_t DIMENSION = std::uniform_int_distribution(1, 128)(gen) * 2; const size_t COUNT = 1000; - IndexMeta meta; - meta.set_meta(IndexMeta::DT_FP32, DIMENSION); - auto converter = IndexFactory::CreateConverter("Int4StreamingConverter"); - ASSERT_TRUE(!!converter); - ASSERT_EQ(0u, converter->init(meta, Params())); - - auto holder = GetHolder(DIMENSION, COUNT, dist); - ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); - auto holder2 = converter->result(); - EXPECT_EQ(COUNT, holder2->count()); - EXPECT_EQ(IndexMeta::DT_INT4, holder2->data_type()); - auto &meta2 = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(meta2.reformer_name()); - ASSERT_TRUE(reformer); - ASSERT_EQ(0u, reformer->init(meta2.reformer_params())); - - ailego::NumericalVector vec(DIMENSION); - for (size_t j = 0; j < DIMENSION; ++j) { - vec[j] = dist(gen); - } - IndexQueryMeta qmeta; - qmeta.set_meta(IndexMeta::DT_FP32, DIMENSION); - IndexQueryMeta qmeta2; - std::string out; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &out, &qmeta2)); - ASSERT_EQ(qmeta2.dimension(), meta2.dimension()); - - auto iter = holder->create_iterator(); - auto iter2 = holder2->create_iterator(); - auto metric = IndexFactory::CreateMetric(meta2.metric_name()); - ASSERT_TRUE(!!metric); - ASSERT_EQ(0, metric->init(meta2, meta2.metric_params())); - auto compute = metric->distance(); - ASSERT_TRUE(compute); - - for (; iter->is_valid(); iter->next(), iter2->next()) { - const float *mf = (const float *)iter->data(); - const int8_t *mi = (const int8_t *)iter2->data(); - const int8_t *qi = reinterpret_cast(&out[0]); - float v1 = - ailego::Distance::SquaredEuclidean(mf, vec.data(), holder->dimension()); - float v2; - compute(mi, qi, holder2->dimension(), &v2); - ASSERT_NEAR(v1, v2, 0.2 * DIMENSION); - std::string out2; - ASSERT_EQ(0, reformer->convert(iter->data(), qmeta, &out2, &qmeta2)); - ASSERT_EQ(out2.size(), holder2->element_size()); - ASSERT_EQ(0, std::memcmp(out2.data(), iter2->data(), out2.size())); + for (const size_t dimension : {2, 4, 8, 14, 32, 128, 256}) { + SCOPED_TRACE(dimension); + IndexMeta meta; + meta.set_meta(IndexMeta::DT_FP32, dimension); + auto converter = IndexFactory::CreateConverter("Int4StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0u, converter->init(meta, Params())); + + auto holder = + std::make_shared>(dimension); + for (size_t i = 0; i < COUNT; ++i) { + ailego::NumericalVector row(dimension); + for (size_t j = 0; j < dimension; ++j) row[j] = dist(gen); + if (dimension == 4 && i == 0) { + row = ailego::NumericalVector{-0.4362f, 1.95854f, 1.87264f, + 1.43786f}; + } + holder->emplace(i, row); + } + ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); + auto holder2 = converter->result(); + EXPECT_EQ(COUNT, holder2->count()); + EXPECT_EQ(IndexMeta::DT_INT4, holder2->data_type()); + auto &meta2 = converter->meta(); + + auto reformer = IndexFactory::CreateReformer(meta2.reformer_name()); + ASSERT_TRUE(reformer); + ASSERT_EQ(0u, reformer->init(meta2.reformer_params())); + + ailego::NumericalVector vec(dimension); + for (size_t j = 0; j < dimension; ++j) vec[j] = dist(gen); + if (dimension == 4) { + vec = ailego::NumericalVector{-0.95011f, 0.08062f, -0.86274f, + 1.42267f}; + } + IndexQueryMeta qmeta(IndexMeta::DT_FP32, dimension); + IndexQueryMeta qmeta2; + std::string out; + ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &out, &qmeta2)); + ASSERT_EQ(qmeta2.dimension(), meta2.dimension()); + std::string restored_query; + ASSERT_EQ(0, reformer->revert(out.data(), qmeta2, &restored_query)); + + auto iter = holder->create_iterator(); + auto iter2 = holder2->create_iterator(); + auto metric = IndexFactory::CreateMetric(meta2.metric_name()); + ASSERT_TRUE(metric); + ASSERT_EQ(0, metric->init(meta2, meta2.metric_params())); + auto compute = metric->distance(); + ASSERT_TRUE(compute); + + for (; iter->is_valid(); iter->next(), iter2->next()) { + SCOPED_TRACE(iter->key()); + std::string restored_row; + ASSERT_EQ(0, reformer->revert(iter2->data(), qmeta2, &restored_row)); + // Validate the metric against decoded vectors. Lossy INT4 rounding does + // not guarantee the former 0.2 * dimension bound against raw vectors. + const float expected = ailego::Distance::SquaredEuclidean( + reinterpret_cast(restored_row.data()), + reinterpret_cast(restored_query.data()), dimension); + float actual; + compute(iter2->data(), out.data(), holder2->dimension(), &actual); + ASSERT_NEAR(expected, actual, 1e-5 * dimension); + + if (dimension == 4 && iter->key() == 0) { + // This fixed pair exceeds the old tolerance with a correct metric. + const float raw = ailego::Distance::SquaredEuclidean( + static_cast(iter->data()), vec.data(), dimension); + ASSERT_GT(std::abs(raw - expected), 0.2 * dimension); + } + + std::string out2; + ASSERT_EQ(0, reformer->convert(iter->data(), qmeta, &out2, &qmeta2)); + ASSERT_EQ(out2.size(), holder2->element_size()); + ASSERT_EQ(0, std::memcmp(out2.data(), iter2->data(), out2.size())); + } } } From 11cd52f45d769e9e3f5cc8d7c3009b1c04229796 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Fri, 9 Oct 2026 12:45:21 +0800 Subject: [PATCH 05/16] chore(examples): remove external flat example --- examples/c++/CMakeLists.txt | 4 - examples/c++/core/external_flat_example.cc | 119 --------------------- 2 files changed, 123 deletions(-) delete mode 100644 examples/c++/core/external_flat_example.cc diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 771502bac0..ecb138c173 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -135,10 +135,6 @@ 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) -add_executable(external-flat-example core/external_flat_example.cc) -target_link_libraries(external-flat-example PRIVATE zvec-core-lib) -list(APPEND ZVEC_EXAMPLE_TARGETS external-flat-example) - # This is also a link-integration check: DiskANN must work through the # all-in-one core library without linking core_knn_diskann separately. add_executable(diskann-core-example core/diskann_main.cc) diff --git a/examples/c++/core/external_flat_example.cc b/examples/c++/core/external_flat_example.cc deleted file mode 100644 index b877e532b5..0000000000 --- a/examples/c++/core/external_flat_example.cc +++ /dev/null @@ -1,119 +0,0 @@ -// 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. - -// External Flat scans logical batches supplied by the source. A disk-backed -// source may read physical blocks and attach page pins to Batch::lease. -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace zvec::core_interface; - -class ExternalRows : public zvec::core::VectorSource { - public: - static constexpr uint32_t kDimension = 4; - struct Row { - uint32_t id; - std::array vector; - }; - std::vector rows; - - const void *get_vector(uint32_t id) const override { - for (const auto &row : rows) { - if (row.id == id) return row.vector.data(); - } - return nullptr; - } - - class Cursor : public Iterator { - public: - explicit Cursor(const ExternalRows &source) : source_(source) {} - int next_batch(uint32_t max_count, Batch *out) override { - if (!out || !max_count) return zvec::core::IndexError_InvalidArgument; - out->clear(); - const size_t end = std::min(source_.rows.size(), position_ + max_count); - for (; position_ < end; ++position_) { - const auto &row = source_.rows[position_]; - out->ids.push_back(row.id); - out->vectors.push_back(row.vector.data()); - } - // No lease is needed: this example keeps rows immutable throughout - // every request. Empty output signals end of scan. - return 0; - } - - private: - const ExternalRows &source_; - size_t position_{0}; - }; - - Iterator::Pointer create_iterator() const override { - return std::make_unique(*this); - } -}; - -int main() { - const std::string path = "external_flat_example.index"; - if (std::filesystem::exists(path)) { - std::cerr << "Remove " << path << " before running this example.\n"; - return 1; - } - ExternalRows source; - source.rows = {{42, {1, 0, 0, 0}}, {7, {0, 1, 0, 0}}, {105, {0, 0, 1, 0}}}; - auto param = FlatIndexParamBuilder() - .with_data_type(DataType::DT_FP32) - .with_dimension(ExternalRows::kDimension) - .with_metric_type(MetricType::kL2sq) - .with_use_external_vector(true) - .build(); - auto index = IndexFactory::CreateAndInitIndex(*param); - if (!index || index->open(path, {StorageOptions::StorageType::kMMAP, true})) - return 1; - // Register IDs only; source owns vector data and subsequent updates. - for (const auto &row : source.rows) { - if (index->add_with_source(VectorData{DenseVector{row.vector.data()}}, - row.id, source)) - return 1; - } - VectorData query{DenseVector{source.rows[0].vector.data()}}; - auto qp = - FlatQueryParamBuilder().with_topk(1).with_fetch_vector(true).build(); - // fetch_vector returns borrowed pointers; source outlives result. - SearchResult result; - if (index->search_with_source(query, qp, source, &result) || - result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42 || - result.doc_list_[0].score() != 0) - return 1; - std::cout << "Full scan: id=42, distance=0\n"; - if (index->close()) return 1; - - // Reopen the persisted ID registry and bind the source on the next request. - index = IndexFactory::CreateAndInitIndex(*param); - if (!index || index->open(path, {StorageOptions::StorageType::kMMAP, false})) - return 1; - qp->bf_pks = std::make_shared>( - std::initializer_list{105, 42}); - if (index->search_with_source(query, qp, source, &result) || - result.doc_list_.size() != 1 || result.doc_list_[0].key() != 42) - return 1; - std::cout << "Candidate search after reopen: id=42\n"; - if (index->close()) return 1; - std::filesystem::remove(path); - return 0; -} From 7d90ca43366cd5a5c85bb8e493b71102d1f089f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JasonXuDeveloper=20-=20=E5=82=91?= Date: Fri, 9 Oct 2026 17:17:17 +1100 Subject: [PATCH 06/16] fix(test): avoid MSVC ARM64 ICE in uniform_uint8_metric_test (#795) Co-authored-by: Claude Opus 5.5 (1M context) --- tests/core/metric/uniform_uint8_metric_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/metric/uniform_uint8_metric_test.cc b/tests/core/metric/uniform_uint8_metric_test.cc index e83dd70a39..b537b9e4ba 100644 --- a/tests/core/metric/uniform_uint8_metric_test.cc +++ b/tests/core/metric/uniform_uint8_metric_test.cc @@ -35,7 +35,7 @@ std::vector EncodeRecord(const std::vector &codes) { int64_t sum_squared = 0; for (size_t i = 0; i < codes.size(); ++i) { encoded[i] = static_cast(static_cast(codes[i]) - 128); - sum_squared += static_cast(codes[i]) * codes[i]; + sum_squared += static_cast(codes[i]) * codes[i]; } const uint32_t tail = static_cast(sum_squared); std::memcpy(encoded.data() + codes.size(), &tail, sizeof(tail)); From c59b8f1b4a68e0f155114b3b28c97c59bf04411a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 12:56:05 +0800 Subject: [PATCH 07/16] fix(flat): preserve snapshot filters and skip external payload copies --- src/core/algorithm/flat/flat_streamer.cc | 5 +- src/core/interface/indexes/flat_index.cc | 10 ++- src/include/zvec/core/interface/index.h | 7 +- .../zvec/core/interface/vector_source.h | 7 +- .../interface/external_flat_index_test.cc | 71 ++++++++++++++++++- 5 files changed, 89 insertions(+), 11 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 5aca246dfd..5a535ad04b 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -525,13 +525,14 @@ int FlatStreamer::search_external( if (!keys) entity_->key_presence(batch.ids, &scratch->key_presence); for (size_t i = 0; i < batch.ids.size(); ++i) { const auto id = batch.ids[i]; - if (id == kInvalidNodeId || !batch.vectors[i]) - return IndexError_ReadData; + if (id == kInvalidNodeId) return IndexError_ReadData; if (!keys && (!scratch->key_presence[i] || (ctx->filter().is_valid() && ctx->filter()(id)))) { ++*ctx->mutable_stats(q)->mutable_filtered_count(); continue; } + // Invisible or unregistered rows need not have a readable vector. + if (!batch.vectors[i]) return IndexError_ReadData; scratch->vector_ptrs.push_back(batch.vectors[i]); scratch->vector_keys.push_back(id); } diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index e9251a77e6..8f3ca303a8 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -156,11 +156,17 @@ int FlatIndex::bind_vector_source(const core::VectorSource &source, int FlatIndex::add_with_source(const VectorData &vector, uint32_t doc_id, const core::VectorSource &source) { if (!is_open_) return core::IndexError_NoReady; + if (is_read_only_) return core::IndexError_Runtime; + const auto *dense = std::get_if(&vector.vector); + if (!dense || !dense->data) return core::IndexError_InvalidArgument; auto &context = acquire_context(); if (!context) return core::IndexError_Runtime; FlatSourceScope scope{context}; int ret = bind_vector_source(source, context); - return ret == 0 ? Index::add(vector, doc_id) : ret; + if (ret != 0) return ret; + // External Flat stores only IDs. Its raw FP32 layout needs no conversion. + return streamer_->add_with_id_impl(doc_id, dense->data, input_vector_meta_, + context); } int FlatIndex::search_with_source( @@ -406,7 +412,7 @@ int FlatIndex::_prepare_for_search( context->set_topk(flat_search_param->topk); context->set_fetch_vector(flat_search_param->fetch_vector); if (flat_search_param->filter && flat_search_param->filter->is_valid()) { - context->set_filter(std::move(*flat_search_param->filter)); + context->set_filter(*flat_search_param->filter); } else { context->reset_filter(); } diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index ed27ab143c..b3c0665e20 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -275,9 +275,10 @@ class ZVEC_CORE_API FlatIndex : public Index { // doc_id is the source ID with either use_id_map setting. int add_with_source(const VectorData &vector, uint32_t doc_id, const core::VectorSource &src) override; - // Full scans enumerate every registered ID once (extra IDs are ignored). - // Candidate searches use get_vectors(). Hold a stable source snapshot for - // the request. With fetch_vector, returned pointers borrow source memory; + // Full scans use the source's visible IDs intersected with registered IDs + // and the request filter; extra IDs are ignored. Candidate searches use + // get_vectors(). Bind the source, scan and filter to one stable snapshot. + // With fetch_vector, returned pointers borrow source memory; // keep that memory valid while using the results, as with external HNSW. // External Flat does not support group-by or binding a source for refinement. int search_with_source(const VectorData &query, diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 455d66c365..07338c7961 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -28,6 +28,8 @@ class ZVEC_CORE_API VectorSource { // Pointers need not be contiguous. External Flat currently expects raw FP32. struct Batch { std::vector ids; + // Rows excluded by Flat's registered-ID set or filter may be null. + // Every row selected for distance computation must have a valid pointer. std::vector vectors; // Optional backing allocation or page pin. Data must remain valid until // the batch is released, or (for scans) the next next_batch() call. @@ -60,8 +62,9 @@ class ZVEC_CORE_API VectorSource { const void **out) const; // Optional full scan. Each call creates an independent cursor. The caller - // holds a stable source snapshot for the whole request, including subsequent - // random reads. A disk source can scan physical blocks behind this interface. + // binds the source, scan domain and visibility filter to the same stable + // snapshot for the whole request, including subsequent random reads. + // A disk source can scan physical blocks behind this interface. virtual Iterator::Pointer create_iterator() const; }; diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 185c0353d2..009f5f3377 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -37,6 +38,7 @@ class BatchSource : public core::VectorSource { std::map> rows; mutable std::atomic scans{0}; size_t page_rows{7}; + std::set null_scan_ids; bool supports_scan{true}, fail_scan{false}, malformed_scan{false}; bool null_scan{false}, fail_read{false}; @@ -75,8 +77,10 @@ class BatchSource : public core::VectorSource { std::memcpy(page->data() + i * kDimension, source_.rows.at(id).data(), kDimension * sizeof(float)); out->ids.push_back(id); - out->vectors.push_back( - source_.null_scan ? nullptr : page->data() + i * kDimension); + out->vectors.push_back(source_.null_scan || + source_.null_scan_ids.count(id) + ? nullptr + : page->data() + i * kDimension); } out->lease = page; previous_ = page; @@ -168,6 +172,69 @@ class ExternalFlatTest : public testing::Test { } }; +class InspectableFlatIndex : public FlatIndex { + public: + using Index::acquire_context; + using Index::init; +}; + +TEST_F(ExternalFlatTest, RegistrationDoesNotCopyVectorPayload) { + BatchSource source(3); + auto index = std::make_shared(); + ASSERT_EQ(0, index->init(*Param(true))); + ASSERT_EQ(0, index->open(test_path("index"), + {StorageOptions::StorageType::kMMAP, true})); + opened_.push_back(index); + auto &context = index->acquire_context(); + ASSERT_NE(nullptr, context); + *context->mutable_features() = "unchanged"; + add_vectors(index.get(), source); + ASSERT_EQ("unchanged", *context->mutable_features()); + EXPECT_EQ(source.rows.size(), index->get_doc_count()); + + // Keep the generic add guards when bypassing its conversion pipeline. + EXPECT_NE( + 0, index->add_with_source(VectorData{DenseVector{nullptr}}, 3, source)); + EXPECT_NE(0, index->add_with_source( + VectorData{SparseVector{0, nullptr, nullptr}}, 3, source)); + const auto vector = Vector(source.rows.at(3)); + EXPECT_NE(0, index->add_with_source(vector, UINT32_MAX, source)); + EXPECT_NE(0, index->add_with_source(vector, 99999, source)); + EXPECT_NE(0, index->add(vector, 3)); + close_index(index); + EXPECT_NE(0, index->add_with_source(vector, 3, source)); +} + +TEST_F(ExternalFlatTest, FullScanAppliesSnapshotFilterBeforeReadingVectors) { + BatchSource source(30); + auto external = open_index(Param(true), test_path("external")); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + // Model an immutable snapshot whose invisible rows have no usable pointer. + const std::set invisible{3, 16, 42}; + source.null_scan_ids = invisible; + source.rows[99999] = std::vector(kDimension, 0.3f); + source.null_scan_ids.insert(99999); // Unregistered rows are also skipped. + std::vector query(kDimension, 0.3f); + auto make_query = [&]() { + auto qp = FlatQueryParamBuilder().with_topk(10).build(); + qp->filter = std::make_shared(); + qp->filter->set([invisible](uint64_t id) { return invisible.count(id); }); + return qp; + }; + SearchResult expected, actual; + ASSERT_EQ(0, embedded->search(Vector(query), make_query(), &expected)); + auto qp = make_query(); + // Reusing the query must preserve its visibility filter across requests. + for (int i = 0; i < 2; ++i) { + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &actual)); + SameResults(expected, actual); + EXPECT_TRUE(qp->filter->is_valid()); + } +} + TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { From fb0124f1df45a29a9ecf5c5591841f048645e605 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 15:53:37 +0800 Subject: [PATCH 08/16] fix(flat): fall back to registered IDs without a source iterator --- src/core/algorithm/flat/flat_streamer.cc | 70 ++++++---- .../algorithm/flat/flat_streamer_entity.h | 18 +++ src/include/zvec/core/interface/index.h | 7 +- .../zvec/core/interface/vector_source.h | 7 +- .../interface/external_flat_index_test.cc | 122 ++++++++++++++++++ 5 files changed, 193 insertions(+), 31 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 3d57aed1cc..45031ab891 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -519,18 +519,19 @@ int FlatStreamer::search_external( // Destroy the outstanding lease before its cursor, including on errors. VectorSource::Iterator::Pointer cursor; VectorSource::Batch batch; + if (!keys) cursor = source->create_iterator(); auto consume = [&]() -> int { if (batch.ids.size() != batch.vectors.size() || batch.ids.size() > batch_limit) return IndexError_ReadData; scratch->vector_ptrs.clear(); scratch->vector_keys.clear(); - if (!keys) entity_->key_presence(batch.ids, &scratch->key_presence); + if (cursor) entity_->key_presence(batch.ids, &scratch->key_presence); for (size_t i = 0; i < batch.ids.size(); ++i) { const auto id = batch.ids[i]; if (id == kInvalidNodeId) return IndexError_ReadData; - if (!keys && (!scratch->key_presence[i] || - (ctx->filter().is_valid() && ctx->filter()(id)))) { + if (cursor && (!scratch->key_presence[i] || + (ctx->filter().is_valid() && ctx->filter()(id)))) { ++*ctx->mutable_stats(q)->mutable_filtered_count(); continue; } @@ -565,10 +566,19 @@ int FlatStreamer::search_external( return 0; }; - if (keys) { + if (cursor) { + while (true) { + batch.clear(); + int ret = cursor->next_batch(batch_limit, &batch); + if (ret != 0) return ret; + ret = consume(); + if (ret != 0) return ret; + if (batch.ids.empty()) break; + } + } else { std::vector ids; ids.reserve(batch_limit); - auto read_candidates = [&]() -> int { + auto read_batch = [&]() -> int { if (ids.empty()) return 0; batch.clear(); batch.ids = ids; @@ -578,28 +588,38 @@ int FlatStreamer::search_external( ids.clear(); return ret; }; - for (auto key : (*keys)[q]) { - if (key >= kInvalidNodeId || !entity_->contains_key(key) || - (ctx->filter().is_valid() && ctx->filter()(key))) - continue; - ids.push_back(static_cast(key)); - if (ids.size() == batch_limit) { - int ret = read_candidates(); - if (ret != 0) return ret; + auto read_candidates = [&](const std::vector &candidates) { + for (auto key : candidates) { + if (key >= kInvalidNodeId || (keys && !entity_->contains_key(key))) + continue; + if (ctx->filter().is_valid() && ctx->filter()(key)) { + if (!keys) ++*ctx->mutable_stats(q)->mutable_filtered_count(); + continue; + } + ids.push_back(static_cast(key)); + if (ids.size() == batch_limit) { + int ret = read_batch(); + if (ret != 0) return ret; + } } - } - int ret = read_candidates(); - if (ret != 0) return ret; - } else { - cursor = source->create_iterator(); - if (!cursor) return IndexError_Unsupported; - while (true) { - batch.clear(); - int ret = cursor->next_batch(batch_limit, &batch); - if (ret != 0) return ret; - ret = consume(); + return read_batch(); + }; + if (keys) { + int ret = read_candidates((*keys)[q]); if (ret != 0) return ret; - if (batch.ids.empty()) break; + } else { + // Sources without a scan cursor can use the existing random reads. + // Fix the scan length and copy only one batch of registered IDs. + const size_t registered_count = entity_->registered_key_count(); + std::vector registered_keys; + for (size_t offset = 0; offset < registered_count; + offset += batch_limit) { + entity_->get_registered_keys( + offset, std::min(batch_limit, registered_count - offset), + ®istered_keys); + int ret = read_candidates(registered_keys); + if (ret != 0) return ret; + } } } batch.clear(); diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index 29aafa090a..4fbe55aa2f 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -189,6 +190,23 @@ class FlatStreamerEntity { } } + size_t registered_key_count() const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + return id_key_vector_.size(); + } + + // Copy a bounded batch so source callbacks never run under the key lock. + void get_registered_keys(size_t offset, size_t count, + std::vector *keys) const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + offset = std::min(offset, id_key_vector_.size()); + count = std::min(count, id_key_vector_.size() - offset); + auto begin = id_key_vector_.begin() + offset; + keys->assign(begin, begin + count); + } + //! Set params void set_segment_size(uint32_t size) { meta_.segment_size = size; diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 438354a169..077a88974f 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -292,9 +292,10 @@ class ZVEC_CORE_API FlatIndex : public Index { // doc_id is the source ID with either use_id_map setting. int add_with_source(const VectorData &vector, uint32_t doc_id, const core::VectorSource &src) override; - // Full scans use the source's visible IDs intersected with registered IDs - // and the request filter; extra IDs are ignored. Candidate searches use - // get_vectors(). Bind the source, scan and filter to one stable snapshot. + // Full scans prefer the source iterator, intersecting its visible IDs with + // registered IDs and the request filter. Without an iterator, registered IDs + // are filtered and read through get_vectors(), as in candidate searches. + // Bind the source, scan and filter to one stable snapshot. // With fetch_vector, returned pointers borrow source memory; // keep that memory valid while using the results, as with external HNSW. // External Flat does not support group-by or binding a source for refinement. diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 07338c7961..e4a7759d30 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -61,9 +61,10 @@ class ZVEC_CORE_API VectorSource { virtual void get_vectors(const uint32_t *ids, uint32_t count, const void **out) const; - // Optional full scan. Each call creates an independent cursor. The caller - // binds the source, scan domain and visibility filter to the same stable - // snapshot for the whole request, including subsequent random reads. + // Optional full scan. Each call creates an independent cursor. Returning + // nullptr makes External Flat read registered IDs through get_vectors(). + // The caller binds the source, scan domain and visibility filter to the same + // stable snapshot for the whole request, including subsequent random reads. // A disk source can scan physical blocks behind this interface. virtual Iterator::Pointer create_iterator() const; }; diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 009f5f3377..3d47e77e02 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -37,6 +37,7 @@ class BatchSource : public core::VectorSource { public: std::map> rows; mutable std::atomic scans{0}; + mutable std::atomic batch_reads{0}; size_t page_rows{7}; std::set null_scan_ids; bool supports_scan{true}, fail_scan{false}, malformed_scan{false}; @@ -57,6 +58,13 @@ class BatchSource : public core::VectorSource { return it == rows.end() ? nullptr : it->second.data(); } + void get_vectors(const uint32_t *ids, uint32_t count, + const void **out) const override { + ++batch_reads; + EXPECT_LE(count, 1024u); + VectorSource::get_vectors(ids, count, out); + } + class Cursor : public Iterator { public: explicit Cursor(const BatchSource &source) : source_(source) { @@ -101,6 +109,22 @@ class BatchSource : public core::VectorSource { } }; +// An existing source that implements only the original random-read API. +class RandomAccessSource : public core::VectorSource { + public: + explicit RandomAccessSource(const BatchSource &source) : source_(source) {} + + const void *get_vector(uint32_t id) const override { + reads.push_back(id); + return source_.get_vector(id); + } + + mutable std::vector reads; + + private: + const BatchSource &source_; +}; + VectorData Vector(const std::vector &row) { return VectorData{DenseVector{row.data()}}; } @@ -258,6 +282,7 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { external->search_with_source(Vector(query), qp, source, &got)); SameResults(want, got); EXPECT_EQ(1, source.scans); + EXPECT_EQ(0, source.batch_reads); close_index(external); ASSERT_EQ(11, got.doc_list_.size()); // External vectors are returned directly in the input layout, without @@ -273,6 +298,83 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { } } +TEST_F(ExternalFlatTest, FullScanWithOnlyGetVector) { + BatchSource data(30); + RandomAccessSource source(data); + auto external = open_index(Param(true), test_path("external")); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(external.get(), data); + add_vectors(embedded.get(), data, false); + // Invisible rows must be excluded before calling the source. + const std::set invisible{3, 16, 42}; + for (auto id : invisible) data.rows.erase(id); + data.rows[99999] = std::vector(kDimension, 0.3f); + auto qp = FlatQueryParamBuilder().with_topk(10).build(); + qp->filter = std::make_shared(); + qp->filter->set([invisible](uint64_t id) { return invisible.count(id); }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(27, source.reads.size()); + for (auto id : source.reads) { + EXPECT_EQ(0, invisible.count(id)); + EXPECT_NE(99999, id); + } + source.reads.clear(); + qp->filter->set([](uint64_t) { return true; }); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); + EXPECT_TRUE(source.reads.empty()); + qp->filter.reset(); + EXPECT_NE(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); +} + +TEST_F(ExternalFlatTest, FallbackMatchesEmbeddedUsesBatchesAndReopens) { + int case_id = 0; + for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { + for (bool id_map : {false, true}) { + BatchSource source(2103); + source.supports_scan = false; + // Break ties in the generated rows when comparing full result lists. + for (auto &[id, row] : source.rows) row[0] += id * 0.0001f; + auto param = Param(true, metric, id_map); + auto suffix = std::to_string(case_id++); + auto path = test_path(suffix + "external"); + auto external = open_index(param, path); + auto embedded = + open_index(Param(false, metric, id_map), test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + source.rows[99999] = std::vector(kDimension, 0.3f); + auto qp = + FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + for (int reopened = 0; reopened < 2; ++reopened) { + source.batch_reads = 0; + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_GE(source.batch_reads.load(), 3u); + for (auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + EXPECT_NE(99999, doc.key()); + } + if (!reopened) { + ASSERT_EQ(0, external->flush()); + close_index(external); + external = open_index(param, path, false); + } + } + EXPECT_EQ(2, source.scans); + } + } +} + TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { BatchSource source; source.supports_scan = false; @@ -358,6 +460,7 @@ TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { source.malformed_scan = failure == 1; source.null_scan = failure == 2; source.supports_scan = failure != 3; + source.fail_read = failure == 3; result.doc_list_.emplace_back(999, 0); EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_TRUE(result.doc_list_.empty()); @@ -365,6 +468,7 @@ TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { } source.fail_scan = source.malformed_scan = source.null_scan = false; source.supports_scan = true; + source.fail_read = false; ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); source.fail_read = true; qp->fetch_vector = true; @@ -382,6 +486,7 @@ TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { BatchSource first(50), second(50, 2.0f); + second.supports_scan = false; auto index = open_index(Param(true), test_path("index")); add_vectors(index.get(), first); std::vector query(kDimension, 0.3f); @@ -407,6 +512,8 @@ TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { for (auto &thread : threads) thread.join(); EXPECT_EQ(61, first.scans); EXPECT_EQ(61, second.scans); + EXPECT_EQ(0, first.batch_reads); + EXPECT_EQ(61, second.batch_reads); } TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { @@ -476,6 +583,11 @@ TEST_F(ExternalFlatTest, BufferPoolReopen) { ASSERT_EQ(1, result.doc_list_.size()); EXPECT_EQ(3, result.doc_list_[0].key()); EXPECT_EQ(0, result.doc_list_[0].score()); + source.supports_scan = false; + ASSERT_EQ(0, index->search_with_source(vector, qp, source, &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); } TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { @@ -509,6 +621,13 @@ TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { EXPECT_EQ(3, context->result(0)[0].key()); EXPECT_EQ(16, context->result(1)[0].key()); EXPECT_EQ(2, source.scans); + source.supports_scan = false; + ASSERT_EQ(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + ASSERT_EQ(1, context->result(0).size()); + ASSERT_EQ(1, context->result(1).size()); + EXPECT_EQ(3, context->result(0)[0].key()); + EXPECT_EQ(16, context->result(1)[0].key()); + EXPECT_EQ(2, source.batch_reads); core::IndexQueryMeta wrong(core::IndexMeta::DT_FP32, kDimension - 1); EXPECT_NE(0, streamer->search_bf_impl(query.data(), wrong, 2, context)); ctx->set_group_params(1, 1); @@ -534,6 +653,9 @@ TEST_F(ExternalFlatTest, EmptyIndexAndEmptyCandidates) { ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_TRUE(result.doc_list_.empty()); source.supports_scan = false; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + EXPECT_EQ(0, source.batch_reads); qp->bf_pks = std::make_shared>(); ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_TRUE(result.doc_list_.empty()); From 04ae8eedd2b45e6ded69d4da5f32f8703cccd6c7 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 17:57:15 +0800 Subject: [PATCH 09/16] feat(flat): accept contiguous external vector batches --- src/core/algorithm/flat/flat_streamer.cc | 42 ++- .../algorithm/flat/flat_streamer_entity.h | 10 +- .../zvec/core/interface/vector_source.h | 36 ++- .../interface/external_flat_index_test.cc | 251 ++++++++++++++++++ 4 files changed, 324 insertions(+), 15 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 45031ab891..9f259b28cc 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "flat_streamer.h" +#include #include #include "flat_streamer_context.h" #include "flat_streamer_dumper.h" @@ -521,14 +522,35 @@ int FlatStreamer::search_external( VectorSource::Batch batch; if (!keys) cursor = source->create_iterator(); auto consume = [&]() -> int { - if (batch.ids.size() != batch.vectors.size() || - batch.ids.size() > batch_limit) + const auto *contiguous = batch.contiguous ? &*batch.contiguous : nullptr; + const size_t batch_size = batch.size(); + if (batch_size > batch_limit) return IndexError_ReadData; + if (contiguous) { + if (!batch.ids.empty() || !batch.vectors.empty()) + return IndexError_ReadData; + if (batch_size && + (!contiguous->data || contiguous->stride < meta_.element_size() || + batch_size - 1 > + (std::numeric_limits::max() - meta_.element_size()) / + contiguous->stride || + (!contiguous->ids && + (contiguous->first_id >= kInvalidNodeId || + batch_size > kInvalidNodeId - contiguous->first_id)))) + return IndexError_ReadData; + } else if (batch.ids.size() != batch.vectors.size()) { return IndexError_ReadData; + } + const auto *ids = contiguous ? contiguous->ids : batch.ids.data(); + const uint32_t first_id = contiguous ? contiguous->first_id : 0; scratch->vector_ptrs.clear(); scratch->vector_keys.clear(); - if (cursor) entity_->key_presence(batch.ids, &scratch->key_presence); - for (size_t i = 0; i < batch.ids.size(); ++i) { - const auto id = batch.ids[i]; + scratch->vector_ptrs.reserve(batch_size); + scratch->vector_keys.reserve(batch_size); + if (cursor) + entity_->key_presence(ids, batch_size, first_id, + &scratch->key_presence); + for (size_t i = 0; i < batch_size; ++i) { + const auto id = ids ? ids[i] : first_id + static_cast(i); if (id == kInvalidNodeId) return IndexError_ReadData; if (cursor && (!scratch->key_presence[i] || (ctx->filter().is_valid() && ctx->filter()(id)))) { @@ -536,8 +558,12 @@ int FlatStreamer::search_external( continue; } // Invisible or unregistered rows need not have a readable vector. - if (!batch.vectors[i]) return IndexError_ReadData; - scratch->vector_ptrs.push_back(batch.vectors[i]); + const void *vector = contiguous + ? static_cast(contiguous->data) + + i * contiguous->stride + : batch.vectors[i]; + if (!vector) return IndexError_ReadData; + scratch->vector_ptrs.push_back(vector); scratch->vector_keys.push_back(id); } const size_t size = scratch->vector_ptrs.size(); @@ -573,7 +599,7 @@ int FlatStreamer::search_external( if (ret != 0) return ret; ret = consume(); if (ret != 0) return ret; - if (batch.ids.empty()) break; + if (batch.empty()) break; } } else { std::vector ids; diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index 4fbe55aa2f..73ca575dee 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -180,13 +180,15 @@ class FlatStreamerEntity { return key_info_map_.find(key) != key_info_map_.end(); } - void key_presence(const std::vector &keys, + // A null keys pointer denotes consecutive IDs starting at first_key. + void key_presence(const uint32_t *keys, size_t count, uint32_t first_key, std::vector *present) const { - present->resize(keys.size()); + present->resize(count); ailego::ReadLock lock(*key_info_map_lock_); std::lock_guard guard(lock); - for (size_t i = 0; i < keys.size(); ++i) { - (*present)[i] = key_info_map_.find(keys[i]) != key_info_map_.end(); + for (size_t i = 0; i < count; ++i) { + const auto key = keys ? keys[i] : first_key + i; + (*present)[i] = key_info_map_.find(key) != key_info_map_.end(); } } diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index e4a7759d30..cd873f85bf 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -14,8 +14,10 @@ #pragma once +#include #include #include +#include #include #include @@ -24,6 +26,19 @@ namespace core { class ZVEC_CORE_API VectorSource { public: + // Borrowed row-major storage. Row i starts at data + i * stride bytes. + // stride must be at least the vector's byte length. Nonempty batches need + // valid backing storage; count includes rows excluded by the query filter. + struct ContiguousBatch { + const void *data{nullptr}; + uint32_t first_id{0}; + uint32_t count{0}; + size_t stride{0}; + // Optional borrowed array of count IDs. If null, row i has first_id + i. + // Otherwise first_id is ignored. IDs use the get_vector() namespace. + const uint32_t *ids{nullptr}; + }; + // Logical row-major vectors in the same ID namespace as get_vector(). // Pointers need not be contiguous. External Flat currently expects raw FP32. struct Batch { @@ -31,13 +46,26 @@ class ZVEC_CORE_API VectorSource { // Rows excluded by Flat's registered-ID set or filter may be null. // Every row selected for distance computation must have a valid pointer. std::vector vectors; - // Optional backing allocation or page pin. Data must remain valid until - // the batch is released, or (for scans) the next next_batch() call. + // Optional backing allocation or page pin. Vectors and any borrowed ID + // array must remain valid until the batch is released, or (for scans) + // the next next_batch() call. std::shared_ptr lease; + // Alternative to ids/vectors, which must both be empty when set. Flat + // adapts this view directly into its reusable distance-computation buffers. + std::optional contiguous; + + size_t size() const { + return contiguous ? contiguous->count : ids.size(); + } + + bool empty() const { + return size() == 0; + } void clear() { vectors.clear(); ids.clear(); + contiguous.reset(); lease.reset(); } }; @@ -47,7 +75,9 @@ class ZVEC_CORE_API VectorSource { using Pointer = std::unique_ptr; virtual ~Iterator() = default; // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an - // IndexError on failure. Enumerate every visible ID exactly once. + // IndexError on failure. Enumerate every visible ID exactly once. Either + // batch representation may include filtered rows; an entirely filtered + // batch still has a nonzero size and must not signal EOF. virtual int next_batch(uint32_t max_count, Batch *out) = 0; }; diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 3d47e77e02..434e556965 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -15,6 +15,8 @@ #include #include #include +#include +#include #include #include #include @@ -125,6 +127,107 @@ class RandomAccessSource : public core::VectorSource { const BatchSource &source_; }; +class ContiguousSource : public BatchSource { + public: + ContiguousSource(uint32_t count, bool explicit_ids, size_t padding = 0, + uint32_t first_id = 3) + : BatchSource(0), + explicit_ids_(explicit_ids), + stride_elements_(kDimension + padding) { + data_.resize(count * stride_elements_, -999.0f); + for (uint32_t i = 0; i < count; ++i) { + auto id = first_id + i * (explicit_ids ? 13 : 1); + row_ids_.push_back(id); + auto &row = rows[id]; + for (uint32_t d = 0; d < kDimension; ++d) { + row.push_back(((i * 19 + d * 7) % 101) / 101.0f + i * 0.001f); + } + std::copy(row.begin(), row.end(), data_.begin() + i * stride_elements_); + } + } + + const void *get_vector(uint32_t id) const override { + auto it = std::lower_bound(row_ids_.begin(), row_ids_.end(), id); + if (it == row_ids_.end() || *it != id) return nullptr; + return data_.data() + (it - row_ids_.begin()) * stride_elements_; + } + + class Cursor : public Iterator { + public: + explicit Cursor(const ContiguousSource &source) : source_(source) {} + ~Cursor() override { + EXPECT_TRUE(previous_.expired()); + } + + int next_batch(uint32_t max_count, Batch *out) override { + EXPECT_TRUE(previous_.expired()); + out->clear(); + if (source_.fail_scan && position_) return core::IndexError_ReadData; + const auto count = static_cast( + std::min({size_t{max_count}, source_.page_rows, + source_.row_ids_.size() - position_})); + // Exercise EOF expressed as an empty contiguous view too. + out->contiguous = ContiguousBatch{}; + if (!count) return 0; + ++source_.scan_batches; + auto &view = *out->contiguous; + view.data = source_.data_.data() + position_ * source_.stride_elements_; + view.first_id = source_.row_ids_[position_]; + view.count = count; + view.stride = source_.stride_elements_ * sizeof(float); + if (source_.explicit_ids_) view.ids = source_.row_ids_.data() + position_; + if (source_.lease_pages) { + auto page = std::make_shared(); + const auto *begin = static_cast(view.data); + page->data.assign(begin, begin + count * source_.stride_elements_); + page->ids.assign(source_.row_ids_.begin() + position_, + source_.row_ids_.begin() + position_ + count); + view.data = page->data.data(); + if (view.ids) view.ids = page->ids.data(); + out->lease = page; + previous_ = page; + } + if (source_.mixed_batches && batch_number_ % 2) { + for (uint32_t i = 0; i < count; ++i) { + out->ids.push_back(view.ids ? view.ids[i] : view.first_id + i); + out->vectors.push_back(static_cast(view.data) + + i * view.stride); + } + out->contiguous.reset(); + } + if (source_.mutate_batch) source_.mutate_batch(out); + position_ += count; + ++batch_number_; + return 0; + } + + private: + struct Page { + std::vector data; + std::vector ids; + }; + const ContiguousSource &source_; + size_t position_{0}; + size_t batch_number_{0}; + std::weak_ptr previous_; + }; + + Iterator::Pointer create_iterator() const override { + ++scans; + return std::make_unique(*this); + } + + mutable std::atomic scan_batches{0}; + bool lease_pages{false}, mixed_batches{false}; + std::function mutate_batch; + std::vector row_ids_; + + private: + bool explicit_ids_; + size_t stride_elements_; + std::vector data_; +}; + VectorData Vector(const std::vector &row) { return VectorData{DenseVector{row.data()}}; } @@ -332,6 +435,154 @@ TEST_F(ExternalFlatTest, FullScanWithOnlyGetVector) { EXPECT_TRUE(got.doc_list_.empty()); } +TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { + int case_id = 0; + for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { + for (bool id_map : {false, true}) { + for (bool explicit_ids : {false, true}) { + for (size_t padding : {0, 3}) { + ContiguousSource source(2103, explicit_ids, padding); + source.page_rows = 1024; + // This row is in the scan but is not registered in either index. + source.rows.erase(source.row_ids_.back()); + auto suffix = std::to_string(case_id++); + auto external = open_index(Param(true, metric, id_map), + test_path(suffix + "ext")); + auto embedded = open_index(Param(false, metric, id_map), + test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + auto qp = FlatQueryParamBuilder() + .with_topk(11) + .with_fetch_vector(true) + .build(); + qp->filter = std::make_shared(); + // The first physical batch is entirely deleted. It must not end + // the scan, and holes in later batches must retain the correct IDs. + qp->filter->set([last_deleted = source.row_ids_[1023]](uint64_t id) { + return id <= last_deleted || id % 7 == 0; + }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(3, source.scan_batches); + EXPECT_EQ(0, source.batch_reads); + for (const auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + EXPECT_NE(source.row_ids_.back(), doc.key()); + } + // Candidate lookup still uses random access, even with this source. + qp->bf_pks = std::make_shared>( + source.row_ids_.begin() + 1024, source.row_ids_.begin() + 1048); + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(1, source.scans); + EXPECT_EQ(1, source.batch_reads); + } + } + } + } +} + +TEST_F(ExternalFlatTest, ContiguousScanLeasesAndMixedBatches) { + ContiguousSource source(30, true, 3); + source.lease_pages = source.mixed_batches = true; + auto external = open_index(Param(true), test_path("external")); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + auto qp = + FlatQueryParamBuilder().with_topk(10).with_fetch_vector(true).build(); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(5, source.scan_batches); + for (const auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + } + qp->filter = std::make_shared(); + qp->filter->set([](uint64_t) { return true; }); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); + EXPECT_EQ(10, source.scan_batches); + qp->filter.reset(); + source.fail_scan = true; + EXPECT_NE(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); +} + +TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { + ContiguousSource source(10, false); + source.lease_pages = true; + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); + auto qp = FlatQueryParamBuilder().with_topk(3).build(); + std::vector query(kDimension, 0.3f); + SearchResult result; + for (int failure = 0; failure < 8; ++failure) { + SCOPED_TRACE(failure); + source.mutate_batch = [failure](core::VectorSource::Batch *batch) { + auto &view = *batch->contiguous; + switch (failure) { + case 0: + batch->ids.push_back(view.first_id); + break; + case 1: + batch->vectors.push_back(view.data); + break; + case 2: + view.data = nullptr; + break; + case 3: + view.stride = kDimension * sizeof(float) - 1; + break; + case 4: + view.stride = std::numeric_limits::max(); + break; + case 5: + view.count = 1025; + break; + case 6: + view.first_id = UINT32_MAX - 1; + break; + case 7: { + static const uint32_t kInvalidIds[7] = {UINT32_MAX}; + view.ids = kInvalidIds; + break; + } + } + }; + result.doc_list_.emplace_back(999, 0); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + } + source.mutate_batch = nullptr; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_EQ(3, result.doc_list_.size()); +} + +TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { + ContiguousSource source(1, false, 0, UINT32_MAX - 1); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + SearchResult result; + auto query = Vector(source.rows.begin()->second); + ASSERT_EQ(0, index->search_with_source(query, qp, source, &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(UINT32_MAX - 1, result.doc_list_[0].key()); + ContiguousSource empty(0, false); + ASSERT_EQ(0, index->search_with_source(query, qp, empty, &result)); + EXPECT_TRUE(result.doc_list_.empty()); +} + TEST_F(ExternalFlatTest, FallbackMatchesEmbeddedUsesBatchesAndReopens) { int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { From 115c21878c1b66239bb0dd83b9a1c8e658768b61 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 19:01:37 +0800 Subject: [PATCH 10/16] refactor(flat): simplify contiguous batches to consecutive IDs --- src/core/algorithm/flat/flat_streamer.cc | 14 +- .../zvec/core/interface/vector_source.h | 16 +- .../interface/external_flat_index_test.cc | 162 ++++++++---------- 3 files changed, 77 insertions(+), 115 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 9f259b28cc..7ba9a99240 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -13,7 +13,6 @@ // limitations under the License. #include "flat_streamer.h" -#include #include #include "flat_streamer_context.h" #include "flat_streamer_dumper.h" @@ -529,18 +528,13 @@ int FlatStreamer::search_external( if (!batch.ids.empty() || !batch.vectors.empty()) return IndexError_ReadData; if (batch_size && - (!contiguous->data || contiguous->stride < meta_.element_size() || - batch_size - 1 > - (std::numeric_limits::max() - meta_.element_size()) / - contiguous->stride || - (!contiguous->ids && - (contiguous->first_id >= kInvalidNodeId || - batch_size > kInvalidNodeId - contiguous->first_id)))) + (!contiguous->data || contiguous->first_id >= kInvalidNodeId || + batch_size > kInvalidNodeId - contiguous->first_id)) return IndexError_ReadData; } else if (batch.ids.size() != batch.vectors.size()) { return IndexError_ReadData; } - const auto *ids = contiguous ? contiguous->ids : batch.ids.data(); + const auto *ids = contiguous ? nullptr : batch.ids.data(); const uint32_t first_id = contiguous ? contiguous->first_id : 0; scratch->vector_ptrs.clear(); scratch->vector_keys.clear(); @@ -560,7 +554,7 @@ int FlatStreamer::search_external( // Invisible or unregistered rows need not have a readable vector. const void *vector = contiguous ? static_cast(contiguous->data) + - i * contiguous->stride + i * meta_.element_size() : batch.vectors[i]; if (!vector) return IndexError_ReadData; scratch->vector_ptrs.push_back(vector); diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index cd873f85bf..b5196c3b2e 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -26,17 +26,14 @@ namespace core { class ZVEC_CORE_API VectorSource { public: - // Borrowed row-major storage. Row i starts at data + i * stride bytes. - // stride must be at least the vector's byte length. Nonempty batches need - // valid backing storage; count includes rows excluded by the query filter. + // Borrowed, tightly packed row-major vectors with consecutive IDs. + // Row i has ID first_id + i and starts at data + i * element_size bytes, + // using the index's vector byte length. No padding or row compaction. + // Nonempty batches need valid backing storage; count includes filtered rows. struct ContiguousBatch { const void *data{nullptr}; uint32_t first_id{0}; uint32_t count{0}; - size_t stride{0}; - // Optional borrowed array of count IDs. If null, row i has first_id + i. - // Otherwise first_id is ignored. IDs use the get_vector() namespace. - const uint32_t *ids{nullptr}; }; // Logical row-major vectors in the same ID namespace as get_vector(). @@ -46,9 +43,8 @@ class ZVEC_CORE_API VectorSource { // Rows excluded by Flat's registered-ID set or filter may be null. // Every row selected for distance computation must have a valid pointer. std::vector vectors; - // Optional backing allocation or page pin. Vectors and any borrowed ID - // array must remain valid until the batch is released, or (for scans) - // the next next_batch() call. + // Optional backing allocation or page pin. Data must remain valid until + // the batch is released, or (for scans) the next next_batch() call. std::shared_ptr lease; // Alternative to ids/vectors, which must both be empty when set. Flat // adapts this view directly into its reusable distance-computation buffers. diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 434e556965..5dca74e2cb 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -129,27 +128,23 @@ class RandomAccessSource : public core::VectorSource { class ContiguousSource : public BatchSource { public: - ContiguousSource(uint32_t count, bool explicit_ids, size_t padding = 0, - uint32_t first_id = 3) - : BatchSource(0), - explicit_ids_(explicit_ids), - stride_elements_(kDimension + padding) { - data_.resize(count * stride_elements_, -999.0f); + explicit ContiguousSource(uint32_t count, uint32_t first_id = 3) + : BatchSource(0), first_id_(first_id) { + data_.resize(count * kDimension); for (uint32_t i = 0; i < count; ++i) { - auto id = first_id + i * (explicit_ids ? 13 : 1); - row_ids_.push_back(id); + auto id = first_id + i; auto &row = rows[id]; for (uint32_t d = 0; d < kDimension; ++d) { row.push_back(((i * 19 + d * 7) % 101) / 101.0f + i * 0.001f); } - std::copy(row.begin(), row.end(), data_.begin() + i * stride_elements_); + std::copy(row.begin(), row.end(), data_.begin() + i * kDimension); } } const void *get_vector(uint32_t id) const override { - auto it = std::lower_bound(row_ids_.begin(), row_ids_.end(), id); - if (it == row_ids_.end() || *it != id) return nullptr; - return data_.data() + (it - row_ids_.begin()) * stride_elements_; + if (id < first_id_ || id - first_id_ >= data_.size() / kDimension) + return nullptr; + return data_.data() + (id - first_id_) * kDimension; } class Cursor : public Iterator { @@ -165,33 +160,28 @@ class ContiguousSource : public BatchSource { if (source_.fail_scan && position_) return core::IndexError_ReadData; const auto count = static_cast( std::min({size_t{max_count}, source_.page_rows, - source_.row_ids_.size() - position_})); + source_.data_.size() / kDimension - position_})); // Exercise EOF expressed as an empty contiguous view too. out->contiguous = ContiguousBatch{}; if (!count) return 0; ++source_.scan_batches; auto &view = *out->contiguous; - view.data = source_.data_.data() + position_ * source_.stride_elements_; - view.first_id = source_.row_ids_[position_]; + view.data = source_.data_.data() + position_ * kDimension; + view.first_id = source_.first_id_ + static_cast(position_); view.count = count; - view.stride = source_.stride_elements_ * sizeof(float); - if (source_.explicit_ids_) view.ids = source_.row_ids_.data() + position_; if (source_.lease_pages) { - auto page = std::make_shared(); const auto *begin = static_cast(view.data); - page->data.assign(begin, begin + count * source_.stride_elements_); - page->ids.assign(source_.row_ids_.begin() + position_, - source_.row_ids_.begin() + position_ + count); - view.data = page->data.data(); - if (view.ids) view.ids = page->ids.data(); + auto page = std::make_shared>( + begin, begin + count * kDimension); + view.data = page->data(); out->lease = page; previous_ = page; } if (source_.mixed_batches && batch_number_ % 2) { for (uint32_t i = 0; i < count; ++i) { - out->ids.push_back(view.ids ? view.ids[i] : view.first_id + i); - out->vectors.push_back(static_cast(view.data) + - i * view.stride); + out->ids.push_back(view.first_id + i); + out->vectors.push_back(static_cast(view.data) + + i * kDimension); } out->contiguous.reset(); } @@ -202,10 +192,6 @@ class ContiguousSource : public BatchSource { } private: - struct Page { - std::vector data; - std::vector ids; - }; const ContiguousSource &source_; size_t position_{0}; size_t batch_number_{0}; @@ -220,11 +206,9 @@ class ContiguousSource : public BatchSource { mutable std::atomic scan_batches{0}; bool lease_pages{false}, mixed_batches{false}; std::function mutate_batch; - std::vector row_ids_; private: - bool explicit_ids_; - size_t stride_elements_; + uint32_t first_id_; std::vector data_; }; @@ -436,61 +420,57 @@ TEST_F(ExternalFlatTest, FullScanWithOnlyGetVector) { } TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { + constexpr uint32_t kFirstId = 3; + constexpr uint32_t kRowCount = 2103; int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { for (bool id_map : {false, true}) { - for (bool explicit_ids : {false, true}) { - for (size_t padding : {0, 3}) { - ContiguousSource source(2103, explicit_ids, padding); - source.page_rows = 1024; - // This row is in the scan but is not registered in either index. - source.rows.erase(source.row_ids_.back()); - auto suffix = std::to_string(case_id++); - auto external = open_index(Param(true, metric, id_map), - test_path(suffix + "ext")); - auto embedded = open_index(Param(false, metric, id_map), - test_path(suffix + "base")); - add_vectors(external.get(), source); - add_vectors(embedded.get(), source, false); - auto qp = FlatQueryParamBuilder() - .with_topk(11) - .with_fetch_vector(true) - .build(); - qp->filter = std::make_shared(); - // The first physical batch is entirely deleted. It must not end - // the scan, and holes in later batches must retain the correct IDs. - qp->filter->set([last_deleted = source.row_ids_[1023]](uint64_t id) { - return id <= last_deleted || id % 7 == 0; - }); - std::vector query(kDimension, 0.3f); - SearchResult want, got; - ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); - ASSERT_EQ( - 0, external->search_with_source(Vector(query), qp, source, &got)); - SameResults(want, got); - EXPECT_EQ(3, source.scan_batches); - EXPECT_EQ(0, source.batch_reads); - for (const auto &doc : got.doc_list_) { - EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); - EXPECT_NE(source.row_ids_.back(), doc.key()); - } - // Candidate lookup still uses random access, even with this source. - qp->bf_pks = std::make_shared>( - source.row_ids_.begin() + 1024, source.row_ids_.begin() + 1048); - ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); - ASSERT_EQ( - 0, external->search_with_source(Vector(query), qp, source, &got)); - SameResults(want, got); - EXPECT_EQ(1, source.scans); - EXPECT_EQ(1, source.batch_reads); - } + ContiguousSource source(kRowCount, kFirstId); + source.page_rows = 1024; + // This row is in the scan but is not registered in either index. + source.rows.erase(kFirstId + kRowCount - 1); + auto suffix = std::to_string(case_id++); + auto external = + open_index(Param(true, metric, id_map), test_path(suffix + "ext")); + auto embedded = + open_index(Param(false, metric, id_map), test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + auto qp = + FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); + qp->filter = std::make_shared(); + // The first physical batch is entirely deleted. It must not end + // the scan, and holes in later batches must retain the correct IDs. + qp->filter->set( + [](uint64_t id) { return id <= kFirstId + 1023 || id % 7 == 0; }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(3, source.scan_batches); + EXPECT_EQ(0, source.batch_reads); + for (const auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + EXPECT_NE(kFirstId + kRowCount - 1, doc.key()); } + // Candidate lookup still uses random access, even with this source. + qp->bf_pks = std::make_shared>(); + for (uint32_t i = 1024; i < 1048; ++i) + qp->bf_pks->push_back(kFirstId + i); + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(1, source.scans); + EXPECT_EQ(1, source.batch_reads); } } } TEST_F(ExternalFlatTest, ContiguousScanLeasesAndMixedBatches) { - ContiguousSource source(30, true, 3); + ContiguousSource source(30); source.lease_pages = source.mixed_batches = true; auto external = open_index(Param(true), test_path("external")); auto embedded = open_index(Param(false), test_path("embedded")); @@ -519,14 +499,14 @@ TEST_F(ExternalFlatTest, ContiguousScanLeasesAndMixedBatches) { } TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { - ContiguousSource source(10, false); + ContiguousSource source(10); source.lease_pages = true; auto index = open_index(Param(true), test_path("index")); add_vectors(index.get(), source); auto qp = FlatQueryParamBuilder().with_topk(3).build(); std::vector query(kDimension, 0.3f); SearchResult result; - for (int failure = 0; failure < 8; ++failure) { + for (int failure = 0; failure < 6; ++failure) { SCOPED_TRACE(failure); source.mutate_batch = [failure](core::VectorSource::Batch *batch) { auto &view = *batch->contiguous; @@ -541,22 +521,14 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { view.data = nullptr; break; case 3: - view.stride = kDimension * sizeof(float) - 1; - break; - case 4: - view.stride = std::numeric_limits::max(); - break; - case 5: view.count = 1025; break; - case 6: + case 4: view.first_id = UINT32_MAX - 1; break; - case 7: { - static const uint32_t kInvalidIds[7] = {UINT32_MAX}; - view.ids = kInvalidIds; + case 5: + view.first_id = UINT32_MAX; break; - } } }; result.doc_list_.emplace_back(999, 0); @@ -569,7 +541,7 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { } TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { - ContiguousSource source(1, false, 0, UINT32_MAX - 1); + ContiguousSource source(1, UINT32_MAX - 1); auto index = open_index(Param(true), test_path("index")); add_vectors(index.get(), source); auto qp = FlatQueryParamBuilder().with_topk(1).build(); @@ -578,7 +550,7 @@ TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { ASSERT_EQ(0, index->search_with_source(query, qp, source, &result)); ASSERT_EQ(1, result.doc_list_.size()); EXPECT_EQ(UINT32_MAX - 1, result.doc_list_[0].key()); - ContiguousSource empty(0, false); + ContiguousSource empty(0); ASSERT_EQ(0, index->search_with_source(query, qp, empty, &result)); EXPECT_TRUE(result.doc_list_.empty()); } From 320f53eb69882c98f5533bdc627184bd47e3a72c Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 19:03:33 +0800 Subject: [PATCH 11/16] docs(flat): trim contiguous batch comments --- src/include/zvec/core/interface/vector_source.h | 12 +++--------- tests/core/interface/external_flat_index_test.cc | 7 ++----- 2 files changed, 5 insertions(+), 14 deletions(-) diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index b5196c3b2e..a97ceb79e6 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -26,10 +26,7 @@ namespace core { class ZVEC_CORE_API VectorSource { public: - // Borrowed, tightly packed row-major vectors with consecutive IDs. - // Row i has ID first_id + i and starts at data + i * element_size bytes, - // using the index's vector byte length. No padding or row compaction. - // Nonempty batches need valid backing storage; count includes filtered rows. + // Packed row-major data; IDs first_id + i, count includes filtered rows. struct ContiguousBatch { const void *data{nullptr}; uint32_t first_id{0}; @@ -46,8 +43,7 @@ class ZVEC_CORE_API VectorSource { // Optional backing allocation or page pin. Data must remain valid until // the batch is released, or (for scans) the next next_batch() call. std::shared_ptr lease; - // Alternative to ids/vectors, which must both be empty when set. Flat - // adapts this view directly into its reusable distance-computation buffers. + // When set, ids and vectors must be empty. std::optional contiguous; size_t size() const { @@ -71,9 +67,7 @@ class ZVEC_CORE_API VectorSource { using Pointer = std::unique_ptr; virtual ~Iterator() = default; // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an - // IndexError on failure. Enumerate every visible ID exactly once. Either - // batch representation may include filtered rows; an entirely filtered - // batch still has a nonzero size and must not signal EOF. + // IndexError on failure. Enumerate every visible ID exactly once. virtual int next_batch(uint32_t max_count, Batch *out) = 0; }; diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 5dca74e2cb..65ce23999c 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -161,7 +161,6 @@ class ContiguousSource : public BatchSource { const auto count = static_cast( std::min({size_t{max_count}, source_.page_rows, source_.data_.size() / kDimension - position_})); - // Exercise EOF expressed as an empty contiguous view too. out->contiguous = ContiguousBatch{}; if (!count) return 0; ++source_.scan_batches; @@ -427,7 +426,7 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { for (bool id_map : {false, true}) { ContiguousSource source(kRowCount, kFirstId); source.page_rows = 1024; - // This row is in the scan but is not registered in either index. + // Keep an unregistered row in the scan. source.rows.erase(kFirstId + kRowCount - 1); auto suffix = std::to_string(case_id++); auto external = @@ -439,8 +438,7 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { auto qp = FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); qp->filter = std::make_shared(); - // The first physical batch is entirely deleted. It must not end - // the scan, and holes in later batches must retain the correct IDs. + // A fully filtered batch must not end the scan. qp->filter->set( [](uint64_t id) { return id <= kFirstId + 1023 || id % 7 == 0; }); std::vector query(kDimension, 0.3f); @@ -455,7 +453,6 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); EXPECT_NE(kFirstId + kRowCount - 1, doc.key()); } - // Candidate lookup still uses random access, even with this source. qp->bf_pks = std::make_shared>(); for (uint32_t i = 1024; i < 1048; ++i) qp->bf_pks->push_back(kFirstId + i); From c8602fccd622879998d186c1f82a1ab638b739e8 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 19:09:03 +0800 Subject: [PATCH 12/16] refactor(flat): separate pointer and contiguous scan batches --- src/core/algorithm/flat/flat_streamer.cc | 54 +++-- .../zvec/core/interface/vector_source.h | 46 ++--- .../interface/external_flat_index_test.cc | 195 ++++++++++-------- 3 files changed, 157 insertions(+), 138 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 7ba9a99240..58226514bc 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "flat_streamer.h" +#include #include #include "flat_streamer_context.h" #include "flat_streamer_dumper.h" @@ -520,21 +521,36 @@ int FlatStreamer::search_external( VectorSource::Iterator::Pointer cursor; VectorSource::Batch batch; if (!keys) cursor = source->create_iterator(); + auto clear_batch = [&]() { + if (auto *pointers = std::get_if(&batch)) { + pointers->ids.clear(); + pointers->vectors.clear(); + pointers->lease.reset(); + } else { + std::get(batch) = {}; + } + }; auto consume = [&]() -> int { - const auto *contiguous = batch.contiguous ? &*batch.contiguous : nullptr; - const size_t batch_size = batch.size(); + const auto *contiguous = + std::get_if(&batch); + const auto *pointers = std::get_if(&batch); + if (!contiguous && !pointers) return IndexError_ReadData; + const size_t batch_size = + contiguous ? contiguous->count : pointers->ids.size(); if (batch_size > batch_limit) return IndexError_ReadData; if (contiguous) { - if (!batch.ids.empty() || !batch.vectors.empty()) - return IndexError_ReadData; if (batch_size && - (!contiguous->data || contiguous->first_id >= kInvalidNodeId || + (!contiguous->data || contiguous->stride < meta_.element_size() || + batch_size - 1 > + (std::numeric_limits::max() - meta_.element_size()) / + contiguous->stride || + contiguous->first_id >= kInvalidNodeId || batch_size > kInvalidNodeId - contiguous->first_id)) return IndexError_ReadData; - } else if (batch.ids.size() != batch.vectors.size()) { + } else if (pointers->ids.size() != pointers->vectors.size()) { return IndexError_ReadData; } - const auto *ids = contiguous ? nullptr : batch.ids.data(); + const auto *ids = pointers ? pointers->ids.data() : nullptr; const uint32_t first_id = contiguous ? contiguous->first_id : 0; scratch->vector_ptrs.clear(); scratch->vector_keys.clear(); @@ -554,8 +570,8 @@ int FlatStreamer::search_external( // Invisible or unregistered rows need not have a readable vector. const void *vector = contiguous ? static_cast(contiguous->data) + - i * meta_.element_size() - : batch.vectors[i]; + i * contiguous->stride + : pointers->vectors[i]; if (!vector) return IndexError_ReadData; scratch->vector_ptrs.push_back(vector); scratch->vector_keys.push_back(id); @@ -588,22 +604,28 @@ int FlatStreamer::search_external( if (cursor) { while (true) { - batch.clear(); + clear_batch(); int ret = cursor->next_batch(batch_limit, &batch); if (ret != 0) return ret; ret = consume(); if (ret != 0) return ret; - if (batch.empty()) break; + const auto *contiguous = + std::get_if(&batch); + if (contiguous + ? contiguous->count == 0 + : std::get(batch).ids.empty()) + break; } } else { std::vector ids; ids.reserve(batch_limit); auto read_batch = [&]() -> int { if (ids.empty()) return 0; - batch.clear(); - batch.ids = ids; - batch.vectors.resize(ids.size()); - source->get_vectors(ids.data(), ids.size(), batch.vectors.data()); + clear_batch(); + auto &pointers = std::get(batch); + pointers.ids = ids; + pointers.vectors.resize(ids.size()); + source->get_vectors(ids.data(), ids.size(), pointers.vectors.data()); int ret = consume(); ids.clear(); return ret; @@ -642,7 +664,7 @@ int FlatStreamer::search_external( } } } - batch.clear(); + clear_batch(); ctx->take_topk_result(q); if (ctx->fetch_vector()) { // Match HNSW's borrowed random-access result contract. Scan pages have diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index a97ceb79e6..9f3cdb6fa6 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -17,7 +17,7 @@ #include #include #include -#include +#include #include #include @@ -26,42 +26,26 @@ namespace core { class ZVEC_CORE_API VectorSource { public: - // Packed row-major data; IDs first_id + i, count includes filtered rows. - struct ContiguousBatch { - const void *data{nullptr}; - uint32_t first_id{0}; - uint32_t count{0}; - }; - - // Logical row-major vectors in the same ID namespace as get_vector(). - // Pointers need not be contiguous. External Flat currently expects raw FP32. - struct Batch { + // Both batch types borrow raw FP32 rows in get_vector()'s ID space. + // Keep memory valid through consumption; lease may pin its backing storage. + struct PointerBatch { std::vector ids; - // Rows excluded by Flat's registered-ID set or filter may be null. - // Every row selected for distance computation must have a valid pointer. + // Only filtered or unregistered rows may have null pointers. std::vector vectors; - // Optional backing allocation or page pin. Data must remain valid until - // the batch is released, or (for scans) the next next_batch() call. std::shared_ptr lease; - // When set, ids and vectors must be empty. - std::optional contiguous; - - size_t size() const { - return contiguous ? contiguous->count : ids.size(); - } - - bool empty() const { - return size() == 0; - } + }; - void clear() { - vectors.clear(); - ids.clear(); - contiguous.reset(); - lease.reset(); - } + // Row i has ID first_id + i; stride is in bytes, at least one vector's size. + struct ContiguousBatch { + const void *data{nullptr}; + uint32_t first_id{0}; + uint32_t count{0}; // Includes filtered rows. + size_t stride{0}; + std::shared_ptr lease; }; + using Batch = std::variant; + class Iterator { public: using Pointer = std::unique_ptr; diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 65ce23999c..7dc1720151 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -75,7 +76,7 @@ class BatchSource : public core::VectorSource { } int next_batch(uint32_t max_count, Batch *out) override { EXPECT_TRUE(previous_.expired()); - out->clear(); + auto &batch = out->emplace(); if (source_.fail_scan && position_) return core::IndexError_ReadData; const size_t count = std::min( {size_t{max_count}, source_.page_rows, ids_.size() - position_}); @@ -85,15 +86,15 @@ class BatchSource : public core::VectorSource { auto id = ids_[position_++]; std::memcpy(page->data() + i * kDimension, source_.rows.at(id).data(), kDimension * sizeof(float)); - out->ids.push_back(id); - out->vectors.push_back(source_.null_scan || - source_.null_scan_ids.count(id) - ? nullptr - : page->data() + i * kDimension); + batch.ids.push_back(id); + batch.vectors.push_back(source_.null_scan || + source_.null_scan_ids.count(id) + ? nullptr + : page->data() + i * kDimension); } - out->lease = page; + batch.lease = page; previous_ = page; - if (source_.malformed_scan) out->vectors.pop_back(); + if (source_.malformed_scan) batch.vectors.pop_back(); return 0; } @@ -128,23 +129,26 @@ class RandomAccessSource : public core::VectorSource { class ContiguousSource : public BatchSource { public: - explicit ContiguousSource(uint32_t count, uint32_t first_id = 3) - : BatchSource(0), first_id_(first_id) { - data_.resize(count * kDimension); + explicit ContiguousSource(uint32_t count, uint32_t first_id = 3, + size_t padding = 0) + : BatchSource(0), + first_id_(first_id), + stride_elements_(kDimension + padding) { + data_.resize(count * stride_elements_, -999.0f); for (uint32_t i = 0; i < count; ++i) { auto id = first_id + i; auto &row = rows[id]; for (uint32_t d = 0; d < kDimension; ++d) { row.push_back(((i * 19 + d * 7) % 101) / 101.0f + i * 0.001f); } - std::copy(row.begin(), row.end(), data_.begin() + i * kDimension); + std::copy(row.begin(), row.end(), data_.begin() + i * stride_elements_); } } const void *get_vector(uint32_t id) const override { - if (id < first_id_ || id - first_id_ >= data_.size() / kDimension) + if (id < first_id_ || id - first_id_ >= data_.size() / stride_elements_) return nullptr; - return data_.data() + (id - first_id_) * kDimension; + return data_.data() + (id - first_id_) * stride_elements_; } class Cursor : public Iterator { @@ -156,35 +160,36 @@ class ContiguousSource : public BatchSource { int next_batch(uint32_t max_count, Batch *out) override { EXPECT_TRUE(previous_.expired()); - out->clear(); + auto &view = out->emplace(); if (source_.fail_scan && position_) return core::IndexError_ReadData; - const auto count = static_cast( - std::min({size_t{max_count}, source_.page_rows, - source_.data_.size() / kDimension - position_})); - out->contiguous = ContiguousBatch{}; + const auto count = static_cast(std::min( + {size_t{max_count}, source_.page_rows, + source_.data_.size() / source_.stride_elements_ - position_})); if (!count) return 0; ++source_.scan_batches; - auto &view = *out->contiguous; - view.data = source_.data_.data() + position_ * kDimension; + view.data = source_.data_.data() + position_ * source_.stride_elements_; view.first_id = source_.first_id_ + static_cast(position_); view.count = count; + view.stride = source_.stride_elements_ * sizeof(float); if (source_.lease_pages) { const auto *begin = static_cast(view.data); auto page = std::make_shared>( - begin, begin + count * kDimension); + begin, begin + count * source_.stride_elements_); view.data = page->data(); - out->lease = page; + view.lease = page; previous_ = page; } + if (source_.mutate_batch) source_.mutate_batch(&view); if (source_.mixed_batches && batch_number_ % 2) { + PointerBatch pointers; for (uint32_t i = 0; i < count; ++i) { - out->ids.push_back(view.first_id + i); - out->vectors.push_back(static_cast(view.data) + - i * kDimension); + pointers.ids.push_back(view.first_id + i); + pointers.vectors.push_back(static_cast(view.data) + + i * view.stride); } - out->contiguous.reset(); + pointers.lease = std::move(view.lease); + *out = std::move(pointers); } - if (source_.mutate_batch) source_.mutate_batch(out); position_ += count; ++batch_number_; return 0; @@ -204,10 +209,11 @@ class ContiguousSource : public BatchSource { mutable std::atomic scan_batches{0}; bool lease_pages{false}, mixed_batches{false}; - std::function mutate_batch; + std::function mutate_batch; private: uint32_t first_id_; + size_t stride_elements_; std::vector data_; }; @@ -424,50 +430,54 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { for (bool id_map : {false, true}) { - ContiguousSource source(kRowCount, kFirstId); - source.page_rows = 1024; - // Keep an unregistered row in the scan. - source.rows.erase(kFirstId + kRowCount - 1); - auto suffix = std::to_string(case_id++); - auto external = - open_index(Param(true, metric, id_map), test_path(suffix + "ext")); - auto embedded = - open_index(Param(false, metric, id_map), test_path(suffix + "base")); - add_vectors(external.get(), source); - add_vectors(embedded.get(), source, false); - auto qp = - FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); - qp->filter = std::make_shared(); - // A fully filtered batch must not end the scan. - qp->filter->set( - [](uint64_t id) { return id <= kFirstId + 1023 || id % 7 == 0; }); - std::vector query(kDimension, 0.3f); - SearchResult want, got; - ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); - ASSERT_EQ(0, - external->search_with_source(Vector(query), qp, source, &got)); - SameResults(want, got); - EXPECT_EQ(3, source.scan_batches); - EXPECT_EQ(0, source.batch_reads); - for (const auto &doc : got.doc_list_) { - EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); - EXPECT_NE(kFirstId + kRowCount - 1, doc.key()); + for (size_t padding : {0, 3}) { + ContiguousSource source(kRowCount, kFirstId, padding); + source.page_rows = 1024; + // Keep an unregistered row in the scan. + source.rows.erase(kFirstId + kRowCount - 1); + auto suffix = std::to_string(case_id++); + auto external = + open_index(Param(true, metric, id_map), test_path(suffix + "ext")); + auto embedded = open_index(Param(false, metric, id_map), + test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + auto qp = FlatQueryParamBuilder() + .with_topk(11) + .with_fetch_vector(true) + .build(); + qp->filter = std::make_shared(); + // A fully filtered batch must not end the scan. + qp->filter->set( + [](uint64_t id) { return id <= kFirstId + 1023 || id % 7 == 0; }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(3, source.scan_batches); + EXPECT_EQ(0, source.batch_reads); + for (const auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + EXPECT_NE(kFirstId + kRowCount - 1, doc.key()); + } + qp->bf_pks = std::make_shared>(); + for (uint32_t i = 1024; i < 1048; ++i) + qp->bf_pks->push_back(kFirstId + i); + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(1, source.scans); + EXPECT_EQ(1, source.batch_reads); } - qp->bf_pks = std::make_shared>(); - for (uint32_t i = 1024; i < 1048; ++i) - qp->bf_pks->push_back(kFirstId + i); - ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); - ASSERT_EQ(0, - external->search_with_source(Vector(query), qp, source, &got)); - SameResults(want, got); - EXPECT_EQ(1, source.scans); - EXPECT_EQ(1, source.batch_reads); } } } TEST_F(ExternalFlatTest, ContiguousScanLeasesAndMixedBatches) { - ContiguousSource source(30); + ContiguousSource source(30, 3, 3); source.lease_pages = source.mixed_batches = true; auto external = open_index(Param(true), test_path("external")); auto embedded = open_index(Param(false), test_path("embedded")); @@ -503,31 +513,34 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { auto qp = FlatQueryParamBuilder().with_topk(3).build(); std::vector query(kDimension, 0.3f); SearchResult result; - for (int failure = 0; failure < 6; ++failure) { + for (int failure = 0; failure < 7; ++failure) { SCOPED_TRACE(failure); - source.mutate_batch = [failure](core::VectorSource::Batch *batch) { - auto &view = *batch->contiguous; - switch (failure) { - case 0: - batch->ids.push_back(view.first_id); - break; - case 1: - batch->vectors.push_back(view.data); - break; - case 2: - view.data = nullptr; - break; - case 3: - view.count = 1025; - break; - case 4: - view.first_id = UINT32_MAX - 1; - break; - case 5: - view.first_id = UINT32_MAX; - break; - } - }; + source.mutate_batch = + [failure](core::VectorSource::ContiguousBatch *batch) { + switch (failure) { + case 0: + batch->data = nullptr; + break; + case 1: + batch->stride = 0; + break; + case 2: + batch->stride = kDimension * sizeof(float) - 1; + break; + case 3: + batch->stride = std::numeric_limits::max(); + break; + case 4: + batch->count = 1025; + break; + case 5: + batch->first_id = UINT32_MAX - 1; + break; + case 6: + batch->first_id = UINT32_MAX; + break; + } + }; result.doc_list_.emplace_back(999, 0); EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_TRUE(result.doc_list_.empty()); From ba9f53e32921d43b04c8d201f900a19f13e9e6ac Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 20:38:54 +0800 Subject: [PATCH 13/16] refactor(core): separate vector scan types and add contiguous source --- src/core/algorithm/flat/flat_streamer.cc | 23 +-- src/core/interface/vector_source.cc | 59 +++++- .../zvec/core/interface/vector_source.h | 72 ++++--- .../interface/external_flat_index_test.cc | 195 +++++++++++++----- 4 files changed, 256 insertions(+), 93 deletions(-) diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 58226514bc..1d7bc62178 100644 --- a/src/core/algorithm/flat/flat_streamer.cc +++ b/src/core/algorithm/flat/flat_streamer.cc @@ -518,22 +518,21 @@ int FlatStreamer::search_external( auto *heap = ctx->result_heap(); heap->clear(); // Destroy the outstanding lease before its cursor, including on errors. - VectorSource::Iterator::Pointer cursor; - VectorSource::Batch batch; + VectorIterator::Pointer cursor; + VectorBatch batch; if (!keys) cursor = source->create_iterator(); auto clear_batch = [&]() { - if (auto *pointers = std::get_if(&batch)) { + if (auto *pointers = std::get_if(&batch)) { pointers->ids.clear(); pointers->vectors.clear(); pointers->lease.reset(); } else { - std::get(batch) = {}; + std::get(batch) = {}; } }; auto consume = [&]() -> int { - const auto *contiguous = - std::get_if(&batch); - const auto *pointers = std::get_if(&batch); + const auto *contiguous = std::get_if(&batch); + const auto *pointers = std::get_if(&batch); if (!contiguous && !pointers) return IndexError_ReadData; const size_t batch_size = contiguous ? contiguous->count : pointers->ids.size(); @@ -609,11 +608,9 @@ int FlatStreamer::search_external( if (ret != 0) return ret; ret = consume(); if (ret != 0) return ret; - const auto *contiguous = - std::get_if(&batch); - if (contiguous - ? contiguous->count == 0 - : std::get(batch).ids.empty()) + const auto *contiguous = std::get_if(&batch); + if (contiguous ? contiguous->count == 0 + : std::get(batch).ids.empty()) break; } } else { @@ -622,7 +619,7 @@ int FlatStreamer::search_external( auto read_batch = [&]() -> int { if (ids.empty()) return 0; clear_batch(); - auto &pointers = std::get(batch); + auto &pointers = std::get(batch); pointers.ids = ids; pointers.vectors.resize(ids.size()); source->get_vectors(ids.data(), ids.size(), pointers.vectors.data()); diff --git a/src/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index 98307cb76b..205bc71264 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -12,16 +12,59 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include +#include #include namespace zvec { namespace core { +namespace { + +bool ValidContiguousView(const ContiguousBatch &view) { + if (!view.count) return true; + constexpr auto kInvalidId = std::numeric_limits::max(); + if (!view.data || !view.stride || view.first_id >= kInvalidId || + view.count > kInvalidId - view.first_id) + return false; + const auto max_span = std::numeric_limits::max() - + reinterpret_cast(view.data); + return view.count <= max_span / view.stride; +} + +class ContiguousVectorIterator final : public VectorIterator { + public: + ContiguousVectorIterator(ContiguousBatch view, bool valid) + : view_(std::move(view)), valid_(valid) {} + + int next_batch(uint32_t max_count, VectorBatch *out) override { + if (!out || !max_count) return IndexError_InvalidArgument; + auto &batch = out->emplace(); + if (!valid_) return IndexError_InvalidArgument; + if (offset_ == view_.count) return 0; + const auto count = std::min(max_count, view_.count - offset_); + batch = { + static_cast(view_.data) + size_t(offset_) * view_.stride, + view_.first_id + offset_, count, view_.stride, view_.lease}; + offset_ += count; + return 0; + } + + private: + const ContiguousBatch view_; + const bool valid_; + uint32_t offset_{0}; +}; + +} // namespace + VectorSource::VectorSource() = default; VectorSource::~VectorSource() = default; -VectorSource::Iterator::Pointer VectorSource::create_iterator() const { +VectorIterator::Pointer VectorSource::create_iterator() const { return nullptr; } @@ -32,5 +75,19 @@ void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, } } +ContiguousVectorSource::ContiguousVectorSource(ContiguousBatch view) + : view_(std::move(view)), valid_(ValidContiguousView(view_)) {} + +const void *ContiguousVectorSource::get_vector(uint32_t id) const { + if (!valid_ || id < view_.first_id || id - view_.first_id >= view_.count) + return nullptr; + return static_cast(view_.data) + + size_t(id - view_.first_id) * view_.stride; +} + +VectorIterator::Pointer ContiguousVectorSource::create_iterator() const { + return std::make_unique(view_, valid_); +} + } // namespace core } // namespace zvec diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 9f3cdb6fa6..0daa6a906b 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -24,37 +24,36 @@ namespace zvec { namespace core { -class ZVEC_CORE_API VectorSource { +// Borrowed raw FP32 rows in get_vector()'s ID space. Leases may pin storage. +struct PointerBatch { + std::vector ids; + // Only filtered or unregistered rows may have null pointers. + std::vector vectors; + std::shared_ptr lease; +}; + +// Row i has ID first_id + i; stride is in bytes, at least one vector's size. +struct ContiguousBatch { + const void *data{nullptr}; + uint32_t first_id{0}; + uint32_t count{0}; // Includes filtered rows. + size_t stride{0}; + std::shared_ptr lease; +}; + +using VectorBatch = std::variant; + +class ZVEC_CORE_API VectorIterator { public: - // Both batch types borrow raw FP32 rows in get_vector()'s ID space. - // Keep memory valid through consumption; lease may pin its backing storage. - struct PointerBatch { - std::vector ids; - // Only filtered or unregistered rows may have null pointers. - std::vector vectors; - std::shared_ptr lease; - }; - - // Row i has ID first_id + i; stride is in bytes, at least one vector's size. - struct ContiguousBatch { - const void *data{nullptr}; - uint32_t first_id{0}; - uint32_t count{0}; // Includes filtered rows. - size_t stride{0}; - std::shared_ptr lease; - }; - - using Batch = std::variant; - - class Iterator { - public: - using Pointer = std::unique_ptr; - virtual ~Iterator() = default; - // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an - // IndexError on failure. Enumerate every visible ID exactly once. - virtual int next_batch(uint32_t max_count, Batch *out) = 0; - }; + using Pointer = std::unique_ptr; + virtual ~VectorIterator() = default; + // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an + // IndexError on failure. Enumerate every visible ID exactly once. + virtual int next_batch(uint32_t max_count, VectorBatch *out) = 0; +}; +class ZVEC_CORE_API VectorSource { + public: VectorSource(); virtual ~VectorSource(); @@ -70,7 +69,20 @@ class ZVEC_CORE_API VectorSource { // The caller binds the source, scan domain and visibility filter to the same // stable snapshot for the whole request, including subsequent random reads. // A disk source can scan physical blocks behind this interface. - virtual Iterator::Pointer create_iterator() const; + virtual VectorIterator::Pointer create_iterator() const; +}; + +class ZVEC_CORE_API ContiguousVectorSource final : public VectorSource { + public: + explicit ContiguousVectorSource(ContiguousBatch view); + + const void *get_vector(uint32_t id) const override; + + VectorIterator::Pointer create_iterator() const override; + + private: + const ContiguousBatch view_; + const bool valid_; }; } // namespace core diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 7dc1720151..bbc586a224 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -67,16 +67,16 @@ class BatchSource : public core::VectorSource { VectorSource::get_vectors(ids, count, out); } - class Cursor : public Iterator { + class Cursor : public core::VectorIterator { public: explicit Cursor(const BatchSource &source) : source_(source) { for (auto it = source.rows.rbegin(); it != source.rows.rend(); ++it) { ids_.push_back(it->first); } } - int next_batch(uint32_t max_count, Batch *out) override { + int next_batch(uint32_t max_count, core::VectorBatch *out) override { EXPECT_TRUE(previous_.expired()); - auto &batch = out->emplace(); + auto &batch = out->emplace(); if (source_.fail_scan && position_) return core::IndexError_ReadData; const size_t count = std::min( {size_t{max_count}, source_.page_rows, ids_.size() - position_}); @@ -105,7 +105,7 @@ class BatchSource : public core::VectorSource { std::weak_ptr previous_; }; - Iterator::Pointer create_iterator() const override { + core::VectorIterator::Pointer create_iterator() const override { ++scans; return supports_scan ? std::make_unique(*this) : nullptr; } @@ -151,16 +151,16 @@ class ContiguousSource : public BatchSource { return data_.data() + (id - first_id_) * stride_elements_; } - class Cursor : public Iterator { + class Cursor : public core::VectorIterator { public: explicit Cursor(const ContiguousSource &source) : source_(source) {} ~Cursor() override { EXPECT_TRUE(previous_.expired()); } - int next_batch(uint32_t max_count, Batch *out) override { + int next_batch(uint32_t max_count, core::VectorBatch *out) override { EXPECT_TRUE(previous_.expired()); - auto &view = out->emplace(); + auto &view = out->emplace(); if (source_.fail_scan && position_) return core::IndexError_ReadData; const auto count = static_cast(std::min( {size_t{max_count}, source_.page_rows, @@ -181,7 +181,7 @@ class ContiguousSource : public BatchSource { } if (source_.mutate_batch) source_.mutate_batch(&view); if (source_.mixed_batches && batch_number_ % 2) { - PointerBatch pointers; + core::PointerBatch pointers; for (uint32_t i = 0; i < count; ++i) { pointers.ids.push_back(view.first_id + i); pointers.vectors.push_back(static_cast(view.data) + @@ -202,14 +202,14 @@ class ContiguousSource : public BatchSource { std::weak_ptr previous_; }; - Iterator::Pointer create_iterator() const override { + core::VectorIterator::Pointer create_iterator() const override { ++scans; return std::make_unique(*this); } mutable std::atomic scan_batches{0}; bool lease_pages{false}, mixed_batches{false}; - std::function mutate_batch; + std::function mutate_batch; private: uint32_t first_id_; @@ -393,6 +393,7 @@ TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { TEST_F(ExternalFlatTest, FullScanWithOnlyGetVector) { BatchSource data(30); RandomAccessSource source(data); + EXPECT_EQ(nullptr, source.create_iterator()); auto external = open_index(Param(true), test_path("external")); auto embedded = open_index(Param(false), test_path("embedded")); add_vectors(external.get(), data); @@ -431,17 +432,24 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { for (bool id_map : {false, true}) { for (size_t padding : {0, 3}) { - ContiguousSource source(kRowCount, kFirstId, padding); - source.page_rows = 1024; + ContiguousSource data(kRowCount, kFirstId, padding); + core::ContiguousVectorSource source( + {data.get_vector(kFirstId), + kFirstId, + kRowCount, + (kDimension + padding) * sizeof(float), + {}}); // Keep an unregistered row in the scan. - source.rows.erase(kFirstId + kRowCount - 1); + data.rows.erase(kFirstId + kRowCount - 1); auto suffix = std::to_string(case_id++); auto external = open_index(Param(true, metric, id_map), test_path(suffix + "ext")); auto embedded = open_index(Param(false, metric, id_map), test_path(suffix + "base")); - add_vectors(external.get(), source); - add_vectors(embedded.get(), source, false); + for (const auto &[id, row] : data.rows) { + ASSERT_EQ(0, external->add_with_source(Vector(row), id, source)); + ASSERT_EQ(0, embedded->add(Vector(row), id)); + } auto qp = FlatQueryParamBuilder() .with_topk(11) .with_fetch_vector(true) @@ -456,8 +464,6 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { ASSERT_EQ( 0, external->search_with_source(Vector(query), qp, source, &got)); SameResults(want, got); - EXPECT_EQ(3, source.scan_batches); - EXPECT_EQ(0, source.batch_reads); for (const auto &doc : got.doc_list_) { EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); EXPECT_NE(kFirstId + kRowCount - 1, doc.key()); @@ -469,8 +475,6 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { ASSERT_EQ( 0, external->search_with_source(Vector(query), qp, source, &got)); SameResults(want, got); - EXPECT_EQ(1, source.scans); - EXPECT_EQ(1, source.batch_reads); } } } @@ -515,32 +519,31 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { SearchResult result; for (int failure = 0; failure < 7; ++failure) { SCOPED_TRACE(failure); - source.mutate_batch = - [failure](core::VectorSource::ContiguousBatch *batch) { - switch (failure) { - case 0: - batch->data = nullptr; - break; - case 1: - batch->stride = 0; - break; - case 2: - batch->stride = kDimension * sizeof(float) - 1; - break; - case 3: - batch->stride = std::numeric_limits::max(); - break; - case 4: - batch->count = 1025; - break; - case 5: - batch->first_id = UINT32_MAX - 1; - break; - case 6: - batch->first_id = UINT32_MAX; - break; - } - }; + source.mutate_batch = [failure](core::ContiguousBatch *batch) { + switch (failure) { + case 0: + batch->data = nullptr; + break; + case 1: + batch->stride = 0; + break; + case 2: + batch->stride = kDimension * sizeof(float) - 1; + break; + case 3: + batch->stride = std::numeric_limits::max(); + break; + case 4: + batch->count = 1025; + break; + case 5: + batch->first_id = UINT32_MAX - 1; + break; + case 6: + batch->first_id = UINT32_MAX; + break; + } + }; result.doc_list_.emplace_back(999, 0); EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); EXPECT_TRUE(result.doc_list_.empty()); @@ -551,20 +554,114 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { } TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { - ContiguousSource source(1, UINT32_MAX - 1); + std::vector row(kDimension, 0.5f); + core::ContiguousVectorSource source( + {row.data(), UINT32_MAX - 1, 1, kDimension * sizeof(float), {}}); auto index = open_index(Param(true), test_path("index")); - add_vectors(index.get(), source); + ASSERT_EQ(0, index->add_with_source(Vector(row), UINT32_MAX - 1, source)); auto qp = FlatQueryParamBuilder().with_topk(1).build(); SearchResult result; - auto query = Vector(source.rows.begin()->second); + auto query = Vector(row); ASSERT_EQ(0, index->search_with_source(query, qp, source, &result)); ASSERT_EQ(1, result.doc_list_.size()); EXPECT_EQ(UINT32_MAX - 1, result.doc_list_[0].key()); - ContiguousSource empty(0); + core::ContiguousVectorSource empty(core::ContiguousBatch{}); ASSERT_EQ(0, index->search_with_source(query, qp, empty, &result)); EXPECT_TRUE(result.doc_list_.empty()); } +TEST(ContiguousVectorSourceTest, RandomReadsAndIndependentCursors) { + const float rows[] = {1, 2, -1, 3, 4, -1, 5, 6, -1}; + core::ContiguousVectorSource source({rows, 100, 3, 3 * sizeof(float), {}}); + EXPECT_EQ(nullptr, source.get_vector(99)); + EXPECT_EQ(rows, source.get_vector(100)); + EXPECT_EQ(rows + 6, source.get_vector(102)); + EXPECT_EQ(nullptr, source.get_vector(103)); + const uint32_t ids[] = {102, 99, 100}; + const void *vectors[3] = {}; + source.get_vectors(ids, 3, vectors); + EXPECT_EQ(rows + 6, vectors[0]); + EXPECT_EQ(nullptr, vectors[1]); + EXPECT_EQ(rows, vectors[2]); + + auto first = source.create_iterator(); + auto second = source.create_iterator(); + core::VectorBatch batch; + ASSERT_NE(nullptr, first); + ASSERT_NE(nullptr, second); + EXPECT_EQ(core::IndexError_InvalidArgument, first->next_batch(0, &batch)); + EXPECT_EQ(core::IndexError_InvalidArgument, first->next_batch(2, nullptr)); + ASSERT_EQ(0, first->next_batch(2, &batch)); + auto view = std::get(batch); + EXPECT_EQ(rows, view.data); + EXPECT_EQ(100, view.first_id); + EXPECT_EQ(2, view.count); + EXPECT_EQ(3 * sizeof(float), view.stride); + ASSERT_EQ(0, first->next_batch(2, &batch)); + view = std::get(batch); + EXPECT_EQ(rows + 6, view.data); + EXPECT_EQ(102, view.first_id); + EXPECT_EQ(1, view.count); + for (int i = 0; i < 2; ++i) { + ASSERT_EQ(0, first->next_batch(2, &batch)); + EXPECT_EQ(0, std::get(batch).count); + } + ASSERT_EQ(0, second->next_batch(10, &batch)); + view = std::get(batch); + EXPECT_EQ(rows, view.data); + EXPECT_EQ(100, view.first_id); + EXPECT_EQ(3, view.count); +} + +TEST(ContiguousVectorSourceTest, SourceCursorAndBatchRetainLease) { + auto storage = std::make_shared>(6, 0.5f); + std::weak_ptr weak = storage; + const auto *data = storage->data(); + auto source = std::make_unique( + core::ContiguousBatch{data, 3, 3, 2 * sizeof(float), storage}); + storage.reset(); + ASSERT_FALSE(weak.expired()); + EXPECT_EQ(data, source->get_vector(3)); + auto cursor = source->create_iterator(); + source.reset(); + ASSERT_FALSE(weak.expired()); + core::VectorBatch batch; + ASSERT_EQ(0, cursor->next_batch(2, &batch)); + cursor.reset(); + ASSERT_FALSE(weak.expired()); + const auto &view = std::get(batch); + EXPECT_EQ(data, view.data); + EXPECT_FLOAT_EQ(0.5f, static_cast(view.data)[0]); + batch = core::PointerBatch{}; + EXPECT_TRUE(weak.expired()); +} + +TEST(ContiguousVectorSourceTest, RejectsInvalidViewsAndSupportsEmptyViews) { + float rows[6] = {}; + const std::vector invalid{ + {nullptr, 3, 3, 2 * sizeof(float), {}}, + {rows, 3, 3, 0, {}}, + {rows, UINT32_MAX, 1, 2 * sizeof(float), {}}, + {rows, UINT32_MAX - 1, 2, 2 * sizeof(float), {}}, + {rows, 3, 3, std::numeric_limits::max(), {}}, + }; + for (const auto &view : invalid) { + core::ContiguousVectorSource source(view); + EXPECT_EQ(nullptr, source.get_vector(view.first_id)); + auto cursor = source.create_iterator(); + ASSERT_NE(nullptr, cursor); + core::VectorBatch batch; + EXPECT_EQ(core::IndexError_InvalidArgument, cursor->next_batch(1, &batch)); + } + core::ContiguousVectorSource empty(core::ContiguousBatch{}); + EXPECT_EQ(nullptr, empty.get_vector(0)); + auto cursor = empty.create_iterator(); + ASSERT_NE(nullptr, cursor); + core::VectorBatch batch; + ASSERT_EQ(0, cursor->next_batch(2, &batch)); + EXPECT_EQ(0, std::get(batch).count); +} + TEST_F(ExternalFlatTest, FallbackMatchesEmbeddedUsesBatchesAndReopens) { int case_id = 0; for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { From 2a12937f2da4036cf863c8abe2e57a6207f1012a Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 20:55:00 +0800 Subject: [PATCH 14/16] refactor(core): let clients provide contiguous vector views --- src/core/interface/vector_source.cc | 31 ++++--- .../zvec/core/interface/vector_source.h | 13 +-- .../interface/external_flat_index_test.cc | 80 ++++++++++++++++--- 3 files changed, 96 insertions(+), 28 deletions(-) diff --git a/src/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index 205bc71264..eb3ae1f239 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -34,10 +34,16 @@ bool ValidContiguousView(const ContiguousBatch &view) { return view.count <= max_span / view.stride; } +const void *GetContiguousVector(const ContiguousBatch &view, uint32_t id) { + if (id < view.first_id || id - view.first_id >= view.count) return nullptr; + return static_cast(view.data) + + size_t(id - view.first_id) * view.stride; +} + class ContiguousVectorIterator final : public VectorIterator { public: - ContiguousVectorIterator(ContiguousBatch view, bool valid) - : view_(std::move(view)), valid_(valid) {} + explicit ContiguousVectorIterator(ContiguousBatch view) + : view_(std::move(view)), valid_(ValidContiguousView(view_)) {} int next_batch(uint32_t max_count, VectorBatch *out) override { if (!out || !max_count) return IndexError_InvalidArgument; @@ -75,18 +81,23 @@ void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, } } -ContiguousVectorSource::ContiguousVectorSource(ContiguousBatch view) - : view_(std::move(view)), valid_(ValidContiguousView(view_)) {} - const void *ContiguousVectorSource::get_vector(uint32_t id) const { - if (!valid_ || id < view_.first_id || id - view_.first_id >= view_.count) - return nullptr; - return static_cast(view_.data) + - size_t(id - view_.first_id) * view_.stride; + const auto view = get_contiguous_view(); + return ValidContiguousView(view) ? GetContiguousVector(view, id) : nullptr; +} + +void ContiguousVectorSource::get_vectors(const uint32_t *ids, uint32_t count, + const void **out) const { + if (!count) return; + const auto view = get_contiguous_view(); + const bool valid = ValidContiguousView(view); + for (uint32_t i = 0; i < count; ++i) { + out[i] = valid ? GetContiguousVector(view, ids[i]) : nullptr; + } } VectorIterator::Pointer ContiguousVectorSource::create_iterator() const { - return std::make_unique(view_, valid_); + return std::make_unique(get_contiguous_view()); } } // namespace core diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 0daa6a906b..09c23b8fd0 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -72,17 +72,18 @@ class ZVEC_CORE_API VectorSource { virtual VectorIterator::Pointer create_iterator() const; }; -class ZVEC_CORE_API ContiguousVectorSource final : public VectorSource { +class ZVEC_CORE_API ContiguousVectorSource : public VectorSource { public: - explicit ContiguousVectorSource(ContiguousBatch view); + // Iterators retain this view and its lease. The caller keeps storage alive + // for random reads and any borrowed search results. + virtual ContiguousBatch get_contiguous_view() const = 0; const void *get_vector(uint32_t id) const override; - VectorIterator::Pointer create_iterator() const override; + void get_vectors(const uint32_t *ids, uint32_t count, + const void **out) const override; - private: - const ContiguousBatch view_; - const bool valid_; + VectorIterator::Pointer create_iterator() const override; }; } // namespace core diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index bbc586a224..76f1cfe17b 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -127,6 +128,20 @@ class RandomAccessSource : public core::VectorSource { const BatchSource &source_; }; +class ContiguousViewSource : public core::ContiguousVectorSource { + public: + explicit ContiguousViewSource(core::ContiguousBatch value) + : view(std::move(value)) {} + + core::ContiguousBatch get_contiguous_view() const override { + ++view_reads; + return view; + } + + core::ContiguousBatch view; + mutable uint32_t view_reads{0}; +}; + class ContiguousSource : public BatchSource { public: explicit ContiguousSource(uint32_t count, uint32_t first_id = 3, @@ -433,12 +448,11 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { for (bool id_map : {false, true}) { for (size_t padding : {0, 3}) { ContiguousSource data(kRowCount, kFirstId, padding); - core::ContiguousVectorSource source( - {data.get_vector(kFirstId), - kFirstId, - kRowCount, - (kDimension + padding) * sizeof(float), - {}}); + ContiguousViewSource source({data.get_vector(kFirstId), + kFirstId, + kRowCount, + (kDimension + padding) * sizeof(float), + {}}); // Keep an unregistered row in the scan. data.rows.erase(kFirstId + kRowCount - 1); auto suffix = std::to_string(case_id++); @@ -555,7 +569,7 @@ TEST_F(ExternalFlatTest, RejectsMalformedContiguousBatches) { TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { std::vector row(kDimension, 0.5f); - core::ContiguousVectorSource source( + ContiguousViewSource source( {row.data(), UINT32_MAX - 1, 1, kDimension * sizeof(float), {}}); auto index = open_index(Param(true), test_path("index")); ASSERT_EQ(0, index->add_with_source(Vector(row), UINT32_MAX - 1, source)); @@ -565,14 +579,14 @@ TEST_F(ExternalFlatTest, ContiguousIdBoundaryAndEmptySource) { ASSERT_EQ(0, index->search_with_source(query, qp, source, &result)); ASSERT_EQ(1, result.doc_list_.size()); EXPECT_EQ(UINT32_MAX - 1, result.doc_list_[0].key()); - core::ContiguousVectorSource empty(core::ContiguousBatch{}); + ContiguousViewSource empty(core::ContiguousBatch{}); ASSERT_EQ(0, index->search_with_source(query, qp, empty, &result)); EXPECT_TRUE(result.doc_list_.empty()); } TEST(ContiguousVectorSourceTest, RandomReadsAndIndependentCursors) { const float rows[] = {1, 2, -1, 3, 4, -1, 5, 6, -1}; - core::ContiguousVectorSource source({rows, 100, 3, 3 * sizeof(float), {}}); + ContiguousViewSource source({rows, 100, 3, 3 * sizeof(float), {}}); EXPECT_EQ(nullptr, source.get_vector(99)); EXPECT_EQ(rows, source.get_vector(100)); EXPECT_EQ(rows + 6, source.get_vector(102)); @@ -580,6 +594,9 @@ TEST(ContiguousVectorSourceTest, RandomReadsAndIndependentCursors) { const uint32_t ids[] = {102, 99, 100}; const void *vectors[3] = {}; source.get_vectors(ids, 3, vectors); + EXPECT_EQ(5, source.view_reads); + source.get_vectors(nullptr, 0, nullptr); + EXPECT_EQ(5, source.view_reads); EXPECT_EQ(rows + 6, vectors[0]); EXPECT_EQ(nullptr, vectors[1]); EXPECT_EQ(rows, vectors[2]); @@ -611,13 +628,47 @@ TEST(ContiguousVectorSourceTest, RandomReadsAndIndependentCursors) { EXPECT_EQ(rows, view.data); EXPECT_EQ(100, view.first_id); EXPECT_EQ(3, view.count); + EXPECT_EQ(7, source.view_reads); +} + +TEST(ContiguousVectorSourceTest, CursorsRetainViewAcrossSourceUpdates) { + const float old_rows[] = {1, 2, -1, 3, 4, -1, 5, 6, -1}; + const float new_rows[] = {7, 8, 9, 10}; + ContiguousViewSource source({old_rows, 100, 3, 3 * sizeof(float), {}}); + auto old_cursor = source.create_iterator(); + core::VectorBatch batch; + ASSERT_EQ(0, old_cursor->next_batch(2, &batch)); + EXPECT_EQ(old_rows, std::get(batch).data); + + source.view = {new_rows, 200, 2, 2 * sizeof(float), {}}; + EXPECT_EQ(nullptr, source.get_vector(100)); + EXPECT_EQ(new_rows + 2, source.get_vector(201)); + auto new_cursor = source.create_iterator(); + ASSERT_EQ(4, source.view_reads); + + ASSERT_EQ(0, old_cursor->next_batch(2, &batch)); + auto view = std::get(batch); + EXPECT_EQ(old_rows + 6, view.data); + EXPECT_EQ(102, view.first_id); + EXPECT_EQ(1, view.count); + EXPECT_EQ(3 * sizeof(float), view.stride); + ASSERT_EQ(0, old_cursor->next_batch(2, &batch)); + EXPECT_EQ(0, std::get(batch).count); + + ASSERT_EQ(0, new_cursor->next_batch(2, &batch)); + view = std::get(batch); + EXPECT_EQ(new_rows, view.data); + EXPECT_EQ(200, view.first_id); + EXPECT_EQ(2, view.count); + EXPECT_EQ(2 * sizeof(float), view.stride); + EXPECT_EQ(4, source.view_reads); } TEST(ContiguousVectorSourceTest, SourceCursorAndBatchRetainLease) { auto storage = std::make_shared>(6, 0.5f); std::weak_ptr weak = storage; const auto *data = storage->data(); - auto source = std::make_unique( + auto source = std::make_unique( core::ContiguousBatch{data, 3, 3, 2 * sizeof(float), storage}); storage.reset(); ASSERT_FALSE(weak.expired()); @@ -646,14 +697,19 @@ TEST(ContiguousVectorSourceTest, RejectsInvalidViewsAndSupportsEmptyViews) { {rows, 3, 3, std::numeric_limits::max(), {}}, }; for (const auto &view : invalid) { - core::ContiguousVectorSource source(view); + ContiguousViewSource source(view); EXPECT_EQ(nullptr, source.get_vector(view.first_id)); + const uint32_t ids[] = {view.first_id, 0}; + const void *vectors[] = {rows, rows}; + source.get_vectors(ids, 2, vectors); + EXPECT_EQ(nullptr, vectors[0]); + EXPECT_EQ(nullptr, vectors[1]); auto cursor = source.create_iterator(); ASSERT_NE(nullptr, cursor); core::VectorBatch batch; EXPECT_EQ(core::IndexError_InvalidArgument, cursor->next_batch(1, &batch)); } - core::ContiguousVectorSource empty(core::ContiguousBatch{}); + ContiguousViewSource empty(core::ContiguousBatch{}); EXPECT_EQ(nullptr, empty.get_vector(0)); auto cursor = empty.create_iterator(); ASSERT_NE(nullptr, cursor); From 44204d7d0afbdce92918a2f024f5a82e8cbc6556 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 27 Oct 2026 10:56:33 +0800 Subject: [PATCH 15/16] test(hnsw): cover contiguous external vector sources --- tests/core/interface/hnsw_turbo_index_test.cc | 73 +++++++++++++++---- 1 file changed, 58 insertions(+), 15 deletions(-) diff --git a/tests/core/interface/hnsw_turbo_index_test.cc b/tests/core/interface/hnsw_turbo_index_test.cc index a1d0a95c2a..35c00a92f5 100644 --- a/tests/core/interface/hnsw_turbo_index_test.cc +++ b/tests/core/interface/hnsw_turbo_index_test.cc @@ -97,6 +97,32 @@ class TestExternalVectorSource final : public zvec::core::VectorSource { const std::vector> *vectors_; }; +class TestContiguousVectorSource final + : public zvec::core::ContiguousVectorSource { + public: + TestContiguousVectorSource(const std::vector> &vectors, + size_t padding) + : stride_elements_(kDimension + padding), + data_(vectors.size() * stride_elements_, -999.0f) { + for (size_t i = 0; i < vectors.size(); ++i) { + std::copy(vectors[i].begin(), vectors[i].end(), + data_.begin() + i * stride_elements_); + } + } + + zvec::core::ContiguousBatch get_contiguous_view() const override { + return {data_.data(), + 0, + static_cast(data_.size() / stride_elements_), + stride_elements_ * sizeof(float), + {}}; + } + + private: + size_t stride_elements_; + std::vector data_; +}; + std::vector> RandomVectors(size_t count = kVectorCount) { std::mt19937 gen(2026); std::uniform_real_distribution dist(-1.0f, 1.0f); @@ -116,6 +142,19 @@ std::vector> RandomVectors(size_t count = kVectorCount) { return vectors; } +std::vector> ExternalVectors(MetricType metric) { + auto vectors = RandomVectors(kGraphVectorCount); + if (metric == MetricType::kCosine) { + for (size_t i = 0; i < vectors.size(); ++i) { + const float scale = static_cast(i % 7 + 1); + for (float &value : vectors[i]) { + value *= scale; + } + } + } + return vectors; +} + HNSWIndexParam::Pointer MakeParam(MetricType metric, QuantizerType quantizer, bool enable_rotate = false) { return HNSWIndexParamBuilder() @@ -294,21 +333,11 @@ void CheckTurboAddSearchReopen(MetricType metric, QuantizerType quantizer, zvec::test_util::RemoveTestFiles(path); } -void CheckExternalTurboAddSearchReopen(MetricType metric, - QuantizerType quantizer, - const char *quantizer_name, - const std::string &path) { +void CheckExternalTurboAddSearchReopen( + MetricType metric, QuantizerType quantizer, const char *quantizer_name, + const std::string &path, const std::vector> &vectors, + const zvec::core::VectorSource &source) { zvec::test_util::RemoveTestFiles(path); - auto vectors = RandomVectors(kGraphVectorCount); - if (metric == MetricType::kCosine) { - for (size_t i = 0; i < vectors.size(); ++i) { - const float scale = static_cast(i % 7 + 1); - for (float &value : vectors[i]) { - value *= scale; - } - } - } - TestExternalVectorSource source(&vectors); auto param = MakeParam(metric, quantizer); param->use_external_vector = true; @@ -641,8 +670,22 @@ TEST_P(HnswTurboIndexTest, AddSearchReopenFetch) { TEST_P(HnswTurboIndexTest, ExternalVectorsUseTurbo) { const auto &[quantizer, metric] = GetParam(); + auto vectors = ExternalVectors(metric.type); + TestExternalVectorSource source(&vectors); CheckExternalTurboAddSearchReopen(metric.type, quantizer.type, quantizer.name, - index_path("external")); + index_path("external"), vectors, source); +} + +TEST_P(HnswTurboIndexTest, ContiguousExternalVectorsUseTurbo) { + const auto &[quantizer, metric] = GetParam(); + auto vectors = ExternalVectors(metric.type); + for (size_t padding : {0, 3}) { + SCOPED_TRACE(padding); + TestContiguousVectorSource source(vectors, padding); + ASSERT_NO_FATAL_FAILURE(CheckExternalTurboAddSearchReopen( + metric.type, quantizer.type, quantizer.name, + index_path("contiguous_external"), vectors, source)); + } } TEST_P(HnswTurboIndexTest, OriginalProviderBuildUsesTurbo) { From 3ca65d8a5601c72ef8d5260d2b52f7583abc5299 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Tue, 27 Oct 2026 14:21:32 +0800 Subject: [PATCH 16/16] fix(tests): explicitly capture Flat filter ID for MSVC --- tests/core/interface/external_flat_index_test.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/core/interface/external_flat_index_test.cc b/tests/core/interface/external_flat_index_test.cc index 76f1cfe17b..7723172219 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -470,8 +470,9 @@ TEST_F(ExternalFlatTest, ContiguousScanMatchesEmbeddedAndSkipsDeletedRows) { .build(); qp->filter = std::make_shared(); // A fully filtered batch must not end the scan. - qp->filter->set( - [](uint64_t id) { return id <= kFirstId + 1023 || id % 7 == 0; }); + qp->filter->set([first_id = kFirstId](uint64_t id) { + return id <= first_id + 1023 || id % 7 == 0; + }); std::vector query(kDimension, 0.3f); SearchResult want, got; ASSERT_EQ(0, embedded->search(Vector(query), qp, &want));