diff --git a/python/tests/test_turbo_hnsw_int8.py b/python/tests/test_turbo_hnsw_int8.py new file mode 100644 index 000000000..5fd0c1f34 --- /dev/null +++ b/python/tests/test_turbo_hnsw_int8.py @@ -0,0 +1,98 @@ +# 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. + +"""Packaged-extension regression tests for the HNSW Turbo INT8 path.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import zvec +from zvec import ( + CollectionOption, + CollectionSchema, + Doc, + HnswIndexParam, + HnswQueryParam, + Query, + VectorSchema, +) +from zvec.typing import DataType, MetricType, QuantizeType + +DIMENSION = 32 +DOC_COUNT = 64 + + +def _normalized_vectors() -> list[list[float]]: + rng = np.random.default_rng(2026) + vectors = rng.standard_normal((DOC_COUNT, DIMENSION)).astype(np.float32) + vectors /= np.linalg.norm(vectors, axis=1, keepdims=True) + return vectors.tolist() + + +@pytest.mark.parametrize( + "metric", [MetricType.COSINE, MetricType.L2], ids=["cosine", "l2"] +) +def test_turbo_int8_hnsw_python_registration_roundtrip(tmp_path, metric): + """Keep Turbo's factory registration reachable from the Python module.""" + path = str(tmp_path / f"turbo_int8_hnsw_{metric.name.lower()}") + vectors = _normalized_vectors() + schema = CollectionSchema( + name="turbo_int8_hnsw", + vectors=[ + VectorSchema( + "dense", + DataType.VECTOR_FP32, + dimension=DIMENSION, + index_param=HnswIndexParam( + metric_type=metric, + m=16, + ef_construction=100, + quantize_type=QuantizeType.INT8, + ), + ) + ], + ) + + collection = zvec.create_and_open( + path=path, + schema=schema, + option=CollectionOption(read_only=False, enable_mmap=True), + ) + try: + docs = [ + Doc(id=str(i), vectors={"dense": vector}) + for i, vector in enumerate(vectors) + ] + for status in collection.insert(docs): + assert status.ok() + collection.optimize() + + query = Query( + field_name="dense", + vector=vectors[17], + param=HnswQueryParam(ef=128), + ) + hits = collection.query(query, topk=5) + assert hits[0].id == "17" + finally: + collection.close() + + reopened = zvec.open(path, option=CollectionOption(read_only=True)) + try: + hits = reopened.query(query, topk=5) + assert hits[0].id == "17" + finally: + reopened.close() diff --git a/src/core/algorithm/hnsw/CMakeLists.txt b/src/core/algorithm/hnsw/CMakeLists.txt index d59dc8ffc..c95cbf105 100644 --- a/src/core/algorithm/hnsw/CMakeLists.txt +++ b/src/core/algorithm/hnsw/CMakeLists.txt @@ -10,7 +10,7 @@ cc_library( NAME core_knn_hnsw STATIC SHARED STRICT ALWAYS_LINK SRCS *.cc - LIBS core_framework core_utility sparsehash + LIBS core_framework core_utility sparsehash zvec_turbo INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm LDFLAGS "${CORE_KNN_HNSW_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" diff --git a/src/core/algorithm/hnsw/hnsw_context.h b/src/core/algorithm/hnsw/hnsw_context.h index 08eaf8a74..39858e5f7 100644 --- a/src/core/algorithm/hnsw/hnsw_context.h +++ b/src/core/algorithm/hnsw/hnsw_context.h @@ -142,6 +142,18 @@ class HnswContext : public IndexContext { return vector_source_; } + //! Keep the unquantized input vector for external-vector graph + //! construction. The streamer receives a quantized datapoint from the + //! interface, but build-time node-to-node comparisons must stay in the + //! external source's input layout. + inline void set_external_build_query(const void *query) { + external_build_query_ = query; + } + + inline const void *external_build_query() const { + return external_build_query_; + } + inline void reset_query_raw(const void *query, const IndexMeta &meta) { dc_.set_dim(meta.dimension()); dc_.reset_query(query); @@ -289,8 +301,9 @@ class HnswContext : public IndexContext { //! distance computation inline void reset_query(const void *query, const IndexMeta &meta) { dc_.set_dim(meta.dimension()); - if (auto query_preprocess_func = index_metric_->get_query_preprocess_func(); - query_preprocess_func != nullptr) { + auto query_preprocess_func = + index_metric_ ? index_metric_->get_query_preprocess_func() : nullptr; + if (query_preprocess_func != nullptr) { size_t dim = meta.dimension(); preprocess_buffer_.resize(dim); memcpy(preprocess_buffer_.data(), query, dim); @@ -444,6 +457,7 @@ class HnswContext : public IndexContext { set_group_params(0, 0); reset_group_by(); set_vector_source(nullptr); + set_external_build_query(nullptr); dc_.set_provider(nullptr); } @@ -616,6 +630,7 @@ class HnswContext : public IndexContext { HnswDistCalculator dc_; IndexMetric::Pointer metric_; const VectorSource *vector_source_{nullptr}; + const void *external_build_query_{nullptr}; size_t vector_data_size_{0}; size_t extra_values_size_{0}; diff --git a/src/core/algorithm/hnsw/hnsw_dist_calculator.h b/src/core/algorithm/hnsw/hnsw_dist_calculator.h index fd2cf47d9..e7ddbb902 100644 --- a/src/core/algorithm/hnsw/hnsw_dist_calculator.h +++ b/src/core/algorithm/hnsw/hnsw_dist_calculator.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include #include @@ -38,8 +39,9 @@ class HnswDistCalculator { HnswDistCalculator(const HnswEntity *entity, const IndexMetric::Pointer &metric, uint32_t dim) : entity_(entity), - distance_(metric->distance()), - batch_distance_(metric->batch_distance()), + distance_(metric ? metric->distance() : IndexMetric::MatrixDistance{}), + batch_distance_(metric ? metric->batch_distance() + : IndexMetric::MatrixBatchDistance{}), query_(nullptr), dim_(dim), compare_cnt_(0) {} @@ -49,8 +51,9 @@ class HnswDistCalculator { const IndexMetric::Pointer &metric, uint32_t dim, const void *query) : entity_(entity), - distance_(metric->distance()), - batch_distance_(metric->batch_distance()), + distance_(metric ? metric->distance() : IndexMetric::MatrixDistance{}), + batch_distance_(metric ? metric->batch_distance() + : IndexMetric::MatrixBatchDistance{}), query_(query), dim_(dim), compare_cnt_(0) {} @@ -59,23 +62,26 @@ class HnswDistCalculator { HnswDistCalculator(const HnswEntity *entity, const IndexMetric::Pointer &metric) : entity_(entity), - distance_(metric->distance()), - batch_distance_(metric->batch_distance()), + distance_(metric ? metric->distance() : IndexMetric::MatrixDistance{}), + batch_distance_(metric ? metric->batch_distance() + : IndexMetric::MatrixBatchDistance{}), query_(nullptr), dim_(0), compare_cnt_(0) {} void update(const HnswEntity *entity, const IndexMetric::Pointer &metric) { entity_ = entity; - distance_ = metric->distance(); - batch_distance_ = metric->batch_distance(); + distance_ = metric ? metric->distance() : IndexMetric::MatrixDistance{}; + batch_distance_ = + metric ? metric->batch_distance() : IndexMetric::MatrixBatchDistance{}; } void update(const HnswEntity *entity, const IndexMetric::Pointer &metric, uint32_t dim) { entity_ = entity; - distance_ = metric->distance(); - batch_distance_ = metric->batch_distance(); + distance_ = metric ? metric->distance() : IndexMetric::MatrixDistance{}; + batch_distance_ = + metric ? metric->batch_distance() : IndexMetric::MatrixBatchDistance{}; dim_ = dim; } @@ -107,6 +113,11 @@ class HnswDistCalculator { float score{0.0f}; + if (ailego_unlikely(!distance_)) { + LOG_ERROR("Distance function is not initialized"); + error_ = true; + return 0.0f; + } distance_(vec_lhs, vec_rhs, dim_, &score); return score; @@ -186,6 +197,12 @@ class HnswDistCalculator { const void **extra_values) { compare_cnt_++; + if (ailego_unlikely(!batch_distance_)) { + LOG_ERROR("Batch distance function is not initialized"); + error_ = true; + std::fill(distances, distances + num, 0.0f); + return; + } batch_distance_(vecs, query_, num, dim_, distances, extra_values); } @@ -198,6 +215,11 @@ class HnswDistCalculator { return 0.0f; } dist_t score = 0; + if (ailego_unlikely(!batch_distance_)) { + LOG_ERROR("Batch distance function is not initialized"); + error_ = true; + return 0.0f; + } if (extra_values != nullptr) { batch_distance_(&feat, query_, 1, dim_, &score, &extra_values); } else { diff --git a/src/core/algorithm/hnsw/hnsw_streamer.cc b/src/core/algorithm/hnsw/hnsw_streamer.cc index c9b11c450..3abcc2c68 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.cc +++ b/src/core/algorithm/hnsw/hnsw_streamer.cc @@ -16,6 +16,7 @@ #include #include #include +#include #include "utility/sparse_utility.h" #include "hnsw_algorithm.h" #include "hnsw_context.h" @@ -25,6 +26,34 @@ namespace zvec { namespace core { +namespace { + +using TurboQuantizer = zvec::turbo::Quantizer; + +// Adapt Turbo's typed operations to the callbacks used by graph traversal. +template +void BindTurboDistances(TurboQuantizer::Pointer quantizer, + IndexMetric::MatrixDistance &distance, + IndexMetric::MatrixBatchDistance &batch_distance) { + distance = [quantizer](const void *lhs, const void *rhs, size_t, float *out) { + *out = ((*quantizer).*Distance)(lhs, rhs); + }; + batch_distance = [quantizer](const void **vectors, const void *query, + size_t count, size_t, float *out, + const void **) { + if constexpr (BatchDistance != nullptr) { + ((*quantizer).*BatchDistance)(vectors, static_cast(count), query, + out); + } else { + for (size_t i = 0; i < count; ++i) { + out[i] = ((*quantizer).*Distance)(vectors[i], query); + } + } + }; +} + +} // namespace + HnswStreamer::HnswStreamer() = default; HnswStreamer::~HnswStreamer() { @@ -190,6 +219,21 @@ int HnswStreamer::init(const IndexMeta &imeta, const ailego::Params ¶ms) { return 0; } +int HnswStreamer::init( + const IndexMeta &imeta, const ailego::Params ¶ms, + const std::shared_ptr &quantizer) { + if (!quantizer) { + return this->init(imeta, params); + } + + quantizer_ = quantizer; + int ret = this->init(imeta, params); + if (ret != 0) { + quantizer_.reset(); + } + return ret; +} + int HnswStreamer::cleanup() { if (state_ == STATE_OPENED) { this->close(); @@ -199,10 +243,16 @@ int HnswStreamer::cleanup() { meta_.clear(); metric_.reset(); + quantizer_.reset(); + add_distance_ = {}; + add_batch_distance_ = {}; + search_distance_ = {}; + search_batch_distance_ = {}; stats_.clear(); provider_.reset(); provider_meta_.clear(); provider_metric_.reset(); + provider_quantizer_.reset(); if (entity_) { entity_->cleanup(); } @@ -308,7 +358,8 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { if (index_meta.dimension() != meta_.dimension() || index_meta.element_size() != meta_.element_size() || index_meta.metric_name() != meta_.metric_name() || - index_meta.data_type() != meta_.data_type()) { + index_meta.data_type() != meta_.data_type() || + index_meta.quantizer_name() != meta_.quantizer_name()) { LOG_ERROR("IndexMeta mismatch from the previous in index"); return IndexError_Mismatch; } @@ -330,106 +381,163 @@ int HnswStreamer::open(IndexStorage::Pointer stg) { } } - metric_ = IndexFactory::CreateMetric(meta_.metric_name()); - if (!metric_) { - LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); - return IndexError_NoExist; - } - ret = metric_->init(meta_, meta_.metric_params()); - if (ret != 0) { - LOG_ERROR("Failed to init metric, ret=%d", ret); - return ret; - } - - if (!metric_->distance()) { - LOG_ERROR("Invalid metric distance"); - return IndexError_InvalidArgument; - } - - if (!metric_->batch_distance()) { - LOG_ERROR("Invalid metric batch distance"); - return IndexError_InvalidArgument; - } - - add_distance_ = metric_->distance(); - add_batch_distance_ = metric_->batch_distance(); - const size_t stored_vector_size = meta_.element_size(); - const size_t stored_extra_values_size = - metric_->extra_values_size_per_vector(); auto valid_vector_layout = [](size_t vector_size, size_t extra_values_size) { return extra_values_size == 0 || extra_values_size < vector_size; }; - if (!valid_vector_layout(stored_vector_size, stored_extra_values_size)) { - LOG_ERROR("Invalid HNSW vector layout, vector_size=%zu extra_size=%zu", - stored_vector_size, stored_extra_values_size); - return IndexError_InvalidArgument; - } - search_distance_ = add_distance_; - search_batch_distance_ = add_batch_distance_; + if (quantizer_) { + if (use_external_vector_) { + // External records and build queries are raw; search queries are encoded. + BindTurboDistances<&TurboQuantizer::calc_distance_input_input, + &TurboQuantizer::calc_distance_input_input_batch>( + quantizer_, add_distance_, add_batch_distance_); + BindTurboDistances<&TurboQuantizer::calc_distance_input_query, + &TurboQuantizer::calc_distance_input_query_batch>( + quantizer_, search_distance_, search_batch_distance_); + } else { + BindTurboDistances<&TurboQuantizer::calc_distance_dp_dp>( + quantizer_, add_distance_, add_batch_distance_); + BindTurboDistances<&TurboQuantizer::calc_distance_dp_query, + &TurboQuantizer::calc_distance_dp_query_batch>( + quantizer_, search_distance_, search_batch_distance_); + } + } else { + metric_ = IndexFactory::CreateMetric(meta_.metric_name()); + if (!metric_) { + LOG_ERROR("Failed to create metric %s", meta_.metric_name().c_str()); + return IndexError_NoExist; + } + ret = metric_->init(meta_, meta_.metric_params()); + if (ret != 0) { + LOG_ERROR("Failed to init metric, ret=%d", ret); + return ret; + } + + if (!metric_->distance() || !metric_->batch_distance()) { + LOG_ERROR("Invalid metric distance"); + return IndexError_InvalidArgument; + } + + add_distance_ = metric_->distance(); + add_batch_distance_ = metric_->batch_distance(); - const auto query_metric = metric_->query_metric(); - if (query_metric && query_metric->distance() && - query_metric->batch_distance()) { - const size_t query_extra_values_size = - query_metric->extra_values_size_per_vector(); - if (query_extra_values_size != stored_extra_values_size) { - LOG_ERROR( - "HNSW query metric layout mismatch, stored_extra_size=%zu " - "query_extra_size=%zu", - stored_extra_values_size, query_extra_values_size); + const size_t stored_vector_size = meta_.element_size(); + const size_t stored_extra_values_size = + metric_->extra_values_size_per_vector(); + if (!valid_vector_layout(stored_vector_size, stored_extra_values_size)) { + LOG_ERROR("Invalid HNSW vector layout, vector_size=%zu extra_size=%zu", + stored_vector_size, stored_extra_values_size); return IndexError_InvalidArgument; } - search_distance_ = query_metric->distance(); - search_batch_distance_ = query_metric->batch_distance(); + + search_distance_ = add_distance_; + search_batch_distance_ = add_batch_distance_; + + const auto query_metric = metric_->query_metric(); + if (query_metric && query_metric->distance() && + query_metric->batch_distance()) { + const size_t query_extra_values_size = + query_metric->extra_values_size_per_vector(); + if (query_extra_values_size != stored_extra_values_size) { + LOG_ERROR( + "HNSW query metric layout mismatch, stored_extra_size=%zu " + "query_extra_size=%zu", + stored_extra_values_size, query_extra_values_size); + return IndexError_InvalidArgument; + } + search_distance_ = query_metric->distance(); + search_batch_distance_ = query_metric->batch_distance(); + } } - //! Create a dedicated build metric when the provider meta differs from - //! the index meta in layout or metric, so build distances run in the - //! original vector space + //! Create a dedicated build distance path for original-vector providers. + //! Plain FP32 providers use Turbo in their original vector space. Check + //! the layout here as core callers can bypass the interface's selection; + //! other providers keep their dtype-compatible legacy build metric. provider_metric_.reset(); + provider_quantizer_.reset(); if (provider_) { - const bool layout_differs = - provider_meta_.data_type() != meta_.data_type() || - provider_meta_.dimension() != meta_.dimension() || - provider_meta_.element_size() != meta_.element_size(); const bool use_index_metric = provider_meta_.metric_name().empty(); - if (layout_differs || !use_index_metric) { - const std::string &metric_name = - use_index_metric ? meta_.metric_name() : provider_meta_.metric_name(); - const ailego::Params &metric_params = - use_index_metric ? meta_.metric_params() - : provider_meta_.metric_params(); - provider_metric_ = IndexFactory::CreateMetric(metric_name); - if (!provider_metric_) { - LOG_ERROR("Failed to create provider metric %s", metric_name.c_str()); + const std::string &provider_metric_name = + use_index_metric ? meta_.metric_name() : provider_meta_.metric_name(); + const bool use_turbo_provider = + quantizer_ && provider_meta_.data_type() == IndexMeta::DT_FP32 && + provider_meta_.unit_size() == sizeof(float) && + provider_meta_.extra_meta_size() == 0 && + provider_meta_.element_size() == + static_cast(provider_meta_.dimension()) * sizeof(float) && + (provider_metric_name == "SquaredEuclidean" || + provider_metric_name == "Cosine" || + provider_metric_name == "InnerProduct"); + if (use_turbo_provider) { + IndexMeta provider_quantizer_meta = provider_meta_; + if (use_index_metric) { + provider_quantizer_meta.set_metric(meta_.metric_name(), 0, + meta_.metric_params()); + } + provider_quantizer_ = IndexFactory::CreateQuantizer("Fp32Quantizer"); + if (!provider_quantizer_) { + LOG_ERROR("Failed to create provider Fp32Quantizer"); return IndexError_NoExist; } - ret = provider_metric_->init(provider_meta_, metric_params); + ret = + provider_quantizer_->init(provider_quantizer_meta, ailego::Params{}); if (ret != 0) { - LOG_ERROR("Failed to init provider metric, ret=%d", ret); + LOG_ERROR("Failed to init provider Fp32Quantizer, ret=%d", ret); return ret; } - if (!provider_metric_->distance() || - !provider_metric_->batch_distance()) { - LOG_ERROR("Invalid provider metric distance"); - return IndexError_InvalidArgument; + BindTurboDistances<&TurboQuantizer::calc_distance_input_input, + &TurboQuantizer::calc_distance_input_input_batch>( + provider_quantizer_, add_distance_, add_batch_distance_); + } else { + const bool layout_differs = + provider_meta_.data_type() != meta_.data_type() || + provider_meta_.dimension() != meta_.dimension() || + provider_meta_.element_size() != meta_.element_size(); + // Turbo search leaves metric_ empty. Always create a separate metric + // when this provider cannot use Turbo, even if its layout matches. + if (quantizer_ || layout_differs || !use_index_metric) { + const std::string &metric_name = use_index_metric + ? meta_.metric_name() + : provider_meta_.metric_name(); + const ailego::Params &metric_params = + use_index_metric ? meta_.metric_params() + : provider_meta_.metric_params(); + provider_metric_ = IndexFactory::CreateMetric(metric_name); + if (!provider_metric_) { + LOG_ERROR("Failed to create provider metric %s", metric_name.c_str()); + return IndexError_NoExist; + } + ret = provider_metric_->init(provider_meta_, metric_params); + if (ret != 0) { + LOG_ERROR("Failed to init provider metric, ret=%d", ret); + return ret; + } + if (!provider_metric_->distance() || + !provider_metric_->batch_distance()) { + LOG_ERROR("Invalid provider metric distance"); + return IndexError_InvalidArgument; + } + add_distance_ = provider_metric_->distance(); + add_batch_distance_ = provider_metric_->batch_distance(); } - add_distance_ = provider_metric_->distance(); - add_batch_distance_ = provider_metric_->batch_distance(); } + // Turbo input vectors have no framework-managed extra values. The legacy + // path keeps validating its metric-specific provider layout. const IndexMetric *add_metric = provider_metric_ ? provider_metric_.get() : metric_.get(); - const size_t provider_vector_size = provider_meta_.element_size(); - const size_t provider_extra_values_size = - add_metric->extra_values_size_per_vector(); - if (!valid_vector_layout(provider_vector_size, - provider_extra_values_size)) { - LOG_ERROR( - "Invalid HNSW provider vector layout, vector_size=%zu " - "extra_size=%zu", - provider_vector_size, provider_extra_values_size); - return IndexError_InvalidArgument; + if (!provider_quantizer_ && add_metric) { + const size_t provider_vector_size = provider_meta_.element_size(); + const size_t provider_extra_values_size = + add_metric->extra_values_size_per_vector(); + if (!valid_vector_layout(provider_vector_size, + provider_extra_values_size)) { + LOG_ERROR( + "Invalid HNSW provider vector layout, vector_size=%zu " + "extra_size=%zu", + provider_vector_size, provider_extra_values_size); + return IndexError_InvalidArgument; + } } } @@ -478,7 +586,9 @@ int HnswStreamer::close() { LOG_INFO("HnswStreamer close"); stats_.clear(); - meta_.set_metric(metric_->name(), 0, metric_->params()); + if (metric_) { + meta_.set_metric(metric_->name(), 0, metric_->params()); + } entity_->set_index_meta(meta_); int ret = entity_->close(); if (ret != 0) { @@ -492,7 +602,9 @@ int HnswStreamer::close() { int HnswStreamer::flush(uint64_t checkpoint) { LOG_INFO("HnswStreamer flush checkpoint=%zu", (size_t)checkpoint); - meta_.set_metric(metric_->name(), 0, metric_->params()); + if (metric_) { + meta_.set_metric(metric_->name(), 0, metric_->params()); + } entity_->set_index_meta(meta_); return entity_->flush(checkpoint); } @@ -588,14 +700,18 @@ void HnswStreamer::bind_add_dist_space(HnswContext *ctx) const { provider_metric_ ? provider_metric_.get() : metric_.get(); const size_t vector_size = provider_ ? provider_meta_.element_size() : meta_.element_size(); + // The quantizer path has no metric, and quantizers keep no extra values + const size_t extra_values_size = + add_metric ? add_metric->extra_values_size_per_vector() : 0; ctx->bind_dist_space(add_distance_, add_batch_distance_, provider_, - vector_size, add_metric->extra_values_size_per_vector()); + vector_size, extra_values_size); } void HnswStreamer::bind_search_dist_space(HnswContext *ctx) const { + const size_t extra_values_size = + metric_ ? metric_->extra_values_size_per_vector() : 0; ctx->bind_dist_space(search_distance_, search_batch_distance_, nullptr, - meta_.element_size(), - metric_->extra_values_size_per_vector()); + meta_.element_size(), extra_values_size); } //! Add a vector with id into index @@ -655,11 +771,18 @@ int HnswStreamer::add_with_id_impl(uint32_t id, const void *query, return ret; } ctx->reset_query_raw(original_query_block.data(), provider_meta_); + } else if (ailego_unlikely(use_external_vector_ && quantizer_)) { + if (ailego_unlikely(ctx->external_build_query() == nullptr)) { + LOG_ERROR("External build query is not set"); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidArgument; + } + ctx->reset_query_raw(ctx->external_build_query(), meta_); } else { ctx->reset_query(query, meta_); } - if (metric_->support_train()) { + if (metric_ && metric_->support_train()) { const std::lock_guard lk(mutex_); ret = metric_->train(query, meta_.dimension()); if (ailego_unlikely(ret != 0)) { @@ -749,11 +872,18 @@ int HnswStreamer::add_impl(uint64_t pkey, const void *query, return ret; } ctx->reset_query_raw(original_query_block.data(), provider_meta_); + } else if (ailego_unlikely(use_external_vector_ && quantizer_)) { + if (ailego_unlikely(ctx->external_build_query() == nullptr)) { + LOG_ERROR("External build query is not set"); + (*stats_.mutable_discarded_count())++; + return IndexError_InvalidArgument; + } + ctx->reset_query_raw(ctx->external_build_query(), meta_); } else { ctx->reset_query(query, meta_); } - if (metric_->support_train()) { + if (metric_ && metric_->support_train()) { const std::lock_guard lk(mutex_); ret = metric_->train(query, meta_.dimension()); if (ailego_unlikely(ret != 0)) { diff --git a/src/core/algorithm/hnsw/hnsw_streamer.h b/src/core/algorithm/hnsw/hnsw_streamer.h index 91ee055fd..f4a198f8c 100644 --- a/src/core/algorithm/hnsw/hnsw_streamer.h +++ b/src/core/algorithm/hnsw/hnsw_streamer.h @@ -90,10 +90,26 @@ class HnswStreamer : public IndexStreamer { return entity_->storage_mode(); } + //! Whether the active search distance path is backed by a turbo quantizer. + bool uses_turbo_distance() const { + return quantizer_ != nullptr; + } + + //! Whether graph construction is backed by a turbo quantizer. Compatible + //! FP32 providers use a dedicated quantizer; other providers use a metric. + bool uses_turbo_build_distance() const { + return quantizer_ != nullptr && + (provider_ == nullptr || provider_quantizer_ != nullptr); + } + protected: //! Initialize Streamer int init(const IndexMeta &imeta, const ailego::Params ¶ms) override; + //! Initialize Streamer with a turbo quantizer for distance computation + int init(const IndexMeta &imeta, const ailego::Params ¶ms, + const std::shared_ptr &quantizer) override; + //! Cleanup Streamer int cleanup() override; @@ -260,6 +276,7 @@ class HnswStreamer : public IndexStreamer { IndexMetric::MatrixBatchDistance add_batch_distance_{}; IndexMetric::MatrixBatchDistance search_batch_distance_{}; + std::shared_ptr quantizer_{}; Stats stats_{}; std::mutex mutex_{}; @@ -267,7 +284,8 @@ class HnswStreamer : public IndexStreamer { // provider of the original vectors used to build graph IndexProvider::Pointer provider_{}; IndexMeta provider_meta_{}; - IndexMetric::Pointer provider_metric_{}; + IndexMetric::Pointer provider_metric_{}; // legacy provider distance path + std::shared_ptr provider_quantizer_{}; size_t max_index_size_{0UL}; size_t chunk_size_{HnswEntity::kDefaultChunkSize}; diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 58ec4acce..1932d4013 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -558,6 +558,10 @@ int Index::open(const std::string &file_path, StorageOptions storage_options) { core::IndexError::What(ret)); return core::IndexError_Runtime; } + ret = prepare_streamer_open(storage_options); + if (ret != 0) { + return ret; + } if (streamer_ == nullptr || streamer_->open(storage_) != 0) { LOG_ERROR("Failed to open streamer, path: %s", file_path.c_str()); return core::IndexError_Runtime; @@ -851,9 +855,7 @@ int Index::_dense_fetch(const uint32_t doc_id, out_vector_buffer.resize(input_vector_meta_.element_size()); if (turbo_quantizer_ != nullptr) { - // The stored record is int8 codes + quantizer tail; dequantize restores - // the original FP32 vector (cosine layouts also denormalize by the - // stored norm). + // Decode the quantizer's stored layout back to the original vector format. if (turbo_quantizer_->dequantize(vector, streamer_vector_meta_, &out_vector_buffer) != 0) { LOG_ERROR("Failed to dequantize vector"); @@ -1184,7 +1186,9 @@ int Index::_collect_dense_result( } } if (turbo_quantizer_) { - if (context->fetch_vector()) { + // External HNSW vectors are already in the caller's input layout. They + // are not stored quantizer codes and therefore must not be dequantized. + if (context->fetch_vector() && !param_.use_external_vector) { int revert_err = 0; auto revert_one = [&](const void *vec, std::vector *out) { if (revert_err) return; diff --git a/src/core/interface/indexes/flat_index.cc b/src/core/interface/indexes/flat_index.cc index 29c5d2e70..1fed9d433 100644 --- a/src/core/interface/indexes/flat_index.cc +++ b/src/core/interface/indexes/flat_index.cc @@ -183,7 +183,7 @@ int FlatIndex::fallback_to_legacy_pipeline() { turbo_quantizer_.reset(); streamer_.reset(); - // Redo the Index::Init() setup down the legacy branch. + // Redo the Index::init() setup down the legacy branch. proxima_index_meta_.clear(); proxima_index_meta_.set_meta(param_.data_type, param_.dimension); proxima_index_meta_.set_meta_type(is_sparse_ diff --git a/src/core/interface/indexes/hnsw_index.cc b/src/core/interface/indexes/hnsw_index.cc index 0b9703af5..3dd81f519 100644 --- a/src/core/interface/indexes/hnsw_index.cc +++ b/src/core/interface/indexes/hnsw_index.cc @@ -14,6 +14,8 @@ #include #include +#include +#include #include #include "algorithm/hnsw/hnsw_context.h" #include "algorithm/hnsw/hnsw_params.h" @@ -23,6 +25,113 @@ namespace zvec::core_interface { +namespace { + +const char *ResolveTurboQuantizerName(const QuantizerParam &quantizer_param, + const HNSWIndexParam &hnsw_param) { + // Turbo quantizers currently consume FP32 inputs and own dense, in-index + // vector storage. External-vector HNSW is also supported: its source stays + // in the FP32 input layout and the streamer quantizes source vectors only + // for distance calculation. + if (hnsw_param.is_sparse || hnsw_param.data_type != DataType::DT_FP32 || + hnsw_param.metric_type == MetricType::kMIPSL2sq) { + return nullptr; + } + + // An original-vector provider is a separate build space. Turbo can keep + // that path in FP32 as long as the provider exposes plain FP32 vectors and + // uses a metric supported by Fp32Quantizer. Other provider layouts retain + // the legacy metric pipeline. + if (hnsw_param.provider) { + const auto &provider_meta = hnsw_param.provider_meta; + const auto &provider_metric = provider_meta.metric_name(); + const bool supported_provider_metric = + provider_metric.empty() || provider_metric == "SquaredEuclidean" || + provider_metric == "Cosine" || provider_metric == "InnerProduct"; + if (provider_meta.data_type() != core::IndexMeta::DT_FP32 || + provider_meta.dimension() != + static_cast(hnsw_param.dimension) || + provider_meta.element_size() != + static_cast(hnsw_param.dimension) * sizeof(float) || + !supported_provider_metric) { + return nullptr; + } + } + + // Rotation is still implemented by the legacy integer converters. + if (quantizer_param.enable_rotate) { + return nullptr; + } + + switch (quantizer_param.type) { + case QuantizerType::kNone: + return "Fp32Quantizer"; + case QuantizerType::kFP16: + return "Fp16Quantizer"; + case QuantizerType::kInt8: + return "Int8Quantizer"; + case QuantizerType::kInt4: + return "Int4Quantizer"; + default: + return nullptr; + } +} + +} // namespace + +int HNSWIndex::prepare_streamer_open(const StorageOptions &options) { + if (!turbo_quantizer_ || options.create_new) { + return 0; + } + core::IndexMeta persisted_meta; + if (core::IndexHelper::DeserializeFromStorage(storage_.get(), + &persisted_meta) != 0 || + !persisted_meta.quantizer_name().empty()) { + return 0; + } + // Reuse normal initialization for old indexes, without selecting Turbo again. + turbo_quantizer_.reset(); + streamer_.reset(); + converter_.reset(); + reformer_.reset(); + metric_.reset(); + proxima_index_meta_.clear(); + use_legacy_pipeline_ = true; + return Index::init(param_); +} + +int HNSWIndex::create_and_init_converter_reformer( + const QuantizerParam &quantizer_param, const BaseIndexParam &index_param) { + const auto &hnsw_param = dynamic_cast(index_param); + const char *quantizer_name = + use_legacy_pipeline_ + ? nullptr + : ResolveTurboQuantizerName(quantizer_param, hnsw_param); + if (quantizer_name != nullptr) { + turbo_quantizer_ = core::IndexFactory::CreateQuantizer(quantizer_name); + if (!turbo_quantizer_) { + LOG_ERROR("Failed to create turbo quantizer %s", quantizer_name); + return core::IndexError_Runtime; + } + if (turbo_quantizer_->init(proxima_index_meta_, ailego::Params{}) != 0) { + LOG_ERROR("Failed to init turbo quantizer %s", quantizer_name); + turbo_quantizer_.reset(); + return core::IndexError_Runtime; + } + + proxima_index_meta_ = turbo_quantizer_->meta(); + proxima_index_meta_.set_quantizer(quantizer_name, 0, ailego::Params{}); + streamer_vector_meta_.set_meta( + proxima_index_meta_.data_type(), proxima_index_meta_.dimension(), + static_cast(turbo_quantizer_->type()), + proxima_index_meta_.extra_meta_size()); + streamer_vector_meta_.set_meta_type(proxima_index_meta_.meta_type()); + return core::IndexError_Success; + } + return Index::create_and_init_converter_reformer(quantizer_param, + index_param); +} + std::string HNSWIndex::storage_mode() const { if (!streamer_) { return ""; @@ -55,6 +164,10 @@ int HNSWIndex::add_with_source(const VectorData &vector_data, } if (auto *ctx = dynamic_cast(context.get())) { ctx->set_vector_source(&src); + if (std::holds_alternative(vector_data.vector)) { + ctx->set_external_build_query( + std::get(vector_data.vector).data); + } } return Index::add(vector_data, doc_id); } @@ -134,8 +247,11 @@ int HNSWIndex::create_and_init_streamer(const BaseIndexParam ¶m) { LOG_ERROR("Failed to create streamer"); return core::IndexError_Runtime; } - if (ailego_unlikely( - streamer_->init(proxima_index_meta_, proxima_index_params_) != 0)) { + int ret = turbo_quantizer_ != nullptr && !is_sparse_ + ? streamer_->init(proxima_index_meta_, proxima_index_params_, + turbo_quantizer_) + : streamer_->init(proxima_index_meta_, proxima_index_params_); + if (ailego_unlikely(ret != 0)) { LOG_ERROR("Failed to init streamer"); return core::IndexError_Runtime; } @@ -173,7 +289,12 @@ int HNSWIndex::_prepare_for_search( context->reset_filter(); } if (hnsw_search_param->radius > 0.0f) { - context->set_threshold(hnsw_search_param->radius); + float threshold = hnsw_search_param->radius; + if (turbo_quantizer_ != nullptr && + turbo_quantizer_->support_score_normalization()) { + turbo_quantizer_->denormalize_score(&threshold); + } + context->set_threshold(threshold); } ailego::Params params; const int real_search_ef = diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.cc b/src/core/mixed_reducer/mixed_streamer_reducer.cc index d36287817..09a513bfa 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.cc +++ b/src/core/mixed_reducer/mixed_streamer_reducer.cc @@ -361,14 +361,16 @@ int MixedStreamerReducer::read_vec(size_t source_streamer_index, bool need_encode = need_revert || reformer == nullptr; if (quantizer != nullptr) { // Quantizer-encoded records can be copied raw only into an identical - // target layout; otherwise (or when a builder consumes original vectors) - // dequantize back to the original format. + // target layout, including the metric: INT4 IP and L2 tails have the + // same size but different meanings. Otherwise (or when a builder consumes + // original vectors), dequantize back to the original format. const auto &source_meta = streamer->meta(); const auto &target_meta = target_streamer_->meta(); const bool same_layout = target_builder_ == nullptr && turbo::QuantizerStorageDataTypeMatches(target_meta, source_meta) && target_meta.quantizer_name() == source_meta.quantizer_name() && + target_meta.metric_name() == source_meta.metric_name() && target_meta.data_type() == source_meta.data_type() && target_meta.dimension() == source_meta.dimension() && target_meta.unit_size() == source_meta.unit_size() && diff --git a/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 5fcc10bcd..db014b0b9 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -229,6 +229,11 @@ class ZVEC_CORE_API Index { const ailego::Params &converter_params = {}); virtual int create_and_init_streamer(const BaseIndexParam ¶m) = 0; + //! Adjust the pipeline after storage opens, before the streamer reads it. + virtual int prepare_streamer_open(const StorageOptions & /*options*/) { + return 0; + } + protected: bool init_context(); core::IndexContext::Pointer &acquire_context(); @@ -246,10 +251,9 @@ class ZVEC_CORE_API Index { core::IndexStreamer::Pointer streamer_{}; core::IndexReformer::Pointer reformer_{}; core::IndexConverter::Pointer converter_{}; // for build() - core::IndexMetric::Pointer metric_{}; // to do normalization - // Turbo quantizer for the FLAT-on-turbo path: quantizes records and - // queries and computes distances via turbo SIMD batch kernels. When set, - // converter_/reformer_/metric_ stay null. + core::IndexMetric::Pointer metric_{}; // legacy distance/score path + // Quantizes records and queries and computes distances through turbo SIMD + // kernels. When set, converter_/reformer_/metric_ stay null. std::shared_ptr turbo_quantizer_{}; size_t context_index_{std::numeric_limits::max()}; @@ -363,6 +367,9 @@ class ZVEC_CORE_API HNSWIndex : public Index { protected: int create_and_init_streamer(const BaseIndexParam ¶m) override; + int create_and_init_converter_reformer( + const QuantizerParam ¶m, const BaseIndexParam &index_param) override; + int _prepare_for_search(const VectorData &query, const BaseIndexQueryParam::Pointer &search_param, core::IndexContext::Pointer &context) override; @@ -370,7 +377,10 @@ class ZVEC_CORE_API HNSWIndex : public Index { const BaseIndexQueryParam::Pointer &search_param) override; private: + int prepare_streamer_open(const StorageOptions &options) override; + HNSWIndexParam param_{}; + bool use_legacy_pipeline_{false}; }; class ZVEC_CORE_API VamanaIndex : public Index { diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc index 91b296b6b..e02813fb0 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.cc @@ -186,6 +186,44 @@ float Fp32Quantizer::calc_distance_dp_dp(const void *dp1, return calc_distance_dp_query(dp1, dp2); } +float Fp32Quantizer::calc_distance_input_query(const void *dp, + const void *query) const { + if (meta_.metric_name() != "Cosine") { + return calc_distance_dp_query(dp, query); + } + return Quantizer::calc_distance_input_query(dp, query); +} + +void Fp32Quantizer::calc_distance_input_query_batch(const void *const *dp_list, + int dp_num, + const void *query, + float *dist_list) const { + if (meta_.metric_name() != "Cosine") { + calc_distance_dp_query_batch(dp_list, dp_num, query, dist_list); + return; + } + Quantizer::calc_distance_input_query_batch(dp_list, dp_num, query, dist_list); +} + +float Fp32Quantizer::calc_distance_input_input(const void *dp1, + const void *dp2) const { + if (meta_.metric_name() != "Cosine") { + return calc_distance_dp_query(dp1, dp2); + } + return Quantizer::calc_distance_input_input(dp1, dp2); +} + +void Fp32Quantizer::calc_distance_input_input_batch(const void *const *dp_list, + int dp_num, + const void *query, + float *dist_list) const { + if (meta_.metric_name() != "Cosine") { + calc_distance_dp_query_batch(dp_list, dp_num, query, dist_list); + return; + } + Quantizer::calc_distance_input_input_batch(dp_list, dp_num, query, dist_list); +} + INDEX_FACTORY_REGISTER_QUANTIZER(Fp32Quantizer); } // namespace turbo diff --git a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h index 24ccb4362..ab690a939 100644 --- a/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h +++ b/src/turbo/quantizer/fp32_quantizer/fp32_quantizer.h @@ -90,6 +90,20 @@ class Fp32Quantizer : public Quantizer { float calc_distance_dp_dp(const void *dp1, const void *dp2) const override; + float calc_distance_input_query(const void *dp, + const void *query) const override; + + void calc_distance_input_query_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + + float calc_distance_input_input(const void *dp1, + const void *dp2) const override; + + void calc_distance_input_input_batch(const void *const *dp_list, int dp_num, + const void *query, + float *dist_list) const override; + int quantize(const void *query, const core::IndexQueryMeta &qmeta, std::string *out, core::IndexQueryMeta *ometa) const override; diff --git a/src/turbo/quantizer/int4_quantizer/int4_quantizer.h b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h index 9f2db8a5f..5b40c6187 100644 --- a/src/turbo/quantizer/int4_quantizer/int4_quantizer.h +++ b/src/turbo/quantizer/int4_quantizer/int4_quantizer.h @@ -110,6 +110,18 @@ class Int4Quantizer : public Quantizer { DistanceImpl distance(const void *query, const core::IndexQueryMeta &qmeta) const override; + void normalize_score(float *score) const override { + *score = -(*score); + } + + void denormalize_score(float *score) const override { + *score = -(*score); + } + + bool support_score_normalization() const override { + return meta_.metric_name() == "InnerProduct"; + } + private: //! Byte length of a quantized vector (packed int4 codes + extra meta). size_t quantized_length() const { diff --git a/src/turbo/quantizer/int8_quantizer/int8_quantizer.h b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h index d4e7617c8..0135337d6 100644 --- a/src/turbo/quantizer/int8_quantizer/int8_quantizer.h +++ b/src/turbo/quantizer/int8_quantizer/int8_quantizer.h @@ -107,6 +107,18 @@ class Int8Quantizer : public Quantizer { DistanceImpl distance(const void *query, const core::IndexQueryMeta &qmeta) const override; + void normalize_score(float *score) const override { + *score = -(*score); + } + + void denormalize_score(float *score) const override { + *score = -(*score); + } + + bool support_score_normalization() const override { + return meta_.metric_name() == "InnerProduct"; + } + private: //! Byte length of a quantized vector (int8 codes + extra meta). size_t quantized_length() const { diff --git a/src/turbo/quantizer/quantizer.cc b/src/turbo/quantizer/quantizer.cc new file mode 100644 index 000000000..860dbc860 --- /dev/null +++ b/src/turbo/quantizer/quantizer.cc @@ -0,0 +1,93 @@ +// 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 "quantizer/quantizer.h" +#include +#include + +namespace zvec::turbo { + +namespace { + +struct InputDistanceScratch { + std::string datapoint; + std::string query; + std::string datapoints; + std::vector datapoint_ptrs; +}; + +InputDistanceScratch &GetInputDistanceScratch() { + thread_local InputDistanceScratch scratch; + return scratch; +} + +void QuantizeDatapoints(const Quantizer &quantizer, const void *const *dp_list, + int dp_num, InputDistanceScratch *scratch) { + const size_t length = quantizer.quantized_datapoint_vector_length(); + scratch->datapoints.resize(length * static_cast(dp_num)); + scratch->datapoint_ptrs.resize(static_cast(dp_num)); + for (int i = 0; i < dp_num; ++i) { + void *output = scratch->datapoints.data() + length * i; + quantizer.quantize_data(dp_list[i], output); + scratch->datapoint_ptrs[i] = output; + } +} + +} // namespace + +float Quantizer::calc_distance_input_query(const void *dp, + const void *query) const { + auto &scratch = GetInputDistanceScratch(); + scratch.datapoint.resize(quantized_datapoint_vector_length()); + quantize_data(dp, scratch.datapoint.data()); + return calc_distance_dp_query(scratch.datapoint.data(), query); +} + +void Quantizer::calc_distance_input_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_num <= 0) { + return; + } + auto &scratch = GetInputDistanceScratch(); + QuantizeDatapoints(*this, dp_list, dp_num, &scratch); + calc_distance_dp_query_batch(scratch.datapoint_ptrs.data(), dp_num, query, + dist_list); +} + +float Quantizer::calc_distance_input_input(const void *dp1, + const void *dp2) const { + auto &scratch = GetInputDistanceScratch(); + scratch.datapoint.resize(quantized_datapoint_vector_length()); + scratch.query.resize(quantized_query_vector_length()); + quantize_data(dp1, scratch.datapoint.data()); + quantize_query(dp2, scratch.query.data()); + return calc_distance_dp_query(scratch.datapoint.data(), scratch.query.data()); +} + +void Quantizer::calc_distance_input_input_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const { + if (dp_num <= 0) { + return; + } + auto &scratch = GetInputDistanceScratch(); + QuantizeDatapoints(*this, dp_list, dp_num, &scratch); + scratch.query.resize(quantized_query_vector_length()); + quantize_query(query, scratch.query.data()); + calc_distance_dp_query_batch(scratch.datapoint_ptrs.data(), dp_num, + scratch.query.data(), dist_list); +} + +} // namespace zvec::turbo diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index f9e7864e8..214e80468 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -158,6 +158,32 @@ class Quantizer { //! Distance between two quantized datapoints virtual float calc_distance_dp_dp(const void *dp1, const void *dp2) const = 0; + //! Distance between an input datapoint and an already-quantized query. + //! + //! This is used by indexes whose vectors live outside the index in the + //! quantizer's input layout. The default implementation quantizes the + //! datapoint before dispatching to the quantized distance kernel. + virtual float calc_distance_input_query(const void *dp, + const void *query) const; + + //! Batched input-datapoint to quantized-query distance. + virtual void calc_distance_input_query_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const; + + //! Distance between two vectors in the quantizer's input layout. + //! + //! The default implementation quantizes both sides, which keeps graph + //! construction in the quantizer pipeline even when an original-vector + //! provider supplies the build vectors. + virtual float calc_distance_input_input(const void *dp1, + const void *dp2) const; + + //! Batched input-datapoint to input-query distance. + virtual void calc_distance_input_input_batch(const void *const *dp_list, + int dp_num, const void *query, + float *dist_list) const; + //! Quantize a query vector for search virtual int quantize(const void * /*query*/, const IndexQueryMeta & /*qmeta*/, std::string * /*out*/, diff --git a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc index d6bae6dfd..0f75a964b 100644 --- a/tests/core/algorithm/hnsw/hnsw_streamer_test.cc +++ b/tests/core/algorithm/hnsw/hnsw_streamer_test.cc @@ -27,8 +27,10 @@ #include #include #include +#include #include #include +#include #include #include "tests/test_util.h" @@ -4584,6 +4586,183 @@ TEST_F(HnswStreamerTest, TestCompareFromOriginalVsBaseline) { EXPECT_GT(topk1_recall_b, 0.90f); } +TEST_F(HnswStreamerTest, TestTurboSearchWithFp16ProviderBuildFallback) { + constexpr size_t kProviderDim = 8; + constexpr size_t kCount = 64; + IndexMeta raw_meta(IndexMeta::DataType::DT_FP32, kProviderDim); + raw_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kProviderDim); + + for (bool explicit_metric : {false, true}) { + for (bool explicit_id : {false, true}) { + SCOPED_TRACE(testing::Message() << "explicit_metric=" << explicit_metric + << ", explicit_id=" << explicit_id); + auto provider = + make_shared>( + kProviderDim); + std::vector> vectors(kCount, + std::vector(kProviderDim)); + for (size_t i = 0; i < kCount; ++i) { + NumericalVector original(kProviderDim); + for (size_t j = 0; j < kProviderDim; ++j) { + vectors[i][j] = i * 0.25f + j * 0.03125f; + original[j] = ailego::FloatHelper::ToFP16(vectors[i][j]); + } + ASSERT_TRUE(provider->emplace(i, std::move(original))); + } + IndexMeta provider_meta(IndexMeta::DataType::DT_FP16, kProviderDim); + if (explicit_metric) { + provider_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + } + + auto quantizer = IndexFactory::CreateQuantizer("Fp32Quantizer"); + ASSERT_NE(nullptr, quantizer); + ASSERT_EQ(0, quantizer->init(raw_meta, ailego::Params())); + auto streamer = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + auto hnsw_streamer = std::dynamic_pointer_cast(streamer); + ASSERT_NE(nullptr, hnsw_streamer); + ASSERT_EQ(0, streamer->set_provider(provider, provider_meta)); + ailego::Params params; + params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + params.set(PARAM_HNSW_STREAMER_EF, 100U); + params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 0U); + ASSERT_EQ(0, streamer->init(raw_meta, params, quantizer)); + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + const std::string path = dir_ + "turbo_fp16_provider_" + + std::to_string(explicit_metric) + "_" + + std::to_string(explicit_id) + ".index"; + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + EXPECT_TRUE(hnsw_streamer->uses_turbo_distance()); + ASSERT_FALSE(hnsw_streamer->uses_turbo_build_distance()); + + auto context = streamer->create_context(); + ASSERT_NE(nullptr, context); + auto *ctx = dynamic_cast(context.get()); + ASSERT_NE(nullptr, ctx); + for (size_t i = 0; i < kCount; ++i) { + ASSERT_EQ(0, explicit_id + ? streamer->add_with_id_impl(i, vectors[i].data(), + query_meta, context) + : streamer->add_impl(i, vectors[i].data(), query_meta, + context)); + } + + // The active build calculator must interpret provider records as FP16, + // even though the search quantizer consumes FP32 records. + EXPECT_FLOAT_EQ(0.5f, + ctx->dist_calculator().dist(uint32_t{0}, uint32_t{1})); + ctx->reset_query_raw(provider->get_vector(3), provider_meta); + const void *candidates[] = {provider->get_vector(0), + provider->get_vector(1), + provider->get_vector(2)}; + float distances[3]; + ctx->dist_calculator().batch_dist(candidates, 3, distances, nullptr); + EXPECT_FLOAT_EQ(4.5f, distances[0]); + EXPECT_FLOAT_EQ(2.0f, distances[1]); + EXPECT_FLOAT_EQ(0.5f, distances[2]); + + context->set_topk(1); + for (size_t probe : {size_t{0}, size_t{31}, kCount - 1}) { + ASSERT_EQ(0, streamer->search_impl(vectors[probe].data(), query_meta, + context)); + ASSERT_EQ(1U, context->result().size()); + EXPECT_EQ(probe, context->result()[0].key()); + EXPECT_FLOAT_EQ(0.0f, context->result()[0].score()); + } + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); + } + } +} + +TEST_F(HnswStreamerTest, TestTurboInt8QuantizerDistance) { + constexpr size_t kTurboDim = 35; + constexpr size_t kCount = 128; + constexpr size_t kTopk = 10; + + IndexMeta raw_meta(IndexMeta::DataType::DT_FP32, kTurboDim); + raw_meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + auto quantizer = IndexFactory::CreateQuantizer("Int8Quantizer"); + ASSERT_NE(nullptr, quantizer); + ASSERT_EQ(0, quantizer->init(raw_meta, ailego::Params())); + + IndexMeta quantized_meta = quantizer->meta(); + quantized_meta.set_quantizer("Int8Quantizer", 0, ailego::Params()); + + ailego::Params params; + params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 16U); + params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + params.set(PARAM_HNSW_STREAMER_EF, 100U); + params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 0U); + + auto streamer = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(quantized_meta, params, quantizer)); + + auto storage = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(dir_ + "turbo_int8.index", true)); + ASSERT_EQ(0, streamer->open(storage)); + + std::mt19937 gen(2026); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector> data(kCount, std::vector(kTurboDim)); + std::vector codes(kCount); + IndexQueryMeta raw_qmeta(IndexMeta::DataType::DT_FP32, kTurboDim); + IndexQueryMeta quantized_qmeta; + auto add_ctx = streamer->create_context(); + ASSERT_NE(nullptr, add_ctx); + for (size_t i = 0; i < kCount; ++i) { + for (float &value : data[i]) { + value = dist(gen); + } + ASSERT_EQ(0, quantizer->quantize(data[i].data(), raw_qmeta, &codes[i], + &quantized_qmeta)); + ASSERT_EQ(0, + streamer->add_impl(i, codes[i].data(), quantized_qmeta, add_ctx)); + } + + const size_t query_index = 37; + auto linear_ctx = streamer->create_context(); + ASSERT_NE(nullptr, linear_ctx); + linear_ctx->set_topk(kTopk); + ASSERT_EQ(0, streamer->search_bf_impl(codes[query_index].data(), + quantized_qmeta, linear_ctx)); + + std::vector> expected(kCount); + for (size_t i = 0; i < kCount; ++i) { + expected[i] = {quantizer->calc_distance_dp_query(codes[i].data(), + codes[query_index].data()), + i}; + } + std::partial_sort(expected.begin(), expected.begin() + kTopk, expected.end()); + + const auto &linear_result = linear_ctx->result(); + ASSERT_EQ(kTopk, linear_result.size()); + for (size_t i = 0; i < kTopk; ++i) { + EXPECT_EQ(expected[i].second, linear_result[i].key()); + EXPECT_NEAR(expected[i].first, linear_result[i].score(), + 1e-5f + std::abs(expected[i].first) * 1e-4f); + } + + auto ann_ctx = streamer->create_context(); + ASSERT_NE(nullptr, ann_ctx); + ann_ctx->set_topk(1); + ASSERT_EQ(0, streamer->search_impl(codes[query_index].data(), quantized_qmeta, + ann_ctx)); + ASSERT_EQ(1U, ann_ctx->result().size()); + EXPECT_EQ(query_index, ann_ctx->result()[0].key()); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); +} + } // namespace core } // namespace zvec diff --git a/tests/core/interface/hnsw_turbo_index_test.cc b/tests/core/interface/hnsw_turbo_index_test.cc new file mode 100644 index 000000000..d6e7cacb8 --- /dev/null +++ b/tests/core/interface/hnsw_turbo_index_test.cc @@ -0,0 +1,1046 @@ +// 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 +#include +#include +#include +#include "algorithm/hnsw/hnsw_context.h" +#include "algorithm/hnsw/hnsw_params.h" +#include "algorithm/hnsw/hnsw_streamer.h" +#include "tests/test_util.h" + +using namespace zvec::core_interface; + +namespace { + +constexpr uint32_t kDimension = 36; +constexpr size_t kVectorCount = 200; +constexpr uint32_t kTopK = 10; +constexpr size_t kGraphVectorCount = + zvec::core::HnswEntity::kDefaultBruteForceThreshold + 1; +constexpr uint32_t kGraphQueryIds[] = {37, 101, kGraphVectorCount - 1}; + +using SearchRowList = std::vector>; + +struct TurboQuantizerCase { + QuantizerType type; + const char *name; + float fetch_tolerance; + float score_tolerance; + const char *test_name; +}; + +struct TurboMetricCase { + MetricType type; + const char *test_name; +}; + +constexpr TurboQuantizerCase kTurboQuantizers[] = { + {QuantizerType::kNone, "Fp32Quantizer", 1e-6f, 1e-6f, "Fp32"}, + {QuantizerType::kFP16, "Fp16Quantizer", 1e-3f, 1e-3f, "Fp16"}, + {QuantizerType::kInt8, "Int8Quantizer", 1e-2f, 1e-2f, "Int8"}, + {QuantizerType::kInt4, "Int4Quantizer", 2e-1f, 5e-2f, "Int4"}, +}; +constexpr TurboMetricCase kTurboMetrics[] = { + {MetricType::kL2sq, "L2"}, + {MetricType::kCosine, "Cosine"}, + {MetricType::kInnerProduct, "InnerProduct"}, +}; + +class HnswTurboIndexTest + : public testing::TestWithParam< + std::tuple> { + protected: + std::string index_path(const char *suffix) const { + const auto &[quantizer, metric] = GetParam(); + return std::string("hnsw_turbo_") + quantizer.test_name + "_" + + metric.test_name + "_" + suffix + ".index"; + } +}; + +class TestExternalVectorSource final : public zvec::core::VectorSource { + public: + explicit TestExternalVectorSource( + const std::vector> *vectors) + : vectors_(vectors) {} + + const void *get_vector(uint32_t node_id) const override { + return (*vectors_)[node_id].data(); + } + + private: + const std::vector> *vectors_; +}; + +std::vector> RandomVectors(size_t count = kVectorCount) { + std::mt19937 gen(2026); + std::uniform_real_distribution dist(-1.0f, 1.0f); + std::vector> vectors(count, + std::vector(kDimension)); + for (auto &vector : vectors) { + float norm = 0.0f; + for (float &value : vector) { + value = dist(gen); + norm += value * value; + } + norm = std::sqrt(norm); + for (float &value : vector) { + value /= norm; + } + } + return vectors; +} + +HNSWIndexParam::Pointer MakeParam(MetricType metric, QuantizerType quantizer, + bool enable_rotate = false) { + return HNSWIndexParamBuilder() + .with_metric_type(metric) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_m(16) + .with_ef_construction(100) + .with_quantizer_param(QuantizerParam(quantizer, enable_rotate)) + .build(); +} + +HNSWIndexParam::Pointer MakeDefaultParam(MetricType metric) { + return HNSWIndexParamBuilder() + .with_metric_type(metric) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_m(16) + .with_ef_construction(100) + .build(); +} + +const char *MetricName(MetricType metric) { + switch (metric) { + case MetricType::kL2sq: + return "SquaredEuclidean"; + case MetricType::kCosine: + return "Cosine"; + case MetricType::kInnerProduct: + return "InnerProduct"; + default: + return ""; + } +} + +SearchRowList SearchRows(Index *index, const std::vector &query, + bool linear, bool fetch_vector = false, + const zvec::core::VectorSource *source = nullptr) { + auto query_param = HNSWQueryParamBuilder() + .with_topk(kTopK) + .with_ef_search(100) + .with_is_linear(linear) + .with_fetch_vector(fetch_vector) + .build(); + VectorData query_data{DenseVector{query.data()}}; + SearchResult result; + EXPECT_EQ(0, source == nullptr + ? index->search(query_data, query_param, &result) + : index->search_with_source(query_data, query_param, *source, + &result)); + SearchRowList rows; + for (const auto &doc : result.doc_list_) { + rows.emplace_back(doc.key(), doc.score()); + } + if (fetch_vector) { + if (source == nullptr) { + EXPECT_EQ(rows.size(), result.reverted_vector_list_.size()); + } else { + EXPECT_TRUE(result.reverted_vector_list_.empty()); + for (const auto &doc : result.doc_list_) { + const auto *fetched = static_cast(doc.vector()); + EXPECT_NE(nullptr, fetched); + if (fetched == nullptr) { + continue; + } + const auto *original = + static_cast(source->get_vector(doc.key())); + for (uint32_t d = 0; d < kDimension; ++d) { + EXPECT_FLOAT_EQ(original[d], fetched[d]); + } + } + } + } + return rows; +} + +void CheckGraphSearchEnabled(Index *index) { + auto context = index->index_searcher()->create_context(); + auto *hnsw_context = dynamic_cast(context.get()); + ASSERT_NE(nullptr, hnsw_context); + // The public interface inherits this threshold. A small data set would + // silently run brute force even with is_linear=false. + ASSERT_GT(index->get_doc_count(), hnsw_context->get_bruteforce_threshold()); +} + +void CheckGraphRecall(const SearchRowList &linear_rows, + const SearchRowList &graph_rows) { + ASSERT_EQ(kTopK, linear_rows.size()); + ASSERT_EQ(kTopK, graph_rows.size()); + size_t matches = 0; + for (const auto &graph_row : graph_rows) { + const auto match = std::find_if( + linear_rows.begin(), linear_rows.end(), + [key = graph_row.first](const auto &row) { return row.first == key; }); + if (match != linear_rows.end()) { + ++matches; + EXPECT_FLOAT_EQ(match->second, graph_row.second); + } + } + EXPECT_GE(matches, 9U) << "Graph search must recover at least 90% of the " + "linear top-10 using the same Turbo distances"; +} + +void AddVectors(Index *index, const std::vector> &vectors) { + for (size_t i = 0; i < vectors.size(); ++i) { + VectorData vector_data{DenseVector{vectors[i].data()}}; + ASSERT_EQ(0, index->add(vector_data, static_cast(i))); + } +} + +void CheckTurboAddSearchReopen(MetricType metric, QuantizerType quantizer, + const char *quantizer_name, + float fetch_tolerance, const std::string &path) { + zvec::test_util::RemoveTestFiles(path); + auto vectors = RandomVectors(kGraphVectorCount); + auto param = MakeParam(metric, quantizer); + + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(quantizer_name, index->index_searcher()->meta().quantizer_name()); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + auto streamer = std::dynamic_pointer_cast( + index->index_searcher()); + ASSERT_NE(nullptr, streamer); + EXPECT_TRUE(streamer->uses_turbo_distance()); + EXPECT_TRUE(streamer->uses_turbo_build_distance()); + AddVectors(index.get(), vectors); + ASSERT_EQ(0, index->train()); + + CheckGraphSearchEnabled(index.get()); + std::vector linear_results; + std::vector graph_results; + for (uint32_t query_id : kGraphQueryIds) { + SCOPED_TRACE(query_id); + auto linear_rows = SearchRows(index.get(), vectors[query_id], true, true); + auto graph_rows = SearchRows(index.get(), vectors[query_id], false, true); + CheckGraphRecall(linear_rows, graph_rows); + ASSERT_FALSE(graph_rows.empty()); + EXPECT_EQ(query_id, graph_rows.front().first); + linear_results.push_back(std::move(linear_rows)); + graph_results.push_back(std::move(graph_rows)); + } + + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(37, &fetched)); + const auto *fetched_vector = reinterpret_cast( + std::get(fetched.vector_buffer).data.data()); + for (uint32_t i = 0; i < kDimension; ++i) { + EXPECT_NEAR(vectors[37][i], fetched_vector[i], fetch_tolerance); + } + + ASSERT_EQ(0, index->close()); + + auto reopened = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, reopened); + ASSERT_EQ(0, + reopened->open(path, {StorageOptions::StorageType::kMMAP, false})); + EXPECT_EQ(quantizer_name, + reopened->index_searcher()->meta().quantizer_name()); + CheckGraphSearchEnabled(reopened.get()); + for (size_t i = 0; i < std::size(kGraphQueryIds); ++i) { + SCOPED_TRACE(kGraphQueryIds[i]); + auto linear_rows = + SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], true, true); + auto graph_rows = + SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], false, true); + EXPECT_EQ(linear_results[i], linear_rows); + EXPECT_EQ(graph_results[i], graph_rows); + CheckGraphRecall(linear_rows, graph_rows); + } + ASSERT_EQ(0, reopened->close()); + zvec::test_util::RemoveTestFiles(path); +} + +void CheckExternalTurboAddSearchReopen(MetricType metric, + QuantizerType quantizer, + const char *quantizer_name, + const std::string &path) { + 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; + + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(quantizer_name, index->index_searcher()->meta().quantizer_name()); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + auto streamer = std::dynamic_pointer_cast( + index->index_searcher()); + ASSERT_NE(nullptr, streamer); + EXPECT_TRUE(streamer->uses_turbo_distance()); + EXPECT_TRUE(streamer->uses_turbo_build_distance()); + + for (size_t i = 0; i < vectors.size(); ++i) { + VectorData vector_data{DenseVector{vectors[i].data()}}; + ASSERT_EQ(0, index->add_with_source(vector_data, static_cast(i), + source)); + } + + CheckGraphSearchEnabled(index.get()); + std::vector linear_results; + std::vector graph_results; + for (uint32_t query_id : kGraphQueryIds) { + SCOPED_TRACE(query_id); + auto linear_rows = + SearchRows(index.get(), vectors[query_id], true, true, &source); + auto graph_rows = + SearchRows(index.get(), vectors[query_id], false, true, &source); + CheckGraphRecall(linear_rows, graph_rows); + ASSERT_FALSE(graph_rows.empty()); + EXPECT_EQ(query_id, graph_rows.front().first); + linear_results.push_back(std::move(linear_rows)); + graph_results.push_back(std::move(graph_rows)); + } + + ASSERT_EQ(0, index->close()); + + auto reopened = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, reopened); + ASSERT_EQ(0, + reopened->open(path, {StorageOptions::StorageType::kMMAP, false})); + CheckGraphSearchEnabled(reopened.get()); + for (size_t i = 0; i < std::size(kGraphQueryIds); ++i) { + SCOPED_TRACE(kGraphQueryIds[i]); + auto linear_rows = SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], + true, true, &source); + auto graph_rows = SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], + false, true, &source); + EXPECT_EQ(linear_results[i], linear_rows); + EXPECT_EQ(graph_results[i], graph_rows); + CheckGraphRecall(linear_rows, graph_rows); + } + ASSERT_EQ(0, reopened->close()); + zvec::test_util::RemoveTestFiles(path); +} + +void CheckOriginalProviderUsesTurbo(MetricType metric, QuantizerType quantizer, + const char *quantizer_name, + const std::string &path) { + 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; + } + } + } + + auto provider = std::make_shared>(kDimension); + for (size_t i = 0; i < vectors.size(); ++i) { + zvec::ailego::NumericalVector vector(kDimension); + for (uint32_t d = 0; d < kDimension; ++d) { + vector[d] = vectors[i][d]; + } + ASSERT_TRUE(provider->emplace(i, vector)); + } + zvec::core::IndexMeta provider_meta(zvec::core::IndexMeta::DT_FP32, + kDimension); + provider_meta.set_metric(MetricName(metric), 0, zvec::ailego::Params{}); + + auto param = MakeParam(metric, quantizer); + param->provider = provider; + param->provider_meta = provider_meta; + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(quantizer_name, index->index_searcher()->meta().quantizer_name()); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + auto streamer = std::dynamic_pointer_cast( + index->index_searcher()); + ASSERT_NE(nullptr, streamer); + EXPECT_TRUE(streamer->uses_turbo_distance()); + EXPECT_TRUE(streamer->uses_turbo_build_distance()); + + AddVectors(index.get(), vectors); + CheckGraphSearchEnabled(index.get()); + std::vector linear_results; + std::vector graph_results; + for (uint32_t query_id : kGraphQueryIds) { + SCOPED_TRACE(query_id); + auto linear_rows = SearchRows(index.get(), vectors[query_id], true); + auto graph_rows = SearchRows(index.get(), vectors[query_id], false); + CheckGraphRecall(linear_rows, graph_rows); + ASSERT_FALSE(graph_rows.empty()); + EXPECT_EQ(query_id, graph_rows.front().first); + linear_results.push_back(std::move(linear_rows)); + graph_results.push_back(std::move(graph_rows)); + } + + ASSERT_EQ(0, index->close()); + + auto reopened = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, reopened); + ASSERT_EQ(0, + reopened->open(path, {StorageOptions::StorageType::kMMAP, false})); + EXPECT_EQ(quantizer_name, + reopened->index_searcher()->meta().quantizer_name()); + CheckGraphSearchEnabled(reopened.get()); + for (size_t i = 0; i < std::size(kGraphQueryIds); ++i) { + SCOPED_TRACE(kGraphQueryIds[i]); + auto linear_rows = + SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], true); + auto graph_rows = + SearchRows(reopened.get(), vectors[kGraphQueryIds[i]], false); + EXPECT_EQ(linear_results[i], linear_rows); + EXPECT_EQ(graph_results[i], graph_rows); + CheckGraphRecall(linear_rows, graph_rows); + } + ASSERT_EQ(0, reopened->close()); + zvec::test_util::RemoveTestFiles(path); +} + +void BuildLegacyFp32Hnsw(const std::string &path, + const std::vector> &vectors) { + namespace core = zvec::core; + core::IndexMeta legacy_meta(core::IndexMeta::DT_FP32, kDimension); + legacy_meta.set_meta_type(core::IndexMeta::MetaType::MT_DENSE); + legacy_meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params()); + ASSERT_TRUE(legacy_meta.quantizer_name().empty()); + + zvec::ailego::Params params; + params.set(core::PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + params.set(core::PARAM_HNSW_STREAMER_SCALING_FACTOR, 16U); + params.set(core::PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + params.set(core::PARAM_HNSW_STREAMER_EF, 100U); + params.set(core::PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); + auto streamer = core::IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(legacy_meta, params)); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(zvec::ailego::Params())); + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + + auto context = streamer->create_context(); + core::IndexQueryMeta qmeta(core::IndexMeta::DT_FP32, kDimension); + for (size_t i = 0; i < vectors.size(); ++i) { + ASSERT_EQ(0, streamer->add_with_id_impl(static_cast(i), + vectors[i].data(), qmeta, context)); + } + ASSERT_EQ(0, streamer->flush(0)); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); +} + +void BuildLegacyInt8Hnsw(const std::string &path, MetricType metric, + const std::vector> &vectors) { + namespace core = zvec::core; + core::IndexMeta raw_meta(core::IndexMeta::DT_FP32, kDimension); + raw_meta.set_meta_type(core::IndexMeta::MetaType::MT_DENSE); + raw_meta.set_metric( + metric == MetricType::kCosine ? "Cosine" : "SquaredEuclidean", 0, + zvec::ailego::Params()); + const char *converter_name = metric == MetricType::kCosine + ? "CosineInt8Converter" + : "Int8StreamingConverter"; + raw_meta.set_converter(converter_name, 0, zvec::ailego::Params()); + auto converter = core::IndexFactory::CreateConverter(converter_name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(raw_meta, zvec::ailego::Params())); + + core::IndexMeta legacy_meta = converter->meta(); + ASSERT_TRUE(legacy_meta.quantizer_name().empty()); + auto reformer = + core::IndexFactory::CreateReformer(legacy_meta.reformer_name()); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(legacy_meta.reformer_params())); + + zvec::ailego::Params params; + params.set(core::PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + params.set(core::PARAM_HNSW_STREAMER_SCALING_FACTOR, 16U); + params.set(core::PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + params.set(core::PARAM_HNSW_STREAMER_EF, 100U); + params.set(core::PARAM_HNSW_STREAMER_GET_VECTOR_ENABLE, true); + auto streamer = core::IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer); + ASSERT_EQ(0, streamer->init(legacy_meta, params)); + auto storage = core::IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(zvec::ailego::Params())); + ASSERT_EQ(0, storage->open(path, true)); + ASSERT_EQ(0, streamer->open(storage)); + + auto context = streamer->create_context(); + core::IndexQueryMeta raw_qmeta(core::IndexMeta::DT_FP32, kDimension); + for (size_t i = 0; i < vectors.size(); ++i) { + std::string converted; + core::IndexQueryMeta converted_meta; + ASSERT_EQ(0, reformer->convert(vectors[i].data(), raw_qmeta, &converted, + &converted_meta)); + ASSERT_EQ(0, streamer->add_with_id_impl(static_cast(i), + converted.data(), converted_meta, + context)); + } + ASSERT_EQ(0, streamer->flush(0)); + ASSERT_EQ(0, streamer->close()); + ASSERT_EQ(0, storage->close()); +} + +} // namespace + +TEST(HnswTurboQuantizerIndex, DefaultUsesFp32TurboQuantizer) { + for (MetricType metric : + {MetricType::kL2sq, MetricType::kCosine, MetricType::kInnerProduct}) { + auto index = IndexFactory::CreateAndInitIndex(*MakeDefaultParam(metric)); + ASSERT_NE(nullptr, index); + EXPECT_EQ("Fp32Quantizer", + index->index_searcher()->meta().quantizer_name()); + } +} + +TEST(HnswTurboQuantizerIndex, InnerProductScoreAndRadiusUseCallerSpace) { + const std::string path{"hnsw_turbo_fp32_ip_radius.index"}; + zvec::test_util::RemoveTestFiles(path); + auto index = IndexFactory::CreateAndInitIndex( + *MakeDefaultParam(MetricType::kInnerProduct)); + ASSERT_NE(nullptr, index); + ASSERT_EQ("Fp32Quantizer", index->index_searcher()->meta().quantizer_name()); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + + std::vector> vectors(3, std::vector(kDimension)); + vectors[0][0] = 1.0f; + vectors[1][0] = 0.75f; + vectors[2][0] = 0.25f; + AddVectors(index.get(), vectors); + + auto query_param = HNSWQueryParamBuilder() + .with_topk(3) + .with_ef_search(100) + .with_is_linear(true) + .with_radius(0.5f) + .build(); + std::vector query(kDimension); + query[0] = 1.0f; + SearchResult result; + ASSERT_EQ(0, index->search(VectorData{DenseVector{query.data()}}, query_param, + &result)); + ASSERT_EQ(2U, result.doc_list_.size()); + EXPECT_EQ(0U, result.doc_list_[0].key()); + EXPECT_FLOAT_EQ(1.0f, result.doc_list_[0].score()); + EXPECT_EQ(1U, result.doc_list_[1].key()); + EXPECT_FLOAT_EQ(0.75f, result.doc_list_[1].score()); + + ASSERT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); +} + +TEST_P(HnswTurboIndexTest, SelectsTurboQuantizer) { + const auto &[quantizer, metric] = GetParam(); + auto index = + IndexFactory::CreateAndInitIndex(*MakeParam(metric.type, quantizer.type)); + ASSERT_NE(nullptr, index); + EXPECT_EQ(quantizer.name, index->index_searcher()->meta().quantizer_name()); +} + +TEST_P(HnswTurboIndexTest, KnownScoresAndRadiusUseCallerSpace) { + const auto &[quantizer, metric] = GetParam(); + const std::string path = index_path("scores"); + zvec::test_util::RemoveTestFiles(path); + auto index = + IndexFactory::CreateAndInitIndex(*MakeParam(metric.type, quantizer.type)); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + + // Non-unit vectors distinguish cosine normalization from inner product. + std::vector> vectors(3, std::vector(kDimension)); + vectors[0][0] = 1.0f; + vectors[1][0] = 0.75f; + vectors[1][1] = 0.25f; + vectors[2][0] = 0.25f; + vectors[2][1] = 0.75f; + AddVectors(index.get(), vectors); + + auto query_param = HNSWQueryParamBuilder() + .with_topk(3) + .with_is_linear(true) + .with_radius(0.5f) + .build(); + SearchResult result; + ASSERT_EQ(0, index->search(VectorData{DenseVector{vectors[0].data()}}, + query_param, &result)); + ASSERT_EQ(2U, result.doc_list_.size()); + EXPECT_EQ(0U, result.doc_list_[0].key()); + EXPECT_EQ(1U, result.doc_list_[1].key()); + float expected_first = 0.0f; + float expected_second = 0.125f; + if (metric.type == MetricType::kCosine) { + expected_second = 1.0f - 0.75f / std::sqrt(0.625f); + } else if (metric.type == MetricType::kInnerProduct) { + expected_first = 1.0f; + expected_second = 0.75f; + } + EXPECT_NEAR(expected_first, result.doc_list_[0].score(), + quantizer.score_tolerance); + EXPECT_NEAR(expected_second, result.doc_list_[1].score(), + quantizer.score_tolerance); + + ASSERT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); +} + +TEST_P(HnswTurboIndexTest, AddSearchReopenFetch) { + const auto &[quantizer, metric] = GetParam(); + CheckTurboAddSearchReopen(metric.type, quantizer.type, quantizer.name, + quantizer.fetch_tolerance, index_path("stored")); +} + +TEST_P(HnswTurboIndexTest, ExternalVectorsUseTurbo) { + const auto &[quantizer, metric] = GetParam(); + CheckExternalTurboAddSearchReopen(metric.type, quantizer.type, quantizer.name, + index_path("external")); +} + +TEST_P(HnswTurboIndexTest, OriginalProviderBuildUsesTurbo) { + const auto &[quantizer, metric] = GetParam(); + CheckOriginalProviderUsesTurbo(metric.type, quantizer.type, quantizer.name, + index_path("provider")); +} + +TEST(HnswTurboQuantizerIndex, UnsupportedCombinationsUseLegacyPipeline) { + auto rotated = IndexFactory::CreateAndInitIndex( + *MakeParam(MetricType::kCosine, QuantizerType::kInt8, true)); + ASSERT_NE(nullptr, rotated); + EXPECT_TRUE(rotated->index_searcher()->meta().quantizer_name().empty()); + + auto mips = IndexFactory::CreateAndInitIndex( + *MakeParam(MetricType::kMIPSL2sq, QuantizerType::kNone)); + ASSERT_NE(nullptr, mips); + EXPECT_TRUE(mips->index_searcher()->meta().quantizer_name().empty()); +} + +TEST_P(HnswTurboIndexTest, MergePreservesTurboLayout) { + const auto &[quantizer, metric] = GetParam(); + const std::string source_path = index_path("merge_source"); + const std::string target_path = index_path("merge_target"); + zvec::test_util::RemoveTestFiles(source_path); + zvec::test_util::RemoveTestFiles(target_path); + auto vectors = RandomVectors(); + auto param = MakeParam(metric.type, quantizer.type); + + auto source = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, source); + ASSERT_EQ( + 0, source->open(source_path, {StorageOptions::StorageType::kMMAP, true})); + AddVectors(source.get(), vectors); + + auto target = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, target); + ASSERT_EQ( + 0, target->open(target_path, {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ(0, target->merge({source}, IndexFilter())); + EXPECT_EQ(kVectorCount, target->get_doc_count()); + EXPECT_EQ(quantizer.name, target->index_searcher()->meta().quantizer_name()); + + auto rows = SearchRows(target.get(), vectors[73], true); + ASSERT_EQ(kTopK, rows.size()); + EXPECT_EQ(73U, rows[0].first); + + ASSERT_EQ(0, target->close()); + ASSERT_EQ(0, source->close()); + zvec::test_util::RemoveTestFiles(source_path); + zvec::test_util::RemoveTestFiles(target_path); +} + +TEST_P(HnswTurboIndexTest, MergeAcrossMetricsReencodesTurboLayout) { + const auto &[quantizer, source_metric] = GetParam(); + auto vectors = RandomVectors(kTopK); + // Exactly representable in INT4: copying the IP tail into an L2 index + // used to produce a negative self-distance for this vector. + vectors[0].assign(kDimension, 0.0f); + vectors[0].front() = vectors[0].back() = 1.0f; + for (size_t i = 1; i < vectors.size(); ++i) { + for (float &value : vectors[i]) { + value *= static_cast(i + 1); + } + } + + auto make_param = [quantizer_type = quantizer.type]( + bool flat, + MetricType metric) -> BaseIndexParam::Pointer { + if (!flat) { + return MakeParam(metric, quantizer_type); + } + return FlatIndexParamBuilder() + .with_metric_type(metric) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_is_sparse(false) + .with_quantizer_param(QuantizerParam(quantizer_type)) + .build(); + }; + + for (bool source_flat : {false, true}) { + const std::string source_path = index_path("cross_metric_source"); + zvec::test_util::RemoveTestFiles(source_path); + auto source = IndexFactory::CreateAndInitIndex( + *make_param(source_flat, source_metric.type)); + ASSERT_NE(nullptr, source); + ASSERT_EQ(0, source->open(source_path, + {StorageOptions::StorageType::kMMAP, true})); + AddVectors(source.get(), vectors); + + // A lossy source cannot recover the initial FP32 values exactly. Build + // the reference from fetched (decoded) values, then encode for the target. + auto decoded = vectors; + for (uint32_t i = 0; i < vectors.size(); ++i) { + VectorDataBuffer fetched; + ASSERT_EQ(0, source->fetch(i, &fetched)); + const auto &bytes = + std::get(fetched.vector_buffer).data; + ASSERT_EQ(kDimension * sizeof(float), bytes.size()); + std::memcpy(decoded[i].data(), bytes.data(), bytes.size()); + } + + for (const auto &target_metric : kTurboMetrics) { + if (target_metric.type == source_metric.type) { + continue; + } + for (bool target_flat : {false, true}) { + SCOPED_TRACE(testing::Message() + << "source_flat=" << source_flat + << " target_flat=" << target_flat + << " target_metric=" << target_metric.test_name); + const std::string target_path = index_path("cross_metric_target"); + const std::string reference_path = index_path("cross_metric_reference"); + zvec::test_util::RemoveTestFiles(target_path); + zvec::test_util::RemoveTestFiles(reference_path); + auto param = make_param(target_flat, target_metric.type); + auto target = IndexFactory::CreateAndInitIndex(*param); + auto reference = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, target); + ASSERT_NE(nullptr, reference); + ASSERT_EQ(0, target->open(target_path, + {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ(0, + reference->open(reference_path, + {StorageOptions::StorageType::kMMAP, true})); + AddVectors(reference.get(), decoded); + ASSERT_EQ(0, target->merge({source}, IndexFilter())); + EXPECT_EQ(vectors.size(), target->get_doc_count()); + + auto search = [&](Index *index, const std::vector &query) { + BaseIndexQueryParam::Pointer query_param; + if (target_flat) { + query_param = FlatQueryParamBuilder().with_topk(kTopK).build(); + } else { + query_param = HNSWQueryParamBuilder() + .with_topk(kTopK) + .with_is_linear(true) + .build(); + } + SearchResult result; + EXPECT_EQ(0, index->search(VectorData{DenseVector{query.data()}}, + query_param, &result)); + SearchRowList rows; + for (const auto &doc : result.doc_list_) { + rows.emplace_back(doc.key(), doc.score()); + } + std::sort(rows.begin(), rows.end()); + return rows; + }; + for (const auto &query : vectors) { + const auto expected = search(reference.get(), query); + const auto actual = search(target.get(), query); + ASSERT_EQ(kTopK, expected.size()); + ASSERT_EQ(expected.size(), actual.size()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(expected[i].first, actual[i].first); + EXPECT_FLOAT_EQ(expected[i].second, actual[i].second); + } + } + ASSERT_EQ(0, target->close()); + ASSERT_EQ(0, reference->close()); + zvec::test_util::RemoveTestFiles(target_path); + zvec::test_util::RemoveTestFiles(reference_path); + } + } + ASSERT_EQ(0, source->close()); + zvec::test_util::RemoveTestFiles(source_path); + } +} + +INSTANTIATE_TEST_SUITE_P( + QuantizersAndMetrics, HnswTurboIndexTest, + testing::Combine(testing::ValuesIn(kTurboQuantizers), + testing::ValuesIn(kTurboMetrics)), + [](const testing::TestParamInfo &info) { + return std::string(std::get<0>(info.param).test_name) + "_" + + std::get<1>(info.param).test_name; + }); + +class HnswLegacyReopenTest : public testing::TestWithParam { + protected: + static void SetUpTestSuite() { + ASSERT_EQ(0, zvec::ailego::MemoryLimitPool::get_instance().init(100 * 1024 * + 1024)); + } +}; + +INSTANTIATE_TEST_SUITE_P( + StorageModes, HnswLegacyReopenTest, + testing::Values( + StorageOptions{StorageOptions::StorageType::kMMAP, false}, + StorageOptions{StorageOptions::StorageType::kMMAP, false, false, true}, + StorageOptions{StorageOptions::StorageType::kBufferPool, false})); + +TEST_P(HnswLegacyReopenTest, LegacyLayoutReopenFallsBack) { + const std::string path{"hnsw_int8_legacy_layout.index"}; + zvec::test_util::RemoveTestFiles(path); + auto vectors = RandomVectors(); + BuildLegacyInt8Hnsw(path, MetricType::kL2sq, vectors); + if (::testing::Test::HasFatalFailure()) { + return; + } + + auto index = IndexFactory::CreateAndInitIndex( + *MakeParam(MetricType::kL2sq, QuantizerType::kInt8)); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, GetParam())); + EXPECT_TRUE(index->index_searcher()->meta().quantizer_name().empty()); + auto rows = SearchRows(index.get(), vectors[37], true); + ASSERT_EQ(kTopK, rows.size()); + EXPECT_EQ(37U, rows[0].first); + + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(37, &fetched)); + const auto *fetched_vector = reinterpret_cast( + std::get(fetched.vector_buffer).data.data()); + for (uint32_t i = 0; i < kDimension; ++i) { + EXPECT_NEAR(vectors[37][i], fetched_vector[i], 5e-2f); + } + ASSERT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); +} + +TEST_P(HnswLegacyReopenTest, LegacyFp32LayoutReopenFallsBack) { + const std::string path{"hnsw_fp32_legacy_layout.index"}; + zvec::test_util::RemoveTestFiles(path); + auto vectors = RandomVectors(); + BuildLegacyFp32Hnsw(path, vectors); + if (::testing::Test::HasFatalFailure()) { + return; + } + + auto index = IndexFactory::CreateAndInitIndex( + *MakeParam(MetricType::kL2sq, QuantizerType::kNone)); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, GetParam())); + EXPECT_TRUE(index->index_searcher()->meta().quantizer_name().empty()); + auto rows = SearchRows(index.get(), vectors[37], true); + ASSERT_EQ(kTopK, rows.size()); + EXPECT_EQ(37U, rows[0].first); + + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(37, &fetched)); + const auto *fetched_vector = reinterpret_cast( + std::get(fetched.vector_buffer).data.data()); + for (uint32_t i = 0; i < kDimension; ++i) { + EXPECT_FLOAT_EQ(vectors[37][i], fetched_vector[i]); + } + ASSERT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); +} + +namespace zvec { +namespace core { +namespace { + +constexpr size_t kDimension = 8; +constexpr size_t kCount = 16; + +class ExternalVectorSource final : public VectorSource { + public: + ExternalVectorSource() : vectors(kCount, std::vector(kDimension)) { + for (size_t i = 0; i < kCount; ++i) { + for (size_t j = 0; j < kDimension; ++j) { + vectors[i][j] = i * 0.125f + j * 0.03125f; + } + } + } + + const void *get_vector(uint32_t node_id) const override { + return vectors[node_id].data(); + } + + std::vector> vectors; +}; + +class HnswExternalCoreCompatibilityTest : public testing::TestWithParam { + protected: + void SetUp() override { + zvec::test_util::RemoveTestPath(directory_); + } + + void TearDown() override { + if (streamer_) { + streamer_->close(); + } + if (storage_) { + storage_->close(); + } + zvec::test_util::RemoveTestPath(directory_); + } + + void open(bool turbo) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + if (turbo) { + quantizer_ = IndexFactory::CreateQuantizer("Int8Quantizer"); + ASSERT_NE(nullptr, quantizer_); + ASSERT_EQ(0, quantizer_->init(meta, ailego::Params())); + meta = quantizer_->meta(); + } + ailego::Params params; + params.set(PARAM_HNSW_STREAMER_USE_EXTERNAL_VECTOR, true); + params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U); + params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100U); + params.set(PARAM_HNSW_STREAMER_EF, 100U); + params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 0U); + streamer_ = IndexFactory::CreateStreamer("HnswStreamer"); + ASSERT_NE(nullptr, streamer_); + ASSERT_EQ(0, turbo ? streamer_->init(meta, params, quantizer_) + : streamer_->init(meta, params)); + storage_ = IndexFactory::CreateStorage("MMapFileStorage"); + ASSERT_NE(nullptr, storage_); + ASSERT_EQ(0, storage_->init(ailego::Params())); + ASSERT_EQ(0, storage_->open(directory_ + "external.index", true)); + ASSERT_EQ(0, streamer_->open(storage_)); + context_ = streamer_->create_context(); + ASSERT_NE(nullptr, context_); + ctx_ = dynamic_cast(context_.get()); + ASSERT_NE(nullptr, ctx_); + ctx_->set_vector_source(&source_); + } + + int add(uint32_t id, const void *query, const IndexQueryMeta &meta) { + return GetParam() ? streamer_->add_with_id_impl(id, query, meta, context_) + : streamer_->add_impl(id, query, meta, context_); + } + + void check_node_count(size_t expected) const { + // Context entities retain their creation-time header. A fresh provider + // snapshots the streamer's current entity, including any orphan nodes. + auto provider = streamer_->create_provider(); + ASSERT_NE(nullptr, provider); + ASSERT_EQ(expected, provider->count()); + } + + const std::string directory_ = "hnsw_streamer_turbo_compat_test_dir/"; + IndexStreamer::Pointer streamer_; + IndexStorage::Pointer storage_; + IndexStreamer::Context::Pointer context_; + HnswContext *ctx_ = nullptr; + turbo::Quantizer::Pointer quantizer_; + ExternalVectorSource source_; +}; + +TEST_P(HnswExternalCoreCompatibilityTest, + LegacyBuildAcceptsQueryWithoutExternalBuildField) { + ASSERT_NO_FATAL_FAILURE(open(false)); + IndexQueryMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + for (uint32_t i = 0; i < kCount; ++i) { + ASSERT_EQ(nullptr, ctx_->external_build_query()); + ASSERT_EQ(0, add(i, source_.get_vector(i), meta)); + } + ASSERT_NO_FATAL_FAILURE(check_node_count(kCount)); + EXPECT_EQ(kCount, streamer_->stats().added_count()); + + context_->set_topk(1); + for (uint32_t probe : {0U, 7U, 15U}) { + ASSERT_EQ( + 0, streamer_->search_impl(source_.get_vector(probe), meta, context_)); + ASSERT_EQ(1U, context_->result().size()); + EXPECT_EQ(probe, context_->result()[0].key()); + EXPECT_FLOAT_EQ(0.0f, context_->result()[0].score()); + } +} + +TEST_P(HnswExternalCoreCompatibilityTest, + TurboRejectsMissingRawQueryBeforeMutatingNodes) { + ASSERT_NO_FATAL_FAILURE(open(true)); + IndexQueryMeta raw_meta(IndexMeta::DataType::DT_FP32, kDimension); + IndexQueryMeta encoded_meta; + std::vector codes(kCount); + for (uint32_t i = 0; i < kCount; ++i) { + ASSERT_EQ(0, quantizer_->quantize(source_.get_vector(i), raw_meta, + &codes[i], &encoded_meta)); + ctx_->set_external_build_query(nullptr); + ASSERT_EQ(IndexError_InvalidArgument, + add(i, codes[i].data(), encoded_meta)); + ASSERT_NO_FATAL_FAILURE(check_node_count(i)); + EXPECT_EQ(i, streamer_->stats().added_count()); + ctx_->set_external_build_query(source_.get_vector(i)); + ASSERT_EQ(0, add(i, codes[i].data(), encoded_meta)); + ASSERT_NO_FATAL_FAILURE(check_node_count(i + 1)); + } + EXPECT_EQ(kCount, streamer_->stats().added_count()); + EXPECT_EQ(kCount, streamer_->stats().discarded_count()); + + context_->set_topk(1); + for (uint32_t probe : {0U, 7U, 15U}) { + ASSERT_EQ( + 0, streamer_->search_impl(codes[probe].data(), encoded_meta, context_)); + ASSERT_EQ(1U, context_->result().size()); + EXPECT_EQ(probe, context_->result()[0].key()); + } +} + +INSTANTIATE_TEST_SUITE_P(AddApis, HnswExternalCoreCompatibilityTest, + testing::Bool()); + +} // namespace +} // namespace core +} // namespace zvec