-
Notifications
You must be signed in to change notification settings - Fork 113
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 (#1747)
Signed-off-by: Heemin Kim <heemin@amazon.com>
- Loading branch information
Showing
35 changed files
with
1,620 additions
and
95 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.