diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 45031ab89..1d7bc6217 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" @@ -517,18 +518,48 @@ 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)) { + pointers->ids.clear(); + pointers->vectors.clear(); + pointers->lease.reset(); + } else { + std::get(batch) = {}; + } + }; auto consume = [&]() -> int { - if (batch.ids.size() != batch.vectors.size() || - batch.ids.size() > batch_limit) + 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_size && + (!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 (pointers->ids.size() != pointers->vectors.size()) { return IndexError_ReadData; + } + 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(); - 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 +567,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 + : pointers->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(); @@ -568,22 +603,26 @@ 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.ids.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; @@ -622,7 +661,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/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/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index 98307cb76..eb3ae1f23 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -12,16 +12,65 @@ // 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; +} + +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: + 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; + 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 +81,24 @@ void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, } } +const void *ContiguousVectorSource::get_vector(uint32_t id) const { + 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(get_contiguous_view()); +} + } // 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 e4a7759d3..09c23b8fd 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -14,43 +14,46 @@ #pragma once +#include #include #include +#include #include #include 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: - // 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; - // 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. - 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; - }; + 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(); @@ -66,7 +69,21 @@ 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 : public VectorSource { + public: + // 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; + + void get_vectors(const uint32_t *ids, uint32_t count, + const void **out) const override; + + 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 3d47e77e0..76f1cfe17 100644 --- a/tests/core/interface/external_flat_index_test.cc +++ b/tests/core/interface/external_flat_index_test.cc @@ -15,9 +15,12 @@ #include #include #include +#include +#include #include #include #include +#include #include #include #include @@ -65,16 +68,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()); - 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_}); @@ -84,15 +87,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; } @@ -103,7 +106,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; } @@ -125,6 +128,110 @@ 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, + 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 * stride_elements_); + } + } + + const void *get_vector(uint32_t id) const override { + if (id < first_id_ || id - first_id_ >= data_.size() / stride_elements_) + return nullptr; + return data_.data() + (id - first_id_) * stride_elements_; + } + + 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, core::VectorBatch *out) override { + EXPECT_TRUE(previous_.expired()); + 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() / source_.stride_elements_ - position_})); + if (!count) return 0; + ++source_.scan_batches; + 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 * source_.stride_elements_); + view.data = page->data(); + view.lease = page; + previous_ = page; + } + if (source_.mutate_batch) source_.mutate_batch(&view); + if (source_.mixed_batches && batch_number_ % 2) { + 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) + + i * view.stride); + } + pointers.lease = std::move(view.lease); + *out = std::move(pointers); + } + position_ += count; + ++batch_number_; + return 0; + } + + private: + const ContiguousSource &source_; + size_t position_{0}; + size_t batch_number_{0}; + std::weak_ptr previous_; + }; + + 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; + + private: + uint32_t first_id_; + size_t stride_elements_; + std::vector data_; +}; + VectorData Vector(const std::vector &row) { return VectorData{DenseVector{row.data()}}; } @@ -301,6 +408,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); @@ -332,6 +440,284 @@ TEST_F(ExternalFlatTest, FullScanWithOnlyGetVector) { EXPECT_TRUE(got.doc_list_.empty()); } +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 (size_t padding : {0, 3}) { + ContiguousSource data(kRowCount, kFirstId, padding); + 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++); + 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")); + 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) + .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); + 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); + } + } + } +} + +TEST_F(ExternalFlatTest, ContiguousScanLeasesAndMixedBatches) { + 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")); + 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); + 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 < 7; ++failure) { + SCOPED_TRACE(failure); + 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()); + } + 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) { + std::vector row(kDimension, 0.5f); + 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)); + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + SearchResult result; + 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()); + 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}; + 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)); + 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(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]); + + 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); + 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( + 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) { + 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)); + } + ContiguousViewSource 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}) { 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) {