-
Notifications
You must be signed in to change notification settings - Fork 126
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add patch to support multi vector in faiss
Signed-off-by: Heemin Kim <[email protected]>
- Loading branch information
Showing
4 changed files
with
305 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
281 changes: 281 additions & 0 deletions
281
jni/patches/faiss/0001-Custom-patch-to-support-multi-vector.patch
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,281 @@ | ||
| From 8f89fbf1cf445a8216b9cc4ee52e7b82e24e906e Mon Sep 17 00:00:00 2001 | ||
| From: Heemin Kim <[email protected]> | ||
| Date: Wed, 6 Dec 2023 16:33:52 -0800 | ||
| Subject: [PATCH] Custom patch to support multi-vector | ||
|
|
||
| Signed-off-by: Heemin Kim <[email protected]> | ||
| --- | ||
| faiss/CMakeLists.txt | 2 + | ||
| faiss/Index.h | 6 ++- | ||
| faiss/IndexIDMap.cpp | 20 ++++++++++ | ||
| faiss/IndexIDMap.h | 1 + | ||
| faiss/impl/HNSW.cpp | 25 ++++++++----- | ||
| faiss/impl/ResultCollector.h | 58 +++++++++++++++++++++++++++++ | ||
| faiss/impl/ResultCollectorFactory.h | 29 +++++++++++++++ | ||
| 7 files changed, 129 insertions(+), 12 deletions(-) | ||
| create mode 100644 faiss/impl/ResultCollector.h | ||
| create mode 100644 faiss/impl/ResultCollectorFactory.h | ||
|
|
||
| diff --git a/faiss/CMakeLists.txt b/faiss/CMakeLists.txt | ||
| index 27701586..af682a05 100644 | ||
| --- a/faiss/CMakeLists.txt | ||
| +++ b/faiss/CMakeLists.txt | ||
| @@ -162,6 +162,8 @@ set(FAISS_HEADERS | ||
| impl/ProductQuantizer.h | ||
| impl/Quantizer.h | ||
| impl/ResidualQuantizer.h | ||
| + impl/ResultCollector.h | ||
| + impl/ResultCollectorFactory.h | ||
| impl/ResultHandler.h | ||
| impl/ScalarQuantizer.h | ||
| impl/ThreadedIndex-inl.h | ||
| diff --git a/faiss/Index.h b/faiss/Index.h | ||
| index 4b4b302b..13eab0c0 100644 | ||
| --- a/faiss/Index.h | ||
| +++ b/faiss/Index.h | ||
| @@ -38,11 +38,12 @@ | ||
|
|
||
| namespace faiss { | ||
|
|
||
| -/// Forward declarations see impl/AuxIndexStructures.h, impl/IDSelector.h and | ||
| -/// impl/DistanceComputer.h | ||
| +/// Forward declarations see impl/AuxIndexStructures.h, impl/IDSelector.h, | ||
| +/// impl/DistanceComputer.h, and impl/ResultCollectorFactory.h | ||
| struct IDSelector; | ||
| struct RangeSearchResult; | ||
| struct DistanceComputer; | ||
| +struct ResultCollectorFactory; | ||
|
|
||
| /** Parent class for the optional search paramenters. | ||
| * | ||
| @@ -52,6 +53,7 @@ struct DistanceComputer; | ||
| struct SearchParameters { | ||
| /// if non-null, only these IDs will be considered during search. | ||
| IDSelector* sel = nullptr; | ||
| + ResultCollectorFactory* col = nullptr; | ||
| /// make sure we can dynamic_cast this | ||
| virtual ~SearchParameters() {} | ||
| }; | ||
| diff --git a/faiss/IndexIDMap.cpp b/faiss/IndexIDMap.cpp | ||
| index 7972bec9..a5c017a9 100644 | ||
| --- a/faiss/IndexIDMap.cpp | ||
| +++ b/faiss/IndexIDMap.cpp | ||
| @@ -102,6 +102,20 @@ struct ScopedSelChange { | ||
| } | ||
| }; | ||
|
|
||
| +/// RAII object to reset the ResultCollectorFactory in the params object | ||
| +struct ScopedColChange { | ||
| + SearchParameters* params = nullptr; | ||
| + void set(SearchParameters* params, const std::vector<int64_t>* id_map) { | ||
| + this->params = params; | ||
| + params->col->id_map = id_map; | ||
| + } | ||
| + ~ScopedColChange() { | ||
| + if (params) { | ||
| + params->col->id_map = nullptr; | ||
| + } | ||
| + } | ||
| +}; | ||
| + | ||
| } // namespace | ||
|
|
||
| template <typename IndexT> | ||
| @@ -114,6 +128,7 @@ void IndexIDMapTemplate<IndexT>::search( | ||
| const SearchParameters* params) const { | ||
| IDSelectorTranslated this_idtrans(this->id_map, nullptr); | ||
| ScopedSelChange sel_change; | ||
| + ScopedColChange col_change; | ||
|
|
||
| if (params && params->sel) { | ||
| auto idtrans = dynamic_cast<const IDSelectorTranslated*>(params->sel); | ||
| @@ -131,6 +146,11 @@ void IndexIDMapTemplate<IndexT>::search( | ||
| sel_change.set(params_non_const, &this_idtrans); | ||
| } | ||
| } | ||
| + | ||
| + if (params && params->col) { | ||
| + auto params_non_const = const_cast<SearchParameters*>(params); | ||
| + col_change.set(params_non_const, &this->id_map); | ||
| + } | ||
| index->search(n, x, k, distances, labels, params); | ||
| idx_t* li = labels; | ||
| #pragma omp parallel for | ||
| diff --git a/faiss/IndexIDMap.h b/faiss/IndexIDMap.h | ||
| index 2d164123..c6a1be73 100644 | ||
| --- a/faiss/IndexIDMap.h | ||
| +++ b/faiss/IndexIDMap.h | ||
| @@ -10,6 +10,7 @@ | ||
| #include <faiss/Index.h> | ||
| #include <faiss/IndexBinary.h> | ||
| #include <faiss/impl/IDSelector.h> | ||
| +#include <faiss/impl/ResultCollectorFactory.h> | ||
|
|
||
| #include <unordered_map> | ||
| #include <vector> | ||
| diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp | ||
| index 9fc201ea..540210a6 100644 | ||
| --- a/faiss/impl/HNSW.cpp | ||
| +++ b/faiss/impl/HNSW.cpp | ||
| @@ -14,6 +14,7 @@ | ||
| #include <faiss/impl/AuxIndexStructures.h> | ||
| #include <faiss/impl/DistanceComputer.h> | ||
| #include <faiss/impl/IDSelector.h> | ||
| +#include <faiss/impl/ResultCollectorFactory.h> | ||
| #include <faiss/utils/prefetch.h> | ||
|
|
||
| #include <faiss/impl/platform_macros.h> | ||
| @@ -530,6 +531,15 @@ int search_from_candidates( | ||
| int level, | ||
| int nres_in = 0, | ||
| const SearchParametersHNSW* params = nullptr) { | ||
| + ResultCollectorFactory defaultFactory; | ||
| + ResultCollectorFactory* collectorFactory; | ||
| + if (params == nullptr || params->col == nullptr) { | ||
| + collectorFactory = &defaultFactory; | ||
| + } else { | ||
| + collectorFactory = params->col; | ||
| + } | ||
| + ResultCollector* collector = collectorFactory->newCollector(); | ||
| + | ||
| int nres = nres_in; | ||
| int ndis = 0; | ||
|
|
||
| @@ -544,11 +554,7 @@ int search_from_candidates( | ||
| float d = candidates.dis[i]; | ||
| FAISS_ASSERT(v1 >= 0); | ||
| if (!sel || sel->is_member(v1)) { | ||
| - if (nres < k) { | ||
| - faiss::maxheap_push(++nres, D, I, d, v1); | ||
| - } else if (d < D[0]) { | ||
| - faiss::maxheap_replace_top(nres, D, I, d, v1); | ||
| - } | ||
| + collector->collect(k, nres, D, I, d, v1); | ||
| } | ||
| vt.set(v1); | ||
| } | ||
| @@ -612,11 +618,7 @@ int search_from_candidates( | ||
|
|
||
| auto add_to_heap = [&](const size_t idx, const float dis) { | ||
| if (!sel || sel->is_member(idx)) { | ||
| - if (nres < k) { | ||
| - faiss::maxheap_push(++nres, D, I, dis, idx); | ||
| - } else if (dis < D[0]) { | ||
| - faiss::maxheap_replace_top(nres, D, I, dis, idx); | ||
| - } | ||
| + collector->collect(k, nres, D, I, dis, idx); | ||
| } | ||
| candidates.push(idx, dis); | ||
| }; | ||
| @@ -660,6 +662,9 @@ int search_from_candidates( | ||
| } | ||
| } | ||
|
|
||
| + collector->finalize(nres, I); | ||
| + collectorFactory->deleteCollector(collector); | ||
| + | ||
| if (level == 0) { | ||
| stats.n1++; | ||
| if (candidates.size() == 0) { | ||
| diff --git a/faiss/impl/ResultCollector.h b/faiss/impl/ResultCollector.h | ||
| new file mode 100644 | ||
| index 00000000..3e4dac34 | ||
| --- /dev/null | ||
| +++ b/faiss/impl/ResultCollector.h | ||
| @@ -0,0 +1,58 @@ | ||
| +/** | ||
| + * Copyright (c) Facebook, Inc. and its affiliates. | ||
| + * | ||
| + * This source code is licensed under the MIT license found in the | ||
| + * LICENSE file in the root directory of this source tree. | ||
| + */ | ||
| + | ||
| +#pragma once | ||
| + | ||
| +#include <unordered_set> | ||
| +#include <vector> | ||
| + | ||
| +#include <faiss/MetricType.h> | ||
| +#include <faiss/utils/Heap.h> | ||
| + | ||
| +/** ResultCollector is intended to define how to collect search result */ | ||
| + | ||
| +namespace faiss { | ||
| + | ||
| +/** Encapsulates a set of ids to handle. */ | ||
| +struct ResultCollector { | ||
| + // For each result, collect method is called to store result | ||
| + virtual void collect( | ||
| + int k, | ||
| + int& nres, | ||
| + float* bh_val, | ||
| + idx_t* bh_ids, | ||
| + float val, | ||
| + idx_t ids) = 0; | ||
| + | ||
| + // This method is called after all result is collected | ||
| + virtual void finalize(idx_t nres, idx_t* bh_ids) = 0; | ||
| + virtual ~ResultCollector() {} | ||
| +}; | ||
| + | ||
| +struct DefaultCollector : ResultCollector { | ||
| + void collect( | ||
| + int k, | ||
| + int& nres, | ||
| + float* bh_val, | ||
| + idx_t* bh_ids, | ||
| + float val, | ||
| + idx_t ids) override { | ||
| + if (nres < k) { | ||
| + faiss::maxheap_push(++nres, bh_val, bh_ids, val, ids); | ||
| + } else if (val < bh_val[0]) { | ||
| + faiss::maxheap_replace_top(nres, bh_val, bh_ids, val, ids); | ||
| + } | ||
| + } | ||
| + | ||
| + void finalize(idx_t nres, idx_t* bh_ids) override { | ||
| + // Do nothing | ||
| + } | ||
| + | ||
| + ~DefaultCollector() override {} | ||
| +}; | ||
| + | ||
| +} // namespace faiss | ||
| diff --git a/faiss/impl/ResultCollectorFactory.h b/faiss/impl/ResultCollectorFactory.h | ||
| new file mode 100644 | ||
| index 00000000..4d903f8d | ||
| --- /dev/null | ||
| +++ b/faiss/impl/ResultCollectorFactory.h | ||
| @@ -0,0 +1,29 @@ | ||
| +/** | ||
| + * Copyright (c) Facebook, Inc. and its affiliates. | ||
| + * | ||
| + * This source code is licensed under the MIT license found in the | ||
| + * LICENSE file in the root directory of this source tree. | ||
| + */ | ||
| + | ||
| +#pragma once | ||
| +#include <faiss/impl/ResultCollector.h> | ||
| +namespace faiss { | ||
| + | ||
| +/** ResultCollector is intended to define how to collect search result */ | ||
| +struct ResultCollectorFactory { | ||
| + DefaultCollector default_collector; | ||
| + const std::vector<int64_t>* id_map; | ||
| + | ||
| + // For each result, collect method is called to store result | ||
| + virtual ResultCollector* newCollector() { | ||
| + return &default_collector; | ||
| + } | ||
| + | ||
| + virtual void deleteCollector(ResultCollector* collector) { | ||
| + // Do nothing | ||
| + } | ||
| + // This method is called after all result is collected | ||
| + virtual ~ResultCollectorFactory() {} | ||
| +}; | ||
| + | ||
| +} // namespace faiss | ||
| -- | ||
| 2.39.3 (Apple Git-145) | ||
|
|