forked from facebookresearch/faiss
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Custom patch to support multi-vector
- Loading branch information
Showing
6 changed files
with
154 additions
and
12 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
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,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 |
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,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* newCollector() { | ||
| 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 deleteCollector(ResultCollector* collector) { | ||
| // Do nothing | ||
| } | ||
|
|
||
| virtual ~ResultCollectorFactory() {} | ||
| }; | ||
|
|
||
| } // namespace faiss |