-
Notifications
You must be signed in to change notification settings - Fork 123
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add the interface for streaming the vectors from java to jni layer wi…
…th initial capacity (#1586) Signed-off-by: Navneet Verma <[email protected]>
- Loading branch information
Showing
19 changed files
with
412 additions
and
73 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| /* | ||
| * 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. | ||
| */ | ||
| #include "jni_util.h" | ||
| #include <jni.h> | ||
| namespace knn_jni { | ||
| namespace commons { | ||
| /** | ||
| * This is utility function that can be used to store data in native memory. This function will allocate memory for | ||
| * the data(rows*columns) with initialCapacity and return the memory address where the data is stored. | ||
| * If you are using this function for first time use memoryAddress = 0 to ensure that a new memory location is created. | ||
| * For subsequent calls you can pass the same memoryAddress. If the data cannot be stored in the memory location | ||
| * will throw Exception. | ||
| * | ||
| * @param memoryAddress The address of the memory location where data will be stored. | ||
| * @param data 2D float array containing data to be stored in native memory. | ||
| * @param initialCapacity The initial capacity of the memory location. | ||
| * @return memory address where the data is stored. | ||
| */ | ||
| jlong storeVectorData(knn_jni::JNIUtilInterface *, JNIEnv *, jlong , jobjectArray, jlong); | ||
|
|
||
| /** | ||
| * Free up the memory allocated for the data stored in memory address. This function should be used with the memory | ||
| * address returned by {@link JNICommons#storeVectorData(long, float[][], long, long)} | ||
| * | ||
| * @param memoryAddress address to be freed. | ||
| */ | ||
| void freeVectorData(jlong); | ||
| } | ||
| } |
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,40 @@ | ||
| /* | ||
| * 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. | ||
| */ | ||
|
|
||
| /* DO NOT EDIT THIS FILE - it is machine generated */ | ||
| #include <jni.h> | ||
| /* Header for class org_opensearch_knn_jni_JNICommons */ | ||
|
|
||
| #ifndef _Included_org_opensearch_knn_jni_JNICommons | ||
| #define _Included_org_opensearch_knn_jni_JNICommons | ||
| #ifdef __cplusplus | ||
| extern "C" { | ||
| #endif | ||
| /* | ||
| * Class: org_opensearch_knn_jni_JNICommons | ||
| * Method: storeVectorData | ||
| * Signature: (J[[FJJ) | ||
| */ | ||
| JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_JNICommons_storeVectorData | ||
| (JNIEnv *, jclass, jlong, jobjectArray, jlong); | ||
|
|
||
| /* | ||
| * Class: org_opensearch_knn_jni_JNICommons | ||
| * Method: freeVectorData | ||
| * Signature: (J)V | ||
| */ | ||
| JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_JNICommons_freeVectorData | ||
| (JNIEnv *, jclass, jlong); | ||
|
|
||
| #ifdef __cplusplus | ||
| } | ||
| #endif | ||
| #endif |
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,41 @@ | ||
| /* | ||
| * 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_COMMONS_H | ||
| #define OPENSEARCH_KNN_COMMONS_H | ||
| #include <jni.h> | ||
|
|
||
| #include <vector> | ||
|
|
||
| #include "jni_util.h" | ||
| #include "commons.h" | ||
|
|
||
| jlong knn_jni::commons::storeVectorData(knn_jni::JNIUtilInterface *jniUtil, JNIEnv *env, jlong memoryAddressJ, | ||
| jobjectArray dataJ, jlong initialCapacityJ) { | ||
| std::vector<float> *vect; | ||
| if ((long) memoryAddressJ == 0) { | ||
| vect = new std::vector<float>(); | ||
| vect->reserve((long)initialCapacityJ); | ||
| } else { | ||
| vect = reinterpret_cast<std::vector<float>*>(memoryAddressJ); | ||
| } | ||
| int dim = jniUtil->GetInnerDimensionOf2dJavaFloatArray(env, dataJ); | ||
| jniUtil->Convert2dJavaObjectArrayAndStoreToFloatVector(env, dataJ, dim, vect); | ||
|
|
||
| return (jlong) vect; | ||
| } | ||
|
|
||
| void knn_jni::commons::freeVectorData(jlong memoryAddressJ) { | ||
| if (memoryAddressJ != 0) { | ||
| auto *vect = reinterpret_cast<std::vector<float>*>(memoryAddressJ); | ||
| delete vect; | ||
| } | ||
| } | ||
| #endif //OPENSEARCH_KNN_COMMONS_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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| /* | ||
| * 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. | ||
| */ | ||
|
|
||
| #include "org_opensearch_knn_jni_JNICommons.h" | ||
|
|
||
| #include <jni.h> | ||
| #include "commons.h" | ||
| #include "jni_util.h" | ||
|
|
||
| static knn_jni::JNIUtil jniUtil; | ||
| static const jint KNN_JNICOMMONS_JNI_VERSION = JNI_VERSION_1_1; | ||
|
|
||
| jint JNI_OnLoad(JavaVM* vm, void* reserved) { | ||
| // Obtain the JNIEnv from the VM and confirm JNI_VERSION | ||
| JNIEnv* env; | ||
| if (vm->GetEnv((void**)&env, KNN_JNICOMMONS_JNI_VERSION) != JNI_OK) { | ||
| return JNI_ERR; | ||
| } | ||
|
|
||
| jniUtil.Initialize(env); | ||
|
|
||
| return KNN_JNICOMMONS_JNI_VERSION; | ||
| } | ||
|
|
||
| void JNI_OnUnload(JavaVM *vm, void *reserved) { | ||
| JNIEnv* env; | ||
| vm->GetEnv((void**)&env, KNN_JNICOMMONS_JNI_VERSION); | ||
| jniUtil.Uninitialize(env); | ||
| } | ||
|
|
||
|
|
||
| JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_JNICommons_storeVectorData(JNIEnv * env, jclass cls, | ||
| jlong memoryAddressJ, jobjectArray dataJ, jlong initialCapacityJ) | ||
|
|
||
| { | ||
| try { | ||
| return knn_jni::commons::storeVectorData(&jniUtil, env, memoryAddressJ, dataJ, initialCapacityJ); | ||
| } catch (...) { | ||
| jniUtil.CatchCppExceptionAndThrowJava(env); | ||
| } | ||
| return (long)memoryAddressJ; | ||
| } | ||
|
|
||
| JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_JNICommons_freeVectorData(JNIEnv * env, jclass cls, | ||
| jlong memoryAddressJ) | ||
| { | ||
| try { | ||
| return knn_jni::commons::freeVectorData(memoryAddressJ); | ||
| } catch (...) { | ||
| jniUtil.CatchCppExceptionAndThrowJava(env); | ||
| } | ||
| } |
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,73 @@ | ||
| /* | ||
| * 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. | ||
| */ | ||
|
|
||
|
|
||
| #include "test_util.h" | ||
| #include <vector> | ||
| #include "gmock/gmock.h" | ||
| #include "gtest/gtest.h" | ||
| #include "jni_util.h" | ||
| #include "commons.h" | ||
|
|
||
| TEST(CommonsTests, BasicAssertions) { | ||
| long dim = 3; | ||
| long totalNumberOfVector = 5; | ||
| std::vector<std::vector<float>> data; | ||
| for(int i = 0 ; i < totalNumberOfVector - 1 ; i++) { | ||
| std::vector<float> vector; | ||
| for(int j = 0 ; j < dim ; j ++) { | ||
| vector.push_back((float)j); | ||
| } | ||
| data.push_back(vector); | ||
| } | ||
| JNIEnv *jniEnv = nullptr; | ||
|
|
||
| testing::NiceMock<test_util::MockJNIUtil> mockJNIUtil; | ||
|
|
||
| jlong memoryAddress = knn_jni::commons::storeVectorData(&mockJNIUtil, jniEnv, (jlong)0, | ||
| reinterpret_cast<jobjectArray>(&data), (jlong)(totalNumberOfVector * dim)); | ||
| ASSERT_NE(memoryAddress, 0); | ||
| auto *vect = reinterpret_cast<std::vector<float>*>(memoryAddress); | ||
| ASSERT_EQ(vect->size(), data.size() * dim); | ||
| ASSERT_EQ(vect->capacity(), totalNumberOfVector * dim); | ||
|
|
||
| // Check by inserting more vectors at same memory location | ||
| jlong oldMemoryAddress = memoryAddress; | ||
| std::vector<std::vector<float>> data2; | ||
| std::vector<float> vector; | ||
| for(int j = 0 ; j < dim ; j ++) { | ||
| vector.push_back((float)j); | ||
| } | ||
| data2.push_back(vector); | ||
| memoryAddress = knn_jni::commons::storeVectorData(&mockJNIUtil, jniEnv, memoryAddress, | ||
| reinterpret_cast<jobjectArray>(&data2), (jlong)(totalNumberOfVector * dim)); | ||
| ASSERT_NE(memoryAddress, 0); | ||
| ASSERT_EQ(memoryAddress, oldMemoryAddress); | ||
| vect = reinterpret_cast<std::vector<float>*>(memoryAddress); | ||
| int currentIndex = 0; | ||
| ASSERT_EQ(vect->size(), totalNumberOfVector*dim); | ||
| ASSERT_EQ(vect->capacity(), totalNumberOfVector * dim); | ||
|
|
||
| // Validate if all vectors data are at correct location | ||
| for(auto & i : data) { | ||
| for(float j : i) { | ||
| ASSERT_FLOAT_EQ(vect->at(currentIndex), j); | ||
| currentIndex++; | ||
| } | ||
| } | ||
|
|
||
| for(auto & i : data2) { | ||
| for(float j : i) { | ||
| ASSERT_FLOAT_EQ(vect->at(currentIndex), j); | ||
| currentIndex++; | ||
| } | ||
| } | ||
| } |
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.