-
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 jni interface to use a binary hnsw index with faiss (#1778)
* Add jni interface to use a binary hnsw index with faiss (#1747) Signed-off-by: Heemin Kim <[email protected]> * Fix memory leak on test code (#1776) Signed-off-by: Heemin Kim <[email protected]> --------- Signed-off-by: Heemin Kim <[email protected]>
- Loading branch information
1 parent
7e95e16
commit 083f668
Showing
36 changed files
with
1,637 additions
and
106 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
| // The OpenSearch Contributors require contributions made to | ||
| // this file be licensed under the Apache-2.0 license or a | ||
| // compatible open source license. | ||
| // | ||
| // Modifications Copyright OpenSearch Contributors. See | ||
| // GitHub history for details. | ||
|
|
||
| /** | ||
| * This file contains classes for index operations which are free of JNI | ||
| */ | ||
|
|
||
| #ifndef OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H | ||
| #define OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H | ||
|
|
||
| #include <jni.h> | ||
| #include "faiss/MetricType.h" | ||
| #include "jni_util.h" | ||
| #include "faiss_methods.h" | ||
| #include <memory> | ||
|
|
||
| namespace knn_jni { | ||
| namespace faiss_wrapper { | ||
|
|
||
|
|
||
| /** | ||
| * A class to provide operations on index | ||
| * This class should evolve to have only cpp object but not jni object | ||
| */ | ||
| class IndexService { | ||
| public: | ||
| IndexService(std::unique_ptr<FaissMethods> faissMethods); | ||
| //TODO Remove dependency on JNIUtilInterface and JNIEnv | ||
| //TODO Reduce the number of parameters | ||
|
|
||
| /** | ||
| * Create index | ||
| * | ||
| * @param jniUtil jni util | ||
| * @param env jni environment | ||
| * @param metric space type for distance calculation | ||
| * @param indexDescription index description to be used by faiss index factory | ||
| * @param dim dimension of vectors | ||
| * @param numIds number of vectors | ||
| * @param threadCount number of thread count to be used while adding data | ||
| * @param vectorsAddress memory address which is holding vector data | ||
| * @param ids a list of document ids for corresponding vectors | ||
| * @param indexPath path to write index | ||
| * @param parameters parameters to be applied to faiss index | ||
| */ | ||
| virtual void createIndex( | ||
| knn_jni::JNIUtilInterface * jniUtil, | ||
| JNIEnv * env, | ||
| faiss::MetricType metric, | ||
| std::string indexDescription, | ||
| int dim, | ||
| int numIds, | ||
| int threadCount, | ||
| int64_t vectorsAddress, | ||
| std::vector<int64_t> ids, | ||
| std::string indexPath, | ||
| std::unordered_map<std::string, jobject> parameters); | ||
| virtual ~IndexService() = default; | ||
| protected: | ||
| std::unique_ptr<FaissMethods> faissMethods; | ||
| }; | ||
|
|
||
| /** | ||
| * A class to provide operations on index | ||
| * This class should evolve to have only cpp object but not jni object | ||
| */ | ||
| class BinaryIndexService : public IndexService { | ||
| public: | ||
| //TODO Remove dependency on JNIUtilInterface and JNIEnv | ||
| //TODO Reduce the number of parameters | ||
| BinaryIndexService(std::unique_ptr<FaissMethods> faissMethods); | ||
| /** | ||
| * Create binary index | ||
| * | ||
| * @param jniUtil jni util | ||
| * @param env jni environment | ||
| * @param metric space type for distance calculation | ||
| * @param indexDescription index description to be used by faiss index factory | ||
| * @param dim dimension of vectors | ||
| * @param numIds number of vectors | ||
| * @param threadCount number of thread count to be used while adding data | ||
| * @param vectorsAddress memory address which is holding vector data | ||
| * @param ids a list of document ids for corresponding vectors | ||
| * @param indexPath path to write index | ||
| * @param parameters parameters to be applied to faiss index | ||
| */ | ||
| virtual void createIndex( | ||
| knn_jni::JNIUtilInterface * jniUtil, | ||
| JNIEnv * env, | ||
| faiss::MetricType metric, | ||
| std::string indexDescription, | ||
| int dim, | ||
| int numIds, | ||
| int threadCount, | ||
| int64_t vectorsAddress, | ||
| std::vector<int64_t> ids, | ||
| std::string indexPath, | ||
| std::unordered_map<std::string, jobject> parameters | ||
| ) override; | ||
| virtual ~BinaryIndexService() = default; | ||
| }; | ||
|
|
||
| } | ||
| } | ||
|
|
||
|
|
||
| #endif //OPENSEARCH_KNN_FAISS_INDEX_SERVICE_H |
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,42 @@ | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
| // | ||
| // The OpenSearch Contributors require contributions made to | ||
| // this file be licensed under the Apache-2.0 license or a | ||
| // compatible open source license. | ||
| // | ||
| // Modifications Copyright OpenSearch Contributors. See | ||
| // GitHub history for details. | ||
|
|
||
| #ifndef OPENSEARCH_KNN_FAISS_METHODS_H | ||
| #define OPENSEARCH_KNN_FAISS_METHODS_H | ||
|
|
||
| #include "faiss/Index.h" | ||
| #include "faiss/IndexBinary.h" | ||
| #include "faiss/IndexIDMap.h" | ||
| #include "faiss/index_io.h" | ||
|
|
||
| namespace knn_jni { | ||
| namespace faiss_wrapper { | ||
|
|
||
| /** | ||
| * A class having wrapped faiss methods | ||
| * | ||
| * This class helps to mock faiss methods during unit test | ||
| */ | ||
| class FaissMethods { | ||
| public: | ||
| FaissMethods() = default; | ||
| virtual faiss::Index* indexFactory(int d, const char* description, faiss::MetricType metric); | ||
| virtual faiss::IndexBinary* indexBinaryFactory(int d, const char* description); | ||
| virtual faiss::IndexIDMapTemplate<faiss::Index>* indexIdMap(faiss::Index* index); | ||
| virtual faiss::IndexIDMapTemplate<faiss::IndexBinary>* indexBinaryIdMap(faiss::IndexBinary* index); | ||
| virtual void writeIndex(const faiss::Index* idx, const char* fname); | ||
| virtual void writeIndexBinary(const faiss::IndexBinary* idx, const char* fname); | ||
| virtual ~FaissMethods() = default; | ||
| }; | ||
|
|
||
| } //namespace faiss_wrapper | ||
| } //namespace knn_jni | ||
|
|
||
|
|
||
| #endif //OPENSEARCH_KNN_FAISS_METHODS_H |
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
Oops, something went wrong.