diff --git a/src/core/algorithm/ivf/ivf_entity.h b/src/core/algorithm/ivf/ivf_entity.h index d27f53e..e6fd4b6 100644 --- a/src/core/algorithm/ivf/ivf_entity.h +++ b/src/core/algorithm/ivf/ivf_entity.h @@ -171,6 +171,19 @@ class IVFEntity { return *static_cast(data); } + //! Retrieve the key-order mapping (sorted rank -> local_id). + //! mapping[rank] is the local_id of the vector with the rank-th smallest + //! key. Returns nullptr if mapping segment is unavailable. + const uint32_t *get_key_order_mapping() const { + if (!mapping_) return nullptr; + const void *data = nullptr; + const size_t size = vector_count() * sizeof(uint32_t); + if (mapping_->read(0, &data, size) != size) { + return nullptr; + } + return static_cast(data); + } + //! Retrieve vector by local id const void *get_vector(size_t id) const; diff --git a/src/core/algorithm/ivf/ivf_index_provider.h b/src/core/algorithm/ivf/ivf_index_provider.h index fda8b33..12f73b1 100644 --- a/src/core/algorithm/ivf/ivf_index_provider.h +++ b/src/core/algorithm/ivf/ivf_index_provider.h @@ -13,6 +13,9 @@ // limitations under the License. #pragma once +#include +#include +#include #include #include "ivf_entity.h" @@ -33,7 +36,7 @@ class IVFIndexProvider : public IndexProvider { public: //! Create a new iterator virtual Iterator::Pointer create_iterator(void) override { - return Iterator::Pointer(new (std::nothrow) Iterator(entity_)); + return Iterator::Pointer(new (std::nothrow) SortedIterator(entity_)); } //! Retrieve count of vectors @@ -67,13 +70,62 @@ class IVFIndexProvider : public IndexProvider { } private: + class SortedIterator : public IndexProvider::Iterator { + public: + SortedIterator(const IVFEntity::Pointer &entity) : entity_(entity) { + count_ = entity_->vector_count(); + mapping_ = entity_->get_key_order_mapping(); + if (!mapping_) { + // Fallback: compute sorting if mapping segment is unavailable + fallback_.resize(count_); + std::iota(fallback_.begin(), fallback_.end(), size_t(0)); + std::sort(fallback_.begin(), fallback_.end(), [&](size_t a, size_t b) { + return entity_->get_key(a) < entity_->get_key(b); + }); + } + } + + //! Retrieve pointer of data + //! NOTICE: the vec feature will be changed after iterating to next, so + //! the caller need to keep a copy of it before iterator to next vector + virtual const void *data(void) const override { + return entity_->get_vector(current_local_id()); + } + + //! Test if the iterator is valid + virtual bool is_valid(void) const override { + return pos_ < count_; + } + + //! Retrieve primary key + virtual uint64_t key(void) const override { + return entity_->get_key(current_local_id()); + } + + //! Next iterator + virtual void next(void) override { + ++pos_; + } + + private: + size_t current_local_id() const { + return mapping_ ? static_cast(mapping_[pos_]) : fallback_[pos_]; + } + + //! Members + IVFEntity::Pointer entity_; + const uint32_t *mapping_{nullptr}; // points into mapping_ segment data + std::vector fallback_; // used only if mapping_ unavailable + size_t count_{0}; + size_t pos_{0}; + }; + + //! Original sequential iterator (kept for potential internal use) class Iterator : public IndexProvider::Iterator { public: Iterator(const IVFEntity::Pointer &entity) : entity_(entity) {} //! Retrieve pointer of data - //! NOTICE: the vec feature will be changed after iterating to next, so - //! the caller need to keep a copy of it before iterator to next vector virtual const void *data(void) const override { return entity_->get_vector(index_); } diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index ecf01a7..f310dc6 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -2588,20 +2588,24 @@ TEST_F(CollectionTest, Feature_Optimize_General) { } TEST_F(CollectionTest, Feature_Optimize_Repeated) { - auto func = [&](QuantizeType quantize_type = QuantizeType::UNDEFINED) { + auto func = [&](QuantizeType quantize_type = QuantizeType::UNDEFINED, + std::string index_type = "HNSW") { FileHelper::RemoveDirectory(col_path); int doc_count = 1000; // create empty collection CollectionSchema::Ptr schema; - if (quantize_type == QuantizeType::UNDEFINED) { - schema = TestHelper::CreateSchemaWithVectorIndex(); - } else { + if (index_type == "HNSW") { schema = TestHelper::CreateSchemaWithVectorIndex( false, "demo", std::make_shared(MetricType::IP, 16, 200, quantize_type)); + } else if (index_type == "IVF") { + schema = TestHelper::CreateSchemaWithVectorIndex( + false, "demo", + std::make_shared(MetricType::IP, 10, 4, false, + quantize_type)); } auto options = CollectionOptions{false, true, 64 * 1024 * 1024}; auto collection = TestHelper::CreateCollectionWithDoc( @@ -2676,6 +2680,10 @@ TEST_F(CollectionTest, Feature_Optimize_Repeated) { check_doc(); std::cout << "check success 2" << std::endl; }; + // unquantized + func(QuantizeType::UNDEFINED, "IVF"); + // quantized + func(QuantizeType::FP16, "IVF"); // unquantized func(); diff --git a/tests/db/index/utils/utils.cc b/tests/db/index/utils/utils.cc index e901ee9..95a72a9 100644 --- a/tests/db/index/utils/utils.cc +++ b/tests/db/index/utils/utils.cc @@ -112,8 +112,15 @@ CollectionSchema::Ptr TestHelper::CreateNormalSchema( "dense_int8", DataType::VECTOR_INT8, 128, false, std::make_shared(MetricType::IP))); + // IVF and HNSW_RABITQ do not support sparse vectors, always use Flat for + // sparse fields in those cases. + auto supports_sparse = [](const IndexParams::Ptr ¶ms) { + auto type = params->type(); + return type != IndexType::IVF && type != IndexType::HNSW_RABITQ; + }; + IndexParams::Ptr sparse_index_params; - if (vector_index_params) { + if (vector_index_params && supports_sparse(vector_index_params)) { sparse_index_params = vector_index_params->clone(); auto v = std::dynamic_pointer_cast(sparse_index_params); // sparse always use IP