Adds ef_search as a query parameter for Lucene, FAISS and NMSLIB (ope…

…nsearch-project#1783) Currently ef_search is set at index level, this change gives the ability to have query time ef-search parameter without manipulating the index settings. query time value supersedes the value from index in FAISS and NMSLIB. For Lucene, max of k and ef_search is used as ef_search value Signed-off-by: Tejas Shah <[email protected]> (cherry picked from commit 989ad7d)
shatejas · Jul 3, 2024 · 17e9d89 · 17e9d89
1 parent 4ab5339
commit 17e9d89
Show file tree

Hide file tree

Showing 63 changed files with 2,928 additions and 646 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.15...2.x)
 ### Features
+* Adds dynamic query parameter ef_search [#1783](https://github.com/opensearch-project/k-NN/pull/1783)
 ### Enhancements
 ### Bug Fixes
 ### Infrastructure

diff --git a/jni/CMakeLists.txt b/jni/CMakeLists.txt
@@ -147,8 +147,10 @@ if ("${WIN32}" STREQUAL "")
         add_executable(
                 jni_test
                 tests/faiss_wrapper_test.cpp
+                tests/faiss_wrapper_unit_test.cpp
                 tests/faiss_util_test.cpp
                 tests/nmslib_wrapper_test.cpp
+                tests/nmslib_wrapper_unit_test.cpp
                 tests/test_util.cpp
                 tests/commons_test.cpp
                 )

diff --git a/jni/cmake/init-nmslib.cmake b/jni/cmake/init-nmslib.cmake
@@ -13,12 +13,13 @@ if (NOT EXISTS ${NMS_REPO_DIR})
 endif ()
 
 # Check if patch exist, this is to skip git apply during CI build. See CI.yml with ubuntu.
-find_path(PATCH_FILE NAMES 0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch PATHS ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib NO_DEFAULT_PATH)
+find_path(PATCH_FILE NAMES 0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch 0002-Adds-ability-to-pass-ef-parameter-in-the-query-for-h.patch PATHS ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib NO_DEFAULT_PATH)
 
 # If it exists, apply patches
 if (EXISTS ${PATCH_FILE})
     message(STATUS "Applying custom patches.")
     execute_process(COMMAND git ${GIT_PATCH_COMMAND} --3way --ignore-space-change --ignore-whitespace ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0001-Initialize-maxlevel-during-add-from-enterpoint-level.patch WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
+    execute_process(COMMAND git ${GIT_PATCH_COMMAND} --3way --ignore-space-change --ignore-whitespace ${CMAKE_CURRENT_SOURCE_DIR}/patches/nmslib/0002-Adds-ability-to-pass-ef-parameter-in-the-query-for-h.patch WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib ERROR_VARIABLE ERROR_MSG RESULT_VARIABLE RESULT_CODE)
 
     if(RESULT_CODE)
         message(FATAL_ERROR "Failed to apply patch:\n${ERROR_MSG}")

diff --git a/jni/include/commons.h b/jni/include/commons.h
@@ -33,5 +33,10 @@ namespace knn_jni {
          * @param memoryAddress address to be freed.
          */
         void freeVectorData(jlong);
+
+        /**
+         * Extracts query time efSearch from method parameters
+         **/
+        int getIntegerMethodParameter(JNIEnv *, knn_jni::JNIUtilInterface *, std::unordered_map<std::string, jobject>, std::string, int);
     }
 }
diff --git a/jni/include/faiss_wrapper.h b/jni/include/faiss_wrapper.h
@@ -45,17 +45,27 @@ namespace knn_jni {
         // Sets the sharedIndexState for an index
         void SetSharedIndexState(jlong indexPointerJ, jlong shareIndexStatePointerJ);
 
-        // Execute a query against the index located in memory at indexPointerJ.
-        //
-        // Return an array of KNNQueryResults
+         /**
+         *  Execute a query against the index located in memory at indexPointerJ
+         *  
+         * Parameters:
+         * methodParamsJ: introduces a map to have additional method parameters
+         * 
+         * Return an array of KNNQueryResults
+        */
         jobjectArray QueryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
-                                jfloatArray queryVectorJ, jint kJ, jintArray parentIdsJ);
-
-        // Execute a query against the index located in memory at indexPointerJ along with Filters
-        //
-        // Return an array of KNNQueryResults
+                                jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ, jintArray parentIdsJ);
+
+        /**
+         *  Execute a query against the index located in memory at indexPointerJ along with Filters
+         *  
+         * Parameters:
+         * methodParamsJ: introduces a map to have additional method parameters
+         * 
+         * Return an array of KNNQueryResults
+        */
         jobjectArray QueryIndex_WithFilter(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
-                                                                jfloatArray queryVectorJ, jint kJ, jlongArray filterIdsJ,
+                                                                jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ, jlongArray filterIdsJ,
                                                                 jint filterIdsTypeJ, jintArray parentIdsJ);
 
         // Free the index located in memory at indexPointerJ

diff --git a/jni/include/nmslib_wrapper.h b/jni/include/nmslib_wrapper.h
@@ -37,7 +37,7 @@ namespace knn_jni {
         //
         // Return an array of KNNQueryResults
         jobjectArray QueryIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jlong indexPointerJ,
-                                jfloatArray queryVectorJ, jint kJ);
+                                jfloatArray queryVectorJ, jint kJ, jobject methodParamsJ);
 
         // Free the index located in memory at indexPointerJ
         void Free(jlong indexPointer);

diff --git a/jni/include/org_opensearch_knn_jni_FaissService.h b/jni/include/org_opensearch_knn_jni_FaissService.h
@@ -69,18 +69,18 @@ JNIEXPORT void JNICALL Java_org_opensearch_knn_jni_FaissService_setSharedIndexSt
 /*
  * Class:     org_opensearch_knn_jni_FaissService
  * Method:    queryIndex
- * Signature: (J[FI[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
+ * Signature: (J[FI[Ljava/util/MapI)[Lorg/opensearch/knn/index/query/KNNQueryResult;
  */
 JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_queryIndex
-  (JNIEnv *, jclass, jlong, jfloatArray, jint, jintArray);
+  (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject, jintArray);
 
 /*
  * Class:     org_opensearch_knn_jni_FaissService
  * Method:    queryIndexWithFilter
- * Signature: (J[FI[JI[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
+ * Signature: (J[FI[JLjava/util/MapI[I)[Lorg/opensearch/knn/index/query/KNNQueryResult;
  */
 JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_FaissService_queryIndexWithFilter
-  (JNIEnv *, jclass, jlong, jfloatArray, jint, jlongArray, jint, jintArray);
+  (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject, jlongArray, jint, jintArray);
 
 /*
  * Class:     org_opensearch_knn_jni_FaissService

diff --git a/jni/include/org_opensearch_knn_jni_NmslibService.h b/jni/include/org_opensearch_knn_jni_NmslibService.h
@@ -40,7 +40,7 @@ JNIEXPORT jlong JNICALL Java_org_opensearch_knn_jni_NmslibService_loadIndex
  * Signature: (J[FI)[Lorg/opensearch/knn/index/query/KNNQueryResult;
  */
 JNIEXPORT jobjectArray JNICALL Java_org_opensearch_knn_jni_NmslibService_queryIndex
-  (JNIEnv *, jclass, jlong, jfloatArray, jint);
+  (JNIEnv *, jclass, jlong, jfloatArray, jint, jobject);
 
 /*
  * Class:     org_opensearch_knn_jni_NmslibService