From 04ae8eedd2b45e6ded69d4da5f32f8703cccd6c7 Mon Sep 17 00:00:00 2001 From: Zefeng Yin Date: Mon, 26 Oct 2026 17:57:15 +0800 Subject: [PATCH 1/7] 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 45031ab89..9f259b28c 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 4fbe55aa2..73ca575de 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 e4a7759d3..cd873f85b 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 3d47e77e0..434e55696 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 2/7] 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 9f259b28c..7ba9a9924 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 cd873f85b..b5196c3b2 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 434e55696..5dca74e2c 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 3/7] 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 b5196c3b2..a97ceb79e 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 5dca74e2c..65ce23999 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 4/7] 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 7ba9a9924..58226514b 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 a97ceb79e..9f3cdb6fa 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 65ce23999..7dc172015 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 5/7] 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 58226514b..1d7bc6217 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 98307cb76..205bc7126 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 9f3cdb6fa..0daa6a906 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 7dc172015..bbc586a22 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 6/7] 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 205bc7126..eb3ae1f23 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 0daa6a906..09c23b8fd 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 bbc586a22..76f1cfe17 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 7/7] 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 a1d0a95c2..35c00a92f 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) {