From ab12d26a04e0aba97d6a805b39646108c9bad745 Mon Sep 17 00:00:00 2001 From: SpadeA Date: Wed, 15 Jul 2026 11:19:50 +0800 Subject: [PATCH 1/4] refactor raw storage; impl bin for lemur Signed-off-by: SpadeA --- include/knowhere/index/emb_list_strategy.h | 42 ++ include/knowhere/index/index_node.h | 50 ++- src/index/emb_list/emb_list_raw_storage.cc | 337 ++++++++++++++++ src/index/emb_list/emb_list_raw_storage.h | 70 ++++ src/index/emb_list/emb_list_strategy_lemur.cc | 335 ++++++++++++++-- .../emb_list/emb_list_strategy_muvera.cc | 9 + .../emb_list/emb_list_strategy_token_ann.cc | 9 + src/index/hnsw/faiss_hnsw.cc | 43 ++- src/index/index.cc | 11 +- src/index/index_node.cc | 359 +++++++++++++----- src/index/index_static.cc | 13 +- src/index/ivf/ivf.cc | 6 +- tests/ut/test_emb_list.cc | 178 +++++++++ tests/ut/test_get_emb_list.cc | 138 ++++++- tests/ut/utils.h | 26 ++ .../faiss/faiss/gpu_metal/MetalDistance.metal | 1 - 16 files changed, 1468 insertions(+), 159 deletions(-) create mode 100644 src/index/emb_list/emb_list_raw_storage.cc create mode 100644 src/index/emb_list/emb_list_raw_storage.h diff --git a/include/knowhere/index/emb_list_strategy.h b/include/knowhere/index/emb_list_strategy.h index dab4787d5..3b76ffac8 100644 --- a/include/knowhere/index/emb_list_strategy.h +++ b/include/knowhere/index/emb_list_strategy.h @@ -14,6 +14,7 @@ #include #include +#include #include #include "knowhere/bitsetview.h" @@ -47,6 +48,37 @@ struct EmbListMetricInfo { bool is_cosine; }; +// Some strategies build ANN data whose physical type differs from the raw vectors. +// This target tells IndexNode whether that ANN data can use the current index node +// or needs a separate index with a different data type. +enum class EmbListAnnIndexTarget { + // Build the ANN dataset directly into this IndexNode's own ANN index. This is used when + // the ANN dataset has the same data type as the current index node, such as TokenANN and MUVERA. + BaseIndex, + // Build the ANN dataset into a separately allocated Index. This is used when + // the strategy's ANN representation has a different data type from the current index node, + // such as bin1 LEMUR producing fp32 learned representations. + SeparateIndex, +}; + +enum class EmbListAnnIndexDataType { + // Valid for BaseIndex: the current IndexNode data type is reused. + // SeparateIndex must request a concrete data type instead. + SameAsBaseIndex, + Fp32, + Fp16, + Bf16, + Int8, + Bin1, +}; + +struct EmbListAnnIndexSpec { + EmbListAnnIndexTarget target; + EmbListAnnIndexDataType data_type; + // nullopt means the ANN index uses the raw sub metric parsed from the outer emb-list metric. + std::optional ann_metric_type; +}; + /** * @brief Parse metric type from config into a shared struct. * @@ -206,6 +238,16 @@ class EmbListStrategy { return false; } + /** + * @brief Describe how the strategy's ANN dataset should be indexed. + * + * BaseIndex means the ANN data is indexed by this IndexNode itself. SeparateIndex means IndexNode must create + * a child Index with the requested data type. nullopt ann_metric_type means the ANN index uses the + * raw sub metric parsed from the outer emb-list metric; a value means the strategy overrides it. + */ + [[nodiscard]] virtual EmbListAnnIndexSpec + AnnIndexSpec(const BaseConfig& config) const = 0; + /** * @brief Execute search with full control over the search flow. * diff --git a/include/knowhere/index/index_node.h b/include/knowhere/index/index_node.h index 118429f3f..8311faec2 100644 --- a/include/knowhere/index/index_node.h +++ b/include/knowhere/index/index_node.h @@ -13,6 +13,7 @@ #define INDEX_NODE_H #include +#include #include #include #include @@ -46,13 +47,11 @@ struct OpContext; class ThreadPool; #endif -namespace faiss { -class IndexFlat; -} // namespace faiss - namespace knowhere { class Interrupt; +class EmbListRawStorage; +struct EmbListSeparateAnnIndexHolder; class IndexNode : public Object { public: @@ -322,6 +321,23 @@ class IndexNode : public Object { virtual int64_t Count() const = 0; + // Returns the ANN index that currently serves vector-level metadata/search. + // EmbList strategies normally use this node itself, but a strategy may build + // a child ANN index when its ANN representation has a different data type. + virtual IndexNode* + AnnIndexNode(); + + virtual const IndexNode* + AnnIndexNode() const; + + virtual int64_t + CountForSearchBitset() const { + if (emb_list_strategy_ != nullptr && emb_list_strategy_->GetDocCount() >= 0) { + return emb_list_strategy_->GetDocCount(); + } + return Count(); + } + virtual std::string Type() const = 0; @@ -500,6 +516,10 @@ class IndexNode : public Object { SearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, milvus::OpContext* op_context = nullptr) const; + virtual expected + SearchEmbListAnnIndex(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + milvus::OpContext* op_context = nullptr) const; + /** * @brief Returns the code size (in bytes) of a single query vector, which varies depending on the data type (e.g., * fp32, bf16, etc). @@ -623,10 +643,10 @@ class IndexNode : public Object { protected: /** - * @brief Compute distances using emb_list_raw_index_ (raw vector storage). + * @brief Compute distances using emb_list raw vector storage. * - * Used by CalcDistByIDs implementations when emb_list_raw_index_ is present - * (MUVERA/LEMUR strategies). The raw index stores original vectors indexed by + * Used by CalcDistByIDs implementations when emb_list raw storage is present + * (MUVERA/LEMUR strategies). The raw storage keeps original vectors indexed by * global vector IDs, so no ID translation is needed. * * @param pool Thread pool for parallel computation @@ -641,14 +661,14 @@ class IndexNode : public Object { std::string el_metric_type_; EmbListStrategyPtr emb_list_strategy_; // emb_list encoding strategy (tokenann/muvera) // Raw vector storage for EmbList strategies (MUVERA/LEMUR) that encode documents - // into different representations for ANN search. Since the base index holds encoded - // vectors (not raw), this IndexFlat stores original vectors for exact distance - // computation during MaxSim reranking. - // Baseline type so that the same shared_ptr can hold either the knowhere - // Jaccard-aware IndexFlat subclass (fresh build path, if ever needed) or - // a plain ::faiss::IndexFlat{,IP,L2} restored by the deserialization - // factory in cppcontrib/knowhere/impl/index_read.cpp. - std::shared_ptr<::faiss::IndexFlat> emb_list_raw_index_; + // into different representations for ANN search. The ANN index may be this node's + // BaseIndex or a SeparateIndex, so raw storage keeps the original vectors available + // for exact distance computation during MaxSim reranking. + std::shared_ptr emb_list_raw_storage_; + // Optional child ANN index used when a strategy's ANN representation has a different + // data type from this IndexNode, for example binary LEMUR producing fp32 vectors. + std::shared_ptr emb_list_separate_ann_index_; + std::string emb_list_raw_metric_type_; #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) struct PrometheusMetrics { diff --git a/src/index/emb_list/emb_list_raw_storage.cc b/src/index/emb_list/emb_list_raw_storage.cc new file mode 100644 index 000000000..9bc5596fc --- /dev/null +++ b/src/index/emb_list/emb_list_raw_storage.cc @@ -0,0 +1,337 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// 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 "index/emb_list/emb_list_raw_storage.h" + +#include +#include + +#include "faiss/cppcontrib/knowhere/IndexBinaryFlat.h" +#include "faiss/cppcontrib/knowhere/IndexFlat.h" +#include "faiss/cppcontrib/knowhere/index_io.h" +#include "io/memory_io.h" +#include "knowhere/comp/index_param.h" +#include "knowhere/context.h" +#include "knowhere/log.h" +#include "knowhere/utils.h" + +namespace knowhere { +namespace { + +inline bool +IsBinarySubMetric(const std::string& metric_type) { + return metric_type == metric::HAMMING || metric_type == metric::JACCARD; +} + +inline faiss::MetricType +ToFaissBinaryMetric(const std::string& metric_type) { + return metric_type == metric::JACCARD ? faiss::METRIC_Jaccard : faiss::METRIC_Hamming; +} + +inline int +PopcountByte(uint8_t value) { + return __builtin_popcount(static_cast(value)); +} + +class EmbListFloatRawStorage final : public EmbListRawStorage { + public: + EmbListFloatRawStorage(int64_t dim, faiss::MetricType metric_type) + : index_(std::make_shared(dim, metric_type)) { + } + + explicit EmbListFloatRawStorage(std::shared_ptr<::faiss::IndexFlat> index) : index_(std::move(index)) { + } + + int64_t + Dim() const override { + return index_->d; + } + + int64_t + Count() const override { + return index_->ntotal; + } + + size_t + CodeSize() const override { + return static_cast(index_->d) * sizeof(float); + } + + Status + Add(const DataSetPtr dataset) override { + index_->add(dataset->GetRows(), static_cast(dataset->GetTensor())); + return Status::success; + } + + Status + ReconstructN(int64_t start, int64_t n, uint8_t* out) const override { + index_->reconstruct_n(start, n, reinterpret_cast(out)); + return Status::success; + } + + Status + Serialize(BinarySet& binset) const override { + MemoryIOWriter writer; + faiss::cppcontrib::knowhere::write_index(index_.get(), &writer); + std::shared_ptr raw_bin(writer.data()); + binset.Append(meta::EMB_LIST_RAW_INDEX, raw_bin, writer.tellg()); + return Status::success; + } + + expected + CalcDistance(const DataSetPtr dataset, const int64_t* labels, size_t labels_len, const std::string& /*metric_type*/, + bool is_cosine, std::shared_ptr pool, milvus::OpContext* op_context) const override { + auto num_queries = dataset->GetRows(); + auto dim = dataset->GetDim(); + auto query_data = dataset->GetTensor(); + auto distances = std::make_unique(num_queries * labels_len); + + try { + std::vector> futs; + futs.reserve(num_queries); + for (int64_t i = 0; i < num_queries; ++i) { + futs.emplace_back(pool->push([&, idx = i]() { + knowhere::checkCancellation(op_context); + std::unique_ptr dist_computer(index_->get_distance_computer()); + + const float* cur_query = static_cast(query_data) + idx * dim; + std::unique_ptr copied_query = nullptr; + if (is_cosine) { + copied_query = CopyAndNormalizeVecs(cur_query, 1, dim); + cur_query = copied_query.get(); + } + + dist_computer->set_query(cur_query); + auto cur_distances = distances.get() + idx * labels_len; + for (size_t j = 0; j < labels_len; ++j) { + cur_distances[j] = (*dist_computer)(labels[j]); + } + })); + } + WaitAllSuccess(futs); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "CalcDistance by float raw storage error: " << e.what(); + return expected::Err(Status::faiss_inner_error, e.what()); + } + + return GenResultDataSet(num_queries, labels_len, std::unique_ptr{}, std::move(distances)); + } + + private: + std::shared_ptr<::faiss::IndexFlat> index_; +}; + +class EmbListBinaryRawStorage final : public EmbListRawStorage { + public: + EmbListBinaryRawStorage(int64_t dim, faiss::MetricType metric_type) + : index_(std::make_shared(dim, metric_type)) { + } + + explicit EmbListBinaryRawStorage(std::shared_ptr index) + : index_(std::move(index)) { + } + + int64_t + Dim() const override { + return index_->d; + } + + int64_t + Count() const override { + return index_->ntotal; + } + + size_t + CodeSize() const override { + return static_cast(index_->code_size); + } + + Status + Add(const DataSetPtr dataset) override { + index_->add(dataset->GetRows(), static_cast(dataset->GetTensor())); + return Status::success; + } + + Status + ReconstructN(int64_t start, int64_t n, uint8_t* out) const override { + index_->reconstruct_n(start, n, out); + return Status::success; + } + + Status + Serialize(BinarySet& binset) const override { + MemoryIOWriter writer; + faiss::cppcontrib::knowhere::write_index_binary(index_.get(), &writer); + std::shared_ptr raw_bin(writer.data()); + binset.Append(meta::EMB_LIST_RAW_INDEX, raw_bin, writer.tellg()); + return Status::success; + } + + expected + CalcDistance(const DataSetPtr dataset, const int64_t* labels, size_t labels_len, const std::string& metric_type, + bool /*is_cosine*/, std::shared_ptr pool, milvus::OpContext* op_context) const override { + auto num_queries = dataset->GetRows(); + auto dim = dataset->GetDim(); + if (dim != index_->d) { + return expected::Err(Status::invalid_args, "binary raw index dim mismatch"); + } + if (dim % 8 != 0) { + return expected::Err(Status::invalid_args, "binary raw distance requires dim multiple of 8"); + } + + const size_t code_size = CodeSize(); + const auto* query_data = static_cast(dataset->GetTensor()); + auto distances = std::make_unique(num_queries * labels_len); + + try { + std::vector label_codes(labels_len * code_size); + for (size_t j = 0; j < labels_len; ++j) { + index_->reconstruct(labels[j], label_codes.data() + j * code_size); + } + + std::vector> futs; + futs.reserve(num_queries); + for (int64_t i = 0; i < num_queries; ++i) { + futs.emplace_back(pool->push([&, idx = i]() { + knowhere::checkCancellation(op_context); + const uint8_t* cur_query = query_data + idx * code_size; + auto cur_distances = distances.get() + idx * labels_len; + + for (size_t j = 0; j < labels_len; ++j) { + const uint8_t* cur_doc = label_codes.data() + j * code_size; + int intersection = 0; + int union_count = 0; + int hamming = 0; + for (size_t b = 0; b < code_size; ++b) { + const uint8_t q = cur_query[b]; + const uint8_t d = cur_doc[b]; + if (metric_type == metric::JACCARD) { + intersection += PopcountByte(static_cast(q & d)); + union_count += PopcountByte(static_cast(q | d)); + } else { + hamming += PopcountByte(static_cast(q ^ d)); + } + } + + if (metric_type == metric::JACCARD) { + cur_distances[j] = + union_count == 0 ? 0.0f : 1.0f - static_cast(intersection) / union_count; + } else { + cur_distances[j] = static_cast(hamming); + } + } + })); + } + WaitAllSuccess(futs); + } catch (const std::exception& e) { + LOG_KNOWHERE_WARNING_ << "CalcDistance by binary raw storage error: " << e.what(); + return expected::Err(Status::faiss_inner_error, e.what()); + } + + return GenResultDataSet(num_queries, labels_len, std::unique_ptr{}, std::move(distances)); + } + + private: + std::shared_ptr index_; +}; + +} // namespace + +expected> +CreateEmbListRawStorageForBuild(const DataSetPtr dataset, const std::string& metric_type) { + auto dim = dataset->GetDim(); + + std::shared_ptr storage; + if (IsBinarySubMetric(metric_type)) { + if (dim % 8 != 0) { + LOG_KNOWHERE_WARNING_ << "Binary emb_list raw storage requires dim to be a multiple of 8, got " << dim; + return expected>::Err( + Status::invalid_args, "binary emb_list raw storage dim must be a multiple of 8"); + } + storage = std::make_shared(dim, ToFaissBinaryMetric(metric_type)); + } else { + faiss::MetricType faiss_metric = faiss::METRIC_INNER_PRODUCT; + if (metric_type == metric::L2) { + faiss_metric = faiss::METRIC_L2; + } + storage = std::make_shared(dim, faiss_metric); + } + + const auto status = storage->Add(dataset); + if (status != Status::success) { + return expected>::Err(status, + "failed to add vectors to emb_list raw storage"); + } + return storage; +} + +expected> +ReadEmbListRawStorageFromBinary(const BinaryPtr& raw_index_bin, const std::string& metric_type) { + MemoryIOReader reader(raw_index_bin->data.get(), raw_index_bin->size); + if (IsBinarySubMetric(metric_type)) { + auto* index = faiss::cppcontrib::knowhere::read_index_binary(&reader); + auto* flat_index = dynamic_cast(index); + if (flat_index == nullptr) { + delete index; + LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX is not an IndexBinaryFlat"; + return expected>::Err(Status::emb_list_inner_error, + "EMB_LIST_RAW_INDEX is not an IndexBinaryFlat"); + } + auto storage = std::make_shared( + std::shared_ptr(flat_index)); + LOG_KNOWHERE_INFO_ << "Loaded binary raw vector storage: " << storage->Count() << " vectors"; + return storage; + } + + auto* index = faiss::cppcontrib::knowhere::read_index(&reader); + auto* flat_index = dynamic_cast<::faiss::IndexFlat*>(index); + if (flat_index == nullptr) { + delete index; + LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX is not an IndexFlat"; + return expected>::Err(Status::emb_list_inner_error, + "EMB_LIST_RAW_INDEX is not an IndexFlat"); + } + auto storage = std::make_shared(std::shared_ptr<::faiss::IndexFlat>(flat_index)); + LOG_KNOWHERE_INFO_ << "Loaded raw vector storage: " << storage->Count() << " vectors"; + return storage; +} + +expected> +ReadEmbListRawStorageFromFile(const std::string& filename, int io_flags, const std::string& metric_type) { + if (IsBinarySubMetric(metric_type)) { + auto* index = faiss::cppcontrib::knowhere::read_index_binary(filename.data(), io_flags); + auto* flat_index = dynamic_cast(index); + if (flat_index == nullptr) { + delete index; + LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX file is not an IndexBinaryFlat"; + return expected>::Err( + Status::emb_list_inner_error, "EMB_LIST_RAW_INDEX file is not an IndexBinaryFlat"); + } + auto storage = std::make_shared( + std::shared_ptr(flat_index)); + LOG_KNOWHERE_INFO_ << "Loaded binary raw vector storage from file: " << storage->Count() << " vectors"; + return storage; + } + + auto* index = faiss::cppcontrib::knowhere::read_index(filename.data(), io_flags); + auto* flat_index = dynamic_cast<::faiss::IndexFlat*>(index); + if (flat_index == nullptr) { + delete index; + LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX file is not an IndexFlat"; + return expected>::Err(Status::emb_list_inner_error, + "EMB_LIST_RAW_INDEX file is not an IndexFlat"); + } + auto storage = std::make_shared(std::shared_ptr<::faiss::IndexFlat>(flat_index)); + LOG_KNOWHERE_INFO_ << "Loaded raw vector storage from file: " << storage->Count() << " vectors"; + return storage; +} + +} // namespace knowhere diff --git a/src/index/emb_list/emb_list_raw_storage.h b/src/index/emb_list/emb_list_raw_storage.h new file mode 100644 index 000000000..ab0a6eb8f --- /dev/null +++ b/src/index/emb_list/emb_list_raw_storage.h @@ -0,0 +1,70 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// 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. + +#pragma once + +#include +#include +#include + +#include "knowhere/binaryset.h" +#include "knowhere/dataset.h" +#include "knowhere/expected.h" + +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) +#include "knowhere/comp/task.h" +#else +class ThreadPool; +#endif + +namespace milvus { +struct OpContext; +} // namespace milvus + +namespace knowhere { + +class EmbListRawStorage { + public: + virtual ~EmbListRawStorage() = default; + + virtual int64_t + Dim() const = 0; + + virtual int64_t + Count() const = 0; + + virtual size_t + CodeSize() const = 0; + + virtual Status + Add(const DataSetPtr dataset) = 0; + + virtual Status + ReconstructN(int64_t start, int64_t n, uint8_t* out) const = 0; + + virtual Status + Serialize(BinarySet& binset) const = 0; + + virtual expected + CalcDistance(const DataSetPtr dataset, const int64_t* labels, size_t labels_len, const std::string& metric_type, + bool is_cosine, std::shared_ptr pool, milvus::OpContext* op_context) const = 0; +}; + +expected> +CreateEmbListRawStorageForBuild(const DataSetPtr dataset, const std::string& metric_type); + +expected> +ReadEmbListRawStorageFromBinary(const BinaryPtr& raw_index_bin, const std::string& metric_type); + +expected> +ReadEmbListRawStorageFromFile(const std::string& filename, int io_flags, const std::string& metric_type); + +} // namespace knowhere diff --git a/src/index/emb_list/emb_list_strategy_lemur.cc b/src/index/emb_list/emb_list_strategy_lemur.cc index 2e396bd2a..201ad6912 100644 --- a/src/index/emb_list/emb_list_strategy_lemur.cc +++ b/src/index/emb_list/emb_list_strategy_lemur.cc @@ -9,6 +9,9 @@ // 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 @@ -16,6 +19,7 @@ #include #include #include +#include #include #include "index/emb_list/cblas_decl.h" @@ -28,7 +32,6 @@ #include "knowhere/object.h" #include "knowhere/thread_pool.h" #include "knowhere/utils.h" -#include "simd/hook.h" #include "simple_mlp.h" #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) @@ -48,6 +51,29 @@ FindMax(const float* data, size_t len) { return max_val; } +static inline bool +IsBinarySubMetric(const std::string& metric_type) { + return metric_type == metric::HAMMING || metric_type == metric::JACCARD; +} + +static inline bool +BinaryBitIsSet(const uint8_t* row, int32_t bit) { + return ((row[bit / 8] >> (bit % 8)) & 1U) != 0; +} + +static void +UnpackBinaryRowsToFloat(const uint8_t* input, size_t rows, int32_t dim, bool signed_mapping, float* output) { + const size_t code_size = static_cast(dim / 8); + for (size_t r = 0; r < rows; ++r) { + const uint8_t* row = input + r * code_size; + float* out = output + r * dim; + for (int32_t d = 0; d < dim; ++d) { + const bool bit = BinaryBitIsSet(row, d); + out[d] = signed_mapping ? (bit ? 1.0f : -1.0f) : (bit ? 1.0f : 0.0f); + } + } +} + /** * @brief LEMUR (Learned Multi-Vector Retrieval) strategy. * @@ -57,7 +83,7 @@ FindMax(const float* data, size_t len) { * * Training: * 1. Sample vectors from corpus as training inputs - * 2. Compute MaxSim between sampled vectors and all documents as labels + * 2. Compute emb-list aggregation labels between sampled vectors and all documents * 3. Train MLP: input(dim) -> hidden(hidden_dim) -> output(num_docs) * 4. W matrix = output_layer weights, shape [num_docs, hidden_dim] * @@ -66,7 +92,7 @@ FindMax(const float* data, size_t len) { * 2. Aggregate query features (sum/mean) * 3. Approximate scoring: query_feat @ W.T * 4. ANN search on W to find candidates - * 5. MaxSim reranking on candidates + * 5. Exact emb-list aggregation reranking on candidates */ class LemurEmbListStrategy : public EmbListStrategy { public: @@ -100,20 +126,31 @@ class LemurEmbListStrategy : public EmbListStrategy { } size_t total_vectors = doc_offset.offset.back(); - // Parse metric for ComputeMaxSimLabels + // Parse metric for training labels. auto metric_or = ParseEmbListMetric(config); if (!metric_or.has_value()) { return expected>::Err(metric_or.error(), metric_or.what()); } - is_l2_ = (metric_or.value().sub_metric_type == metric::L2); + const auto metric_info = metric_or.value(); + const auto& sub_metric_type = metric_info.sub_metric_type; + is_binary_ = IsBinarySubMetric(sub_metric_type); + binary_metric_type_ = is_binary_ ? sub_metric_type : ""; + is_l2_ = (sub_metric_type == metric::L2); + if (is_binary_ && original_dim_ % 8 != 0) { + return expected>::Err(Status::invalid_args, + "LEMUR binary requires dim to be a multiple of 8"); + } LOG_KNOWHERE_INFO_ << "LEMUR PrepareDataForBuild: num_docs=" << num_docs_ << ", total_vectors=" << total_vectors << ", original_dim=" << original_dim_ << ", hidden_dim=" << hidden_dim_ - << ", num_train_samples=" << num_train_samples_ << ", epochs=" << num_epochs_; + << ", num_train_samples=" << num_train_samples_ << ", epochs=" << num_epochs_ + << ", binary=" << is_binary_ << ", agg=" << metric_info.el_metric_type; // 2. Store doc_offset for reranking emb_list_offset_ = std::make_shared(doc_offset.offset); - const float* raw_data = static_cast(dataset->GetTensor()); + const float* raw_data = is_binary_ ? nullptr : static_cast(dataset->GetTensor()); + const uint8_t* raw_binary_data = is_binary_ ? static_cast(dataset->GetTensor()) : nullptr; + const size_t binary_code_size = is_binary_ ? static_cast(original_dim_ / 8) : 0; // 3. Sample training vectors (reservoir sampling: O(actual_samples) memory). std::mt19937 rng(seed_); @@ -145,22 +182,39 @@ class LemurEmbListStrategy : public EmbListStrategy { std::vector X_train(actual_samples * original_dim_); for (size_t i = 0; i < actual_samples; ++i) { - std::memcpy(X_train.data() + i * original_dim_, raw_data + sample_indices[i] * original_dim_, - original_dim_ * sizeof(float)); + if (is_binary_) { + const uint8_t* sample = raw_binary_data + static_cast(sample_indices[i]) * binary_code_size; + UnpackBinaryRowsToFloat(sample, 1, original_dim_, binary_metric_type_ == metric::HAMMING, + X_train.data() + i * original_dim_); + } else { + std::memcpy(X_train.data() + i * original_dim_, raw_data + sample_indices[i] * original_dim_, + original_dim_ * sizeof(float)); + } } LOG_KNOWHERE_INFO_ << "LEMUR: Sampled " << actual_samples << " vectors for training"; - // 4. Compute training labels (MaxSim for each sample vector against each document) + // 4. Compute training labels for each sampled vector against each document. auto label_start = std::chrono::high_resolution_clock::now(); std::vector y_train(actual_samples * num_docs_); - ComputeMaxSimLabels(X_train.data(), actual_samples, raw_data, doc_offset, y_train.data()); + if (is_binary_) { + ComputeBinaryMaxSimLabels(sample_indices.data(), actual_samples, raw_binary_data, doc_offset, + y_train.data()); + } else { + ComputeMaxSimLabels(X_train.data(), actual_samples, raw_data, doc_offset, y_train.data()); + } + std::string empty_doc_label_error; + auto fill_label_status = + FillEmptyDocLabelsWithRowWiseWorst(doc_offset, actual_samples, y_train.data(), empty_doc_label_error); + if (fill_label_status != Status::success) { + return expected>::Err(fill_label_status, empty_doc_label_error); + } auto label_end = std::chrono::high_resolution_clock::now(); double label_ms = std::chrono::duration(label_end - label_start).count(); - LOG_KNOWHERE_INFO_ << "LEMUR: Computed MaxSim labels in " << label_ms << " ms"; + LOG_KNOWHERE_INFO_ << "LEMUR: Computed " << metric_info.el_metric_type << " labels in " << label_ms << " ms"; - // 5. Save raw labels for OLS (original LEMUR uses raw MaxSim, not normalized) + // 5. Save raw labels for OLS (original LEMUR uses raw aggregation labels, not normalized) std::vector y_train_raw = y_train; // 6. Normalize labels (z-score normalization for stable MLP training) @@ -199,7 +253,7 @@ class LemurEmbListStrategy : public EmbListStrategy { LOG_KNOWHERE_INFO_ << "LEMUR: MLP training completed in " << train_ms << " ms, final_loss=" << final_loss; // 8. Compute W using pseudoinverse (fit_corpus step from original LEMUR) - // W = pinv(Z) @ Y, where Z = features of sampled vectors, Y = MaxSim labels + // W = pinv(Z) @ Y, where Z = features of sampled vectors, Y = raw aggregation labels // This is a least-squares solution: find W such that Z @ W^T ≈ Y auto ols_start = std::chrono::high_resolution_clock::now(); final_hidden_dim_ = mlp_->FinalHiddenDim(); @@ -225,7 +279,7 @@ class LemurEmbListStrategy : public EmbListStrategy { ZtZ[i * final_hidden_dim_ + i] += lambda; } - // Compute Z^T @ Y using BLAS (Y is raw MaxSim labels, matching original LEMUR) + // Compute Z^T @ Y using BLAS. cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans, final_hidden_dim_, num_docs_, actual_samples, 1.0f, Z.data(), final_hidden_dim_, y_train_raw.data(), num_docs_, 0.0f, ZtY.data(), num_docs_); @@ -294,6 +348,23 @@ class LemurEmbListStrategy : public EmbListStrategy { return true; } + [[nodiscard]] EmbListAnnIndexSpec + AnnIndexSpec(const BaseConfig& config) const override { + auto metric_or = ParseEmbListMetric(config); + if (metric_or.has_value() && IsBinarySubMetric(metric_or.value().sub_metric_type)) { + return EmbListAnnIndexSpec{ + .target = EmbListAnnIndexTarget::SeparateIndex, + .data_type = EmbListAnnIndexDataType::Fp32, + .ann_metric_type = metric::IP, + }; + } + return EmbListAnnIndexSpec{ + .target = EmbListAnnIndexTarget::BaseIndex, + .data_type = EmbListAnnIndexDataType::SameAsBaseIndex, + .ann_metric_type = metric::IP, + }; + } + expected Search(const DataSetPtr query_dataset, const EmbListOffset& query_offset, const IndexNode* index, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { @@ -310,9 +381,16 @@ class LemurEmbListStrategy : public EmbListStrategy { return expected::Err(metric_or.error(), metric_or.what()); } auto& mi = metric_or.value(); + const bool is_binary_query = IsBinarySubMetric(mi.sub_metric_type); + if (is_binary_query && original_dim_ % 8 != 0) { + return expected::Err(Status::invalid_args, + "LEMUR binary search requires dim to be a multiple of 8"); + } auto num_query_docs = query_offset.num_el(); - const float* query_data = static_cast(query_dataset->GetTensor()); + const float* query_data = nullptr; + const uint8_t* query_binary_data = nullptr; + const size_t binary_code_size = is_binary_query ? static_cast(original_dim_ / 8) : 0; LOG_KNOWHERE_DEBUG_ << "LEMUR Search: num_query_docs=" << num_query_docs << ", k=" << k; @@ -340,6 +418,17 @@ class LemurEmbListStrategy : public EmbListStrategy { if (total_query_tokens > static_cast(std::numeric_limits::max())) { return expected::Err(Status::emb_list_inner_error, "total query tokens exceeds int32 limit"); } + std::vector query_float_data; + if (is_binary_query) { + query_binary_data = static_cast(query_dataset->GetTensor()); + query_float_data.resize(total_query_tokens * original_dim_); + UnpackBinaryRowsToFloat(query_binary_data, total_query_tokens, original_dim_, + mi.sub_metric_type == metric::HAMMING, query_float_data.data()); + query_data = query_float_data.data(); + } else { + query_data = static_cast(query_dataset->GetTensor()); + } + std::vector all_feats(total_query_tokens * final_hidden_dim_); mlp_->ExtractFeatures(query_data, static_cast(total_query_tokens), all_feats.data()); @@ -360,8 +449,8 @@ class LemurEmbListStrategy : public EmbListStrategy { // 4. ANN search on W auto query_feat_dataset = GenDataSet(num_query_docs, final_hidden_dim_, query_feats.data()); config.k = ann_k; - config.metric_type = mi.sub_metric_type; - auto ann_result = index->Search(query_feat_dataset, std::move(cfg), bitset, op_context); + config.metric_type = metric::IP; + auto ann_result = index->SearchEmbListAnnIndex(query_feat_dataset, std::move(cfg), bitset, op_context); if (!ann_result.has_value()) { LOG_KNOWHERE_ERROR_ << "LEMUR ANN search failed: " << ann_result.what(); return expected::Err(Status::emb_list_inner_error, "ANN search failed"); @@ -381,7 +470,6 @@ class LemurEmbListStrategy : public EmbListStrategy { auto ids = std::make_unique(num_query_docs * k); auto dists = std::make_unique(num_query_docs * k); const auto* ann_dists = ann_result.value()->GetDistance(); - for (size_t q = 0; q < num_query_docs; ++q) { for (int32_t i = 0; i < k; ++i) { if (i < ann_k) { @@ -389,8 +477,7 @@ class LemurEmbListStrategy : public EmbListStrategy { dists[q * k + i] = ann_dists[q * ann_k + i]; } else { ids[q * k + i] = -1; - dists[q * k + i] = mi.larger_is_closer ? -std::numeric_limits::infinity() - : std::numeric_limits::infinity(); + dists[q * k + i] = -std::numeric_limits::infinity(); } } } @@ -398,7 +485,7 @@ class LemurEmbListStrategy : public EmbListStrategy { std::move(dists)); } - // 6. MaxSim reranking via CalcDistByIDs + // 6. Exact emb-list aggregation reranking via CalcDistByIDs auto ids = std::make_unique(num_query_docs * k); auto dists = std::make_unique(num_query_docs * k); @@ -431,7 +518,12 @@ class LemurEmbListStrategy : public EmbListStrategy { total_candidates += candidate_docs.size(); // Build query dataset for this query document - auto bf_query_dataset = GenDataSet(nq, original_dim_, query_data + q_vec_start * original_dim_); + DataSetPtr bf_query_dataset; + if (is_binary_query) { + bf_query_dataset = GenDataSet(nq, original_dim_, query_binary_data + q_vec_start * binary_code_size); + } else { + bf_query_dataset = GenDataSet(nq, original_dim_, query_data + q_vec_start * original_dim_); + } auto status = RerankByCalcDistByIDs(candidate_docs, bf_query_dataset, nq, k, mi.larger_is_closer, mi.is_cosine, @@ -712,6 +804,197 @@ class LemurEmbListStrategy : public EmbListStrategy { } private: + Status + FillEmptyDocLabelsWithRowWiseWorst(const EmbListOffset& offset, size_t num_samples, float* labels, + std::string& error_msg) const { + const size_t num_docs = offset.num_el(); + std::vector empty_docs; + for (size_t doc_id = 0; doc_id < num_docs; ++doc_id) { + if (offset.offset[doc_id] == offset.offset[doc_id + 1]) { + empty_docs.push_back(doc_id); + } + } + + if (empty_docs.empty()) { + return Status::success; + } + if (empty_docs.size() == num_docs) { + error_msg = "LEMUR training labels cannot be computed when all emb_list documents are empty"; + return Status::invalid_args; + } + + std::vector non_empty_docs; + non_empty_docs.reserve(num_docs - empty_docs.size()); + for (size_t doc_id = 0; doc_id < num_docs; ++doc_id) { + if (offset.offset[doc_id] != offset.offset[doc_id + 1]) { + non_empty_docs.push_back(doc_id); + } + } + + // LEMUR labels are all larger-is-better, so the finite minimum is the row-wise worst label. + for (size_t sample_id = 0; sample_id < num_samples; ++sample_id) { + float worst_label = std::numeric_limits::infinity(); + float* label_row = labels + sample_id * num_docs; + for (const auto doc_id : non_empty_docs) { + const float label = label_row[doc_id]; + if (!std::isfinite(label)) { + error_msg = "LEMUR training produced non-finite label for non-empty emb_list document"; + return Status::emb_list_inner_error; + } + worst_label = std::min(worst_label, label); + } + + for (const auto doc_id : empty_docs) { + label_row[doc_id] = worst_label; + } + } + + return Status::success; + } + + /** + * @brief Compute MaxSim labels for binary raw tokens. + * + * Hamming uses signed binary dot as the training target: + * dot({-1,+1}, {-1,+1}) = dim - 2 * hamming_distance. + * Jaccard uses max Jaccard similarity. Exact rerank still uses the true binary distance. + */ + void + ComputeBinaryMaxSimLabels(const int64_t* sample_indices, size_t num_samples, const uint8_t* raw_data, + const EmbListOffset& offset, float* labels) const { + auto pool = ThreadPool::GetGlobalSearchThreadPool(); + const size_t num_docs = offset.num_el(); + const size_t total_vecs = offset.offset.back(); + const size_t code_size = static_cast(original_dim_ / 8); + + struct Chunk { + size_t doc_start; + size_t doc_end; + size_t vec_start; + size_t vec_end; + }; + + constexpr size_t kTargetChunkVecs = 50000; + std::vector chunks; + { + size_t cur_doc = 0; + while (cur_doc < num_docs) { + size_t chunk_doc_start = cur_doc; + size_t chunk_vec_start = offset.offset[cur_doc]; + size_t chunk_vec_end = chunk_vec_start; + + while (cur_doc < num_docs) { + chunk_vec_end = offset.offset[cur_doc + 1]; + cur_doc++; + if (chunk_vec_end - chunk_vec_start >= kTargetChunkVecs) { + break; + } + } + chunks.push_back({.doc_start = chunk_doc_start, + .doc_end = cur_doc, + .vec_start = chunk_vec_start, + .vec_end = chunk_vec_end}); + } + } + + std::vector raw_popcounts; + if (binary_metric_type_ == metric::JACCARD) { + raw_popcounts.resize(total_vecs); + for (size_t v = 0; v < total_vecs; ++v) { + raw_popcounts[v] = faiss::cppcontrib::knowhere::popcnt(raw_data + v * code_size, code_size); + } + } + + constexpr size_t kSampleBatchSize = 512; + std::vector> futs; + futs.reserve(chunks.size()); + for (size_t sample_start = 0; sample_start < num_samples; sample_start += kSampleBatchSize) { + size_t sample_end = std::min(sample_start + kSampleBatchSize, num_samples); + size_t batch_samples = sample_end - sample_start; + + std::vector sample_codes(batch_samples * code_size); + for (size_t s = 0; s < batch_samples; ++s) { + const uint8_t* sample = raw_data + static_cast(sample_indices[sample_start + s]) * code_size; + std::memcpy(sample_codes.data() + s * code_size, sample, code_size); + } + + std::vector sample_popcounts; + if (binary_metric_type_ == metric::JACCARD) { + sample_popcounts.resize(batch_samples); + for (size_t s = 0; s < batch_samples; ++s) { + sample_popcounts[s] = + faiss::cppcontrib::knowhere::popcnt(sample_codes.data() + s * code_size, code_size); + } + } + + futs.clear(); + + for (const auto& chunk : chunks) { + futs.emplace_back(pool->push([&, chunk, batch_samples, sample_start]() { + const size_t chunk_vecs = chunk.vec_end - chunk.vec_start; + const uint8_t* chunk_data = raw_data + chunk.vec_start * code_size; + + if (binary_metric_type_ == metric::JACCARD) { + std::vector local_distances(batch_samples * chunk_vecs); + faiss::cppcontrib::knowhere::all_jaccard_distances(sample_codes.data(), chunk_data, code_size, + batch_samples, chunk_vecs, + local_distances.data(), nullptr); + + for (size_t s = 0; s < batch_samples; ++s) { + const float* distance_row = local_distances.data() + s * chunk_vecs; + float* label_row = labels + (sample_start + s) * num_docs; + const bool sample_empty = sample_popcounts[s] == 0; + + for (size_t d = chunk.doc_start; d < chunk.doc_end; ++d) { + const size_t doc_vec_start = offset.offset[d] - chunk.vec_start; + const size_t doc_vec_end = offset.offset[d + 1] - chunk.vec_start; + float best = -std::numeric_limits::infinity(); + + for (size_t v = doc_vec_start; v < doc_vec_end; ++v) { + const bool token_empty = raw_popcounts[chunk.vec_start + v] == 0; + const float score = (sample_empty && token_empty) ? 1.0f : 1.0f - distance_row[v]; + best = std::max(best, score); + } + label_row[d] = best; + } + } + } else { + std::vector local_distances(batch_samples * chunk_vecs); + // hammings() has correct row-major output only on its explicit specialized code sizes. + if (code_size == 8 || code_size == 16 || code_size == 32 || code_size == 64) { + faiss::cppcontrib::knowhere::hammings(sample_codes.data(), chunk_data, batch_samples, + chunk_vecs, code_size, local_distances.data()); + } else { + faiss::cppcontrib::knowhere::all_hamming_distances(sample_codes.data(), chunk_data, + code_size, batch_samples, chunk_vecs, + local_distances.data(), nullptr); + } + + for (size_t s = 0; s < batch_samples; ++s) { + const int32_t* distance_row = local_distances.data() + s * chunk_vecs; + float* label_row = labels + (sample_start + s) * num_docs; + + for (size_t d = chunk.doc_start; d < chunk.doc_end; ++d) { + const size_t doc_vec_start = offset.offset[d] - chunk.vec_start; + const size_t doc_vec_end = offset.offset[d + 1] - chunk.vec_start; + int32_t best = std::numeric_limits::max(); + + for (size_t v = doc_vec_start; v < doc_vec_end; ++v) { + best = std::min(best, distance_row[v]); + } + label_row[d] = best == std::numeric_limits::max() + ? -std::numeric_limits::infinity() + : static_cast(original_dim_ - 2 * best); + } + } + } + })); + } + + WaitAllSuccess(futs); + } + } + /** * @brief Compute MaxSim labels for training (chunk-parallel, fused sgemm+reduce). * @@ -824,7 +1107,9 @@ class LemurEmbListStrategy : public EmbListStrategy { for (size_t d = chunk.doc_start; d < chunk.doc_end; ++d) { size_t doc_vec_start = offset.offset[d] - chunk.vec_start; size_t doc_vec_end = offset.offset[d + 1] - chunk.vec_start; - label_row[d] = FindMax(ip_row + doc_vec_start, doc_vec_end - doc_vec_start); + label_row[d] = doc_vec_start == doc_vec_end + ? -std::numeric_limits::infinity() + : FindMax(ip_row + doc_vec_start, doc_vec_end - doc_vec_start); } } })); @@ -848,6 +1133,8 @@ class LemurEmbListStrategy : public EmbListStrategy { int32_t original_dim_ = 0; int64_t num_docs_ = 0; bool is_l2_ = false; + bool is_binary_ = false; + std::string binary_metric_type_; // MLP model std::unique_ptr mlp_; diff --git a/src/index/emb_list/emb_list_strategy_muvera.cc b/src/index/emb_list/emb_list_strategy_muvera.cc index e77adb8f1..6881a17ae 100644 --- a/src/index/emb_list/emb_list_strategy_muvera.cc +++ b/src/index/emb_list/emb_list_strategy_muvera.cc @@ -156,6 +156,15 @@ class MuveraEmbListStrategy : public EmbListStrategy { return true; } + [[nodiscard]] EmbListAnnIndexSpec + AnnIndexSpec(const BaseConfig& /*config*/) const override { + return EmbListAnnIndexSpec{ + .target = EmbListAnnIndexTarget::BaseIndex, + .data_type = EmbListAnnIndexDataType::SameAsBaseIndex, + .ann_metric_type = std::nullopt, + }; + } + expected Search(const DataSetPtr query_dataset, const EmbListOffset& query_offset, const IndexNode* index, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { diff --git a/src/index/emb_list/emb_list_strategy_token_ann.cc b/src/index/emb_list/emb_list_strategy_token_ann.cc index c6d14e91c..45d40a9e0 100644 --- a/src/index/emb_list/emb_list_strategy_token_ann.cc +++ b/src/index/emb_list/emb_list_strategy_token_ann.cc @@ -47,6 +47,15 @@ class TokenANNEmbListStrategy : public EmbListStrategy { return true; // needs vector_id -> doc_id mapping for bitset filtering } + [[nodiscard]] EmbListAnnIndexSpec + AnnIndexSpec(const BaseConfig& /*config*/) const override { + return EmbListAnnIndexSpec{ + .target = EmbListAnnIndexTarget::BaseIndex, + .data_type = EmbListAnnIndexDataType::SameAsBaseIndex, + .ann_metric_type = std::nullopt, + }; + } + expected Search(const DataSetPtr query_dataset, const EmbListOffset& query_offset, const IndexNode* index, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 63101aa21..6d9d33068 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -70,7 +70,7 @@ namespace knowhere { // class BaseFaissIndexNode : public IndexNode { public: - BaseFaissIndexNode(const int32_t& /*version*/, const Object& object) { + BaseFaissIndexNode(const int32_t& version, const Object& object) : IndexNode(version) { build_pool = ThreadPool::GetGlobalBuildThreadPool(); search_pool = ThreadPool::GetGlobalSearchThreadPool(); } @@ -1507,8 +1507,8 @@ class BaseFaissRegularIndexHNSWNode : public BaseFaissRegularIndexNode { expected CalcDistByIDs(const DataSetPtr dataset, const BitsetView& bitset_, const int64_t* labels, const size_t labels_len, const bool is_cosine, milvus::OpContext* op_context) const override { - // When emb_list_raw_index_ exists (MUVERA/LEMUR), use it for exact distance computation - if (emb_list_raw_index_) { + // When emb_list raw storage exists (MUVERA/LEMUR), use it for exact distance computation. + if (emb_list_raw_storage_) { return CalcDistByRawIndex(dataset, labels, labels_len, is_cosine, search_pool, op_context); } @@ -2280,6 +2280,33 @@ class HNSWIndexNodeWithFallback : public IndexNode { } } + int64_t + CountForSearchBitset() const override { + if (use_base_index) { + return base_index->CountForSearchBitset(); + } else { + return fallback_search_index->CountForSearchBitset(); + } + } + + IndexNode* + AnnIndexNode() override { + if (use_base_index) { + return base_index->AnnIndexNode(); + } else { + return fallback_search_index->AnnIndexNode(); + } + } + + const IndexNode* + AnnIndexNode() const override { + if (use_base_index) { + return static_cast(base_index.get())->AnnIndexNode(); + } else { + return static_cast(fallback_search_index.get())->AnnIndexNode(); + } + } + int64_t Size() const override { if (use_base_index) { @@ -2449,6 +2476,16 @@ class HNSWIndexNodeWithFallback : public IndexNode { } } + expected + SearchEmbListAnnIndex(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + if (use_base_index) { + return base_index->SearchEmbListAnnIndex(dataset, std::move(config), bitset, op_context); + } else { + return fallback_search_index->SearchEmbListAnnIndex(dataset, std::move(config), bitset, op_context); + } + } + expected RangeSearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { diff --git a/src/index/index.cc b/src/index/index.cc index 70f615494..e4586eb1c 100644 --- a/src/index/index.cc +++ b/src/index/index.cc @@ -143,9 +143,10 @@ Index::Search(const DataSetPtr dataset, const Json& json, const BitsetView& b // when index is mutable, it could happen that data count larger than bitset size, see // https://github.com/zilliztech/knowhere/issues/70 // so something must be wrong at caller side when passed bitset size larger than data count - if (bitset_.size() > static_cast(this->Count())) { + const auto count_for_bitset = this->node->CountForSearchBitset(); + if (bitset_.size() > static_cast(count_for_bitset)) { msg = fmt::format("bitset size should be <= data count, but we get bitset size: {}, data count: {}", - bitset_.size(), this->Count()); + bitset_.size(), count_for_bitset); LOG_KNOWHERE_ERROR_ << msg; return expected::Err(Status::invalid_args, msg); } @@ -409,19 +410,19 @@ Index::DeserializeFromFile(const std::string& filename, const Json& json) noe template inline int64_t Index::Dim() const noexcept { - return GuardedCall([&]() { return this->node->Dim(); }); + return GuardedCall([&]() { return this->node->AnnIndexNode()->Dim(); }); } template inline int64_t Index::Size() const noexcept { - return GuardedCall([&]() { return this->node->Size(); }); + return GuardedCall([&]() { return this->node->AnnIndexNode()->Size(); }); } template inline int64_t Index::Count() const noexcept { - return GuardedCall([&]() { return this->node->Count(); }); + return GuardedCall([&]() { return this->node->AnnIndexNode()->Count(); }); } template diff --git a/src/index/index_node.cc b/src/index/index_node.cc index a5d41e3f8..283333fa3 100644 --- a/src/index/index_node.cc +++ b/src/index/index_node.cc @@ -12,14 +12,18 @@ #include "knowhere/index/index_node.h" #include +#include #include +#include #include #include +#include -#include "faiss/cppcontrib/knowhere/IndexFlat.h" #include "faiss/cppcontrib/knowhere/index_io.h" +#include "index/emb_list/emb_list_raw_storage.h" #include "io/memory_io.h" #include "knowhere/context.h" +#include "knowhere/index/index_factory.h" #include "knowhere/log.h" #include "knowhere/range_util.h" #include "knowhere/utils.h" @@ -32,6 +36,104 @@ namespace knowhere { +struct EmbListSeparateAnnIndexHolder { + explicit EmbListSeparateAnnIndexHolder(Index&& index_in) : index(std::move(index_in)) { + } + + Index index; +}; + +IndexNode* +IndexNode::AnnIndexNode() { + if (emb_list_separate_ann_index_) { + return emb_list_separate_ann_index_->index.Node(); + } + return this; +} + +const IndexNode* +IndexNode::AnnIndexNode() const { + if (emb_list_separate_ann_index_) { + return emb_list_separate_ann_index_->index.Node(); + } + return this; +} + +namespace { + +inline const char* +AnnIndexTargetName(EmbListAnnIndexTarget target) { + switch (target) { + case EmbListAnnIndexTarget::BaseIndex: + return "BaseIndex"; + case EmbListAnnIndexTarget::SeparateIndex: + return "SeparateIndex"; + } + return "Unknown"; +} + +inline const char* +AnnIndexDataTypeName(EmbListAnnIndexDataType data_type) { + switch (data_type) { + case EmbListAnnIndexDataType::SameAsBaseIndex: + return "SameAsBaseIndex"; + case EmbListAnnIndexDataType::Fp32: + return "Fp32"; + case EmbListAnnIndexDataType::Fp16: + return "Fp16"; + case EmbListAnnIndexDataType::Bf16: + return "Bf16"; + case EmbListAnnIndexDataType::Int8: + return "Int8"; + case EmbListAnnIndexDataType::Bin1: + return "Bin1"; + } + return "Unknown"; +} + +expected> +CreateEmbListSeparateAnnIndex(EmbListAnnIndexDataType data_type, const std::string& index_type, int32_t version) { + switch (data_type) { + case EmbListAnnIndexDataType::Fp32: + return IndexFactory::Instance().Create(index_type, version); + case EmbListAnnIndexDataType::SameAsBaseIndex: + return expected>::Err( + Status::not_implemented, + "emb_list separate ANN index does not support SameAsBaseIndex; use BaseIndex target instead"); + case EmbListAnnIndexDataType::Fp16: + return expected>::Err(Status::not_implemented, + "emb_list separate fp16 ANN index is not implemented"); + case EmbListAnnIndexDataType::Bf16: + return expected>::Err(Status::not_implemented, + "emb_list separate bf16 ANN index is not implemented"); + case EmbListAnnIndexDataType::Int8: + return expected>::Err(Status::not_implemented, + "emb_list separate int8 ANN index is not implemented"); + case EmbListAnnIndexDataType::Bin1: + return expected>::Err(Status::not_implemented, + "emb_list separate binary ANN index is not implemented"); + } + return expected>::Err(Status::not_implemented, "unknown emb_list separate ANN index type"); +} + +class ScopedMetricTypeOverride { + public: + ScopedMetricTypeOverride(BaseConfig& config, std::string metric_type) + : config_(config), original_metric_type_(config.metric_type) { + config_.metric_type = std::move(metric_type); + } + + ~ScopedMetricTypeOverride() { + config_.metric_type = std::move(original_metric_type_); + } + + private: + BaseConfig& config_; + std::optional original_metric_type_; +}; + +} // namespace + // NOLINTBEGIN(google-default-arguments) expected IndexNode::RangeSearch(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, @@ -293,6 +395,12 @@ IndexNode::SearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr return SearchEmbList(dataset, std::move(config), bitset, op_context); } +expected +IndexNode::SearchEmbListAnnIndex(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + milvus::OpContext* op_context) const { + return AnnIndexNode()->Search(dataset, std::move(config), bitset, op_context); +} + expected IndexNode::RangeSearchEmbListIfNeed(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const { @@ -332,9 +440,9 @@ IndexNode::GetEmbListByIds(const DataSetPtr dataset, const std::string& metric_t "GetEmbListByIds: invalid metric type " + metric_type); } - // Raw data can come from emb_list_raw_index_ (MUVERA/LEMUR) or base index (TokenANN) - bool use_raw_index = (emb_list_raw_index_ != nullptr); - if (!use_raw_index && !HasRawData(sub_metric.value())) { + // Raw data can come from emb_list_raw_storage_ (MUVERA/LEMUR) or base index (TokenANN). + bool use_raw_storage = (emb_list_raw_storage_ != nullptr); + if (!use_raw_storage && !HasRawData(sub_metric.value())) { return expected::Err( Status::not_implemented, "GetEmbListByIds requires raw data support, but the index does not store raw vectors"); @@ -342,7 +450,7 @@ IndexNode::GetEmbListByIds(const DataSetPtr dataset, const std::string& metric_t auto num_el_ids = dataset->GetRows(); auto el_ids = dataset->GetIds(); - auto dim = use_raw_index ? emb_list_raw_index_->d : Dim(); + auto dim = use_raw_storage ? emb_list_raw_storage_->Dim() : Dim(); // Build the output offset array std::vector out_offsets(num_el_ids + 1); @@ -370,16 +478,20 @@ IndexNode::GetEmbListByIds(const DataSetPtr dataset, const std::string& metric_t const void* tensor = nullptr; - if (use_raw_index) { - // MUVERA/LEMUR: vectors are contiguous per el in emb_list_raw_index_, use reconstruct_n - auto data = std::make_unique(total_vecs * dim); - float* ptr = data.get(); + if (use_raw_storage) { + // MUVERA/LEMUR raw vectors are contiguous per el in emb_list_raw_storage_, use reconstruct_n. + const size_t code_size = emb_list_raw_storage_->CodeSize(); + auto data = std::make_unique(total_vecs * code_size); + auto* ptr = reinterpret_cast(data.get()); for (int64_t i = 0; i < num_el_ids; i++) { auto start = static_cast(emb_list_offset_->offset[el_ids[i]]); auto len = static_cast(out_offsets[i + 1] - out_offsets[i]); if (len > 0) { - emb_list_raw_index_->reconstruct_n(start, len, ptr); - ptr += len * dim; + const auto status = emb_list_raw_storage_->ReconstructN(start, len, ptr); + if (status != Status::success) { + return expected::Err(status, "failed to reconstruct vectors from emb_list raw storage"); + } + ptr += len * code_size; } } tensor = data.release(); @@ -436,6 +548,7 @@ IndexNode::BuildEmbList(const DataSetPtr dataset, std::shared_ptr cfg, c } el_metric_type_ = metric_info_or.value().el_metric_type; auto sub_metric_type = metric_info_or.value().sub_metric_type; + emb_list_raw_metric_type_ = sub_metric_type; // 2. Create document offset structure EmbListOffset doc_offset(lims, num_rows); @@ -449,38 +562,68 @@ IndexNode::BuildEmbList(const DataSetPtr dataset, std::shared_ptr cfg, c } emb_list_strategy_ = std::move(strategy_or.value()); - // 4. Prepare data for build (strategy may transform data, e.g., FDE encoding in MUVERA) + // 4. Prepare the ANN dataset and ask the strategy how it should be indexed. auto build_data_or = emb_list_strategy_->PrepareDataForBuild(dataset, doc_offset, config); if (!build_data_or.has_value()) { - LOG_KNOWHERE_WARNING_ << "Failed to prepare data for build"; + LOG_KNOWHERE_WARNING_ << "Failed to prepare emb_list ANN build data"; return build_data_or.error(); } + // Some strategies build an ANN index over a representation whose data type differs from the + // raw vectors, e.g. binary LEMUR keeps raw data as bin1 but builds an fp32 learned ANN index. + // Ask the strategy where that ANN dataset should be indexed, which data type it uses, and + // whether it needs a metric different from the outer emb-list raw sub metric. + const auto ann_index_spec = emb_list_strategy_->AnnIndexSpec(config); - // Override metric_type to sub_metric for base index build - config.metric_type = sub_metric_type; + const auto ann_index_target = ann_index_spec.target; + const auto ann_index_data_type = ann_index_spec.data_type; + const auto ann_metric_type = ann_index_spec.ann_metric_type.value_or(sub_metric_type); // 5. Build underlying index (if strategy provides data) LOG_KNOWHERE_INFO_ << "Build EmbList-Index with strategy: " << strategy_type << ", metric type: " << el_metric_type_ - << ", sub metric type: " << sub_metric_type; + << ", sub metric type: " << sub_metric_type << ", ann metric type: " << ann_metric_type + << ", ann index target: " << AnnIndexTargetName(ann_index_target) + << ", ann index data type: " << AnnIndexDataTypeName(ann_index_data_type); if (build_data_or.value().has_value()) { - RETURN_IF_ERROR(Build(build_data_or.value().value(), cfg, use_knowhere_build_pool)); + // The underlying ANN index consumes a plain metric, not the outer emb-list metric. + // Keep this override scoped because cfg is shared with later emb-list build steps. + ScopedMetricTypeOverride metric_guard(config, ann_metric_type); + switch (ann_index_target) { + case EmbListAnnIndexTarget::BaseIndex: + // BaseIndex means the ANN dataset is compatible with this IndexNode's own data type. + // No child index is allocated; the current node builds its normal ANN index directly. + RETURN_IF_ERROR(Build(build_data_or.value().value(), cfg, use_knowhere_build_pool)); + break; + case EmbListAnnIndexTarget::SeparateIndex: { + // SeparateIndex is used when the strategy emits a representation whose data type differs from + // this IndexNode, for example binary LEMUR producing fp32 learned vectors. + auto separate_ann_index_or = + CreateEmbListSeparateAnnIndex(ann_index_data_type, Type(), version_.VersionNumber()); + if (!separate_ann_index_or.has_value()) { + LOG_KNOWHERE_WARNING_ + << "Failed to create separate ANN index for emb_list strategy: " << strategy_type + << ", index type: " << Type() + << ", ann index data type: " << AnnIndexDataTypeName(ann_index_data_type); + return separate_ann_index_or.error(); + } + auto separate_ann_index = + std::make_shared(std::move(separate_ann_index_or.value())); + RETURN_IF_ERROR(separate_ann_index->index.Node()->Build(build_data_or.value().value(), cfg, + use_knowhere_build_pool)); + emb_list_separate_ann_index_ = std::move(separate_ann_index); + break; + } + } } // 6. Create raw vector storage if strategy needs it if (emb_list_strategy_->NeedsRawVectorStorage()) { - auto original_dim = dataset->GetDim(); - auto total_vectors = dataset->GetRows(); - const float* raw_data = static_cast(dataset->GetTensor()); - - faiss::MetricType faiss_metric = faiss::METRIC_INNER_PRODUCT; - if (sub_metric_type == metric::L2) { - faiss_metric = faiss::METRIC_L2; + auto raw_storage_or = CreateEmbListRawStorageForBuild(dataset, sub_metric_type); + if (!raw_storage_or.has_value()) { + return raw_storage_or.error(); } - - emb_list_raw_index_ = std::make_shared(original_dim, faiss_metric); - emb_list_raw_index_->add(total_vectors, raw_data); - - LOG_KNOWHERE_INFO_ << "Created raw vector storage: " << total_vectors << " vectors, dim=" << original_dim; + emb_list_raw_storage_ = std::move(raw_storage_or.value()); + LOG_KNOWHERE_INFO_ << "Created raw vector storage: " << emb_list_raw_storage_->Count() + << " vectors, dim=" << emb_list_raw_storage_->Dim() << ", metric=" << sub_metric_type; } // 7. Strategy post-build hook @@ -536,18 +679,18 @@ IndexNode::SerializeEmbList(BinarySet& binset) const { binset.Append(meta::EMB_LIST_META, meta_data, writer.tellg()); // 3. Raw vector index as separate key (large, needs mmap in file path) - if (emb_list_raw_index_) { - MemoryIOWriter writer; - faiss::cppcontrib::knowhere::write_index(emb_list_raw_index_.get(), &writer); - std::shared_ptr raw_bin(writer.data()); - binset.Append(meta::EMB_LIST_RAW_INDEX, raw_bin, writer.tellg()); + if (emb_list_raw_storage_) { + RETURN_IF_ERROR(emb_list_raw_storage_->Serialize(binset)); } } catch (const std::exception& e) { LOG_KNOWHERE_WARNING_ << "serialize emb_list error: " << e.what(); return Status::emb_list_inner_error; } - // 4. Serialize base index + // 4. Serialize ANN index. BaseIndex is this node; SeparateIndex is the optional child index. + if (emb_list_separate_ann_index_) { + return emb_list_separate_ann_index_->index.Node()->Serialize(binset); + } return Serialize(binset); } @@ -561,13 +704,10 @@ IndexNode::DeserializeEmbListFromBinarySet(const BinarySet& binset, std::shared_ return metric_info_or.error(); } el_metric_type_ = metric_info_or.value().el_metric_type; - - // 2. Deserialize base index from BinarySet (override metric_type for base index) - cfg.metric_type = metric_info_or.value().sub_metric_type; - RETURN_IF_ERROR(Deserialize(binset, config)); + emb_list_raw_metric_type_ = metric_info_or.value().sub_metric_type; try { - // 3. Read EMB_LIST_META and parse strategy type + strategy blob + // 2. Read EMB_LIST_META and parse strategy type + strategy blob auto meta_bin = binset.GetByName(meta::EMB_LIST_META); if (!meta_bin) { LOG_KNOWHERE_WARNING_ << "EMB_LIST_META not found in binary set"; @@ -577,7 +717,7 @@ IndexNode::DeserializeEmbListFromBinarySet(const BinarySet& binset, std::shared_ auto [strategy_type, strategy_blob, strategy_blob_size] = ParseEmbListMetaHeader(meta_bin->data.get(), meta_bin->size); - // 4. Create strategy and deserialize strategy-specific data + // 3. Create strategy and deserialize strategy-specific data auto strategy_or = CreateEmbListStrategy(strategy_type, cfg); if (!strategy_or.has_value()) { LOG_KNOWHERE_WARNING_ << "Failed to create emb_list strategy: " << strategy_type; @@ -588,19 +728,43 @@ IndexNode::DeserializeEmbListFromBinarySet(const BinarySet& binset, std::shared_ LOG_KNOWHERE_INFO_ << "Deserialize emb_list with strategy: " << strategy_type; RETURN_IF_ERROR(emb_list_strategy_->Deserialize(strategy_blob, strategy_blob_size, cfg)); + // 4. Deserialize ANN index using the same strategy spec selected at build time. + const auto ann_index_spec = emb_list_strategy_->AnnIndexSpec(cfg); + const auto ann_metric_type = ann_index_spec.ann_metric_type.value_or(emb_list_raw_metric_type_); + { + ScopedMetricTypeOverride metric_guard(cfg, ann_metric_type); + switch (ann_index_spec.target) { + case EmbListAnnIndexTarget::BaseIndex: + // BaseIndex was serialized as this IndexNode's own ANN index. + RETURN_IF_ERROR(Deserialize(binset, config)); + break; + case EmbListAnnIndexTarget::SeparateIndex: { + auto separate_ann_index_or = + CreateEmbListSeparateAnnIndex(ann_index_spec.data_type, Type(), version_.VersionNumber()); + if (!separate_ann_index_or.has_value()) { + LOG_KNOWHERE_WARNING_ + << "Failed to create separate ANN index for emb_list strategy: " << strategy_type + << ", index type: " << Type() + << ", ann index data type: " << AnnIndexDataTypeName(ann_index_spec.data_type); + return separate_ann_index_or.error(); + } + auto separate_ann_index = + std::make_shared(std::move(separate_ann_index_or.value())); + RETURN_IF_ERROR(separate_ann_index->index.Node()->Deserialize(binset, config)); + emb_list_separate_ann_index_ = std::move(separate_ann_index); + break; + } + } + } + // 5. Deserialize raw vector index from BinarySet (if present) auto raw_index_bin = binset.GetByName(meta::EMB_LIST_RAW_INDEX); if (raw_index_bin) { - MemoryIOReader reader(raw_index_bin->data.get(), raw_index_bin->size); - auto* index = faiss::cppcontrib::knowhere::read_index(&reader); - auto* flat_index = dynamic_cast<::faiss::IndexFlat*>(index); - if (flat_index == nullptr) { - delete index; - LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX is not an IndexFlat"; - return Status::emb_list_inner_error; + auto raw_storage_or = ReadEmbListRawStorageFromBinary(raw_index_bin, emb_list_raw_metric_type_); + if (!raw_storage_or.has_value()) { + return raw_storage_or.error(); } - emb_list_raw_index_.reset(flat_index); - LOG_KNOWHERE_INFO_ << "Loaded raw vector index: " << emb_list_raw_index_->ntotal << " vectors"; + emb_list_raw_storage_ = std::move(raw_storage_or.value()); } else if (emb_list_strategy_->NeedsRawVectorStorage()) { LOG_KNOWHERE_WARNING_ << "Strategy requires raw vector storage but EMB_LIST_RAW_INDEX not found"; return Status::emb_list_inner_error; @@ -629,12 +793,9 @@ IndexNode::DeserializeEmbListFromFile(const std::string& filename, std::shared_p return metric_info_or.error(); } el_metric_type_ = metric_info_or.value().el_metric_type; + emb_list_raw_metric_type_ = metric_info_or.value().sub_metric_type; - // 2. Deserialize base index from file (override metric_type for base index) - cfg.metric_type = metric_info_or.value().sub_metric_type; - RETURN_IF_ERROR(DeserializeFromFile(filename, config)); - - // 3. Read meta file and parse strategy type + strategy blob + // 2. Read meta file and parse strategy type + strategy blob if (!cfg.emb_list_meta_file_path.has_value() || cfg.emb_list_meta_file_path.value().empty()) { LOG_KNOWHERE_WARNING_ << "emb_list_meta_file is empty, but metric type is emb_list"; return Status::emb_list_inner_error; @@ -668,7 +829,7 @@ IndexNode::DeserializeEmbListFromFile(const std::string& filename, std::shared_p auto [strategy_type, strategy_blob, strategy_blob_size] = ParseEmbListMetaHeader(file_data.get(), file_size); try { - // 4. Create strategy and deserialize strategy-specific data + // 3. Create strategy and deserialize strategy-specific data auto strategy_or = CreateEmbListStrategy(strategy_type, cfg); if (!strategy_or.has_value()) { LOG_KNOWHERE_WARNING_ << "Failed to create emb_list strategy: " << strategy_type; @@ -679,6 +840,35 @@ IndexNode::DeserializeEmbListFromFile(const std::string& filename, std::shared_p LOG_KNOWHERE_INFO_ << "Deserialize emb_list from file with strategy: " << strategy_type; RETURN_IF_ERROR(emb_list_strategy_->Deserialize(strategy_blob, strategy_blob_size, cfg)); + // 4. Deserialize ANN index from file using the same strategy spec selected at build time. + const auto ann_index_spec = emb_list_strategy_->AnnIndexSpec(cfg); + const auto ann_metric_type = ann_index_spec.ann_metric_type.value_or(emb_list_raw_metric_type_); + { + ScopedMetricTypeOverride metric_guard(cfg, ann_metric_type); + switch (ann_index_spec.target) { + case EmbListAnnIndexTarget::BaseIndex: + // BaseIndex was serialized as this IndexNode's own ANN index. + RETURN_IF_ERROR(DeserializeFromFile(filename, config)); + break; + case EmbListAnnIndexTarget::SeparateIndex: { + auto separate_ann_index_or = + CreateEmbListSeparateAnnIndex(ann_index_spec.data_type, Type(), version_.VersionNumber()); + if (!separate_ann_index_or.has_value()) { + LOG_KNOWHERE_WARNING_ + << "Failed to create separate ANN index for emb_list strategy: " << strategy_type + << ", index type: " << Type() + << ", ann index data type: " << AnnIndexDataTypeName(ann_index_spec.data_type); + return separate_ann_index_or.error(); + } + auto separate_ann_index = + std::make_shared(std::move(separate_ann_index_or.value())); + RETURN_IF_ERROR(separate_ann_index->index.Node()->DeserializeFromFile(filename, config)); + emb_list_separate_ann_index_ = std::move(separate_ann_index); + break; + } + } + } + // 5. Load raw vector index from separate file (if strategy needs it) if (emb_list_strategy_->NeedsRawVectorStorage()) { if (!cfg.emb_list_raw_index_file_path.has_value() || cfg.emb_list_raw_index_file_path.value().empty()) { @@ -692,17 +882,13 @@ IndexNode::DeserializeEmbListFromFile(const std::string& filename, std::shared_p io_flags |= faiss::cppcontrib::knowhere::IO_FLAG_MMAP_IFC; } - auto raw_index_file = cfg.emb_list_raw_index_file_path.value(); - auto* index = faiss::cppcontrib::knowhere::read_index(raw_index_file.data(), io_flags); - auto* flat_index = dynamic_cast<::faiss::IndexFlat*>(index); - if (flat_index == nullptr) { - delete index; - LOG_KNOWHERE_WARNING_ << "EMB_LIST_RAW_INDEX file is not an IndexFlat"; - return Status::emb_list_inner_error; + auto raw_storage_or = ReadEmbListRawStorageFromFile(cfg.emb_list_raw_index_file_path.value(), io_flags, + emb_list_raw_metric_type_); + if (!raw_storage_or.has_value()) { + return raw_storage_or.error(); } - emb_list_raw_index_.reset(flat_index); - LOG_KNOWHERE_INFO_ << "Loaded raw vector index from file: " << emb_list_raw_index_->ntotal - << " vectors, mmap=" << cfg.enable_mmap.value(); + emb_list_raw_storage_ = std::move(raw_storage_or.value()); + LOG_KNOWHERE_INFO_ << "Loaded raw vector storage from file, mmap=" << cfg.enable_mmap.value(); } // 6. Set ID mapping if needed @@ -721,44 +907,11 @@ IndexNode::DeserializeEmbListFromFile(const std::string& filename, std::shared_p expected IndexNode::CalcDistByRawIndex(const DataSetPtr dataset, const int64_t* labels, size_t labels_len, bool is_cosine, std::shared_ptr pool, milvus::OpContext* op_context) const { - if (!emb_list_raw_index_) { - return expected::Err(Status::emb_list_inner_error, "emb_list_raw_index not initialized"); - } - - auto num_queries = dataset->GetRows(); - auto dim = dataset->GetDim(); - auto query_data = dataset->GetTensor(); - auto distances = std::make_unique(num_queries * labels_len); - - try { - std::vector> futs; - futs.reserve(num_queries); - for (int64_t i = 0; i < num_queries; ++i) { - futs.emplace_back(pool->push([&, idx = i]() { - knowhere::checkCancellation(op_context); - std::unique_ptr dist_computer(emb_list_raw_index_->get_distance_computer()); - - const float* cur_query = static_cast(query_data) + idx * dim; - std::unique_ptr copied_query = nullptr; - if (is_cosine) { - copied_query = CopyAndNormalizeVecs(cur_query, 1, dim); - cur_query = copied_query.get(); - } - - dist_computer->set_query(cur_query); - auto cur_distances = distances.get() + idx * labels_len; - for (size_t j = 0; j < labels_len; ++j) { - cur_distances[j] = (*dist_computer)(labels[j]); - } - })); - } - WaitAllSuccess(futs); - } catch (const std::exception& e) { - LOG_KNOWHERE_WARNING_ << "CalcDistByRawIndex error: " << e.what(); - return expected::Err(Status::faiss_inner_error, e.what()); + if (!emb_list_raw_storage_) { + return expected::Err(Status::emb_list_inner_error, "emb_list raw storage not initialized"); } - - return GenResultDataSet(num_queries, labels_len, std::unique_ptr{}, std::move(distances)); + return emb_list_raw_storage_->CalcDistance(dataset, labels, labels_len, emb_list_raw_metric_type_, is_cosine, + std::move(pool), op_context); } } // namespace knowhere diff --git a/src/index/index_static.cc b/src/index/index_static.cc index 8dbfc2846..2f2d2bf87 100644 --- a/src/index/index_static.cc +++ b/src/index/index_static.cc @@ -13,6 +13,7 @@ #include +#include "knowhere/index/emb_list_strategy.h" #include "knowhere/operands.h" namespace knowhere { @@ -68,10 +69,18 @@ IndexStaticFaced::ConfigCheck(const IndexType& indexType, const IndexV if constexpr (!std::is_same_v) { auto strategy = cfg->emb_list_strategy.value_or(""); - if (strategy == meta::EMB_LIST_STRATEGY_MUVERA || strategy == meta::EMB_LIST_STRATEGY_LEMUR) { - msg = "MUVERA/LEMUR strategies only support fp32 data type, got '" + strategy + "'"; + if (strategy == meta::EMB_LIST_STRATEGY_MUVERA) { + msg = "MUVERA strategy only supports fp32 data type"; return Status::invalid_args; } + if (strategy == meta::EMB_LIST_STRATEGY_LEMUR) { + if constexpr (!std::is_same_v) { + msg = + "LEMUR strategy for this data type is not implemented; supported data types today are fp32 " + "and bin1"; + return Status::not_implemented; + } + } } if (Instance().staticConfigCheckMap.find(indexType) != Instance().staticConfigCheckMap.end()) { diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc index 6c76d1351..fdad212c2 100644 --- a/src/index/ivf/ivf.cc +++ b/src/index/ivf/ivf.cc @@ -1222,9 +1222,9 @@ expected IvfIndexNode::CalcDistByIDs(const DataSetPtr dataset, const BitsetView& bitset, const int64_t* labels, const size_t labels_len, const bool is_cosine, milvus::OpContext* op_context) const { - // When emb_list_raw_index_ exists (MUVERA/LEMUR), base index holds encoded vectors, - // so use the raw index for exact distance computation during reranking. - if (emb_list_raw_index_) { + // When emb_list raw storage exists (MUVERA/LEMUR), base index holds encoded vectors, + // so use raw storage for exact distance computation during reranking. + if (emb_list_raw_storage_) { return CalcDistByRawIndex(dataset, labels, labels_len, is_cosine, search_pool_, op_context); } diff --git a/tests/ut/test_emb_list.cc b/tests/ut/test_emb_list.cc index bb7040d6a..2d5f32e35 100644 --- a/tests/ut/test_emb_list.cc +++ b/tests/ut/test_emb_list.cc @@ -1782,6 +1782,118 @@ TEST_CASE("Search for EMBList Indices (Binary)", "Benchmark and validation on bi } } + SECTION("HNSW FLAT LEMUR") { + const std::string& index_type = knowhere::IndexEnum::INDEX_HNSW; + + for (size_t distance_type = 0; distance_type < DISTANCE_TYPES.size(); distance_type++) { + const int32_t dim = 32; + const int32_t nb = 256; + const uint64_t query_rng_seed = get_params_hash({100, static_cast(distance_type), dim}); + auto query_ds_ptr = GenQueryEmbListBinDataSet(NQ, dim, query_rng_seed); + + std::vector params = {100, static_cast(distance_type), dim, nb}; + const uint64_t rng_seed = get_params_hash(params); + auto default_ds_ptr = GenEmbListBinDataSetWithSomeEmpty(nb, dim, rng_seed, each_el_len); + const auto* lims = default_ds_ptr->Get(knowhere::meta::EMB_LIST_OFFSET); + knowhere::EmbListOffset doc_offset(lims, default_ds_ptr->GetRows()); + + knowhere::Json conf = default_conf; + conf[knowhere::meta::INDEX_TYPE] = index_type; + conf[knowhere::meta::METRIC_TYPE] = DISTANCE_TYPES[distance_type]; + conf[knowhere::meta::DIM] = dim; + conf[knowhere::meta::ROWS] = nb; + conf[knowhere::indexparam::RETRIEVAL_ANN_RATIO] = 10.0f; + conf["emb_list_strategy"] = "lemur"; + conf["lemur_hidden_dim"] = 8; + conf["lemur_num_train_samples"] = 1000; + conf["lemur_num_epochs"] = 1; + conf["lemur_batch_size"] = 16; + conf["lemur_learning_rate"] = 0.001f; + conf["lemur_seed"] = 42; + conf["lemur_num_layers"] = 1; + + auto golden_result = + knowhere::BruteForce::Search(default_ds_ptr, query_ds_ptr, conf, nullptr); + REQUIRE(golden_result.has_value()); + + auto version = GenTestEmbListVersionList(); + auto index = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + auto build_status = index.Build(default_ds_ptr, conf); + REQUIRE(build_status == knowhere::Status::success); + REQUIRE(index.Count() == static_cast(doc_offset.num_el())); + REQUIRE(index.Dim() == 8); + REQUIRE(index.Size() > 0); + + auto result = index.Search(query_ds_ptr, conf, nullptr); + REQUIRE(result.has_value()); + auto recall = GetKNNRecall(*golden_result.value(), *result.value()); + REQUIRE(recall >= 0.75f); + + knowhere::BinarySet binset; + REQUIRE(index.Serialize(binset) == knowhere::Status::success); + + auto index_loaded = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(index_loaded.Deserialize(binset, conf) == knowhere::Status::success); + REQUIRE(index_loaded.Count() == static_cast(doc_offset.num_el())); + REQUIRE(index_loaded.Dim() == 8); + REQUIRE(index_loaded.Size() > 0); + + auto result_loaded = index_loaded.Search(query_ds_ptr, conf, nullptr); + REQUIRE(result_loaded.has_value()); + auto recall_loaded = GetKNNRecall(*golden_result.value(), *result_loaded.value()); + REQUIRE(recall_loaded >= 0.75f); + + std::string file_suffix = std::to_string(distance_type); + std::string base_index_file = "/tmp/test_emb_list_bin_lemur_file_" + file_suffix + ".index"; + std::string meta_file = "/tmp/test_emb_list_bin_lemur_file_" + file_suffix + "_meta.bin"; + std::string raw_index_file = "/tmp/test_emb_list_bin_lemur_file_" + file_suffix + "_raw.index"; + { + auto hnsw_bin = binset.GetByName(knowhere::IndexEnum::INDEX_HNSW); + REQUIRE(hnsw_bin != nullptr); + std::ofstream base_out(base_index_file, std::ios::binary); + base_out.write(reinterpret_cast(hnsw_bin->data.get()), hnsw_bin->size); + + auto meta_bin = binset.GetByName(knowhere::meta::EMB_LIST_META); + REQUIRE(meta_bin != nullptr); + std::ofstream meta_out(meta_file, std::ios::binary); + meta_out.write(reinterpret_cast(meta_bin->data.get()), meta_bin->size); + + auto raw_bin = binset.GetByName(knowhere::meta::EMB_LIST_RAW_INDEX); + REQUIRE(raw_bin != nullptr); + std::ofstream raw_out(raw_index_file, std::ios::binary); + raw_out.write(reinterpret_cast(raw_bin->data.get()), raw_bin->size); + } + + for (bool enable_mmap : {false, true}) { + knowhere::Json load_conf = conf; + load_conf["emb_list_meta_file_path"] = meta_file; + load_conf["emb_list_raw_index_file_path"] = raw_index_file; + load_conf["enable_mmap"] = enable_mmap; + + auto index_file_loaded = + knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(index_file_loaded.DeserializeFromFile(base_index_file, load_conf) == knowhere::Status::success); + REQUIRE(index_file_loaded.Count() == static_cast(doc_offset.num_el())); + REQUIRE(index_file_loaded.Dim() == 8); + REQUIRE(index_file_loaded.Size() > 0); + + auto result_file_loaded = index_file_loaded.Search(query_ds_ptr, conf, nullptr); + REQUIRE(result_file_loaded.has_value()); + + const auto* ids = result.value()->GetIds(); + const auto* file_loaded_ids = result_file_loaded.value()->GetIds(); + auto num_q = result.value()->GetRows(); + for (int64_t i = 0; i < num_q * TOPK; ++i) { + REQUIRE(ids[i] == file_loaded_ids[i]); + } + } + + std::remove(base_index_file.c_str()); + std::remove(meta_file.c_str()); + std::remove(raw_index_file.c_str()); + } + } + #ifdef KNOWHERE_WITH_CARDINAL SECTION("CARDINAL_TIERED") { static const int64_t mb = 1024 * 1024; @@ -1976,6 +2088,60 @@ TEST_CASE("Test with some empty emb list", "[empty_emb_list]") { } } } + + SECTION("HNSW FLAT LEMUR") { + const std::string& index_type = knowhere::IndexEnum::INDEX_HNSW; + + for (size_t distance_type = 0; distance_type < DISTANCE_TYPES.size(); distance_type++) { + for (const int32_t dim : DIMS) { + const uint64_t query_rng_seed = get_params_hash({(int)distance_type, dim, 2}); + auto query_ds_ptr = GenQueryEmbListDataSet(NQ, dim, query_rng_seed); + + for (const int32_t nb : NBS) { + knowhere::Json conf = default_conf; + conf[knowhere::meta::METRIC_TYPE] = DISTANCE_TYPES[distance_type]; + conf[knowhere::meta::DIM] = dim; + conf[knowhere::meta::ROWS] = nb; + conf[knowhere::meta::INDEX_TYPE] = index_type; + conf["emb_list_strategy"] = "lemur"; + conf["lemur_hidden_dim"] = 8; + conf["lemur_num_train_samples"] = 128; + conf["lemur_num_epochs"] = 1; + conf["lemur_batch_size"] = 16; + conf["lemur_learning_rate"] = 0.001f; + conf["lemur_seed"] = 42; + conf["lemur_num_layers"] = 1; + + std::vector params = {(int)distance_type, dim, nb, 2}; + const uint64_t rng_seed = get_params_hash(params); + int num_el = int(nb / each_el_len) + 1; + auto default_ds_ptr = GenEmbListDataSetWithSomeEmpty(nb, dim, rng_seed, each_el_len); + + for (const float bitset_rate : BITSET_RATES) { + printf("bitset_rate: %f\n", bitset_rate); + const std::vector bitset_data = + GenerateBitsetByPartition(num_el, 1.0f - bitset_rate, 1); + knowhere::BitsetView bitset_view = nullptr; + if (bitset_rate != 0.0f) { + bitset_view = knowhere::BitsetView(bitset_data.data(), num_el); + } + + auto golden_result = knowhere::BruteForce::Search(default_ds_ptr, query_ds_ptr, + conf, bitset_view); + REQUIRE(golden_result.has_value()); + + printf( + "\nProcessing EMBList HNSW,Flat LEMUR fp32 (empty) for %s distance, dim=%d, nrows=%d, " + "%d%% points filtered out\n", + DISTANCE_TYPES[distance_type].c_str(), dim, nb, int(bitset_rate * 100)); + auto index_file = test_emb_list_index( + default_ds_ptr, query_ds_ptr, golden_result.value(), params, conf, false, bitset_view); + std::remove(index_file.c_str()); + } + } + } + } + } } template @@ -3423,6 +3589,18 @@ TEST_CASE("EmbList Serialization", "Strategy and IndexNode serialization/deseria auto result = strategy_or.value()->PrepareDataForBuild(ds, doc_offset, cfg); REQUIRE(result.has_value()); } + + // Case 4: all docs are empty; no finite training target can be derived + { + auto ds = knowhere::GenDataSet(0, DIM, nullptr); + knowhere::EmbListOffset doc_offset(std::vector{0, 0, 0}); + + auto strategy_or = knowhere::CreateEmbListStrategy("lemur", cfg); + REQUIRE(strategy_or.has_value()); + auto result = strategy_or.value()->PrepareDataForBuild(ds, doc_offset, cfg); + REQUIRE(!result.has_value()); + REQUIRE(result.error() == knowhere::Status::invalid_args); + } } SECTION("File-based: LEMUR DeserializeFromFile with mmap") { diff --git a/tests/ut/test_get_emb_list.cc b/tests/ut/test_get_emb_list.cc index 3d4353c38..067d88143 100644 --- a/tests/ut/test_get_emb_list.cc +++ b/tests/ut/test_get_emb_list.cc @@ -9,7 +9,11 @@ // 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 "catch2/catch_test_macros.hpp" #include "catch2/generators/catch_generators.hpp" @@ -590,7 +594,7 @@ TEST_CASE("Test GetEmbListByIds with MUVERA + HNSW_FLAT", "[GetEmbListByIds]") { REQUIRE(result_offsets[0] == 0); REQUIRE(result_offsets[1] == each_el_len); - // MUVERA stores raw vectors in emb_list_raw_index_, verify data matches original + // MUVERA stores raw vectors in emb_list raw storage, verify data matches original auto result_data = (const float*)result_ds->GetTensor(); size_t orig_vec_start = original_offsets[el_id]; for (int64_t v = 0; v < each_el_len; v++) { @@ -689,7 +693,7 @@ TEST_CASE("Test GetEmbListByIds with LEMUR + HNSW_FLAT", "[GetEmbListByIds]") { REQUIRE(result_offsets[0] == 0); REQUIRE(result_offsets[1] == each_el_len); - // LEMUR stores raw vectors in emb_list_raw_index_, verify data matches original + // LEMUR stores raw vectors in emb_list raw storage, verify data matches original auto result_data = (const float*)result_ds->GetTensor(); size_t orig_vec_start = original_offsets[el_id]; for (int64_t v = 0; v < each_el_len; v++) { @@ -729,7 +733,110 @@ TEST_CASE("Test GetEmbListByIds with LEMUR + HNSW_FLAT", "[GetEmbListByIds]") { } } -TEST_CASE("Test MUVERA/LEMUR reject non-fp32 data types", "[GetEmbListByIds]") { +TEST_CASE("Test GetEmbListByIds with binary LEMUR after DeserializeFromFile", "[GetEmbListByIds]") { + const int64_t dim = 32; + const int64_t each_el_len = 10; + const int64_t nb = 256; + const int64_t num_el = (nb + each_el_len - 1) / each_el_len; + const size_t code_size = static_cast(dim / 8); + + auto version = GenTestEmbListVersionList(); + + knowhere::Json conf; + conf[knowhere::meta::DIM] = dim; + conf[knowhere::meta::METRIC_TYPE] = knowhere::metric::MAX_SIM_HAMMING; + conf[knowhere::meta::TOPK] = 10; + conf[knowhere::meta::ROWS] = nb; + conf[knowhere::meta::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW; + conf[knowhere::indexparam::HNSW_M] = 16; + conf[knowhere::indexparam::EFCONSTRUCTION] = 96; + conf[knowhere::indexparam::EF] = 64; + conf[knowhere::indexparam::RETRIEVAL_ANN_RATIO] = 3.0f; + conf["emb_list_strategy"] = "lemur"; + conf["lemur_hidden_dim"] = 8; + conf["lemur_num_train_samples"] = 128; + conf["lemur_num_epochs"] = 1; + conf["lemur_batch_size"] = 16; + conf["lemur_learning_rate"] = 0.001f; + conf["lemur_seed"] = 42; + conf["lemur_num_layers"] = 1; + + auto train_ds = GenEmbListBinDataSet(nb, dim, 42, each_el_len); + const auto* original_data = static_cast(train_ds->GetTensor()); + auto original_offsets = train_ds->Get(knowhere::meta::EMB_LIST_OFFSET); + + auto idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = idx.Build(train_ds, conf); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + REQUIRE(idx.Serialize(bs) == knowhere::Status::success); + + std::string base_index_file = "/tmp/test_get_emb_list_bin_lemur_file.index"; + std::string meta_file = "/tmp/test_get_emb_list_bin_lemur_file_meta.bin"; + std::string raw_index_file = "/tmp/test_get_emb_list_bin_lemur_file_raw.index"; + { + auto hnsw_bin = bs.GetByName(knowhere::IndexEnum::INDEX_HNSW); + REQUIRE(hnsw_bin != nullptr); + std::ofstream base_out(base_index_file, std::ios::binary); + base_out.write(reinterpret_cast(hnsw_bin->data.get()), hnsw_bin->size); + + auto meta_bin = bs.GetByName(knowhere::meta::EMB_LIST_META); + REQUIRE(meta_bin != nullptr); + std::ofstream meta_out(meta_file, std::ios::binary); + meta_out.write(reinterpret_cast(meta_bin->data.get()), meta_bin->size); + + auto raw_bin = bs.GetByName(knowhere::meta::EMB_LIST_RAW_INDEX); + REQUIRE(raw_bin != nullptr); + std::ofstream raw_out(raw_index_file, std::ios::binary); + raw_out.write(reinterpret_cast(raw_bin->data.get()), raw_bin->size); + } + + std::vector el_ids = {0, 5, num_el - 1}; + auto ids_ds = knowhere::GenIdsDataSet(el_ids.size(), el_ids.data()); + + for (bool enable_mmap : {false, true}) { + knowhere::Json load_conf = conf; + load_conf["emb_list_meta_file_path"] = meta_file; + load_conf["emb_list_raw_index_file_path"] = raw_index_file; + load_conf["enable_mmap"] = enable_mmap; + + auto idx_loaded = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(idx_loaded.DeserializeFromFile(base_index_file, load_conf) == knowhere::Status::success); + + auto result = idx_loaded.GetEmbListByIds(ids_ds, knowhere::metric::MAX_SIM_HAMMING); + REQUIRE_HAS_VALUE(result); + + auto result_ds = result.value(); + REQUIRE(result_ds->GetRows() == static_cast(el_ids.size())); + REQUIRE(result_ds->GetDim() == dim); + + auto result_offsets = result_ds->Get(knowhere::meta::EMB_LIST_OFFSET); + REQUIRE(result_offsets != nullptr); + REQUIRE(result_offsets[0] == 0); + + const auto* result_data = static_cast(result_ds->GetTensor()); + for (size_t i = 0; i < el_ids.size(); i++) { + const auto el_id = el_ids[i]; + const size_t orig_vec_start = original_offsets[el_id]; + const size_t el_len = original_offsets[el_id + 1] - original_offsets[el_id]; + REQUIRE(result_offsets[i + 1] - result_offsets[i] == el_len); + + for (size_t byte = 0; byte < el_len * code_size; byte++) { + REQUIRE(result_data[result_offsets[i] * code_size + byte] == + original_data[orig_vec_start * code_size + byte]); + } + } + } + + std::remove(base_index_file.c_str()); + std::remove(meta_file.c_str()); + std::remove(raw_index_file.c_str()); +} + +TEST_CASE("Test MUVERA fp32 restriction and LEMUR binary ConfigCheck", "[GetEmbListByIds]") { auto version = knowhere::Version::GetCurrentVersion().VersionNumber(); std::string msg; @@ -765,6 +872,20 @@ TEST_CASE("Test MUVERA/LEMUR reject non-fp32 data types", "[GetEmbListByIds]") { REQUIRE(msg.find("fp32") != std::string::npos); } + SECTION("MUVERA with bin1 hamming should fail ConfigCheck") { + conf[knowhere::meta::DIM] = 32; + conf[knowhere::meta::METRIC_TYPE] = knowhere::metric::MAX_SIM_HAMMING; + conf["emb_list_strategy"] = "muvera"; + conf["muvera_num_projections"] = 3; + conf["muvera_num_repeats"] = 5; + conf["muvera_seed"] = 42; + + auto status = knowhere::IndexStaticFaced::ConfigCheck(knowhere::IndexEnum::INDEX_HNSW, version, + conf, msg); + REQUIRE(status != knowhere::Status::success); + REQUIRE(msg.find("fp32") != std::string::npos); + } + SECTION("LEMUR with fp16 should fail ConfigCheck") { conf["emb_list_strategy"] = "lemur"; conf["lemur_hidden_dim"] = 32; @@ -775,6 +896,17 @@ TEST_CASE("Test MUVERA/LEMUR reject non-fp32 data types", "[GetEmbListByIds]") { REQUIRE(msg.find("fp32") != std::string::npos); } + SECTION("LEMUR with bin1 hamming should pass ConfigCheck") { + conf[knowhere::meta::DIM] = 32; + conf[knowhere::meta::METRIC_TYPE] = knowhere::metric::MAX_SIM_HAMMING; + conf["emb_list_strategy"] = "lemur"; + conf["lemur_hidden_dim"] = 32; + + auto status = knowhere::IndexStaticFaced::ConfigCheck(knowhere::IndexEnum::INDEX_HNSW, version, + conf, msg); + REQUIRE(status == knowhere::Status::success); + } + SECTION("MUVERA with fp32 should pass ConfigCheck") { conf["emb_list_strategy"] = "muvera"; conf["muvera_num_projections"] = 3; diff --git a/tests/ut/utils.h b/tests/ut/utils.h index 030db4803..3b73b3465 100644 --- a/tests/ut/utils.h +++ b/tests/ut/utils.h @@ -677,6 +677,32 @@ GenEmbListBinDataSet(int rows, int dim, int seed = 42, int each_el_len = 10) { return ds; } +inline knowhere::DataSetPtr +GenEmbListBinDataSetWithSomeEmpty(int rows, int dim, int seed = 42, int each_el_len = 10) { + std::mt19937 rng(seed); + std::uniform_int_distribution<> distrib(0.0, 100.0); + int uint8_dim = dim / 8; + uint64_t total_size = rows * uint8_dim; + uint8_t* ts = new uint8_t[total_size]; + for (uint64_t i = 0; i < total_size; ++i) ts[i] = (uint8_t)distrib(rng); + auto ds = knowhere::GenDataSet(rows, dim, ts); + auto ptr = std::make_unique(size_t(rows / each_el_len) + 2); + size_t i = 0; + for (; i * each_el_len < (size_t)rows; i += 1) { + if (i % 2 == 1) { + // make some empty emb lists, like [0, 0, 20, 20, 40, 40, ...] + ptr[i] = ptr[i - 1]; + } else { + ptr[i] = i * each_el_len; + } + } + ptr[i] = (size_t)rows; + const size_t* ptr_const = ptr.release(); + ds->Set(knowhere::meta::EMB_LIST_OFFSET, ptr_const); + ds->SetIsOwner(true); + return ds; +} + inline knowhere::DataSetPtr GenQueryEmbListDataSet(int rows, int dim, const uint64_t seed = 42) { std::mt19937 rng(seed); diff --git a/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal b/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal index d2f96de47..3fd487d17 100644 --- a/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal +++ b/thirdparty/faiss/faiss/gpu_metal/MetalDistance.metal @@ -1076,4 +1076,3 @@ inline float sq6_decode_component(device const uchar* code, uint i) { } return (float(bits) + 0.5f) / 63.0f; } - From b4d6b72132179fd1e840270e336c3385d1d6ed5e Mon Sep 17 00:00:00 2001 From: SpadeA Date: Mon, 20 Jul 2026 11:48:14 +0800 Subject: [PATCH 2/4] fix clang Signed-off-by: SpadeA --- src/index/emb_list/emb_list_raw_storage.cc | 12 ++++++------ src/index/emb_list/emb_list_raw_storage.h | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/index/emb_list/emb_list_raw_storage.cc b/src/index/emb_list/emb_list_raw_storage.cc index 9bc5596fc..e45b9eedf 100644 --- a/src/index/emb_list/emb_list_raw_storage.cc +++ b/src/index/emb_list/emb_list_raw_storage.cc @@ -50,17 +50,17 @@ class EmbListFloatRawStorage final : public EmbListRawStorage { explicit EmbListFloatRawStorage(std::shared_ptr<::faiss::IndexFlat> index) : index_(std::move(index)) { } - int64_t + [[nodiscard]] int64_t Dim() const override { return index_->d; } - int64_t + [[nodiscard]] int64_t Count() const override { return index_->ntotal; } - size_t + [[nodiscard]] size_t CodeSize() const override { return static_cast(index_->d) * sizeof(float); } @@ -139,17 +139,17 @@ class EmbListBinaryRawStorage final : public EmbListRawStorage { : index_(std::move(index)) { } - int64_t + [[nodiscard]] int64_t Dim() const override { return index_->d; } - int64_t + [[nodiscard]] int64_t Count() const override { return index_->ntotal; } - size_t + [[nodiscard]] size_t CodeSize() const override { return static_cast(index_->code_size); } diff --git a/src/index/emb_list/emb_list_raw_storage.h b/src/index/emb_list/emb_list_raw_storage.h index ab0a6eb8f..5307f0d7b 100644 --- a/src/index/emb_list/emb_list_raw_storage.h +++ b/src/index/emb_list/emb_list_raw_storage.h @@ -35,13 +35,13 @@ class EmbListRawStorage { public: virtual ~EmbListRawStorage() = default; - virtual int64_t + [[nodiscard]] virtual int64_t Dim() const = 0; - virtual int64_t + [[nodiscard]] virtual int64_t Count() const = 0; - virtual size_t + [[nodiscard]] virtual size_t CodeSize() const = 0; virtual Status From 3fc2955a1bba58ae7fb8a94bce87eeffef1f8588 Mon Sep 17 00:00:00 2001 From: SpadeA Date: Tue, 4 Aug 2026 17:41:09 +0800 Subject: [PATCH 3/4] fix test Signed-off-by: SpadeA --- tests/ut/test_emb_list.cc | 2 +- tests/ut/test_get_emb_list.cc | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/ut/test_emb_list.cc b/tests/ut/test_emb_list.cc index 2d5f32e35..ed999c959 100644 --- a/tests/ut/test_emb_list.cc +++ b/tests/ut/test_emb_list.cc @@ -2105,7 +2105,7 @@ TEST_CASE("Test with some empty emb list", "[empty_emb_list]") { conf[knowhere::meta::INDEX_TYPE] = index_type; conf["emb_list_strategy"] = "lemur"; conf["lemur_hidden_dim"] = 8; - conf["lemur_num_train_samples"] = 128; + conf["lemur_num_train_samples"] = 1000; conf["lemur_num_epochs"] = 1; conf["lemur_batch_size"] = 16; conf["lemur_learning_rate"] = 0.001f; diff --git a/tests/ut/test_get_emb_list.cc b/tests/ut/test_get_emb_list.cc index 067d88143..a516d52c9 100644 --- a/tests/ut/test_get_emb_list.cc +++ b/tests/ut/test_get_emb_list.cc @@ -754,7 +754,7 @@ TEST_CASE("Test GetEmbListByIds with binary LEMUR after DeserializeFromFile", "[ conf[knowhere::indexparam::RETRIEVAL_ANN_RATIO] = 3.0f; conf["emb_list_strategy"] = "lemur"; conf["lemur_hidden_dim"] = 8; - conf["lemur_num_train_samples"] = 128; + conf["lemur_num_train_samples"] = 1000; conf["lemur_num_epochs"] = 1; conf["lemur_batch_size"] = 16; conf["lemur_learning_rate"] = 0.001f; From 0ff1fbaa242f99e85ac4a05a066b5ee6073367b6 Mon Sep 17 00:00:00 2001 From: SpadeA Date: Wed, 5 Aug 2026 17:35:24 +0800 Subject: [PATCH 4/4] fix test Signed-off-by: SpadeA --- src/index/emb_list/simple_mlp.h | 35 ++++++++++++++------------------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/src/index/emb_list/simple_mlp.h b/src/index/emb_list/simple_mlp.h index a926b6ad8..54616d43b 100644 --- a/src/index/emb_list/simple_mlp.h +++ b/src/index/emb_list/simple_mlp.h @@ -20,38 +20,33 @@ #include #include +#ifdef _OPENMP +#include +#endif + #include "index/emb_list/cblas_decl.h" #include "knowhere/log.h" -// OpenBLAS thread control (only available when linked against OpenBLAS) -#ifdef __APPLE__ -// Apple Accelerate does not provide openblas thread control; make it a no-op. -#else -extern "C" { -void -openblas_set_num_threads(int num_threads); -int -openblas_get_num_threads(void); -} -#endif - namespace knowhere { -// RAII guard for BLAS thread control. On platforms without OpenBLAS (e.g. macOS -// Accelerate), this is a no-op. +// OpenBLAS built with USE_OPENMP follows the current OpenMP task's thread +// setting. Keep this guard task-local instead of changing OpenBLAS's +// process-global thread state for every operation. class ScopedBLASThreads { int old_; public: - explicit ScopedBLASThreads(int n) : old_(0) { -#ifndef __APPLE__ - old_ = openblas_get_num_threads(); - openblas_set_num_threads(n); + explicit ScopedBLASThreads(int n) : old_(1) { +#ifdef _OPENMP + old_ = omp_get_max_threads(); + omp_set_num_threads(n); +#else + (void)n; #endif } ~ScopedBLASThreads() { -#ifndef __APPLE__ - openblas_set_num_threads(old_); +#ifdef _OPENMP + omp_set_num_threads(old_); #endif } ScopedBLASThreads(const ScopedBLASThreads&) = delete;