diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index 5572e16bf..45708e6a2 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -7,6 +7,8 @@ #include +#include + #include #include #include @@ -577,6 +579,7 @@ void hnsw_search( } } size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0; + size_t n_rabitq_1bit = 0, n_rabitq_refine = 0; idx_t check_period = InterruptCallback::get_period_hint( hnsw.max_level * index->d * efSearch); @@ -602,7 +605,9 @@ void hnsw_search( omp_capture_exception(ex, [&] { interrupt = true; }); } -#pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided) +#pragma omp for reduction( \ + + : n1, n2, ndis, nhops, n_rabitq_1bit, n_rabitq_refine) \ + schedule(guided) for (idx_t i = i0; i < i1; i++) { if (interrupt.load(std::memory_order_relaxed)) { continue; @@ -610,6 +615,10 @@ void hnsw_search( try { res->begin(i); dis->set_query(x + i * index->d); + auto* rq = dynamic_cast(dis.get()); + if (rq) { + rq->stats.reset(); + } HNSWStats stats = hnsw.search(*dis, index, *res, *vt, params); @@ -617,6 +626,10 @@ void hnsw_search( n2 += stats.n2; ndis += stats.ndis; nhops += stats.nhops; + if (rq) { + n_rabitq_1bit += rq->stats.n_1bit; + n_rabitq_refine += rq->stats.n_refine; + } res->end(); vt->advance(); } catch (...) { @@ -629,6 +642,7 @@ void hnsw_search( } hnsw_stats.combine({n1, n2, ndis, nhops}); + rabitq_stats.add({n_rabitq_1bit, n_rabitq_refine}); } } // anonymous namespace @@ -799,6 +813,7 @@ void IndexHNSW::search_level_0( { std::unique_ptr qdis; HNSWStats search_stats; + RaBitQStats rq_search_stats; VisitedTable* vt = nullptr; std::unique_ptr res; try { @@ -818,6 +833,11 @@ void IndexHNSW::search_level_0( try { res->begin(i); qdis->set_query(x + i * d); + auto* rq = + dynamic_cast(qdis.get()); + if (rq) { + rq->stats.reset(); + } hnsw.search_level_0( *qdis.get(), @@ -829,6 +849,9 @@ void IndexHNSW::search_level_0( search_stats, *vt, params); + if (rq) { + rq_search_stats.add(rq->stats); + } res->end(); vt->advance(); } catch (...) { @@ -838,6 +861,7 @@ void IndexHNSW::search_level_0( #pragma omp critical { hnsw_stats.combine(search_stats); + rabitq_stats.add(rq_search_stats); } } omp_rethrow_if_exception(ex); @@ -1061,7 +1085,7 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama( // Enable Panorama search mode. // This is not ideal, but is still more simple than making a subclass of // HNSW and overriding the search logic. - hnsw.is_panorama = true; + hnsw.search_method = HNSW::SM_PANORAMA; } void IndexHNSWFlatPanorama::add(idx_t n, const float* x) { @@ -1128,6 +1152,38 @@ IndexHNSWSQ::IndexHNSWSQ( IndexHNSWSQ::IndexHNSWSQ() = default; +/************************************************************** + * IndexHNSWRaBitQ implementation + **************************************************************/ + +IndexHNSWRaBitQ::IndexHNSWRaBitQ() = default; + +namespace { + +IndexRaBitQ* make_hnsw_rabitq_storage( + int d, + uint8_t nb_bits, + MetricType metric) { + FAISS_THROW_IF_NOT_MSG( + metric == METRIC_L2, "IndexHNSWRaBitQ supports only the L2 metric"); + return new IndexRaBitQ(d, metric, nb_bits); +} + +} // namespace + +IndexHNSWRaBitQ::IndexHNSWRaBitQ( + int d, + int M, + uint8_t nb_bits, + MetricType metric) + : IndexHNSW(make_hnsw_rabitq_storage(d, nb_bits, metric), M) { + own_fields = true; + is_trained = storage->is_trained; + // 1-bit codes store plain SignBitFactors with no f_error, so there is no + // bound to prune with and the staged path does not apply. + hnsw.search_method = nb_bits >= 2 ? HNSW::SM_RABITQ : HNSW::SM_DEFAULT; +} + /************************************************************** * IndexHNSW2Level implementation **************************************************************/ diff --git a/faiss/IndexHNSW.h b/faiss/IndexHNSW.h index 405e14d68..42d2406e2 100644 --- a/faiss/IndexHNSW.h +++ b/faiss/IndexHNSW.h @@ -216,6 +216,36 @@ struct IndexHNSWSQ : IndexHNSW { MetricType metric = METRIC_L2); }; +/** HNSW index whose storage is RaBitQ-compressed. + * + * The graph is built directly from the compressed codes. Neighbor-diversity + * pruning uses a fast 1-bit symmetric L2 estimator, including when the storage + * contains multi-bit codes. + * + * With nb_bits >= 2 the codes carry a per-vector error factor, so search uses + * the staged search method: a 1-bit estimate for every neighbor and the full + * multi-bit distance only for candidates the error bound cannot rule out. + * nb_bits = 1 has no error factor and uses ordinary HNSW search. + */ +struct IndexHNSWRaBitQ : IndexHNSW { + IndexHNSWRaBitQ(); + IndexHNSWRaBitQ( + int d, + int M, + uint8_t nb_bits = 1, + MetricType metric = METRIC_L2); + + IndexHNSWRaBitQ& operator=(const IndexHNSWRaBitQ&) = delete; + + private: + // clone_index() replaces the shallow-copied storage with a deep copy + // before returning it. Keep ordinary C++ copies from sharing ownership. + IndexHNSWRaBitQ(const IndexHNSWRaBitQ&) = default; +#ifndef SWIG + friend IndexHNSW* clone_IndexHNSW(const IndexHNSW* index); +#endif +}; + /** 2-level code structure with fast random access */ struct IndexHNSW2Level : IndexHNSW { diff --git a/faiss/clone_index.cpp b/faiss/clone_index.cpp index 1fd08482f..76dbdd829 100644 --- a/faiss/clone_index.cpp +++ b/faiss/clone_index.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -139,6 +140,7 @@ IndexIDMap* clone_IndexIDMap(const IndexIDMap* im) { IndexHNSW* clone_IndexHNSW(const IndexHNSW* ihnsw) { TRYCLONE(IndexHNSW2Level, ihnsw) TRYCLONE(IndexHNSWFlatPanorama, ihnsw) + TRYCLONE(IndexHNSWRaBitQ, ihnsw) TRYCLONE(IndexHNSWFlat, ihnsw) TRYCLONE(IndexHNSWPQ, ihnsw) TRYCLONE(IndexHNSWSQ, ihnsw) @@ -294,6 +296,7 @@ Index* Cloner::clone_Index(const Index* index) { TRYCLONE(IndexEDEN, index) TRYCLONE(IndexScalarQuantizer, index) + TRYCLONE(IndexRaBitQ, index) TRYCLONE(MultiIndexQuantizer, index) if (const IndexIVF* ivf = dynamic_cast(index)) { diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp index 4b07c10b7..da6ea14bc 100644 --- a/faiss/impl/HNSW.cpp +++ b/faiss/impl/HNSW.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -1116,13 +1117,82 @@ inline void extract_search_params( } } -/** Templated body of `search_from_candidates` — instantiated once per - * VisitedTable subclass × comparator. +struct DefaultCandidateDistanceEvaluator { + DistanceComputer& qdis; + + template + size_t evaluate( + const storage_idx_t* ids, + int count, + GetThreshold&& /* get_threshold */, + AddResult&& add_result) { + if (count == 4) { + float distances[4]; + qdis.distances_batch_4( + ids[0], + ids[1], + ids[2], + ids[3], + distances[0], + distances[1], + distances[2], + distances[3]); + for (int i = 0; i < 4; ++i) { + add_result(ids[i], distances[i]); + } + } else { + for (int i = 0; i < count; ++i) { + add_result(ids[i], qdis(ids[i])); + } + } + return count; + } +}; + +/** RaBitQ's staged evaluator plugs into the same graph traversal as the + * default distance computer. Each candidate first receives a 1-bit estimate; + * the full code is evaluated only when its error bound can still beat the + * current result threshold. */ -template +struct RaBitQCandidateDistanceEvaluator { + RaBitQDistanceComputer& rq; + const bool is_similarity; + + template + size_t evaluate( + const storage_idx_t* ids, + int count, + GetThreshold&& get_threshold, + AddResult&& add_result) { + size_t ndis = 0; + for (int i = 0; i < count; ++i) { + const uint8_t* code = + rq.codes + static_cast(ids[i]) * rq.code_size; + const float estimate = rq.distance_to_code_1bit(code); + rq.stats.n_1bit++; + ndis++; + + float distance = estimate; + if (rq.should_refine( + code, estimate, get_threshold(), is_similarity)) { + distance = rq.distance_to_code_full(code); + rq.stats.n_refine++; + ndis++; + } + add_result(ids[i], distance); + } + return ndis; + } +}; + +/** Templated body shared by the default and staged RaBitQ searches. The + * evaluator changes only how a group of unvisited neighbors receives its + * distance; traversal, result handling and stopping conditions stay identical. + */ +template int search_from_candidates_fixVT( const HNSW& hnsw, - DistanceComputer& qdis, + DistanceEvaluator& evaluator, ResultHandler& res, MinimaxHeapT>& candidates, VTType& vt, @@ -1131,7 +1201,7 @@ int search_from_candidates_fixVT( int nres_in, const SearchParameters* params) { int nres = nres_in; - int ndis = 0; + size_t ndis = 0; bool do_dis_check; int efSearch; @@ -1189,7 +1259,7 @@ int search_from_candidates_fixVT( } int counter = 0; - size_t saved_j[4]; + storage_idx_t saved_j[4]; threshold = res.threshold; @@ -1212,32 +1282,18 @@ int search_from_candidates_fixVT( counter += vt.set(v1) ? 1 : 0; if (counter == 4) { - float dis[4]; - qdis.distances_batch_4( - saved_j[0], - saved_j[1], - saved_j[2], - saved_j[3], - dis[0], - dis[1], - dis[2], - dis[3]); - - for (size_t id4 = 0; id4 < 4; id4++) { - add_to_heap(saved_j[id4], dis[id4]); - } - - ndis += 4; - + ndis += evaluator.evaluate( + saved_j, + counter, + [&] { return threshold; }, + add_to_heap); counter = 0; } } - for (int icnt = 0; icnt < counter; icnt++) { - float dis = qdis(saved_j[icnt]); - add_to_heap(saved_j[icnt], dis); - - ndis += 1; + if (counter > 0) { + ndis += evaluator.evaluate( + saved_j, counter, [&] { return threshold; }, add_to_heap); } nstep++; @@ -1258,12 +1314,12 @@ int search_from_candidates_fixVT( return nres; } -/// Dispatches the VisitedTable concrete type for a given C, then calls -/// the templated `search_from_candidates_fixVT`. -template -int search_from_candidates_dispatch( +/// Dispatches the VisitedTable concrete type while keeping the distance +/// evaluator statically selected outside the graph traversal loop. +template +int search_from_candidates_evaluator_dispatch( const HNSW& hnsw, - DistanceComputer& qdis, + DistanceEvaluator& evaluator, ResultHandler& res, MinimaxHeapT>& candidates, VisitedTable& vt, @@ -1274,7 +1330,7 @@ int search_from_candidates_dispatch( auto call = [&](VTType& vt_concrete) -> int { return search_from_candidates_fixVT( hnsw, - qdis, + evaluator, res, candidates, vt_concrete, @@ -1290,6 +1346,59 @@ int search_from_candidates_dispatch( return call(vts); } +template +int search_from_candidates_dispatch( + const HNSW& hnsw, + DistanceComputer& qdis, + ResultHandler& res, + MinimaxHeapT>& candidates, + VisitedTable& vt, + HNSWStats& stats, + int level, + int nres_in, + const SearchParameters* params) { + DefaultCandidateDistanceEvaluator evaluator{qdis}; + return search_from_candidates_evaluator_dispatch( + hnsw, + evaluator, + res, + candidates, + vt, + stats, + level, + nres_in, + params); +} + +template +int search_from_candidates_rabitq_dispatch( + const HNSW& hnsw, + DistanceComputer& qdis, + ResultHandler& res, + MinimaxHeapT>& candidates, + VisitedTable& vt, + HNSWStats& stats, + int level, + int nres_in, + const SearchParameters* params) { + auto* rq = dynamic_cast(&qdis); + FAISS_THROW_IF_NOT_MSG( + rq, "staged RaBitQ search requires an IndexRaBitQ storage"); + FAISS_THROW_IF_NOT_MSG( + rq->nb_bits >= 2, "staged RaBitQ search requires nb_bits >= 2"); + RaBitQCandidateDistanceEvaluator evaluator{*rq, hnsw.is_similarity}; + return search_from_candidates_evaluator_dispatch( + hnsw, + evaluator, + res, + candidates, + vt, + stats, + level, + nres_in, + params); +} + } // namespace /** Do a BFS on the candidates list. Public dispatcher: only handles the @@ -1716,37 +1825,49 @@ HNSWStats search_impl( } int ef = std::max(cur_efSearch, k); - if (bounded_queue) { // this is the most common branch, for now we only - // support Panorama search in this branch + if (bounded_queue) { MinimaxHeapT> candidates(ef); candidates.push(nearest, d_nearest); - if (!hnsw.is_panorama) { - search_from_candidates_dispatch( - hnsw, qdis, res, candidates, vt, stats, 0, 0, params); - } else { - // Panorama is L2-specific and is only valid for C_distance. - // The public dispatch ensures we never reach this code path - // with C != C_distance, but assert in debug builds. - if constexpr (std::is_same_v) { - hnsw_detail::search_from_candidates_panorama( - hnsw, - index, - qdis, - res, - candidates, - vt, - stats, - 0, - 0, - params); - } else { - FAISS_THROW_MSG( - "Panorama search is not supported with is_similarity=true"); - } + switch (hnsw.search_method) { + case HNSW::SM_DEFAULT: + search_from_candidates_dispatch( + hnsw, qdis, res, candidates, vt, stats, 0, 0, params); + break; + case HNSW::SM_RABITQ: + search_from_candidates_rabitq_dispatch( + hnsw, qdis, res, candidates, vt, stats, 0, 0, params); + break; + case HNSW::SM_PANORAMA: + if constexpr (std::is_same_v) { + hnsw_detail::search_from_candidates_panorama( + hnsw, + index, + qdis, + res, + candidates, + vt, + stats, + 0, + 0, + params); + } else { + FAISS_THROW_MSG( + "Panorama search does not support similarity metrics"); + } + break; + default: + FAISS_THROW_MSG("invalid HNSW search method"); } } else { + FAISS_THROW_IF_NOT_MSG( + hnsw.search_method != HNSW::SM_RABITQ, + "staged RaBitQ search requires bounded_queue=true"); + FAISS_THROW_IF_NOT_MSG( + hnsw.search_method == HNSW::SM_DEFAULT || + hnsw.search_method == HNSW::SM_PANORAMA, + "invalid HNSW search method"); auto call = [&](VTType& vt_concrete) { return search_from_candidate_unbounded_fixVT( hnsw, @@ -1804,6 +1925,50 @@ void search_level_0_impl( int k = extract_k_from_ResultHandler(res); + auto search_candidates = [&](MinimaxHeapT>& candidates, + int nres) { + switch (hnsw.search_method) { + case HNSW::SM_DEFAULT: + return search_from_candidates_dispatch( + hnsw, + qdis, + res, + candidates, + vt, + search_stats, + 0, + nres, + params); + case HNSW::SM_RABITQ: + return search_from_candidates_rabitq_dispatch( + hnsw, + qdis, + res, + candidates, + vt, + search_stats, + 0, + nres, + params); + case HNSW::SM_PANORAMA: + // Preserve the pre-enum behavior: search_level_0 did not use + // Panorama's progressive traversal and fell back to the + // ordinary HNSW candidate search. + return search_from_candidates_dispatch( + hnsw, + qdis, + res, + candidates, + vt, + search_stats, + 0, + nres, + params); + default: + FAISS_THROW_MSG("invalid HNSW search method"); + } + }; + if (search_type == 1) { int nres = 0; @@ -1823,16 +1988,7 @@ void search_level_0_impl( candidates.push(cj, nearest_d[j]); - nres = search_from_candidates_dispatch( - hnsw, - qdis, - res, - candidates, - vt, - search_stats, - 0, - nres, - params); + nres = search_candidates(candidates, nres); nres = std::min(nres, candidates_size); } } else if (search_type == 2) { @@ -1849,8 +2005,7 @@ void search_level_0_impl( candidates.push(cj, nearest_d[j]); } - search_from_candidates_dispatch( - hnsw, qdis, res, candidates, vt, search_stats, 0, 0, params); + search_candidates(candidates, 0); } } diff --git a/faiss/impl/HNSW.h b/faiss/impl/HNSW.h index 958424d9f..dc442c802 100644 --- a/faiss/impl/HNSW.h +++ b/faiss/impl/HNSW.h @@ -154,8 +154,14 @@ struct HNSW { /// use bounded queue during exploration bool search_bounded_queue = true; - /// use Panorama progressive pruning in search - bool is_panorama = false; + /// Specialized level-0 search implementations. This state is derived from + /// the owning IndexHNSW subtype after construction or deserialization. + enum Search_method_t : uint8_t { + SM_DEFAULT, + SM_PANORAMA, + SM_RABITQ, + }; + Search_method_t search_method = SM_DEFAULT; /// distance comparison semantics: when true, distances are treated as /// similarity scores (larger is better). Default false matches the diff --git a/faiss/impl/RaBitQuantizer.cpp b/faiss/impl/RaBitQuantizer.cpp index 90c524f6f..7e91375c0 100644 --- a/faiss/impl/RaBitQuantizer.cpp +++ b/faiss/impl/RaBitQuantizer.cpp @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -24,6 +25,8 @@ namespace faiss { +RaBitQStats rabitq_stats; + // Import shared utilities from RaBitQUtils using rabitq_utils::ExtraBitsFactors; using rabitq_utils::QueryFactorsData; @@ -224,6 +227,51 @@ void RaBitQuantizer::decode_core( } } +template +float symmetric_dis_1bit(const RaBitQDistanceComputer& dc, idx_t i, idx_t j) { + FAISS_THROW_IF_NOT_MSG( + dc.metric_type == MetricType::METRIC_L2, + "RaBitQ symmetric distance supports only L2"); + FAISS_ASSERT(i >= 0 && j >= 0); + FAISS_ASSERT(dc.codes != nullptr); + + const size_t sign_bytes = (dc.d + 7) / 8; + const uint8_t* code_i = dc.codes + static_cast(i) * dc.code_size; + const uint8_t* code_j = dc.codes + static_cast(j) * dc.code_size; + const auto* factors_i = + reinterpret_cast(code_i + sign_bytes); + const auto* factors_j = + reinterpret_cast(code_j + sign_bytes); + + const uint64_t xor_popcount = + rabitq::bitwise_xor_dot_product(code_i, code_j, sign_bytes, 1); + const float sign_dot = + static_cast(dc.d) - 2.0f * static_cast(xor_popcount); + + // The L2-optimal reconstruction of residual r is alpha * sign(r), where + // alpha = ||r||_1 / d. The stored factors give + // alpha_i * alpha_j = ||r_i||^2 * ||r_j||^2 / + // (d * dp_multiplier_i * dp_multiplier_j). + float cross_term = 0.0f; + if (factors_i->dp_multiplier != 0.0f && factors_j->dp_multiplier != 0.0f) { + // Dividing each norm first avoids overflowing the product of two + // squared norms even when the final distance is representable. + const float scaled_norm_i = + factors_i->or_minus_c_l2sqr / factors_i->dp_multiplier; + const float scaled_norm_j = + factors_j->or_minus_c_l2sqr / factors_j->dp_multiplier; + cross_term = (scaled_norm_i * (sign_dot / static_cast(dc.d))) * + scaled_norm_j; + } + const float distance = factors_i->or_minus_c_l2sqr + + factors_j->or_minus_c_l2sqr - 2.0f * cross_term; + return std::max(0.0f, distance); +} + +float RaBitQDistanceComputer::symmetric_dis(idx_t i, idx_t j) { + return symmetric_dis_1bit(*this, i, j); +} + namespace { // Distance computers templatized on SIMDLevel to avoid per-call dynamic @@ -240,6 +288,10 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer { RaBitQDistanceComputerNotQ() = default; + float symmetric_dis(idx_t i, idx_t j) final { + return symmetric_dis_1bit(*this, i, j); + } + // Compute distance using only 1-bit codes (fast) float distance_to_code_1bit_impl( const uint8_t* binary_data, @@ -455,6 +507,10 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer { RaBitQDistanceComputerQ() = default; + float symmetric_dis(idx_t i, idx_t j) final { + return symmetric_dis_1bit(*this, i, j); + } + // Compute distance using only 1-bit codes (fast) float distance_to_code_1bit_impl( const uint8_t* binary_data, diff --git a/faiss/impl/RaBitQuantizer.h b/faiss/impl/RaBitQuantizer.h index 7766df35c..c9af5e094 100644 --- a/faiss/impl/RaBitQuantizer.h +++ b/faiss/impl/RaBitQuantizer.h @@ -13,6 +13,8 @@ #include #include #include +#include +#include namespace faiss { @@ -20,6 +22,31 @@ struct IDSelector; template struct ResultHandlerUnordered; +/** Statistics for staged RaBitQ distance evaluation. + * + * These counters belong to the quantizer rather than to a particular ANN + * graph. They are currently populated by HNSW's staged search path. + */ +struct RaBitQStats { + uint64_t n_1bit = 0; + uint64_t n_refine = 0; + + void reset() { + n_1bit = n_refine = 0; + } + + void add(const RaBitQStats& other) { + n_1bit += other.n_1bit; + n_refine += other.n_refine; + } + + double refine_ratio() const { + return n_1bit ? double(n_refine) / double(n_1bit) : 0.0; + } +}; + +FAISS_API extern RaBitQStats rabitq_stats; + // the reference implementation of the https://arxiv.org/pdf/2405.12497 // Jianyang Gao, Cheng Long, "RaBitQ: Quantizing High-Dimensional Vectors // with a Theoretical Error Bound for Approximate Nearest Neighbor Search". @@ -122,10 +149,10 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer { // Used with f_error to compute error bounds for two-stage filtering float g_error = 0.0f; - float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override { - // Not used for RaBitQ - FAISS_THROW_MSG("Not implemented"); - } + /// Per-distance-computer counters, aggregated by the owning index. + RaBitQStats stats; + + float symmetric_dis(idx_t i, idx_t j) override; // Compute 1-bit distance estimate (fast) virtual float distance_to_code_1bit(const uint8_t* code) = 0; @@ -133,6 +160,22 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer { // Compute full multi-bit distance (accurate) virtual float distance_to_code_full(const uint8_t* code) = 0; + /// Apply the RaBitQ error bound without exposing the packed code layout to + /// ANN consumers. Precondition: nb_bits >= 2. +#ifndef SWIG + FAISS_ALWAYS_INLINE bool should_refine( + const uint8_t* code, + float estimate, + float threshold, + bool is_similarity) const { + const auto* factors = + reinterpret_cast( + code + (d + 7) / 8); + return rabitq_utils::should_refine_candidate( + estimate, factors->f_error, g_error, threshold, is_similarity); + } +#endif + virtual void set_centroid(const float* centroid_in) { centroid = centroid_in; } diff --git a/faiss/impl/ResultHandler.h b/faiss/impl/ResultHandler.h index 5dfd7bc30..a30f18786 100644 --- a/faiss/impl/ResultHandler.h +++ b/faiss/impl/ResultHandler.h @@ -123,14 +123,20 @@ struct SingleQueryBlockResultHandler : BlockResultHandler { using ResultHandlerT::threshold; explicit SingleResultHandler(SingleQueryBlockResultHandler& hr) - : the_handler(hr.the_handler) {} + : the_handler(hr.the_handler) { + threshold = the_handler.threshold; + } /// begin results for query # i - void begin(const size_t /* qid */) {} + void begin(const size_t /* qid */) { + threshold = the_handler.threshold; + } /// add one result for query i bool add_result(T dis, TI idx) final { - return the_handler.add_result(dis, idx); + bool updated = the_handler.add_result(dis, idx); + threshold = the_handler.threshold; + return updated; } /// series of results for query i is done diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index 5847843fa..134ef9d53 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -2507,7 +2507,7 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { } else if ( h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") || h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHc2") || - h == fourcc("IHfP") || h == fourcc("IH00")) { + h == fourcc("IHfP") || h == fourcc("IHNr") || h == fourcc("IH00")) { std::unique_ptr idxhnsw; if (h == fourcc("IH00")) { idxhnsw = std::make_unique(); @@ -2525,6 +2525,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxhnsw = std::make_unique(); } else if (h == fourcc("IHc2")) { idxhnsw = std::make_unique(); + } else if (h == fourcc("IHNr")) { + idxhnsw = std::make_unique(); } read_index_header(*idxhnsw, f); if (h == fourcc("IHfP")) { @@ -2559,7 +2561,8 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { "HNSW levels size %zu != index ntotal %" PRId64, idxhnsw->hnsw.levels.size(), idxhnsw->ntotal); - idxhnsw->hnsw.is_panorama = (h == fourcc("IHfP")); + idxhnsw->hnsw.search_method = + h == fourcc("IHfP") ? HNSW::SM_PANORAMA : HNSW::SM_DEFAULT; // `HNSW::is_similarity` is intentionally not serialized, so we // re-derive it here from the persisted metric type. Without this, // a saved IP/similarity index would come back configured as a @@ -2581,6 +2584,57 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { idxhnsw->storage->d, idxhnsw->d); } + if (h == fourcc("IHNr")) { + auto* idx_rabitq = dynamic_cast(idxhnsw.get()); + FAISS_THROW_IF_NOT_MSG( + idx_rabitq, "IHNr must deserialize to an IndexHNSWRaBitQ"); + FAISS_THROW_IF_NOT_MSG( + idxhnsw->metric_type == METRIC_L2, + "IndexHNSWRaBitQ supports only the L2 metric"); + bool staged; + READ1_BOOL(staged); + idxhnsw->hnsw.search_method = + staged ? HNSW::SM_RABITQ : HNSW::SM_DEFAULT; + if (idxhnsw->storage) { + auto* rq = dynamic_cast(idxhnsw->storage); + FAISS_THROW_IF_NOT_MSG( + rq, "IndexHNSWRaBitQ storage must be an IndexRaBitQ"); + FAISS_THROW_IF_NOT_MSG( + rq->metric_type == idxhnsw->metric_type && + rq->rabitq.metric_type == idxhnsw->metric_type, + "IndexHNSWRaBitQ storage metric mismatch"); + FAISS_THROW_IF_NOT_MSG( + rq->is_trained == idxhnsw->is_trained, + "IndexHNSWRaBitQ storage training state mismatch"); + FAISS_THROW_IF_NOT_FMT( + rq->rabitq.nb_bits >= 1 && rq->rabitq.nb_bits <= 9, + "invalid RaBitQ nb_bits=%zu", + rq->rabitq.nb_bits); + const size_t expected_code_size = + rq->rabitq.compute_code_size(rq->d, rq->rabitq.nb_bits); + validate_code_size_match( + rq->rabitq.code_size, + expected_code_size, + "IndexHNSWRaBitQ quantizer"); + validate_code_size_match( + rq->code_size, + expected_code_size, + "IndexHNSWRaBitQ storage"); + FAISS_THROW_IF_NOT( + rq->codes.size() == + mul_no_overflow( + static_cast(rq->ntotal), + rq->code_size, + "IndexHNSWRaBitQ codes")); + FAISS_THROW_IF_NOT_MSG( + !rq->is_trained || + rq->center.size() == static_cast(rq->d), + "IndexHNSWRaBitQ center size mismatch"); + FAISS_THROW_IF_NOT_MSG( + staged == (rq->rabitq.nb_bits >= 2), + "IndexHNSWRaBitQ staged-search metadata mismatch"); + } + } if (h == fourcc("IHN2")) { FAISS_THROW_IF_NOT_MSG( idxhnsw->storage, @@ -3402,7 +3456,7 @@ std::unique_ptr read_index_binary_up(IOReader* f, int io_flags) { auto idxhnsw = std::make_unique(); read_index_binary_header(*idxhnsw, f); read_HNSW(idxhnsw->hnsw, f); - idxhnsw->hnsw.is_panorama = false; + idxhnsw->hnsw.search_method = HNSW::SM_DEFAULT; FAISS_THROW_IF_NOT_FMT( idxhnsw->hnsw.levels.size() == (size_t)idxhnsw->ntotal, "IndexBinaryHNSW HNSW levels size %zu != ntotal %" PRId64, @@ -3426,7 +3480,7 @@ std::unique_ptr read_index_binary_up(IOReader* f, int io_flags) { READ1_BOOL(idxhnsw->base_level_only); READ1(idxhnsw->num_base_level_search_entrypoints); read_HNSW(idxhnsw->hnsw, f); - idxhnsw->hnsw.is_panorama = false; + idxhnsw->hnsw.search_method = HNSW::SM_DEFAULT; FAISS_THROW_IF_NOT_FMT( idxhnsw->hnsw.levels.size() == (size_t)idxhnsw->ntotal, "IndexBinaryHNSWCagra HNSW levels size %zu != ntotal %" PRId64, diff --git a/faiss/impl/index_write.cpp b/faiss/impl/index_write.cpp index 873fc0f9c..a0b7abc6b 100644 --- a/faiss/impl/index_write.cpp +++ b/faiss/impl/index_write.cpp @@ -873,12 +873,22 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { : dynamic_cast(idx) ? fourcc("IHNs") : dynamic_cast(idx) ? fourcc("IHN2") : dynamic_cast(idx) ? fourcc("IHc2") + : dynamic_cast(idx) ? fourcc("IHNr") : typeid(*idx) == typeid(IndexHNSW) ? fourcc("IH00") : 0; FAISS_THROW_IF_NOT_FMT( h != 0, "don't know how to serialize this IndexHNSW subtype: %s", typeid(*idx).name()); + const IndexRaBitQ* storage_rabitq = nullptr; + if (h == fourcc("IHNr")) { + storage_rabitq = dynamic_cast(idxhnsw->storage); + FAISS_THROW_IF_NOT_MSG( + storage_rabitq || + ((io_flags & IO_FLAG_SKIP_STORAGE) && + idxhnsw->storage == nullptr), + "IndexHNSWRaBitQ requires IndexRaBitQ storage"); + } WRITE1(h); write_index_header(idxhnsw, f); if (h == fourcc("IHfP")) { @@ -901,6 +911,16 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) { } else { write_index(idxhnsw->storage, f); } + if (h == fourcc("IHNr")) { + // The staged flag is graph-traversal state, so it has to live here: + // with IO_FLAG_SKIP_STORAGE there is no storage to derive it from. + // Storage-owned settings are not duplicated in this payload; + // IndexRaBitQ serializes whatever it owns. + const bool staged = storage_rabitq + ? storage_rabitq->rabitq.nb_bits >= 2 + : idxhnsw->hnsw.search_method == HNSW::SM_RABITQ; + WRITE1(staged); + } } else if (const IndexNSG* idxnsg = dynamic_cast(idx)) { uint32_t h = dynamic_cast(idx) ? fourcc("INSf") : dynamic_cast(idx) ? fourcc("INSp") diff --git a/faiss/index_factory.cpp b/faiss/index_factory.cpp index 73d1517bb..cdbd0bea0 100644 --- a/faiss/index_factory.cpp +++ b/faiss/index_factory.cpp @@ -575,6 +575,13 @@ IndexHNSW* parse_IndexHNSW( if (match(sq_pattern)) { return new IndexHNSWSQ(d, sq_types[sm[1].str()], hnsw_M, mt); } + // Keep the bare RaBitQ token consistent with Flat and IVF: it means 1 bit. + // Use an explicit width such as RaBitQ4 to enable staged refinement. + if (match("RaBitQ([1-9])?")) { + // the capture is the bare digit, so no substr offset here + int nb_bits = mres_to_int(sm[1], 1); + return new IndexHNSWRaBitQ(d, hnsw_M, nb_bits, mt); + } if (match("([0-9]+)\\+PQ([0-9]+)?")) { int ncent = mres_to_int(sm[1]); int pq_m = mres_to_int(sm[2]); diff --git a/faiss/python/__init__.pyi b/faiss/python/__init__.pyi index 746d42c8d..f0c2c8658 100644 --- a/faiss/python/__init__.pyi +++ b/faiss/python/__init__.pyi @@ -1939,12 +1939,23 @@ class HNSWStats: n3: int ndis: int nreorder: int +class RaBitQStats: + n_1bit: int + n_refine: int + + def reset(self) -> None: ... + def add(self, other: RaBitQStats) -> None: ... + def refine_ratio(self) -> float: ... class HNSW: + SM_DEFAULT: int + SM_PANORAMA: int + SM_RABITQ: int max_level: int entry_point: int efConstruction: int efSearch: int + search_method: int hnsw_stats: HNSWStats assign_probas: Float32Vector cum_nneighbor_per_level: Int32Vector @@ -2008,6 +2019,15 @@ class IndexHNSWSQ(IndexHNSW): metric: MetricType = METRIC_L2, ) -> None: ... +class IndexHNSWRaBitQ(IndexHNSW): + def __init__( + self, + d: int, + M: int, + nb_bits: int = 1, + metric: MetricType = METRIC_L2, + ) -> None: ... + class IndexHNSW2Level(IndexHNSW): q1: Any # Level1Quantizer pq: ProductQuantizer # from Index2Layer diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index b4d31de0a..b1eb71602 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -914,6 +914,7 @@ typedef int cudaDataType_t; DOWNCAST ( IndexHNSWPQ ) DOWNCAST ( IndexHNSWSQ ) DOWNCAST ( IndexHNSWCagra ) + DOWNCAST ( IndexHNSWRaBitQ ) DOWNCAST ( IndexHNSW ) DOWNCAST ( IndexHNSW2Level ) DOWNCAST ( IndexNNDescentFlat ) diff --git a/tests/test_custom_result_handler.cpp b/tests/test_custom_result_handler.cpp index 306d468a2..a8beebec3 100644 --- a/tests/test_custom_result_handler.cpp +++ b/tests/test_custom_result_handler.cpp @@ -80,8 +80,28 @@ struct CollectAllResultHandler : ResultHandler { } }; +struct ThresholdUpdatingResultHandler : ResultHandler { + bool add_result(float distance, idx_t /* i */) override { + threshold = distance; + return true; + } +}; + namespace { +TEST(TestSingleQueryBlockResultHandler, SynchronizesWrappedThreshold) { + ThresholdUpdatingResultHandler handler; + handler.threshold = 10.0f; + SingleQueryBlockResultHandler, false> block(handler); + SingleQueryBlockResultHandler, false>:: + SingleResultHandler single(block); + + single.begin(0); + EXPECT_FLOAT_EQ(single.threshold, 10.0f); + single.add_result(4.0f, 0); + EXPECT_FLOAT_EQ(single.threshold, 4.0f); +} + // dimension of the vectors to index constexpr int d = 64; diff --git a/tests/test_hnsw_rabitq.py b/tests/test_hnsw_rabitq.py new file mode 100644 index 000000000..f338a1863 --- /dev/null +++ b/tests/test_hnsw_rabitq.py @@ -0,0 +1,205 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import faiss +import numpy as np + + +class TestHNSWRaBitQ(unittest.TestCase): + def make_data(self, d=32, nt=400, nb=600, nq=20): + rs = np.random.RandomState(123) + xt = rs.randn(nt, d).astype("float32") + xb = rs.randn(nb, d).astype("float32") + xq = rs.randn(nq, d).astype("float32") + return xt, xb, xq + + def make_index(self, nb_bits=3): + xt, xb, xq = self.make_data() + index = faiss.IndexHNSWRaBitQ( + xt.shape[1], 8, nb_bits, faiss.METRIC_L2 + ) + index.hnsw.efConstruction = 40 + index.hnsw.efSearch = 64 + index.train(xt) + index.add(xb) + return index, xb, xq + + @staticmethod + def recall_at_k(actual, expected): + matches = sum( + len(set(actual[i]).intersection(expected[i])) + for i in range(actual.shape[0]) + ) + return matches / actual.size + + def test_staged_search_quality(self): + index, xb, xq = self.make_index(nb_bits=3) + + faiss.cvar.rabitq_stats.reset() + D, I = index.search(xq, 10) + stats = faiss.cvar.rabitq_stats + self.assertGreater(stats.n_1bit, 0) + self.assertGreater(stats.n_refine, 0) + self.assertLess(stats.n_refine, stats.n_1bit) + + exact = faiss.IndexFlatL2(xb.shape[1]) + exact.add(xb) + _, Iexact = exact.search(xq, 10) + self.assertGreaterEqual(self.recall_at_k(I, Iexact), 0.8) + self.assertTrue(np.all(I >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + + def test_symmetric_distance_uses_sign_plane(self): + d = 35 # Exercise padding bits in the final code byte. + rs = np.random.RandomState(456) + xt = np.full((80, d), 2.0, dtype="float32") + residuals = rs.randn(3, d).astype("float32") * 1e10 + residuals[2] = 0 + xb = (residuals + 2.0).astype("float32") + residuals = xb - 2.0 + + norms = np.sum(residuals * residuals, axis=1) + alphas = np.sum(np.abs(residuals), axis=1) / d + signs = np.where(residuals > 0, 1.0, -1.0) + pairs = ((0, 1), (0, 2)) + expected = [ + norms[i] + + norms[j] + - 2 * (alphas[i] * alphas[j] * np.dot(signs[i], signs[j])) + for i, j in pairs + ] + + distances = [] + for nb_bits in (1, 4): + storage = faiss.IndexRaBitQ(d, faiss.METRIC_L2, nb_bits) + storage.train(xt) + storage.add(xb) + dc = storage.get_FlatCodesDistanceComputer() + distances.append([dc.symmetric_dis(i, j) for i, j in pairs]) + + for row in distances: + np.testing.assert_allclose(row, expected, rtol=2e-5, atol=2e-5) + np.testing.assert_array_equal(distances[0], distances[1]) + + def test_clone_and_io_remain_mutable(self): + index, xb, xq = self.make_index(nb_bits=3) + storage = faiss.downcast_index(index.storage) + Dref, Iref = index.search(xq, 10) + + cloned = faiss.clone_index(index) + self.assertIsInstance(cloned, faiss.IndexHNSWRaBitQ) + Dclone, Iclone = cloned.search(xq, 10) + np.testing.assert_array_equal(Iclone, Iref) + np.testing.assert_array_equal(Dclone, Dref) + cloned.add(xb[:1]) + self.assertEqual(cloned.ntotal, len(xb) + 1) + self.assertEqual(index.ntotal, len(xb)) + self.assertEqual(storage.ntotal, len(xb)) + + loaded = faiss.deserialize_index(faiss.serialize_index(index)) + self.assertIsInstance(loaded, faiss.IndexHNSWRaBitQ) + loaded_storage = faiss.downcast_index(loaded.storage) + self.assertEqual(loaded_storage.rabitq.nb_bits, 3) + Dloaded, Iloaded = loaded.search(xq, 10) + np.testing.assert_array_equal(Iloaded, Iref) + np.testing.assert_array_equal(Dloaded, Dref) + loaded.add(xb[:1]) + self.assertEqual(loaded.ntotal, len(xb) + 1) + + def test_io_and_reset(self): + index, xb, xq = self.make_index(nb_bits=4) + cloned = faiss.clone_index(index) + cloned.add(xb[:1]) + self.assertEqual(cloned.ntotal, len(xb) + 1) + + loaded = faiss.deserialize_index(faiss.serialize_index(index)) + loaded.add(xb[:1]) + self.assertEqual(loaded.ntotal, len(xb) + 1) + + loaded.reset() + loaded.add(xb) + self.assertEqual(loaded.ntotal, len(xb)) + D, I = loaded.search(xq, 10) + self.assertTrue(np.all(I >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + + def test_one_bit_fallback(self): + index, _, xq = self.make_index(nb_bits=1) + self.assertEqual(index.hnsw.search_method, 0) + faiss.cvar.rabitq_stats.reset() + D, I = index.search(xq, 10) + self.assertTrue(np.all(I >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + self.assertEqual(faiss.cvar.rabitq_stats.n_1bit, 0) + self.assertEqual(faiss.cvar.rabitq_stats.n_refine, 0) + + loaded = faiss.deserialize_index(faiss.serialize_index(index)) + loaded_storage = faiss.downcast_index(loaded.storage) + self.assertEqual(loaded_storage.rabitq.nb_bits, 1) + self.assertEqual(loaded.hnsw.search_method, 0) + Dloaded, Iloaded = loaded.search(xq, 10) + np.testing.assert_array_equal(Dloaded, D) + np.testing.assert_array_equal(Iloaded, I) + + def test_staged_search_requires_bounded_queue(self): + index, _, xq = self.make_index(nb_bits=3) + index.hnsw.search_bounded_queue = False + with self.assertRaises(RuntimeError): + index.search(xq, 10) + + def test_permute_then_add(self): + index, xb, xq = self.make_index(nb_bits=3) + perm = np.arange(len(xb) - 1, -1, -1, dtype=np.int64) + index.permute_entries(perm) + index.add(xb[:2]) + self.assertEqual(index.ntotal, len(xb) + 2) + D, I = index.search(xq, 10) + self.assertTrue(np.all(I >= 0)) + self.assertTrue(np.all(np.isfinite(D))) + + def test_factory(self): + default_index = faiss.index_factory(32, "HNSW8,RaBitQ") + self.assertIsInstance(default_index, faiss.IndexHNSWRaBitQ) + self.assertFalse(default_index.is_trained) + default_storage = faiss.downcast_index(default_index.storage) + self.assertIsInstance(default_storage, faiss.IndexRaBitQ) + self.assertFalse(default_storage.is_trained) + self.assertEqual(default_storage.rabitq.nb_bits, 1) + self.assertEqual(default_index.hnsw.search_method, 0) + with self.assertRaises(RuntimeError): + default_index.add(np.zeros((1, 32), dtype="float32")) + + four_bit = faiss.index_factory(32, "HNSW8,RaBitQ4") + four_bit_storage = faiss.downcast_index(four_bit.storage) + self.assertEqual(four_bit_storage.rabitq.nb_bits, 4) + self.assertNotEqual(four_bit.hnsw.search_method, 0) + + def test_skip_storage_can_be_reserialized(self): + index, _, _ = self.make_index(nb_bits=3) + data = faiss.serialize_index(index, faiss.IO_FLAG_SKIP_STORAGE) + metadata_only = faiss.deserialize_index(data) + self.assertIsNone(metadata_only.storage) + self.assertNotEqual(metadata_only.hnsw.search_method, 0) + + data_again = faiss.serialize_index( + metadata_only, faiss.IO_FLAG_SKIP_STORAGE + ) + metadata_again = faiss.deserialize_index(data_again) + self.assertIsNone(metadata_again.storage) + self.assertEqual( + metadata_again.hnsw.search_method, + metadata_only.hnsw.search_method, + ) + + def test_unsupported_metric_throws(self): + for metric in (faiss.METRIC_INNER_PRODUCT, faiss.METRIC_L1): + with self.assertRaises(RuntimeError): + faiss.IndexHNSWRaBitQ(32, 8, 1, metric) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_io.py b/tests/test_io.py index 98c6b8bdb..feaab0681 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -775,8 +775,8 @@ class Test_IO_HNSW(unittest.TestCase): index = faiss.IndexHNSWFlat(d, 16) index.train(self.xt) index.add(self.xb) - # Note: RaBitQ lacks symmetric_distance, so it can be used for - # `search()`, but not `add()`. + # Build the graph with Flat storage, then verify that swapping in an + # equivalent RaBitQ storage preserves useful search results. I_flat = index.storage.assign(self.xq, 10) I_hnsw_flat = index.assign(self.xq, 10) index.storage = faiss.IndexRaBitQ(d)