Add the interface for streaming the vectors from java to jni layer wi…

…th initial capacity

Signed-off-by: Navneet Verma <[email protected]>

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(

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 <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);
+    }
+}

jni/include/jni_util.h

@@ -69,6 +69,9 @@ namespace knn_jni {
         virtual std::vector<float> Convert2dJavaObjectArrayToCppFloatVector(JNIEnv *env, jobjectArray array2dJ,
                                                                             int dim) = 0;
 
+        virtual void Convert2dJavaObjectArrayAndStoreToFloatVector(JNIEnv *env, jobjectArray array2dJ,
+                                                                   int dim, std::vector<float> *vect ) = 0;
+
         virtual std::vector<int64_t> 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<float> *vect);
 
     private:
         std::unordered_map<std::string, jclass> cachedClasses;

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

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 <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