From b36a28778a493df46a7f40b03cc5b426e94a4ba1 Mon Sep 17 00:00:00 2001 From: Junqiu Lei Date: Tue, 23 Apr 2024 17:18:56 -0700 Subject: [PATCH] Custom patch to support range search params Signed-off-by: Junqiu Lei --- faiss/IndexIDMap.cpp | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp index 3f375e7bb6..11f3a84775 100644 --- a/faiss/IndexIDMap.cpp +++ b/faiss/IndexIDMap.cpp @@ -176,11 +176,31 @@ void IndexIDMapTemplate::range_search( RangeSearchResult* result, const SearchParameters* params) const { if (params) { - SearchParameters internal_search_parameters; - IDSelectorTranslated id_selector_translated(id_map, params->sel); - internal_search_parameters.sel = &id_selector_translated; + IDSelectorTranslated this_idtrans(this->id_map, nullptr); + ScopedSelChange sel_change; + IDGrouperTranslated this_idgrptrans(this->id_map, nullptr); + ScopedGrpChange grp_change; + + if (params->sel) { + auto idtrans = dynamic_cast(params->sel); + + if (!idtrans) { + auto params_non_const = const_cast(params); + this_idtrans.sel = params->sel; + sel_change.set(params_non_const, &this_idtrans); + } + } + + if (params->grp) { + auto idtrans = dynamic_cast(params->grp); - index->range_search(n, x, radius, result, &internal_search_parameters); + if (!idtrans) { + auto params_non_const = const_cast(params); + this_idgrptrans.grp = params->grp; + grp_change.set(params_non_const, &this_idgrptrans); + } + } + index->range_search(n, x, radius, result, params); } else { index->range_search(n, x, radius, result); }