Add patch to support multi vector in faiss #1358
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -38,6 +38,15 @@ jobs: | |
| with: | ||
| submodules: true | ||
|
|
||
| # Git functionality in CMAKE file does not work with given ubuntu image. Therefore, handling it here. | ||
| - name: Apply Git Patch | ||
| # Deleting file at the end to skip `git apply` inside CMAKE file | ||
| run: | | ||
| cd jni/external/faiss | ||
| git apply --ignore-space-change --ignore-whitespace --3way ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch | ||
| rm ../../patches/faiss/0001-Custom-patch-to-support-multi-vector.patch | ||
|
There was a problem hiding this comment. Commented in line #42. |
||
| working-directory: ${{ github.workspace }} | ||
|
|
||
| - name: Setup Java ${{ matrix.java }} | ||
| uses: actions/setup-java@v1 | ||
| with: | ||
|
|
||
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,301 @@ | ||
| From 70c9b5d6408bb85fd1503002d29f9428fe0d48dc 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 | 24 ++++++++++ | ||
| faiss/impl/HNSW.cpp | 27 +++++++---- | ||
| faiss/impl/ResultCollector.h | 74 +++++++++++++++++++++++++++++ | ||
| faiss/impl/ResultCollectorFactory.h | 33 +++++++++++++ | ||
| 6 files changed, 154 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..0f82a17c 100644 | ||
| --- a/faiss/IndexIDMap.cpp | ||
| +++ b/faiss/IndexIDMap.cpp | ||
| @@ -18,6 +18,7 @@ | ||
| #include <faiss/impl/FaissAssert.h> | ||
| #include <faiss/utils/Heap.h> | ||
| #include <faiss/utils/WorkerThread.h> | ||
| +#include <faiss/impl/ResultCollectorFactory.h> | ||
|
|
||
| namespace faiss { | ||
|
|
||
| @@ -102,6 +103,24 @@ struct ScopedSelChange { | ||
| } | ||
| }; | ||
|
|
||
| +// RAII object to reset the id_map parameter in ResultCollectorFactory object | ||
| +// This object make sure to reset the id_map parameter in ResultCollectorFactory | ||
| +// once the program exist current method scope. | ||
| +struct ScopedColChange { | ||
| + ResultCollectorFactory* collector_factory = nullptr; | ||
| + void set( | ||
| + ResultCollectorFactory* collector_factory, | ||
| + const std::vector<int64_t>* id_map) { | ||
| + this->collector_factory = collector_factory; | ||
| + collector_factory->id_map = id_map; | ||
| + } | ||
| + ~ScopedColChange() { | ||
| + if (collector_factory) { | ||
| + collector_factory->id_map = nullptr; | ||
| + } | ||
| + } | ||
| +}; | ||
| + | ||
| } // namespace | ||
|
|
||
| template <typename IndexT> | ||
| @@ -114,6 +133,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 +151,10 @@ void IndexIDMapTemplate<IndexT>::search( | ||
| sel_change.set(params_non_const, &this_idtrans); | ||
| } | ||
| } | ||
| + | ||
| + if (params && params->col && !params->col->id_map) { | ||
| + col_change.set(params->col, &this->id_map); | ||
| + } | ||
| index->search(n, x, k, distances, labels, params); | ||
| idx_t* li = labels; | ||
| #pragma omp parallel for | ||
| diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp | ||
| index 9fc201ea..5b5900d1 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->new_collector(); | ||
| + | ||
| 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,11 @@ int search_from_candidates( | ||
| } | ||
| } | ||
|
|
||
| + // Completed collection of result. Run post processor. | ||
| + collector->post_process(nres, I); | ||
| + // Collector completed its task. Release all resource of the collector. | ||
| + collectorFactory->delete_collector(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..a0489fd6 | ||
| --- /dev/null | ||
| +++ b/faiss/impl/ResultCollector.h | ||
| @@ -0,0 +1,74 @@ | ||
| +/** | ||
| + * 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 | ||
| + * For each single search result, collect method will be called. | ||
| + * After every results are collected, post_process method is called at the end. | ||
| + */ | ||
| + | ||
| +namespace faiss { | ||
| + | ||
| +/** Encapsulates a set of ids to handle. */ | ||
| +struct ResultCollector { | ||
| + /** | ||
| + * For each result, collect method is called to store result | ||
| + * @param k number of vectors to search | ||
| + * @param nres number of results in queue | ||
| + * @param bh_val search result, distances from query | ||
| + * @param bh_ids search result, ids of vectors | ||
| + * @param val distance from query for current vector | ||
| + * @param ids id of current vector | ||
| + */ | ||
| + 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 post_process(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); | ||
| + } | ||
| + } | ||
| + | ||
| + // This method is called once all result is collected so that final post | ||
| + // processing can be done For example, if the result is collected using | ||
| + // group id, the group id can be converted back to its original id inside | ||
| + // this method | ||
| + void post_process(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..70918266 | ||
| --- /dev/null | ||
| +++ b/faiss/impl/ResultCollectorFactory.h | ||
| @@ -0,0 +1,33 @@ | ||
| +/** | ||
| + * 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 { | ||
| + | ||
| +/** ResultCollectorFactory to create a ResultCollector object */ | ||
| +struct ResultCollectorFactory { | ||
| + DefaultCollector default_collector; | ||
| + const std::vector<int64_t>* id_map; | ||
| + | ||
| + // Create a new ResultCollector object | ||
| + virtual ResultCollector* new_collector() { | ||
| + return &default_collector; | ||
| + } | ||
| + | ||
| + // For default case, the factory share single object and no need to delete | ||
| + // the object. For other case, the factory can create a new object which | ||
| + // need to be deleted later. We have deleteCollector method to handle both | ||
| + // case as factory class knows how to release resource that it created | ||
| + virtual void delete_collector(ResultCollector* collector) { | ||
| + // Do nothing | ||
| + } | ||
| + | ||
| + virtual ~ResultCollectorFactory() {} | ||
| +}; | ||
| + | ||
| +} // namespace faiss | ||
| -- | ||
| 2.39.3 (Apple Git-145) | ||
|
|
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What is meant by this? If we do this, everytime we add a patch we will need to update this.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Correct. Git command in CMAKE file like
git submodule update --init --orgit applydoes not work for ubuntu. Need to talk with @peterzhuamazon to see if we can fix it by using different ubuntu image.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@heemin32 what was the conclusion for this?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The code is needed as of now to make it work. Will talk to @peterzhuamazon and create an issue in opensearch build repo to resolve the problem outside of this PR.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is that discussion blocking this PR?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It won't.