diff --git a/src/core/algorithm/flat/flat_streamer.cc b/src/core/algorithm/flat/flat_streamer.cc index 2b70d619d7..1d7bc62178 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" @@ -109,6 +110,21 @@ int FlatStreamer::init(const IndexMeta &imeta, params.get(PARAM_FLAT_USE_ID_MAP, &use_key_info_map_); params.get(PARAM_FLAT_USE_CONTIGUOUS_MEMORY, &use_contiguous_memory_); + use_external_vector_ = false; + params.get(PARAM_FLAT_USE_EXTERNAL_VECTOR, &use_external_vector_); + if (use_external_vector_) { + if (use_contiguous_memory_ || meta_.major_order() != IndexMeta::MO_ROW || + meta_.data_type() != IndexMeta::DT_FP32 || + meta_.meta_type() != IndexMeta::MT_DENSE || meta_.dimension() == 0 || + meta_.element_size() != meta_.dimension() * sizeof(float) || + (meta_.metric_name() != "SquaredEuclidean" && + meta_.metric_name() != "InnerProduct") || + (quantizer_ && meta_.quantizer_name() != "Fp32Quantizer")) { + return IndexError_Unsupported; + } + use_key_info_map_ = true; + } + state_ = STATE_INITED; return 0; @@ -182,6 +198,7 @@ int FlatStreamer::open(IndexStorage::Pointer stg) { entity_->enable_filter_same_key(true); entity_->set_linear_list_count(1); entity_->set_use_key_info_map(use_key_info_map_); + entity_->set_use_external_vector(use_external_vector_); entity_->set_quantizer(quantizer_); *entity_->mutable_meta() = meta_; @@ -236,6 +253,7 @@ int FlatStreamer::flush(uint64_t checkpoint) { template int FlatStreamer::dump(const IndexDumper::Pointer &dumper) { + if (use_external_vector_) return IndexError_Unsupported; std::string searcher_name = "FlatSearcher"; if constexpr (BATCH_SIZE == 16) { searcher_name = "FlatSearcher16"; @@ -288,6 +306,10 @@ int FlatStreamer::add_impl(uint64_t pkey, const void *query, (*stats_.mutable_discarded_count())++; return IndexError_Cast; } + if (use_external_vector_) { + if (pkey >= kInvalidNodeId) return IndexError_InvalidArgument; + return add_with_id_impl(static_cast(pkey), query, qmeta, context); + } READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); if (!dump_lock.try_lock()) { @@ -337,6 +359,12 @@ int FlatStreamer::add_with_id_impl(uint32_t id, const void *query, (*stats_.mutable_discarded_count())++; return IndexError_Cast; } + if (use_external_vector_) { + if (state_ != STATE_OPENED || !ctx->vector_source()) + return IndexError_NoReady; + if (id == kInvalidNodeId) return IndexError_InvalidArgument; + if (!ctx->vector_source()->get_vector(id)) return IndexError_NoExist; + } READ_LOCK_GUARD_DEFER(dump_mutex_, dump_lock); if (!dump_lock.try_lock()) { @@ -360,6 +388,10 @@ int FlatStreamer::search_bf_impl(const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, nullptr, context); + } if (state_ != STATE_OPENED) return IndexError_NoReady; ailego_assert(query && count && !!context); ailego_assert(quantizer_ || metric_->is_matched(meta_, qmeta)); @@ -408,6 +440,10 @@ int FlatStreamer::search_bf_by_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) { + if (state_ != STATE_OPENED) return IndexError_NoReady; + return search_external(query, qmeta, count, &p_keys, context); + } if (state_ != STATE_OPENED) return IndexError_NoReady; if (count == 0 || count > p_keys.size()) return IndexError_InvalidArgument; ailego_assert(query && count && !!context); @@ -452,12 +488,203 @@ int FlatStreamer::search_bf_by_p_keys_impl( return 0; } +template +int FlatStreamer::search_external( + const void *query, const IndexQueryMeta &qmeta, uint32_t count, + const std::vector> *keys, + Context::Pointer &context) const { + auto *ctx = dynamic_cast *>(context.get()); + if (!ctx || !query || !count || (keys && keys->size() < count) || + qmeta.meta_type() != meta_.meta_type() || + qmeta.data_type() != meta_.data_type() || + qmeta.dimension() != meta_.dimension() || + qmeta.element_size() != meta_.element_size()) { + return IndexError_InvalidArgument; + } + if (ctx->magic() != magic_) ctx->reset(this); + const auto *source = ctx->vector_source(); + if (!source || !ctx->threshold_is_valid()) return IndexError_NoReady; + if (ctx->group_by_search()) return IndexError_Unsupported; + ctx->reset_results(count); + // Bound both pointer scratch space and per-request pinned vector bytes. + const uint32_t batch_limit = static_cast(std::min( + 1024, + std::max(BATCH_SIZE, read_block_size_ / meta_.element_size()))); + auto *scratch = ctx->search_scratch(); + + for (uint32_t q = 0; q < count; ++q) { + const void *current_query = + static_cast(query) + q * qmeta.element_size(); + auto *heap = ctx->result_heap(); + heap->clear(); + // Destroy the outstanding lease before its cursor, including on errors. + 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 { + 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(); + 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)))) { + ++*ctx->mutable_stats(q)->mutable_filtered_count(); + continue; + } + // Invisible or unregistered rows need not have a readable vector. + 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(); + if (!size) return 0; + scratch->distances.resize(size); + if (quantizer_) { + quantizer_->calc_distance_dp_query_batch( + scratch->vector_ptrs.data(), static_cast(size), current_query, + scratch->distances.data()); + } else if (auto distance = metric_->batch_distance()) { + distance(scratch->vector_ptrs.data(), current_query, size, + meta_.dimension(), scratch->distances.data(), nullptr); + } else { + auto scalar_distance = metric_->distance(); + for (size_t i = 0; i < size; ++i) { + scalar_distance(current_query, scratch->vector_ptrs[i], + meta_.dimension(), &scratch->distances[i]); + } + } + *ctx->mutable_stats(q)->mutable_dist_calced_count() += size; + for (size_t i = 0; i < size; ++i) { + const auto id = scratch->vector_keys[i]; + const auto distance = scratch->distances[i]; + heap->emplace(id, distance); + } + return 0; + }; + + if (cursor) { + while (true) { + clear_batch(); + int ret = cursor->next_batch(batch_limit, &batch); + 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()) + break; + } + } else { + std::vector ids; + ids.reserve(batch_limit); + auto read_batch = [&]() -> int { + if (ids.empty()) return 0; + 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; + }; + auto read_candidates = [&](const std::vector &candidates) { + for (auto key : candidates) { + if (key >= kInvalidNodeId || (keys && !entity_->contains_key(key))) + continue; + if (ctx->filter().is_valid() && ctx->filter()(key)) { + if (!keys) ++*ctx->mutable_stats(q)->mutable_filtered_count(); + continue; + } + ids.push_back(static_cast(key)); + if (ids.size() == batch_limit) { + int ret = read_batch(); + if (ret != 0) return ret; + } + } + return read_batch(); + }; + if (keys) { + int ret = read_candidates((*keys)[q]); + if (ret != 0) return ret; + } else { + // Sources without a scan cursor can use the existing random reads. + // Fix the scan length and copy only one batch of registered IDs. + const size_t registered_count = entity_->registered_key_count(); + std::vector registered_keys; + for (size_t offset = 0; offset < registered_count; + offset += batch_limit) { + entity_->get_registered_keys( + offset, std::min(batch_limit, registered_count - offset), + ®istered_keys); + int ret = read_candidates(registered_keys); + if (ret != 0) return ret; + } + } + } + clear_batch(); + ctx->take_topk_result(q); + if (ctx->fetch_vector()) { + // Match HNSW's borrowed random-access result contract. Scan pages have + // already been released, so result pointers must come from the source. + for (auto &doc : *ctx->mutable_result(q)) { + const void *vector = source->get_vector(doc.key()); + if (!vector) return IndexError_NoExist; + IndexStorage::MemoryBlock block; + block.reset(const_cast(vector)); + doc = IndexDocument(doc.key(), doc.score(), doc.index(), block); + } + } + } + return 0; +} + template int FlatStreamer::search_by_p_keys_fast( const void *query, const std::vector &keys, int64_t *output_ids, float *output_scores, size_t topk, const IndexQueryMeta &qmeta, Context::UPointer &context) const { if (state_ != STATE_OPENED) return IndexError_NoReady; + if (use_external_vector_) return IndexError_NotImplemented; if (!query || !output_ids || topk == 0 || !context) { return IndexError_InvalidArgument; } @@ -487,6 +714,7 @@ template int FlatStreamer::group_by_search_impl( const void *query, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) return IndexError_Unsupported; FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); if (!bf_context) { @@ -535,6 +763,7 @@ int FlatStreamer::group_by_search_p_keys_impl( const void *query, const std::vector> &p_keys, const IndexQueryMeta &qmeta, uint32_t count, Context::Pointer &context) const { + if (use_external_vector_) return IndexError_Unsupported; if (count == 0 || count > p_keys.size()) return IndexError_InvalidArgument; FlatStreamerContext *bf_context = dynamic_cast *>(context.get()); diff --git a/src/core/algorithm/flat/flat_streamer.h b/src/core/algorithm/flat/flat_streamer.h index b5f28480ba..e42c9f4a5c 100644 --- a/src/core/algorithm/flat/flat_streamer.h +++ b/src/core/algorithm/flat/flat_streamer.h @@ -17,6 +17,7 @@ #include #include #include +#include #include "flat_streamer_entity.h" #include "flat_utility.h" @@ -139,6 +140,10 @@ class FlatStreamer : public IndexStreamer { return *entity_; } + bool use_external_vector() const { + return use_external_vector_; + } + //! Retrieve the turbo quantizer const std::shared_ptr &quantizer() const { return quantizer_; @@ -184,6 +189,11 @@ class FlatStreamer : public IndexStreamer { } private: + int search_external(const void *query, const IndexQueryMeta &qmeta, + uint32_t count, + const std::vector> *keys, + Context::Pointer &context) const; + //! Constants static constexpr uint32_t kDefaultBlockVecCount = 32u; static constexpr uint32_t kDefaultSegmentSize = 4 * 1024 * 1024u; @@ -206,6 +216,7 @@ class FlatStreamer : public IndexStreamer { bool use_key_info_map_{true}; uint32_t read_block_size_{0}; bool use_contiguous_memory_{false}; + bool use_external_vector_{false}; std::shared_ptr quantizer_{}; std::unique_ptr entity_{}; }; diff --git a/src/core/algorithm/flat/flat_streamer_context.h b/src/core/algorithm/flat/flat_streamer_context.h index 3080448718..bafd90a15d 100644 --- a/src/core/algorithm/flat/flat_streamer_context.h +++ b/src/core/algorithm/flat/flat_streamer_context.h @@ -128,7 +128,7 @@ class FlatStreamerContext : public IndexStreamer::Context { } key_t key = result_heap_[i].key(); - if (fetch_vector_) { + if (fetch_vector_ && !owner_->use_external_vector()) { IndexStorage::MemoryBlock block; owner_->entity().get_vector_by_key(key, block); results_[idx].emplace_back(key, score, key, block); @@ -224,7 +224,16 @@ class FlatStreamerContext : public IndexStreamer::Context { group_topk_heaps_.clear(); } + void set_vector_source(const VectorSource *source) { + vector_source_ = source; + } + + const VectorSource *vector_source() const { + return vector_source_; + } + void reset() override { + vector_source_ = nullptr; for (auto &it : results_) { it.clear(); } @@ -235,7 +244,9 @@ class FlatStreamerContext : public IndexStreamer::Context { //! Reset the context void reset(const FlatStreamer *owner) { + const auto *source = vector_source_; this->reset(); + vector_source_ = source; magic_ = owner->magic(); feature_size_ = owner->meta().element_size(); @@ -270,6 +281,7 @@ class FlatStreamerContext : public IndexStreamer::Context { private: const FlatStreamer *owner_{nullptr}; + const VectorSource *vector_source_{nullptr}; FlatSearchScratch search_scratch_{}; std::vector stats_vec_{}; uint32_t magic_{0}; diff --git a/src/core/algorithm/flat/flat_streamer_entity.cc b/src/core/algorithm/flat/flat_streamer_entity.cc index ff345a6b97..4215057b90 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.cc +++ b/src/core/algorithm/flat/flat_streamer_entity.cc @@ -156,7 +156,7 @@ int FlatStreamerEntity::open(IndexStorage::Pointer storage, vec_cols_ = index_meta_.element_size() / vec_unit_size_; meta_.header.block_size = ailego_align(sizeof(BlockHeader) + sizeof(DeletionMap) + - (index_meta_.element_size() + sizeof(uint64_t)) * + (stored_vector_size() + sizeof(uint64_t)) * meta_.header.block_vector_count, 32); @@ -354,6 +354,7 @@ int FlatStreamerEntity::search(const void *query, const IndexFilter &filter, IndexContext::Stats *context_stats, FlatSearchScratch * /*scratch*/, size_t /*batch_size*/) const { + if (use_external_vector_) return IndexError_Unsupported; IndexStorage::MemoryBlock head_block; this->get_head_block(head_block); const BlockLocation *bl = @@ -400,6 +401,7 @@ int FlatStreamerEntity::search_by_p_keys(const void *query, IndexDocumentHeap *heap, FlatSearchScratch * /*scratch*/, size_t /*batch_size*/) const { + if (use_external_vector_) return IndexError_Unsupported; for (uint64_t key : p_keys) { if (filter.is_valid() && filter(key)) { continue; @@ -504,6 +506,7 @@ FlatStreamerEntity::Pointer FlatStreamerEntity::clone() const { LOG_ERROR("Failed to New FlatStreamerEntity object"); return nullptr; } + entity->use_external_vector_ = this->use_external_vector_; entity->index_meta_ = this->index_meta_; entity->storage_ = this->storage_; // entity->reformer_ = this->reformer_; @@ -802,6 +805,7 @@ const void *FlatContiguousStreamerEntity::get_vector_ptr( } const void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const { + if (use_external_vector_) return nullptr; // Keep the legacy pointer valid until the next call on this thread. static thread_local LegacyReadBuffer buffer; IndexStorage::MemoryBlock block; @@ -812,6 +816,7 @@ const void *FlatStreamerEntity::get_vector_by_key(uint64_t key) const { int FlatStreamerEntity::get_vector_by_key( const uint64_t key, IndexStorage::MemoryBlock &block) const { + if (use_external_vector_) return IndexError_Unsupported; VectorLocation loc{}; key_info_map_lock_->lock_shared(); if (use_key_info_map_) { @@ -843,6 +848,7 @@ int FlatStreamerEntity::get_vector_by_key( int FlatStreamerEntity::get_vectors_by_key( const uint64_t *keys, uint32_t count, std::vector &blocks) const { + if (use_external_vector_) return IndexError_Unsupported; std::vector locations(count); key_info_map_lock_->lock_shared(); if (use_key_info_map_) { @@ -899,6 +905,8 @@ int FlatStreamerEntity::get_vectors_by_key( } IndexProvider::Iterator::Pointer FlatStreamerEntity::creater_iterator() const { + // An external scan requires a request-bound source. + if (use_external_vector_) return nullptr; auto entity = this->clone(); if (!entity) { LOG_ERROR("Failed to clone entity"); @@ -1046,6 +1054,10 @@ int FlatStreamerEntity::load_linear_meta(IndexStorage::Pointer storage) { return IndexError_InvalidFormat; } auto *mt = reinterpret_cast(data_block.data()); + if (mt->header.reserved_[0] != meta_.header.reserved_[0]) { + LOG_ERROR("Unmatched Flat external-vector storage mode"); + return IndexError_Mismatch; + } if (mt->header.block_vector_count != meta_.header.block_vector_count) { LOG_ERROR("Unmatched BlockVecCount Setting, Index %u vs Setting %u", mt->header.block_vector_count, meta_.header.block_vector_count); @@ -1379,7 +1391,7 @@ int FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key, size_t vector_off = get_block_vector_offset(block.block_index, header->vector_count); - if (segment->write(vector_off, data, size) != size) { + if (!use_external_vector_ && segment->write(vector_off, data, size) != size) { LOG_ERROR("Failed to write vector, off=%zu size=%zu", vector_off, size); return IndexError_WriteData; } @@ -1413,6 +1425,11 @@ int FlatStreamerEntity::add_to_block(const BlockLocation &block, uint64_t key, int FlatStreamerEntity::add_vector_with_id(const uint32_t id, const void *query, const uint32_t size) { + if (use_external_vector_) { + // Keep sparse external IDs compact. Updates belong to the source. + int ret = add(id, query, size); + return ret == IndexError_Duplicate ? 0 : ret; + } std::lock_guard lock(mutex_); // if (filter_same_key_) { // key_info_map_lock_->lock_shared(); diff --git a/src/core/algorithm/flat/flat_streamer_entity.h b/src/core/algorithm/flat/flat_streamer_entity.h index 7659953e9d..73ca575dee 100644 --- a/src/core/algorithm/flat/flat_streamer_entity.h +++ b/src/core/algorithm/flat/flat_streamer_entity.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -39,6 +40,7 @@ struct FlatSearchScratch { std::vector vector_keys{}; std::vector distances{}; std::vector query_buffer{}; + std::vector key_presence{}; }; /*! Flat Streamer Entity @@ -166,6 +168,47 @@ class FlatStreamerEntity { LOG_DEBUG("use_key_info_map_: %d", (int)use_key_info_map_); } + void set_use_external_vector(bool enabled) { + use_external_vector_ = enabled; + // Persist the physical layout independently of the logical vector meta. + meta_.header.reserved_[0] = enabled ? 1 : 0; + } + + bool contains_key(uint64_t key) const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + return key_info_map_.find(key) != key_info_map_.end(); + } + + // 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(count); + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + 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(); + } + } + + size_t registered_key_count() const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + return id_key_vector_.size(); + } + + // Copy a bounded batch so source callbacks never run under the key lock. + void get_registered_keys(size_t offset, size_t count, + std::vector *keys) const { + ailego::ReadLock lock(*key_info_map_lock_); + std::lock_guard guard(lock); + offset = std::min(offset, id_key_vector_.size()); + count = std::min(count, id_key_vector_.size() - offset); + auto begin = id_key_vector_.begin() + offset; + keys->assign(begin, begin + count); + } + //! Set params void set_segment_size(uint32_t size) { meta_.segment_size = size; @@ -379,6 +422,10 @@ class FlatStreamerEntity { size_t size); private: + size_t stored_vector_size() const { + return use_external_vector_ ? 0 : index_meta_.element_size(); + } + size_t get_block_offset(uint32_t block_index) const { return block_index * linear_block_size(); } @@ -395,14 +442,14 @@ class FlatStreamerEntity { size_t get_block_key_offset(uint32_t block_index, uint32_t vector_index) const { return get_block_offset(block_index) + - block_vector_count() * index_meta_.element_size() + + block_vector_count() * stored_vector_size() + sizeof(uint64_t) * vector_index; } size_t get_block_vector_offset(uint32_t block_index, uint32_t vector_index) const { return this->get_block_offset(block_index) + - vector_index * index_meta_.element_size(); + vector_index * stored_vector_size(); } //! Get header block of an linear list @@ -488,6 +535,7 @@ class FlatStreamerEntity { // Distance loops repeatedly read this metadata. Keep it off the cache line // written by segment-cache locking in concurrent queries. alignas(AILEGO_CACHE_LINE_SIZE) IndexMeta index_meta_{}; + bool use_external_vector_{false}; IndexStorage::Pointer storage_{}; IndexMetric::MatrixDistance row_distance_{}, column_distance_{}; IndexMetric::MatrixBatchDistance batch_distance_{}; diff --git a/src/core/algorithm/flat/flat_utility.h b/src/core/algorithm/flat/flat_utility.h index 4020c8c9a5..31b8c0869d 100644 --- a/src/core/algorithm/flat/flat_utility.h +++ b/src/core/algorithm/flat/flat_utility.h @@ -34,6 +34,8 @@ static const std::string FLAT_SEGMENT_FEATURES_SEG_ID("flat.features"); static const std::string FLAT_SEGMENT_MAPPING_SEG_ID("flat.mapping"); // index params +static const std::string PARAM_FLAT_USE_EXTERNAL_VECTOR( + "proxima.flat.use_external_vector"); static const std::string PARAM_FLAT_COLUMN_MAJOR_ORDER( "proxima.flat.column_major_order"); static const std::string PARAM_FLAT_BATCH_SIZE("proxima.flat.batch_size"); diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 1932d40135..dc95a8b938 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -1418,6 +1418,16 @@ int Index::merge(const std::vector &indexes, if (indexes.empty()) { return core::IndexError_Success; } + // External Flat providers contain IDs only; merge has no bound source. + if (param_.index_type == IndexType::kFlat && param_.use_external_vector) { + return core::IndexError_Unsupported; + } + for (const auto &index : indexes) { + if (index && index->param_.index_type == IndexType::kFlat && + index->param_.use_external_vector) { + return core::IndexError_Unsupported; + } + } // ivf need builder auto reducer = core::IndexFactory::CreateStreamerReducer("MixedStreamerReducer"); diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index 1fed9d4333..8f3ca303a8 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -18,12 +18,20 @@ #include #include #include "algorithm/flat/flat_index_format.h" +#include "algorithm/flat/flat_streamer_context.h" #include "algorithm/flat/flat_utility.h" namespace zvec::core_interface { namespace { +struct FlatSourceScope { + core::IndexContext::Pointer &context; + ~FlatSourceScope() { + if (context) context->reset(); + } +}; + //! Read the IndexMeta persisted in the flat linear meta segment without //! initializing a streamer. Returns non-zero when the storage or segment //! cannot be read (e.g. the index file does not exist yet). @@ -134,6 +142,46 @@ ailego::Params MakeTurboQuantizerParams(const QuantizerParam &quantizer_param, } // namespace +int FlatIndex::bind_vector_source(const core::VectorSource &source, + core::IndexContext::Pointer &context) { + if (!param_.use_external_vector) return core::IndexError_Unsupported; + auto *ctx = dynamic_cast *>(context.get()); + auto *streamer = dynamic_cast *>(streamer_.get()); + if (!ctx || !streamer) return core::IndexError_Cast; + if (ctx->magic() != streamer->magic()) ctx->reset(streamer); + ctx->set_vector_source(&source); + return 0; +} + +int FlatIndex::add_with_source(const VectorData &vector, uint32_t doc_id, + const core::VectorSource &source) { + if (!is_open_) return core::IndexError_NoReady; + if (is_read_only_) return core::IndexError_Runtime; + const auto *dense = std::get_if(&vector.vector); + if (!dense || !dense->data) return core::IndexError_InvalidArgument; + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + FlatSourceScope scope{context}; + int ret = bind_vector_source(source, context); + if (ret != 0) return ret; + // External Flat stores only IDs. Its raw FP32 layout needs no conversion. + return streamer_->add_with_id_impl(doc_id, dense->data, input_vector_meta_, + context); +} + +int FlatIndex::search_with_source( + const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, + const core::VectorSource &source, SearchResult *result) { + if (!is_open_) return core::IndexError_NoReady; + if (!search_param || !result) return core::IndexError_InvalidArgument; + *result = SearchResult{}; + auto &context = acquire_context(); + if (!context) return core::IndexError_Runtime; + FlatSourceScope scope{context}; + int ret = bind_vector_source(source, context); + return ret == 0 ? Index::search(query, search_param, result) : ret; +} + int FlatIndex::open(const std::string &file_path, StorageOptions storage_options) { // Restore the persisted encoding while keeping the configured structural @@ -220,6 +268,21 @@ int FlatIndex::fallback_to_legacy_pipeline() { int FlatIndex::create_and_init_converter_reformer( const QuantizerParam &quantizer_param, const BaseIndexParam &index_param) { const auto &flat_param = dynamic_cast(index_param); + if (flat_param.use_external_vector && + (flat_param.is_sparse || flat_param.data_type != DataType::DT_FP32 || + flat_param.dimension <= 0 || flat_param.dimension > MAX_DIMENSION || + (flat_param.metric_type != MetricType::kL2sq && + flat_param.metric_type != MetricType::kInnerProduct) || + flat_param.major_order == IndexMeta::MO_COLUMN || + flat_param.use_contiguous_memory || quantizer_param.enable_rotate || + quantizer_param.type != QuantizerType::kNone || + flat_param.preprocess_param.type != PreprocessorType::kNone || + (flat_param.storage_data_type != DataType::DT_UNDEFINED && + flat_param.storage_data_type != DataType::DT_FP32))) { + LOG_ERROR("External Flat requires raw row-major FP32 L2/IP vectors"); + return core::IndexError_Unsupported; + } + // Prefer the turbo quantizer path (quantized records + SIMD batch distance // kernels funneled through the streamer entity) whenever a turbo quantizer // can express the configuration; only the remaining combinations fall @@ -299,6 +362,8 @@ int FlatIndex::create_and_init_streamer(const BaseIndexParam ¶m) { proxima_index_params_.set(core::PARAM_FLAT_COLUMN_MAJOR_ORDER, param_.major_order == IndexMeta::MO_COLUMN); proxima_index_params_.set(core::PARAM_FLAT_USE_ID_MAP, param_.use_id_map); + proxima_index_params_.set(core::PARAM_FLAT_USE_EXTERNAL_VECTOR, + param_.use_external_vector); proxima_index_params_.set(core::PARAM_FLAT_USE_CONTIGUOUS_MEMORY, param_.use_contiguous_memory); if (is_sparse_) { @@ -340,10 +405,14 @@ int FlatIndex::_prepare_for_search( return core::IndexError_Runtime; } + if (param_.use_external_vector && flat_search_param->group_by_param) { + return core::IndexError_Unsupported; + } + context->set_topk(flat_search_param->topk); context->set_fetch_vector(flat_search_param->fetch_vector); if (flat_search_param->filter && flat_search_param->filter->is_valid()) { - context->set_filter(std::move(*flat_search_param->filter)); + context->set_filter(*flat_search_param->filter); } else { context->reset_filter(); } diff --git a/src/core/interface/vector_source.cc b/src/core/interface/vector_source.cc index ea19016e11..eb3ae1f239 100644 --- a/src/core/interface/vector_source.cc +++ b/src/core/interface/vector_source.cc @@ -12,15 +12,68 @@ // 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; +VectorIterator::Pointer VectorSource::create_iterator() const { + return nullptr; +} + void VectorSource::get_vectors(const uint32_t *ids, uint32_t count, const void **out) const { for (uint32_t i = 0; i < count; ++i) { @@ -28,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/index.h b/src/include/zvec/core/interface/index.h index eb2093c22d..077a88974f 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -286,6 +286,23 @@ class ZVEC_CORE_API Index { class ZVEC_CORE_API FlatIndex : public Index { public: FlatIndex() = default; + + // External Flat supports raw row-major FP32 L2/IP. Register IDs after + // writing the source; vector bytes and subsequent updates stay in the source. + // doc_id is the source ID with either use_id_map setting. + int add_with_source(const VectorData &vector, uint32_t doc_id, + const core::VectorSource &src) override; + // Full scans prefer the source iterator, intersecting its visible IDs with + // registered IDs and the request filter. Without an iterator, registered IDs + // are filtered and read through get_vectors(), as in candidate searches. + // Bind the source, scan and filter to one stable snapshot. + // With fetch_vector, returned pointers borrow source memory; + // keep that memory valid while using the results, as with external HNSW. + // External Flat does not support group-by or binding a source for refinement. + int search_with_source(const VectorData &query, + const BaseIndexQueryParam::Pointer &search_param, + const core::VectorSource &src, + SearchResult *result) override; // FlatIndex(const FlatIndexParam ¶m) : param_(param) {} // FlatIndex(FlatIndexParam &¶m) : param(std::move(param)) {} @@ -321,6 +338,9 @@ class ZVEC_CORE_API FlatIndex : public Index { int create_and_init_legacy_converter_reformer( const QuantizerParam ¶m, const BaseIndexParam &index_param); + int bind_vector_source(const core::VectorSource &source, + core::IndexContext::Pointer &context); + FlatIndexParam param_{}; }; diff --git a/src/include/zvec/core/interface/vector_source.h b/src/include/zvec/core/interface/vector_source.h index 66b4da5dd2..09c23b8fd0 100644 --- a/src/include/zvec/core/interface/vector_source.h +++ b/src/include/zvec/core/interface/vector_source.h @@ -14,21 +14,76 @@ #pragma once +#include #include +#include +#include +#include #include namespace zvec { namespace core { +// 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: + 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(); + // Borrowed random-access pointers. Keep all get_vectors() results valid + // through consumption; fetch_vector results also borrow this memory. virtual const void *get_vector(uint32_t node_id) const = 0; virtual void get_vectors(const uint32_t *ids, uint32_t count, const void **out) const; + + // Optional full scan. Each call creates an independent cursor. Returning + // nullptr makes External Flat read registered IDs through get_vectors(). + // The caller binds the source, scan domain and visibility filter to the same + // stable snapshot for the whole request, including subsequent random reads. + // A disk source can scan physical blocks behind this interface. + virtual 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 new file mode 100644 index 0000000000..7723172219 --- /dev/null +++ b/tests/core/interface/external_flat_index_test.cc @@ -0,0 +1,1050 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "algorithm/flat/flat_streamer_context.h" +#include "tests/test_util.h" + +using namespace zvec::core_interface; +namespace core = zvec::core; + +namespace { +constexpr uint32_t kDimension = 8; + +// Each batch owns a temporary page. Release must precede advancing the cursor. +// Reverse scan order and noncontiguous IDs catch assumptions about row offsets. +class BatchSource : public core::VectorSource { + public: + std::map> rows; + mutable std::atomic scans{0}; + mutable std::atomic batch_reads{0}; + size_t page_rows{7}; + std::set null_scan_ids; + bool supports_scan{true}, fail_scan{false}, malformed_scan{false}; + bool null_scan{false}, fail_read{false}; + + explicit BatchSource(uint32_t count = 89, float offset = 0) { + for (uint32_t i = 0; i < count; ++i) { + auto &row = rows[i * 13 + 3]; + for (uint32_t d = 0; d < kDimension; ++d) { + row.push_back(((i * 19 + d * 7) % 101) / 101.0f + offset); + } + } + } + + const void *get_vector(uint32_t id) const override { + if (fail_read) return nullptr; + auto it = rows.find(id); + return it == rows.end() ? nullptr : it->second.data(); + } + + void get_vectors(const uint32_t *ids, uint32_t count, + const void **out) const override { + ++batch_reads; + EXPECT_LE(count, 1024u); + VectorSource::get_vectors(ids, count, out); + } + + class Cursor : public 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, core::VectorBatch *out) override { + EXPECT_TRUE(previous_.expired()); + 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_}); + if (!count) return 0; + auto page = std::make_shared>(count * kDimension); + for (size_t i = 0; i < count; ++i) { + auto id = ids_[position_++]; + std::memcpy(page->data() + i * kDimension, source_.rows.at(id).data(), + kDimension * sizeof(float)); + batch.ids.push_back(id); + batch.vectors.push_back(source_.null_scan || + source_.null_scan_ids.count(id) + ? nullptr + : page->data() + i * kDimension); + } + batch.lease = page; + previous_ = page; + if (source_.malformed_scan) batch.vectors.pop_back(); + return 0; + } + + private: + const BatchSource &source_; + std::vector ids_; + size_t position_{0}; + std::weak_ptr previous_; + }; + + core::VectorIterator::Pointer create_iterator() const override { + ++scans; + return supports_scan ? std::make_unique(*this) : nullptr; + } +}; + +// An existing source that implements only the original random-read API. +class RandomAccessSource : public core::VectorSource { + public: + explicit RandomAccessSource(const BatchSource &source) : source_(source) {} + + const void *get_vector(uint32_t id) const override { + reads.push_back(id); + return source_.get_vector(id); + } + + mutable std::vector reads; + + private: + const BatchSource &source_; +}; + +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()}}; +} + +FlatIndexParam::Pointer Param(bool external, + MetricType metric = MetricType::kL2sq, + bool id_map = true) { + return FlatIndexParamBuilder() + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_metric_type(metric) + .with_use_external_vector(external) + .with_use_id_map(id_map) + .build(); +} + +void SameResults(const SearchResult &want, const SearchResult &got) { + ASSERT_EQ(want.doc_list_.size(), got.doc_list_.size()); + for (size_t i = 0; i < want.doc_list_.size(); ++i) { + EXPECT_EQ(want.doc_list_[i].key(), got.doc_list_[i].key()); + EXPECT_NEAR(want.doc_list_[i].score(), got.doc_list_[i].score(), 1e-5); + } +} + +class ExternalFlatTest : public testing::Test { + protected: + std::vector opened_; + std::vector paths_; + + std::string test_path(const std::string &suffix) { + std::string path = + "test_external_flat_" + + std::string( + testing::UnitTest::GetInstance()->current_test_info()->name()) + + suffix; + paths_.push_back(path); + zvec::test_util::RemoveTestFiles(path); + return path; + } + + Index::Pointer open_index(const BaseIndexParam::Pointer ¶m, + const std::string &path, bool create = true) { + auto index = IndexFactory::CreateAndInitIndex(*param); + EXPECT_NE(nullptr, index); + if (!index) return index; + int ret = index->open(path, {StorageOptions::StorageType::kMMAP, create}); + EXPECT_EQ(0, ret); + if (!ret) opened_.push_back(index); + return index; + } + + void add_vectors(Index *index, const BatchSource &source, + bool external = true) { + for (auto &[id, row] : source.rows) { + ASSERT_EQ(0, external ? index->add_with_source(Vector(row), id, source) + : index->add(Vector(row), id)); + } + } + + void close_index(const Index::Pointer &index) { + EXPECT_EQ(0, index->close()); + opened_.erase(std::remove(opened_.begin(), opened_.end(), index), + opened_.end()); + } + + void TearDown() override { + for (auto &index : opened_) index->close(); + for (auto &path : paths_) zvec::test_util::RemoveTestFiles(path); + } +}; + +class InspectableFlatIndex : public FlatIndex { + public: + using Index::acquire_context; + using Index::init; +}; + +TEST_F(ExternalFlatTest, RegistrationDoesNotCopyVectorPayload) { + BatchSource source(3); + auto index = std::make_shared(); + ASSERT_EQ(0, index->init(*Param(true))); + ASSERT_EQ(0, index->open(test_path("index"), + {StorageOptions::StorageType::kMMAP, true})); + opened_.push_back(index); + auto &context = index->acquire_context(); + ASSERT_NE(nullptr, context); + *context->mutable_features() = "unchanged"; + add_vectors(index.get(), source); + ASSERT_EQ("unchanged", *context->mutable_features()); + EXPECT_EQ(source.rows.size(), index->get_doc_count()); + + // Keep the generic add guards when bypassing its conversion pipeline. + EXPECT_NE( + 0, index->add_with_source(VectorData{DenseVector{nullptr}}, 3, source)); + EXPECT_NE(0, index->add_with_source( + VectorData{SparseVector{0, nullptr, nullptr}}, 3, source)); + const auto vector = Vector(source.rows.at(3)); + EXPECT_NE(0, index->add_with_source(vector, UINT32_MAX, source)); + EXPECT_NE(0, index->add_with_source(vector, 99999, source)); + EXPECT_NE(0, index->add(vector, 3)); + close_index(index); + EXPECT_NE(0, index->add_with_source(vector, 3, source)); +} + +TEST_F(ExternalFlatTest, FullScanAppliesSnapshotFilterBeforeReadingVectors) { + BatchSource source(30); + auto external = open_index(Param(true), test_path("external")); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + // Model an immutable snapshot whose invisible rows have no usable pointer. + const std::set invisible{3, 16, 42}; + source.null_scan_ids = invisible; + source.rows[99999] = std::vector(kDimension, 0.3f); + source.null_scan_ids.insert(99999); // Unregistered rows are also skipped. + std::vector query(kDimension, 0.3f); + auto make_query = [&]() { + auto qp = FlatQueryParamBuilder().with_topk(10).build(); + qp->filter = std::make_shared(); + qp->filter->set([invisible](uint64_t id) { return invisible.count(id); }); + return qp; + }; + SearchResult expected, actual; + ASSERT_EQ(0, embedded->search(Vector(query), make_query(), &expected)); + auto qp = make_query(); + // Reusing the query must preserve its visibility filter across requests. + for (int i = 0; i < 2; ++i) { + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &actual)); + SameResults(expected, actual); + EXPECT_TRUE(qp->filter->is_valid()); + } +} + +TEST_F(ExternalFlatTest, ScanMatchesEmbeddedAndBorrowsReturnedVectors) { + int case_id = 0; + for (auto metric : {MetricType::kL2sq, MetricType::kInnerProduct}) { + for (bool id_map : {false, true}) { + BatchSource source; + 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); + ASSERT_EQ(source.rows.size(), external->get_doc_count()); + source.rows[99999] = std::vector(kDimension, 0.3f); + std::vector query(kDimension, 0.3f); + auto qp = + FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, + external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(1, source.scans); + EXPECT_EQ(0, source.batch_reads); + close_index(external); + ASSERT_EQ(11, got.doc_list_.size()); + // External vectors are returned directly in the input layout, without + // creating a redundant dequantized copy in reverted_vector_list_. + EXPECT_TRUE(got.reverted_vector_list_.empty()); + for (size_t i = 0; i < got.doc_list_.size(); ++i) { + const auto &doc = got.doc_list_[i]; + // Returned vectors borrow the stable random-access source, not the + // transient scan pages. The source remains alive after index close. + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + } + } + } +} + +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); + add_vectors(embedded.get(), data, false); + // Invisible rows must be excluded before calling the source. + const std::set invisible{3, 16, 42}; + for (auto id : invisible) data.rows.erase(id); + data.rows[99999] = std::vector(kDimension, 0.3f); + auto qp = FlatQueryParamBuilder().with_topk(10).build(); + qp->filter = std::make_shared(); + qp->filter->set([invisible](uint64_t id) { return invisible.count(id); }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_EQ(27, source.reads.size()); + for (auto id : source.reads) { + EXPECT_EQ(0, invisible.count(id)); + EXPECT_NE(99999, id); + } + source.reads.clear(); + qp->filter->set([](uint64_t) { return true; }); + ASSERT_EQ(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); + EXPECT_TRUE(source.reads.empty()); + qp->filter.reset(); + EXPECT_NE(0, external->search_with_source(Vector(query), qp, source, &got)); + EXPECT_TRUE(got.doc_list_.empty()); +} + +TEST_F(ExternalFlatTest, 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([first_id = kFirstId](uint64_t id) { + return id <= first_id + 1023 || id % 7 == 0; + }); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + 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}) { + for (bool id_map : {false, true}) { + BatchSource source(2103); + source.supports_scan = false; + // Break ties in the generated rows when comparing full result lists. + for (auto &[id, row] : source.rows) row[0] += id * 0.0001f; + auto param = Param(true, metric, id_map); + auto suffix = std::to_string(case_id++); + auto path = test_path(suffix + "external"); + auto external = open_index(param, path); + auto embedded = + open_index(Param(false, metric, id_map), test_path(suffix + "base")); + add_vectors(external.get(), source); + add_vectors(embedded.get(), source, false); + source.rows[99999] = std::vector(kDimension, 0.3f); + auto qp = + FlatQueryParamBuilder().with_topk(11).with_fetch_vector(true).build(); + std::vector query(kDimension, 0.3f); + SearchResult want, got; + ASSERT_EQ(0, embedded->search(Vector(query), qp, &want)); + for (int reopened = 0; reopened < 2; ++reopened) { + source.batch_reads = 0; + ASSERT_EQ( + 0, external->search_with_source(Vector(query), qp, source, &got)); + SameResults(want, got); + EXPECT_GE(source.batch_reads.load(), 3u); + for (auto &doc : got.doc_list_) { + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + EXPECT_NE(99999, doc.key()); + } + if (!reopened) { + ASSERT_EQ(0, external->flush()); + close_index(external); + external = open_index(param, path, false); + } + } + EXPECT_EQ(2, source.scans); + } + } +} + +TEST_F(ExternalFlatTest, CandidateLookupFilterRadiusAndFetch) { + BatchSource source; + source.supports_scan = false; + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); + auto qp = + FlatQueryParamBuilder().with_topk(5).with_fetch_vector(true).build(); + qp->radius = 0.5f; + qp->filter = std::make_shared(); + qp->filter->set([](uint64_t id) { return id == 3; }); + qp->bf_pks = std::make_shared>(); + for (auto &[id, row] : source.rows) qp->bf_pks->push_back(id); + qp->bf_pks->push_back(999999); + qp->bf_pks->push_back(1ull << 40); + std::vector query(kDimension, 0.3f); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_EQ(0, source.scans); + ASSERT_FALSE(result.doc_list_.empty()); + for (auto &doc : result.doc_list_) { + EXPECT_NE(3, doc.key()); + EXPECT_LE(doc.score(), qp->radius); + EXPECT_EQ(source.get_vector(doc.key()), doc.vector()); + } + VectorDataBuffer buffer; + EXPECT_NE(0, index->fetch(3, &buffer)); + EXPECT_NE(0, index->add(Vector(query), 3)); + EXPECT_NE(0, index->search(Vector(query), qp, &result)); +} + +TEST_F(ExternalFlatTest, ReopenUpdatesAndRejectsWrongStorageMode) { + BatchSource source; + auto path = test_path("index"); + auto index = open_index(Param(true, MetricType::kL2sq, false), path); + add_vectors(index.get(), source); + auto *streamer = + dynamic_cast *>(index->index_searcher().get()); + ASSERT_NE(nullptr, streamer); + EXPECT_EQ(ailego_align(sizeof(core::BlockHeader) + sizeof(core::DeletionMap) + + 32 * sizeof(uint64_t), + 32), + streamer->entity().linear_block_size()); + source.rows[3] = std::vector(kDimension, 12.0f); + ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(3)), 3, source)); + EXPECT_EQ(89, index->get_doc_count()); + ASSERT_EQ(0, index->flush()); + close_index(index); + index = open_index(Param(true, MetricType::kL2sq, false), path, false); + ASSERT_EQ(89, index->get_doc_count()); + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(source.rows.at(3)), qp, source, + &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); + source.rows[1000000] = std::vector(kDimension, 24.0f); + ASSERT_EQ(0, index->add_with_source(Vector(source.rows.at(1000000)), 1000000, + source)); + EXPECT_EQ(90, index->get_doc_count()); + close_index(index); + auto wrong = IndexFactory::CreateAndInitIndex(*Param(false)); + ASSERT_NE(nullptr, wrong); + EXPECT_NE(0, wrong->open(path, {StorageOptions::StorageType::kMMAP, false})); + wrong.reset(); + auto embedded_path = test_path("embedded"); + auto embedded = open_index(Param(false), embedded_path); + close_index(embedded); + wrong = IndexFactory::CreateAndInitIndex(*Param(true)); + EXPECT_NE(0, wrong->open(embedded_path, + {StorageOptions::StorageType::kMMAP, false})); +} + +TEST_F(ExternalFlatTest, RejectsReadFailuresAndClearsRequestSource) { + BatchSource source; + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); + auto qp = FlatQueryParamBuilder().with_topk(3).build(); + std::vector query(kDimension, 0.3f); + SearchResult result; + for (int failure = 0; failure < 4; ++failure) { + source.fail_scan = failure == 0; + source.malformed_scan = failure == 1; + source.null_scan = failure == 2; + source.supports_scan = failure != 3; + source.fail_read = failure == 3; + result.doc_list_.emplace_back(999, 0); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + EXPECT_NE(0, index->search(Vector(query), qp, &result)); + } + source.fail_scan = source.malformed_scan = source.null_scan = false; + source.supports_scan = true; + source.fail_read = false; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + source.fail_read = true; + qp->fetch_vector = true; + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + qp->fetch_vector = false; + qp->bf_pks = std::make_shared>( + std::initializer_list{3}); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + source.fail_read = false; + source.rows.erase(3); + EXPECT_NE(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_NE(0, index->add_with_source(Vector(query), 3, source)); +} + +TEST_F(ExternalFlatTest, ConcurrentRequestsKeepTheirSourceAndCursor) { + BatchSource first(50), second(50, 2.0f); + second.supports_scan = false; + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), first); + std::vector query(kDimension, 0.3f); + SearchResult expected[2]; + BatchSource *sources[] = {&first, &second}; + auto qp = FlatQueryParamBuilder().with_topk(5).build(); + for (int i = 0; i < 2; ++i) { + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, *sources[i], + &expected[i])); + } + std::vector threads; + for (int i = 0; i < 4; ++i) { + threads.emplace_back([&, i] { + auto local_qp = FlatQueryParamBuilder().with_topk(5).build(); + for (int j = 0; j < 30; ++j) { + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), local_qp, + *sources[i % 2], &result)); + SameResults(expected[i % 2], result); + } + }); + } + for (auto &thread : threads) thread.join(); + EXPECT_EQ(61, first.scans); + EXPECT_EQ(61, second.scans); + EXPECT_EQ(0, first.batch_reads); + EXPECT_EQ(61, second.batch_reads); +} + +TEST_F(ExternalFlatTest, RejectsUnsupportedFormatsAndUnboundOperations) { + for (int i = 0; i < 6; ++i) { + auto param = Param(true); + if (i == 0) param->metric_type = MetricType::kCosine; + if (i == 1) + param->quantizer_param = + std::make_shared(QuantizerType::kInt8); + if (i == 2) param->storage_data_type = DataType::DT_FP16; + if (i == 3) param->major_order = core::IndexMeta::MO_COLUMN; + if (i == 4) param->use_contiguous_memory = true; + if (i == 5) param->is_sparse = true; + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)); + } + BatchSource source(1); + auto index = open_index(Param(true), test_path("index")); + add_vectors(index.get(), source); + EXPECT_NE(0, index->index_searcher()->dump(nullptr)); + auto provider = index->create_index_provider(); + ASSERT_NE(nullptr, provider); + EXPECT_EQ(1, provider->count()); + EXPECT_EQ(nullptr, provider->create_iterator()); + EXPECT_EQ(nullptr, provider->get_vector(3)); + EXPECT_NE(0, index->merge({index}, IndexFilter{})); + EXPECT_EQ(0, index->merge({}, IndexFilter{})); + auto embedded = open_index(Param(false), test_path("embedded")); + add_vectors(embedded.get(), source, false); + EXPECT_NE(0, embedded->merge({index}, IndexFilter{})); + + auto qp = FlatQueryParamBuilder().with_topk(1).build(); + qp->group_by_param = std::make_shared(); + qp->group_by_param->group_count = 1; + qp->group_by_param->group_topk = 1; + qp->group_by_param->group_by = [](uint64_t) { return std::string("group"); }; + SearchResult result; + const auto query = Vector(source.rows.at(3)); + EXPECT_EQ(core::IndexError_Unsupported, + index->search_with_source(query, qp, source, &result)); + EXPECT_TRUE(result.group_doc_list_.empty()); + qp->group_by_param.reset(); + // The existing refine entry point has no source binding for external Flat. + qp->refiner_param = std::make_shared(); + qp->refiner_param->scale_factor_ = 1; + qp->refiner_param->reference_index = index; + EXPECT_NE(0, embedded->search(query, qp, &result)); +} + +TEST_F(ExternalFlatTest, BufferPoolReopen) { + BatchSource source(40); + auto path = test_path("index"); + auto index = open_index(Param(true), path); + add_vectors(index.get(), source); + close_index(index); + index = IndexFactory::CreateAndInitIndex(*Param(true)); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kBufferPool, + false, true})); + opened_.push_back(index); + auto qp = + FlatQueryParamBuilder().with_topk(1).with_fetch_vector(true).build(); + SearchResult result; + auto vector = Vector(source.rows.at(3)); + EXPECT_NE(0, index->add_with_source(vector, 3, source)); + EXPECT_NE(0, index->search(vector, qp, &result)); + ASSERT_EQ(0, index->search_with_source(vector, qp, source, &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); + source.supports_scan = false; + ASSERT_EQ(0, index->search_with_source(vector, qp, source, &result)); + ASSERT_EQ(1, result.doc_list_.size()); + EXPECT_EQ(3, result.doc_list_[0].key()); + EXPECT_EQ(0, result.doc_list_[0].score()); +} + +TEST_F(ExternalFlatTest, LowLevelBatchQueriesAndMetadataValidation) { + BatchSource source(20); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_EQ(0, storage->init(zvec::ailego::Params{})); + ASSERT_EQ(0, storage->open(test_path("index"), true)); + core::IndexMeta meta; + meta.set_meta(core::IndexMeta::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params{}); + zvec::ailego::Params params; + params.set(core::PARAM_FLAT_USE_EXTERNAL_VECTOR, true); + auto streamer = core::IndexFactory::CreateStreamer("FlatStreamer16"); + ASSERT_EQ(0, streamer->init(meta, params)); + ASSERT_EQ(0, streamer->open(storage)); + auto context = streamer->create_context(); + auto *ctx = dynamic_cast *>(context.get()); + ASSERT_NE(nullptr, ctx); + ctx->set_vector_source(&source); + ctx->set_topk(1); + core::IndexQueryMeta qmeta(core::IndexMeta::DT_FP32, kDimension); + for (auto &[id, row] : source.rows) { + ASSERT_EQ(0, streamer->add_impl(id, row.data(), qmeta, context)); + } + auto query = source.rows.at(3); + query.insert(query.end(), source.rows.at(16).begin(), + source.rows.at(16).end()); + ASSERT_EQ(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + ASSERT_EQ(1, context->result(0).size()); + ASSERT_EQ(1, context->result(1).size()); + EXPECT_EQ(3, context->result(0)[0].key()); + EXPECT_EQ(16, context->result(1)[0].key()); + EXPECT_EQ(2, source.scans); + source.supports_scan = false; + ASSERT_EQ(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + ASSERT_EQ(1, context->result(0).size()); + ASSERT_EQ(1, context->result(1).size()); + EXPECT_EQ(3, context->result(0)[0].key()); + EXPECT_EQ(16, context->result(1)[0].key()); + EXPECT_EQ(2, source.batch_reads); + core::IndexQueryMeta wrong(core::IndexMeta::DT_FP32, kDimension - 1); + EXPECT_NE(0, streamer->search_bf_impl(query.data(), wrong, 2, context)); + ctx->set_group_params(1, 1); + EXPECT_EQ(core::IndexError_Unsupported, + streamer->search_bf_impl(query.data(), qmeta, 2, context)); + EXPECT_EQ(core::IndexError_Unsupported, + streamer->search_bf_by_p_keys_impl(query.data(), {{3}, {16}}, qmeta, + 2, context)); + ctx->set_group_params(0, 0); + context->reset(); + EXPECT_NE(0, streamer->search_bf_impl(query.data(), qmeta, 2, context)); + EXPECT_EQ(0, streamer->close()); + context.reset(); + EXPECT_EQ(0, storage->close()); +} + +TEST_F(ExternalFlatTest, EmptyIndexAndEmptyCandidates) { + BatchSource source(0); + auto index = open_index(Param(true), test_path("index")); + auto qp = FlatQueryParamBuilder().with_topk(3).build(); + std::vector query(kDimension, 0.3f); + SearchResult result; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + source.supports_scan = false; + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); + EXPECT_EQ(0, source.batch_reads); + qp->bf_pks = std::make_shared>(); + ASSERT_EQ(0, index->search_with_source(Vector(query), qp, source, &result)); + EXPECT_TRUE(result.doc_list_.empty()); +} +} // namespace diff --git a/tests/core/interface/hnsw_turbo_index_test.cc b/tests/core/interface/hnsw_turbo_index_test.cc index a1d0a95c2a..35c00a92f5 100644 --- a/tests/core/interface/hnsw_turbo_index_test.cc +++ b/tests/core/interface/hnsw_turbo_index_test.cc @@ -97,6 +97,32 @@ class TestExternalVectorSource final : public zvec::core::VectorSource { const std::vector> *vectors_; }; +class TestContiguousVectorSource final + : public zvec::core::ContiguousVectorSource { + public: + TestContiguousVectorSource(const std::vector> &vectors, + size_t padding) + : stride_elements_(kDimension + padding), + data_(vectors.size() * stride_elements_, -999.0f) { + for (size_t i = 0; i < vectors.size(); ++i) { + std::copy(vectors[i].begin(), vectors[i].end(), + data_.begin() + i * stride_elements_); + } + } + + zvec::core::ContiguousBatch get_contiguous_view() const override { + return {data_.data(), + 0, + static_cast(data_.size() / stride_elements_), + stride_elements_ * sizeof(float), + {}}; + } + + private: + size_t stride_elements_; + std::vector data_; +}; + std::vector> RandomVectors(size_t count = kVectorCount) { std::mt19937 gen(2026); std::uniform_real_distribution dist(-1.0f, 1.0f); @@ -116,6 +142,19 @@ std::vector> RandomVectors(size_t count = kVectorCount) { return vectors; } +std::vector> ExternalVectors(MetricType metric) { + auto vectors = RandomVectors(kGraphVectorCount); + if (metric == MetricType::kCosine) { + for (size_t i = 0; i < vectors.size(); ++i) { + const float scale = static_cast(i % 7 + 1); + for (float &value : vectors[i]) { + value *= scale; + } + } + } + return vectors; +} + HNSWIndexParam::Pointer MakeParam(MetricType metric, QuantizerType quantizer, bool enable_rotate = false) { return HNSWIndexParamBuilder() @@ -294,21 +333,11 @@ void CheckTurboAddSearchReopen(MetricType metric, QuantizerType quantizer, zvec::test_util::RemoveTestFiles(path); } -void CheckExternalTurboAddSearchReopen(MetricType metric, - QuantizerType quantizer, - const char *quantizer_name, - const std::string &path) { +void CheckExternalTurboAddSearchReopen( + MetricType metric, QuantizerType quantizer, const char *quantizer_name, + const std::string &path, const std::vector> &vectors, + const zvec::core::VectorSource &source) { zvec::test_util::RemoveTestFiles(path); - auto vectors = RandomVectors(kGraphVectorCount); - if (metric == MetricType::kCosine) { - for (size_t i = 0; i < vectors.size(); ++i) { - const float scale = static_cast(i % 7 + 1); - for (float &value : vectors[i]) { - value *= scale; - } - } - } - TestExternalVectorSource source(&vectors); auto param = MakeParam(metric, quantizer); param->use_external_vector = true; @@ -641,8 +670,22 @@ TEST_P(HnswTurboIndexTest, AddSearchReopenFetch) { TEST_P(HnswTurboIndexTest, ExternalVectorsUseTurbo) { const auto &[quantizer, metric] = GetParam(); + auto vectors = ExternalVectors(metric.type); + TestExternalVectorSource source(&vectors); CheckExternalTurboAddSearchReopen(metric.type, quantizer.type, quantizer.name, - index_path("external")); + index_path("external"), vectors, source); +} + +TEST_P(HnswTurboIndexTest, ContiguousExternalVectorsUseTurbo) { + const auto &[quantizer, metric] = GetParam(); + auto vectors = ExternalVectors(metric.type); + for (size_t padding : {0, 3}) { + SCOPED_TRACE(padding); + TestContiguousVectorSource source(vectors, padding); + ASSERT_NO_FATAL_FAILURE(CheckExternalTurboAddSearchReopen( + metric.type, quantizer.type, quantizer.name, + index_path("contiguous_external"), vectors, source)); + } } TEST_P(HnswTurboIndexTest, OriginalProviderBuildUsesTurbo) { diff --git a/tests/core/metric/quantized_integer_metric_test.cc b/tests/core/metric/quantized_integer_metric_test.cc index 57541fcd41..09ec8437ed 100644 --- a/tests/core/metric/quantized_integer_metric_test.cc +++ b/tests/core/metric/quantized_integer_metric_test.cc @@ -344,62 +344,86 @@ TEST(QuantizedIntegerMetric, TestInt8SquaredEuclideanMetric) { } TEST(QuantizedIntegerMetric, TestInt4SquaredEuclidean) { - std::random_device rd; - std::mt19937 gen(rd()); + std::mt19937 gen(42); std::uniform_real_distribution dist(-1.0, 2.0); - - const size_t DIMENSION = std::uniform_int_distribution(1, 128)(gen) * 2; const size_t COUNT = 1000; - IndexMeta meta; - meta.set_meta(IndexMeta::DT_FP32, DIMENSION); - auto converter = IndexFactory::CreateConverter("Int4StreamingConverter"); - ASSERT_TRUE(!!converter); - ASSERT_EQ(0u, converter->init(meta, Params())); - - auto holder = GetHolder(DIMENSION, COUNT, dist); - ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); - auto holder2 = converter->result(); - EXPECT_EQ(COUNT, holder2->count()); - EXPECT_EQ(IndexMeta::DT_INT4, holder2->data_type()); - auto &meta2 = converter->meta(); - - auto reformer = IndexFactory::CreateReformer(meta2.reformer_name()); - ASSERT_TRUE(reformer); - ASSERT_EQ(0u, reformer->init(meta2.reformer_params())); - - ailego::NumericalVector vec(DIMENSION); - for (size_t j = 0; j < DIMENSION; ++j) { - vec[j] = dist(gen); - } - IndexQueryMeta qmeta; - qmeta.set_meta(IndexMeta::DT_FP32, DIMENSION); - IndexQueryMeta qmeta2; - std::string out; - ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &out, &qmeta2)); - ASSERT_EQ(qmeta2.dimension(), meta2.dimension()); - - auto iter = holder->create_iterator(); - auto iter2 = holder2->create_iterator(); - auto metric = IndexFactory::CreateMetric(meta2.metric_name()); - ASSERT_TRUE(!!metric); - ASSERT_EQ(0, metric->init(meta2, meta2.metric_params())); - auto compute = metric->distance(); - ASSERT_TRUE(compute); - - for (; iter->is_valid(); iter->next(), iter2->next()) { - const float *mf = (const float *)iter->data(); - const int8_t *mi = (const int8_t *)iter2->data(); - const int8_t *qi = reinterpret_cast(&out[0]); - float v1 = - ailego::Distance::SquaredEuclidean(mf, vec.data(), holder->dimension()); - float v2; - compute(mi, qi, holder2->dimension(), &v2); - ASSERT_NEAR(v1, v2, 0.2 * DIMENSION); - std::string out2; - ASSERT_EQ(0, reformer->convert(iter->data(), qmeta, &out2, &qmeta2)); - ASSERT_EQ(out2.size(), holder2->element_size()); - ASSERT_EQ(0, std::memcmp(out2.data(), iter2->data(), out2.size())); + for (const size_t dimension : {2, 4, 8, 14, 32, 128, 256}) { + SCOPED_TRACE(dimension); + IndexMeta meta; + meta.set_meta(IndexMeta::DT_FP32, dimension); + auto converter = IndexFactory::CreateConverter("Int4StreamingConverter"); + ASSERT_TRUE(converter); + ASSERT_EQ(0u, converter->init(meta, Params())); + + auto holder = + std::make_shared>(dimension); + for (size_t i = 0; i < COUNT; ++i) { + ailego::NumericalVector row(dimension); + for (size_t j = 0; j < dimension; ++j) row[j] = dist(gen); + if (dimension == 4 && i == 0) { + row = ailego::NumericalVector{-0.4362f, 1.95854f, 1.87264f, + 1.43786f}; + } + holder->emplace(i, row); + } + ASSERT_EQ(0u, IndexConverter::TrainAndTransform(converter, holder)); + auto holder2 = converter->result(); + EXPECT_EQ(COUNT, holder2->count()); + EXPECT_EQ(IndexMeta::DT_INT4, holder2->data_type()); + auto &meta2 = converter->meta(); + + auto reformer = IndexFactory::CreateReformer(meta2.reformer_name()); + ASSERT_TRUE(reformer); + ASSERT_EQ(0u, reformer->init(meta2.reformer_params())); + + ailego::NumericalVector vec(dimension); + for (size_t j = 0; j < dimension; ++j) vec[j] = dist(gen); + if (dimension == 4) { + vec = ailego::NumericalVector{-0.95011f, 0.08062f, -0.86274f, + 1.42267f}; + } + IndexQueryMeta qmeta(IndexMeta::DT_FP32, dimension); + IndexQueryMeta qmeta2; + std::string out; + ASSERT_EQ(0, reformer->transform(vec.data(), qmeta, &out, &qmeta2)); + ASSERT_EQ(qmeta2.dimension(), meta2.dimension()); + std::string restored_query; + ASSERT_EQ(0, reformer->revert(out.data(), qmeta2, &restored_query)); + + auto iter = holder->create_iterator(); + auto iter2 = holder2->create_iterator(); + auto metric = IndexFactory::CreateMetric(meta2.metric_name()); + ASSERT_TRUE(metric); + ASSERT_EQ(0, metric->init(meta2, meta2.metric_params())); + auto compute = metric->distance(); + ASSERT_TRUE(compute); + + for (; iter->is_valid(); iter->next(), iter2->next()) { + SCOPED_TRACE(iter->key()); + std::string restored_row; + ASSERT_EQ(0, reformer->revert(iter2->data(), qmeta2, &restored_row)); + // Validate the metric against decoded vectors. Lossy INT4 rounding does + // not guarantee the former 0.2 * dimension bound against raw vectors. + const float expected = ailego::Distance::SquaredEuclidean( + reinterpret_cast(restored_row.data()), + reinterpret_cast(restored_query.data()), dimension); + float actual; + compute(iter2->data(), out.data(), holder2->dimension(), &actual); + ASSERT_NEAR(expected, actual, 1e-5 * dimension); + + if (dimension == 4 && iter->key() == 0) { + // This fixed pair exceeds the old tolerance with a correct metric. + const float raw = ailego::Distance::SquaredEuclidean( + static_cast(iter->data()), vec.data(), dimension); + ASSERT_GT(std::abs(raw - expected), 0.2 * dimension); + } + + std::string out2; + ASSERT_EQ(0, reformer->convert(iter->data(), qmeta, &out2, &qmeta2)); + ASSERT_EQ(out2.size(), holder2->element_size()); + ASSERT_EQ(0, std::memcmp(out2.data(), iter2->data(), out2.size())); + } } }