From fccc5a965bbe8d315017a4d69884240f5bbc401d Mon Sep 17 00:00:00 2001 From: Navneet Verma Date: Tue, 2 Apr 2024 11:25:33 -0700 Subject: [PATCH] Add the interface for streaming the vectors from java to jni layer with initial capacity (#1586) Signed-off-by: Navneet Verma --- jni/CMakeLists.txt | 3 +- jni/include/commons.h | 37 ++++++++++ jni/include/jni_util.h | 4 + .../org_opensearch_knn_jni_FaissService.h | 8 -- .../org_opensearch_knn_jni_JNICommons.h | 40 ++++++++++ jni/src/commons.cpp | 41 +++++++++++ jni/src/jni_util.cpp | 11 ++- .../org_opensearch_knn_jni_FaissService.cpp | 19 ----- jni/src/org_opensearch_knn_jni_JNICommons.cpp | 60 +++++++++++++++ jni/tests/commons_test.cpp | 73 +++++++++++++++++++ jni/tests/test_util.cpp | 8 ++ jni/tests/test_util.h | 2 + .../knn/TransferVectorsBenchmarks.java | 24 ++++-- .../opensearch/knn/common/KNNConstants.java | 3 + .../org/opensearch/knn/jni/FaissService.java | 19 ++--- .../org/opensearch/knn/jni/JNICommons.java | 62 ++++++++++++++++ .../org/opensearch/knn/jni/JNIService.java | 30 ++------ .../plugin-metadata/plugin-security.policy | 1 + .../opensearch/knn/jni/JNICommonsTest.java | 40 ++++++++++ 19 files changed, 412 insertions(+), 73 deletions(-) create mode 100644 jni/include/commons.h create mode 100644 jni/include/org_opensearch_knn_jni_JNICommons.h create mode 100644 jni/src/commons.cpp create mode 100644 jni/src/org_opensearch_knn_jni_JNICommons.cpp create mode 100644 jni/tests/commons_test.cpp create mode 100644 src/main/java/org/opensearch/knn/jni/JNICommons.java create mode 100644 src/test/java/org/opensearch/knn/jni/JNICommonsTest.java diff --git a/jni/CMakeLists.txt b/jni/CMakeLists.txt index 60321ed1b..4f32c87b9 100644 --- a/jni/CMakeLists.txt +++ b/jni/CMakeLists.txt @@ -61,7 +61,7 @@ endif() # ---------------------------------------------------------------------------- # ---------------------------------- COMMON ---------------------------------- -add_library(${TARGET_LIB_COMMON} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/jni_util.cpp) +add_library(${TARGET_LIB_COMMON} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/jni_util.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/org_opensearch_knn_jni_JNICommons.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/commons.cpp) target_include_directories(${TARGET_LIB_COMMON} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE}) set_target_properties(${TARGET_LIB_COMMON} PROPERTIES SUFFIX ${LIB_EXT}) set_target_properties(${TARGET_LIB_COMMON} PROPERTIES POSITION_INDEPENDENT_CODE ON) @@ -236,6 +236,7 @@ if ("${WIN32}" STREQUAL "") tests/faiss_util_test.cpp tests/nmslib_wrapper_test.cpp tests/test_util.cpp + tests/commons_test.cpp ) target_link_libraries( diff --git a/jni/include/commons.h b/jni/include/commons.h new file mode 100644 index 000000000..05367a693 --- /dev/null +++ b/jni/include/commons.h @@ -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 +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); + } +} diff --git a/jni/include/jni_util.h b/jni/include/jni_util.h index 52b08a202..b3d55f1c1 100644 --- a/jni/include/jni_util.h +++ b/jni/include/jni_util.h @@ -69,6 +69,9 @@ namespace knn_jni { virtual std::vector Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim) = 0; + virtual void Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ, + int dim, std::vector *vect ) = 0; + virtual std::vector ConvertJavaIntArrayToCppIntVector(JNIEnv *env, jintArray arrayJ) = 0; // -------------------------------------------------------------------------- @@ -164,6 +167,7 @@ namespace knn_jni { void ReleaseLongArrayElements(JNIEnv *env, jlongArray array, jlong *elems, jint mode); void SetObjectArrayElement(JNIEnv *env, jobjectArray array, jsize index, jobject val); void SetByteArrayRegion(JNIEnv *env, jbyteArray array, jsize start, jsize len, const jbyte * buf); + void Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim, std::vector *vect); private: std::unordered_map cachedClasses; diff --git a/jni/include/org_opensearch_knn_jni_FaissService.h b/jni/include/org_opensearch_knn_jni_FaissService.h index ec1f46bc3..64a858f84 100644 --- a/jni/include/org_opensearch_knn_jni_FaissService.h +++ b/jni/include/org_opensearch_knn_jni_FaissService.h @@ -122,14 +122,6 @@ JNIEXPORT jbyteArray JNICALL Java_org_opensearch_knn_jni_FaissService_trainIndex JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors (JNIEnv *, jclass, jlong, jobjectArray); -/* - * Class: org_opensearch_knn_jni_FaissService - * Method: transferVectorsV2 - * Signature: (J[[F)J - */ -JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV2 - (JNIEnv *, jclass, jlong, jobjectArray); - /* * Class: org_opensearch_knn_jni_FaissService * Method: freeVectors diff --git a/jni/include/org_opensearch_knn_jni_JNICommons.h b/jni/include/org_opensearch_knn_jni_JNICommons.h new file mode 100644 index 000000000..d0758d7c8 --- /dev/null +++ b/jni/include/org_opensearch_knn_jni_JNICommons.h @@ -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 +/* 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 diff --git a/jni/src/commons.cpp b/jni/src/commons.cpp new file mode 100644 index 000000000..3c03ac49d --- /dev/null +++ b/jni/src/commons.cpp @@ -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 + +#include + +#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 *vect; + if ((long) memoryAddressJ == 0) { + vect = new std::vector(); + vect->reserve((long)initialCapacityJ); + } else { + vect = reinterpret_cast*>(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*>(memoryAddressJ); + delete vect; + } +} +#endif //OPENSEARCH_KNN_COMMONS_H \ No newline at end of file diff --git a/jni/src/jni_util.cpp b/jni/src/jni_util.cpp index a0c1d5733..a1faa4894 100644 --- a/jni/src/jni_util.cpp +++ b/jni/src/jni_util.cpp @@ -223,6 +223,13 @@ int knn_jni::JNIUtil::ConvertJavaObjectToCppInteger(JNIEnv *env, jobject objectJ std::vector knn_jni::JNIUtil::Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ, int dim) { + std::vector vect; + Convert2dJavaObjectArrayAndStoreToFloatVector(env, array2dJ, dim, &vect); + return vect; +} + +void knn_jni::JNIUtil::Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ, + int dim, std::vector *vect) { if (array2dJ == nullptr) { throw std::runtime_error("Array cannot be null"); @@ -231,7 +238,6 @@ std::vector knn_jni::JNIUtil::Convert2dJavaObjectArrayToCppFloatVector(JN int numVectors = env->GetArrayLength(array2dJ); this->HasExceptionInStack(env); - std::vector floatVectorCpp; for (int i = 0; i < numVectors; ++i) { auto vectorArray = (jfloatArray)env->GetObjectArrayElement(array2dJ, i); this->HasExceptionInStack(env, "Unable to get object array element"); @@ -247,13 +253,12 @@ std::vector knn_jni::JNIUtil::Convert2dJavaObjectArrayToCppFloatVector(JN } for(int j = 0; j < dim; ++j) { - floatVectorCpp.push_back(vector[j]); + vect->push_back(vector[j]); } env->ReleaseFloatArrayElements(vectorArray, vector, JNI_ABORT); } this->HasExceptionInStack(env); env->DeleteLocalRef(array2dJ); - return floatVectorCpp; } std::vector knn_jni::JNIUtil::ConvertJavaIntArrayToCppIntVector(JNIEnv *env, jintArray arrayJ) { diff --git a/jni/src/org_opensearch_knn_jni_FaissService.cpp b/jni/src/org_opensearch_knn_jni_FaissService.cpp index 3d9624c25..ae9d461ff 100644 --- a/jni/src/org_opensearch_knn_jni_FaissService.cpp +++ b/jni/src/org_opensearch_knn_jni_FaissService.cpp @@ -13,7 +13,6 @@ #include -#include #include #include "faiss_wrapper.h" @@ -191,24 +190,6 @@ JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectors return (jlong) vect; } -JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_FaissService_transferVectorsV2(JNIEnv * env, jclass cls, -jlong vectorsPointerJ, - jobjectArray vectorsJ) -{ - std::vector *vect; - if ((long) vectorsPointerJ == 0) { - vect = new std::vector; - } else { - vect = reinterpret_cast*>(vectorsPointerJ); - } - - int dim = jniUtil.GetInnerDimensionOf2dJavaFloatArray(env, vectorsJ); - auto dataset = jniUtil.Convert2dJavaObjectArrayToCppFloatVector(env, vectorsJ, dim); - vect->insert(vect->end(), dataset.begin(), dataset.end()); - - return (jlong) vect; -} - JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_freeVectors(JNIEnv * env, jclass cls, jlong vectorsPointerJ) { diff --git a/jni/src/org_opensearch_knn_jni_JNICommons.cpp b/jni/src/org_opensearch_knn_jni_JNICommons.cpp new file mode 100644 index 000000000..ccdd11882 --- /dev/null +++ b/jni/src/org_opensearch_knn_jni_JNICommons.cpp @@ -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 +#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); + } +} diff --git a/jni/tests/commons_test.cpp b/jni/tests/commons_test.cpp new file mode 100644 index 000000000..09323f0fb --- /dev/null +++ b/jni/tests/commons_test.cpp @@ -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 +#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> data; + for(int i = 0 ; i < totalNumberOfVector - 1 ; i++) { + std::vector vector; + for(int j = 0 ; j < dim ; j ++) { + vector.push_back((float)j); + } + data.push_back(vector); + } + JNIEnv *jniEnv = nullptr; + + testing::NiceMock mockJNIUtil; + + jlong memoryAddress = knn_jni::commons::storeVectorData(&mockJNIUtil, jniEnv, (jlong)0, + reinterpret_cast(&data), (jlong)(totalNumberOfVector * dim)); + ASSERT_NE(memoryAddress, 0); + auto *vect = reinterpret_cast*>(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> data2; + std::vector 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(&data2), (jlong)(totalNumberOfVector * dim)); + ASSERT_NE(memoryAddress, 0); + ASSERT_EQ(memoryAddress, oldMemoryAddress); + vect = reinterpret_cast*>(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++; + } + } +} diff --git a/jni/tests/test_util.cpp b/jni/tests/test_util.cpp index 89b19f9aa..92532b9e2 100644 --- a/jni/tests/test_util.cpp +++ b/jni/tests/test_util.cpp @@ -45,6 +45,14 @@ test_util::MockJNIUtil::MockJNIUtil() { return data; }); + ON_CALL(*this, Convert2dJavaObjectArrayAndStoreToFloatVector) + .WillByDefault([this](JNIEnv *env, jobjectArray array2dJ, int dim, std::vector* data) { + for (const auto &v : + (*reinterpret_cast> *>(array2dJ))) + for (auto item : v) data->push_back(item); + }); + + // arrayJ is re-interpreted as std::vector * ON_CALL(*this, ConvertJavaIntArrayToCppIntVector) .WillByDefault([this](JNIEnv *env, jintArray arrayJ) { diff --git a/jni/tests/test_util.h b/jni/tests/test_util.h index 1e32ad3c3..8e73a8ab0 100644 --- a/jni/tests/test_util.h +++ b/jni/tests/test_util.h @@ -44,6 +44,8 @@ namespace test_util { // TODO: Figure out why this cant use "new" MOCK_METHOD MOCK_METHOD(std::vector, Convert2dJavaObjectArrayToCppFloatVector, (JNIEnv * env, jobjectArray array2dJ, int dim)); + MOCK_METHOD(void, Convert2dJavaObjectArrayAndStoreToFloatVector, + (JNIEnv * env, jobjectArray array2dJ, int dim, std::vector*vect)); MOCK_METHOD(std::vector, ConvertJavaIntArrayToCppIntVector, (JNIEnv * env, jintArray arrayJ)); MOCK_METHOD2(ConvertJavaMapToCppMap, diff --git a/micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java b/micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java index ad1076484..2bce54ee6 100644 --- a/micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java +++ b/micro-benchmarks/src/main/java/org/opensearch/knn/TransferVectorsBenchmarks.java @@ -23,7 +23,7 @@ import org.openjdk.jmh.annotations.Setup; import org.openjdk.jmh.annotations.State; import org.openjdk.jmh.annotations.Warmup; -import org.opensearch.knn.jni.JNIService; +import org.opensearch.knn.jni.JNICommons; import java.util.ArrayList; import java.util.List; @@ -42,9 +42,9 @@ @State(Scope.Benchmark) public class TransferVectorsBenchmarks { private static final Random random = new Random(1212121212); - private static final int TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED = 1000000; + private static final long TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED = 1000000; - @Param({ "128", "256", "384", "512" }) + @Param({ "128", "256", "384", "512", "960", "1024", "1536" }) private int dimension; @Param({ "100000", "500000", "1000000" }) @@ -61,20 +61,30 @@ public void setup() { } @Benchmark - public void transferVectors() { + public void transferVectors_withCapacity() { long vectorsAddress = 0; List vectorToTransfer = new ArrayList<>(); + long startingIndex = 0; for (float[] floats : vectorList) { if (vectorToTransfer.size() == vectorsPerTransfer) { - vectorsAddress = JNIService.transferVectorsV2(vectorsAddress, vectorToTransfer.toArray(new float[][] {})); + vectorsAddress = JNICommons.storeVectorData( + vectorsAddress, + vectorToTransfer.toArray(new float[][] {}), + dimension * TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED + ); + startingIndex += vectorsPerTransfer; vectorToTransfer = new ArrayList<>(); } vectorToTransfer.add(floats); } if (!vectorToTransfer.isEmpty()) { - vectorsAddress = JNIService.transferVectorsV2(vectorsAddress, vectorToTransfer.toArray(new float[][] {})); + vectorsAddress = JNICommons.storeVectorData( + vectorsAddress, + vectorToTransfer.toArray(new float[][] {}), + dimension * TOTAL_NUMBER_OF_VECTOR_TO_BE_TRANSFERRED + ); } - JNIService.freeVectors(vectorsAddress); + JNICommons.freeVectorData(vectorsAddress); } private float[] generateRandomVector(int dimensions) { diff --git a/src/main/java/org/opensearch/knn/common/KNNConstants.java b/src/main/java/org/opensearch/knn/common/KNNConstants.java index 34805b7e5..98622ee85 100644 --- a/src/main/java/org/opensearch/knn/common/KNNConstants.java +++ b/src/main/java/org/opensearch/knn/common/KNNConstants.java @@ -72,6 +72,7 @@ public class KNNConstants { // nmslib specific constants public static final String NMSLIB_NAME = "nmslib"; + public static final String COMMONS_NAME = "common"; public static final String SPACE_TYPE = "spaceType"; // used as field info key public static final String HNSW_ALGO_M = "M"; public static final String HNSW_ALGO_EF_CONSTRUCTION = "efConstruction"; @@ -121,6 +122,8 @@ public class KNNConstants { public static final String FAISS_AVX2_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + FAISS_NAME + "_avx2"; public static final String NMSLIB_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + NMSLIB_NAME; + public static final String COMMON_JNI_LIBRARY_NAME = JNI_LIBRARY_PREFIX + COMMONS_NAME; + // API Constants public static final String CLEAR_CACHE = "clear_cache"; diff --git a/src/main/java/org/opensearch/knn/jni/FaissService.java b/src/main/java/org/opensearch/knn/jni/FaissService.java index 4b5045359..517945968 100644 --- a/src/main/java/org/opensearch/knn/jni/FaissService.java +++ b/src/main/java/org/opensearch/knn/jni/FaissService.java @@ -173,33 +173,26 @@ public static native KNNQueryResult[] queryIndexWithFilter( public static native byte[] trainIndex(Map indexParameters, int dimension, long trainVectorsPointer); /** + *

+ * The function is deprecated. Use {@link JNICommons#storeVectorData(long, float[][], long)} + *

* Transfer vectors from Java to native * * @param vectorsPointer pointer to vectors in native memory. Should be 0 to create vector as well * @param trainingData data to be transferred * @return pointer to native memory location of training data */ + @Deprecated(since = "2.14.0", forRemoval = true) public static native long transferVectors(long vectorsPointer, float[][] trainingData); /** - * Transfer vectors from Java to native layer. This is the version 2 of transfer vector functionality. The - * difference between this and the version 1 is, this version puts vectors at the end rather than in front. - * Keeping this name as V2 for now, will come up with better name going forward. *

- * TODO: Rename the function - *
- * TODO: Make this function native function and use a common cpp file to host these functions. + * The function is deprecated. Use {@link JNICommons#freeVectorData(long)} *

- * @param vectorsPointer pointer to vectors in native memory. Should be 0 to create vector as well - * @param data data to be transferred - * @return pointer to native memory location for data - */ - public static native long transferVectorsV2(long vectorsPointer, float[][] data); - - /** * Free vectors from memory * * @param vectorsPointer to be freed */ + @Deprecated(since = "2.14.0", forRemoval = true) public static native void freeVectors(long vectorsPointer); } diff --git a/src/main/java/org/opensearch/knn/jni/JNICommons.java b/src/main/java/org/opensearch/knn/jni/JNICommons.java new file mode 100644 index 000000000..90ad70c3d --- /dev/null +++ b/src/main/java/org/opensearch/knn/jni/JNICommons.java @@ -0,0 +1,62 @@ +/* + * 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. + */ + +package org.opensearch.knn.jni; + +import org.opensearch.knn.common.KNNConstants; + +import java.security.AccessController; +import java.security.PrivilegedAction; + +/** + * Common class for providing the JNI related functionality to various JNIServices. + */ +public class JNICommons { + + static { + AccessController.doPrivileged((PrivilegedAction) () -> { + System.loadLibrary(KNNConstants.COMMON_JNI_LIBRARY_NAME); + return null; + }); + } + + /** + * 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. + * + *

+ * The function is not threadsafe. If multiple threads are trying to insert on same memory location, then it can + * lead to data corruption. + *

+ * + * @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. + */ + public static native long storeVectorData(long memoryAddress, float[][] data, long initialCapacity); + + /** + * 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)} + * + *

+ * The function is not threadsafe. If multiple threads are trying to free up same memory location, then it can + * lead to errors. + *

+ * + * @param memoryAddress address to be freed. + */ + public static native void freeVectorData(long memoryAddress); +} diff --git a/src/main/java/org/opensearch/knn/jni/JNIService.java b/src/main/java/org/opensearch/knn/jni/JNIService.java index 80b56b173..a85089517 100644 --- a/src/main/java/org/opensearch/knn/jni/JNIService.java +++ b/src/main/java/org/opensearch/knn/jni/JNIService.java @@ -235,44 +235,30 @@ public static byte[] trainIndex(Map indexParameters, int dimensi } /** + *

+ * The function is deprecated. Use {@link JNICommons#storeVectorData(long, float[][], long)} + *

* Transfer vectors from Java to native * * @param vectorsPointer pointer to vectors in native memory. Should be 0 to create vector as well * @param trainingData data to be transferred * @return pointer to native memory location of training data */ + @Deprecated(since = "2.14.0", forRemoval = true) public static long transferVectors(long vectorsPointer, float[][] trainingData) { return FaissService.transferVectors(vectorsPointer, trainingData); } /** + *

+ * The function is deprecated. Use {@link JNICommons#freeVectorData(long)} + *

* Free vectors from memory * * @param vectorsPointer to be freed */ + @Deprecated(since = "2.14.0", forRemoval = true) public static void freeVectors(long vectorsPointer) { FaissService.freeVectors(vectorsPointer); } - - /** - * Experimental: Transfer vectors from Java to native layer. This is the version 2 of transfer vector - * functionality. The difference between this and the version 1 is, this version puts vectors at the end rather - * than in front. Keeping this name as V2 for now, will come up with better name going forward. - *

- * This is not a production ready function for now. Adding this to ensure that we are able to run atleast 1 - * micro-benchmarks. - *

- *

- * TODO: Rename the function - *
- * TODO: Make this function native function and use a common cpp file to host these functions. - *

- * @param vectorsPointer pointer to vectors in native memory. Should be 0 to create vector as well - * @param data data to be transferred - * @return pointer to native memory location for data - * - */ - public static long transferVectorsV2(long vectorsPointer, float[][] data) { - return FaissService.transferVectorsV2(vectorsPointer, data); - } } diff --git a/src/main/plugin-metadata/plugin-security.policy b/src/main/plugin-metadata/plugin-security.policy index 91624613c..d5ab0be21 100644 --- a/src/main/plugin-metadata/plugin-security.policy +++ b/src/main/plugin-metadata/plugin-security.policy @@ -1,6 +1,7 @@ grant { permission java.lang.RuntimePermission "loadLibrary.opensearchknn_nmslib"; permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss"; + permission java.lang.RuntimePermission "loadLibrary.opensearchknn_common"; permission java.lang.RuntimePermission "loadLibrary.opensearchknn_faiss_avx2"; permission java.net.SocketPermission "*", "connect,resolve"; permission java.lang.RuntimePermission "accessDeclaredMembers"; diff --git a/src/test/java/org/opensearch/knn/jni/JNICommonsTest.java b/src/test/java/org/opensearch/knn/jni/JNICommonsTest.java new file mode 100644 index 000000000..bf27458b0 --- /dev/null +++ b/src/test/java/org/opensearch/knn/jni/JNICommonsTest.java @@ -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. + */ + +package org.opensearch.knn.jni; + +import org.opensearch.knn.KNNTestCase; + +public class JNICommonsTest extends KNNTestCase { + + public void testStoreVectorData_whenVaildInputThenSuccess() { + float[][] data = new float[2][2]; + for (int i = 0; i < 2; i++) { + for (int j = 0; j < 2; j++) { + data[i][j] = i + j; + } + } + long memoryAddress = JNICommons.storeVectorData(0, data, 8); + assertTrue(memoryAddress > 0); + assertEquals(memoryAddress, JNICommons.storeVectorData(memoryAddress, data, 8)); + } + + public void testFreeVectorData_whenValidInput_ThenSuccess() { + float[][] data = new float[2][2]; + for (int i = 0; i < 2; i++) { + for (int j = 0; j < 2; j++) { + data[i][j] = i + j; + } + } + long memoryAddress = JNICommons.storeVectorData(0, data, 8); + JNICommons.freeVectorData(memoryAddress); + } +}