Adding support for FAISS JNI (#285)

.github/workflows/test-workflow.yml

@@ -9,6 +9,7 @@ on:
     branches:
       - master
       - opendistro-*
+      - faiss*
 jobs:
   build:
     strategy:

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "jni/external/nmslib"]
 	path = jni/external/nmslib
 	url = https://github.com/nmslib/nmslib.git
+[submodule "jniFaiss/external/faiss"]
+	path = jniFaiss/external/faiss
+	url = https://github.com/facebookresearch/faiss

build.gradle

@@ -137,31 +137,41 @@ loggerUsageCheck.enabled = false
 def es_tmp_dir = rootProject.file('build/private/es_tmp').absoluteFile
 es_tmp_dir.mkdirs()
 
-task cmakeJniLib(type:Exec) {
+task cmakeJniNmsLib(type:Exec) {
     workingDir 'jni'
     commandLine 'cmake', '.'
 }
-
-task buildJniLib(type:Exec) {
-    dependsOn cmakeJniLib
+task cmakeJniFaissLib(type:Exec) {
+    workingDir 'jniFaiss'
+    commandLine 'cmake', '.'
+}
+task buildJniNmsLib(type:Exec) {
+    dependsOn cmakeJniNmsLib
     workingDir 'jni'
     commandLine 'make'
 }
+task buildJniFaissLib(type:Exec) {
+    dependsOn cmakeJniFaissLib
+    workingDir 'jniFaiss'
+    commandLine 'make'
+}
 
 test {
-    dependsOn buildJniLib
+    dependsOn buildJniNmsLib
+    dependsOn buildJniFaissLib
     systemProperty 'tests.security.manager', 'false'
-    systemProperty "java.library.path", "$rootDir/jni/release"
+    systemProperty "java.library.path", "$rootDir/jniFaiss/release:$rootDir/jni/release"
 }
 
 def _numNodes = findProperty('numNodes') as Integer ?: 1
 integTest {
     if (integTestDependOnJniLib) {
-        dependsOn buildJniLib
+        dependsOn buildJniNmsLib
+        dependsOn buildJniFaissLib
     }
     systemProperty 'tests.security.manager', 'false'
     systemProperty 'java.io.tmpdir', es_tmp_dir.absolutePath
-    systemProperty "java.library.path", "$rootDir/jni/release"
+    systemProperty "java.library.path", "$rootDir/jniFaiss/release:$rootDir/jni/release"
     // allows integration test classes to access test resource from project root path
     systemProperty('project.root', project.rootDir.absolutePath)
 
@@ -204,12 +214,13 @@ testClusters.integTest {
                 debugPort += 1
             }
         }
-        systemProperty("java.library.path", "$rootDir/jni/release")
+        systemProperty("java.library.path", "$rootDir/jniFaiss/release:$rootDir/jni/release")
 }
 
 run {
     useCluster project.testClusters.integTest
-    dependsOn buildJniLib
+    dependsOn buildJniNmsLib
+    dependsOn buildJniFaissLib
     doFirst {
         // There seems to be an issue when running multi node run or integ tasks with unicast_hosts
         // not being written, the waitForAllConditions ensures it's written

jni/CMakeLists.txt

@@ -15,7 +15,7 @@
 
 cmake_minimum_required(VERSION 2.8)
 
-project(KNNIndexV2_0_11)
+project(KNNIndex_NMSLIB)
 
 # Corner case. For CMake 2.8, there is no option to specify set(CMAKE_CXX_STANDARD 11). Instead, the flag manually needs
 # to be set.
@@ -27,7 +27,7 @@ else()
 endif()
 
 # Target Library to be built
-set(KNN_INDEX KNNIndexV2_0_11)
+set(KNN_INDEX KNNIndex_NMSLIB_V2_0_11)
 set(KNN_PACKAGE_NAME opendistro-knnlib)
 
 # Check if similarity search exists
@@ -55,7 +55,7 @@ else()
 endif()
 
 # Compile the library
-add_library(${KNN_INDEX} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex.cpp)
+add_library(${KNN_INDEX} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex.cpp)
 target_link_libraries(${KNN_INDEX} NonMetricSpaceLib)
 target_include_directories(${KNN_INDEX} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/nmslib/similarity_search/include)
 

jni/src/com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex.cpp → jni/src/com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex.cpp

@@ -13,7 +13,7 @@
  *   permissions and limitations under the License.
  */
 
-#include "com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex.h"
+#include "com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex.h"
 
 #include "init.h"
 #include "index.h"
@@ -85,7 +85,7 @@ void catch_cpp_exception_and_throw_java(JNIEnv* env)
     }
 }
 
-JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex_saveIndex(JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jobjectArray algoParams, jstring spaceType)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex_saveIndex(JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jobjectArray algoParams, jstring spaceType)
 {
     Space<float>* space = NULL;
     ObjectVector dataset;
@@ -144,7 +144,7 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v201
     }
 }
 
-JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex_queryIndex(JNIEnv* env, jclass cls, jlong indexPointer, jfloatArray queryVector, jint k)
+JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex_queryIndex(JNIEnv* env, jclass cls, jlong indexPointer, jfloatArray queryVector, jint k)
 {
     try {
         IndexWrapper *indexWrapper = reinterpret_cast<IndexWrapper*>(indexPointer);
@@ -175,7 +175,7 @@ JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_in
     return NULL;
 }
 
-JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex_init(JNIEnv* env, jclass cls,  jstring indexPath, jobjectArray algoParams, jstring spaceType)
+JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex_init(JNIEnv* env, jclass cls,  jstring indexPath, jobjectArray algoParams, jstring spaceType)
 {
     IndexWrapper *indexWrapper = NULL;
     try {
@@ -215,7 +215,7 @@ JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v20
     return NULL;
 }
 
-JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex_gc(JNIEnv* env, jclass cls,  jlong indexPointer)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex_gc(JNIEnv* env, jclass cls,  jlong indexPointer)
 {
     try {
         IndexWrapper *indexWrapper = reinterpret_cast<IndexWrapper*>(indexPointer);
@@ -228,7 +228,7 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v201
     }
 }
 
-JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v2011_KNNIndex_initLibrary(JNIEnv *, jclass)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_nmslib_v2011_KNNNmsLibIndex_initLibrary(JNIEnv *, jclass)
 {
     initLibrary();
 

jniFaiss/CMakeLists.txt

@@ -0,0 +1,108 @@
+cmake_minimum_required(VERSION 2.8)
+
+project(KNNIndex_FAISS)
+
+# Corner case. For CMake 2.8, there is no option to specify set(CMAKE_CXX_STANDARD 11). Instead, the flag manually needs
+# to be set.
+if (CMAKE_VERSION VERSION_LESS "3.1")
+    set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
+else()
+    set(CMAKE_CXX_STANDARD 11)
+    set(CMAKE_CXX_STANDARD_REQUIRED True)
+endif()
+
+# Target Library to be built
+set(KNN_INDEX KNNIndex_FAISS_V1_6_5)
+set(KNN_PACKAGE_NAME opendistro-knnlib)
+
+# --- FAISS BEGIN ---
+if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
+    if(CMAKE_C_COMPILER_ID MATCHES "Clang\$")
+        set(OpenMP_C_FLAGS "-Xpreprocessor -fopenmp")
+        set(OpenMP_C_LIB_NAMES "omp")
+        set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
+    endif()
+
+    if(CMAKE_CXX_COMPILER_ID MATCHES "Clang\$")
+        set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp -I/usr/local/opt/libomp/include")
+        set(OpenMP_CXX_LIB_NAMES "omp")
+        set(OpenMP_omp_LIBRARY /usr/local/opt/libomp/lib/libomp.dylib)
+    endif()
+endif()
+
+find_package(OpenMP REQUIRED)
+find_package(ZLIB REQUIRED)
+find_package(BLAS REQUIRED)
+find_package(LAPACK REQUIRED)
+
+# Check if faiss search exists
+find_path(FAISS_REPO_DIR NAMES faiss PATHS ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss)
+
+# If not, pull the updated submodule
+if (NOT EXISTS ${FAISS_REPO_DIR})
+    message(STATUS "Could not find faiss. Pulling updated submodule.")
+    execute_process(COMMAND git submodule update --init -- external/faiss WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+endif ()
+
+# Add the subdirectory so it is possible to use its targets
+set(FAISS_ENABLE_GPU OFF)
+set(FAISS_ENABLE_PYTHON OFF)
+add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/external/faiss EXCLUDE_FROM_ALL)
+# --- FAISS END
+
+# Set OS specific variables
+if (${CMAKE_SYSTEM_NAME} STREQUAL Darwin)
+    set(CMAKE_MACOSX_RPATH 1)
+    set(JVM_OS_TYPE darwin)
+    set(LIB_EXT .jnilib)
+elseif(${CMAKE_SYSTEM_NAME} STREQUAL Linux)
+    set(JVM_OS_TYPE linux)
+    set(LIB_EXT .so)
+else()
+    message( FATAL_ERROR "Unable to run on system: ${CMAKE_SYSTEM_NAME}")
+endif()
+
+# Compile the library
+include_directories($ENV{JAVA_HOME}/include )
+add_library(${KNN_INDEX} SHARED ${CMAKE_CURRENT_SOURCE_DIR}/src/com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex.cpp)
+target_link_libraries(${KNN_INDEX} faiss ${BLAS_LIBRARIES} ${LAPACK_LIBRARIES} OpenMP::OpenMP_CXX)
+target_include_directories(${KNN_INDEX} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include $ENV{JAVA_HOME}/include $ENV{JAVA_HOME}/include/${JVM_OS_TYPE} ${CMAKE_CURRENT_SOURCE_DIR}/external/faiss/faiss )
+
+set_target_properties(${KNN_INDEX} PROPERTIES SUFFIX ${LIB_EXT})
+set_target_properties(${KNN_INDEX} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+set_target_properties(${KNN_INDEX} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
+
+# Installation rules for shared library
+install(TARGETS ${KNN_INDEX}
+        LIBRARY DESTINATION lib
+        COMPONENT library)
+
+# CPack section to build artifacts
+set(KNN_MAINTAINER "OpenDistro for Elasticsearch Team <[email protected]>")
+set(ODFE_DOWNLOAD_URL "https://opendistro.github.io/elasticsearch/downloads")
+set(CPACK_PACKAGE_NAME ${KNN_PACKAGE_NAME})
+set(CPACK_PACKAGE_VERSION 1.12.0.0)
+set(CMAKE_INSTALL_PREFIX /usr)
+set(CPACK_GENERATOR "RPM;DEB")
+SET(CPACK_OUTPUT_FILE_PREFIX packages)
+set(CPACK_PACKAGE_RELEASE 1)
+set(CPACK_PACKAGE_DESCRIPTION_SUMMARY "KNN JNI library built off of Faiss for OpenDistro for Elasticsearch. Reference documentation can be found at https://opendistro.github.io/for-elasticsearch-docs/.")
+set(CPACK_PACKAGE_VENDOR "Amazon")
+set(CPACK_PACKAGE_CONTACT "Maintainer: ${KNN_MAINTAINER}")
+set(CPACK_PACKAGING_INSTALL_PREFIX ${CMAKE_INSTALL_PREFIX})
+set(CPACK_PACKAGE_FILE_NAME "${CPACK_PACKAGE_NAME}-${CPACK_PACKAGE_VERSION}-${CPACK_PACKAGE_RELEASE}_${JVM_OS_TYPE}.${CMAKE_SYSTEM_PROCESSOR}")
+
+# RPM Specific variables
+set(CPACK_RPM_PACKAGE_RELEASE ${CPACK_PACKAGE_RELEASE})
+set(CPACK_RPM_PACKAGE_URL ${ODFE_DOWNLOAD_URL})
+set(CPACK_RPM_PACKAGE_DESCRIPTION "Open Distro for Elasticsearch KNN Faiss JNI Library")
+set(CPACK_RPM_PACKAGE_LICENSE "ASL-2.0")
+
+# DEB Specific variables
+set(CPACK_DEBIAN_PACKAGE_HOMEPAGE ${ODFE_DOWNLOAD_URL})
+set(CPACK_DEBIAN_PACKAGE_MAINTAINER ${KNN_MAINTAINER})
+set(CPACK_DEBIAN_PACKAGE_SOURCE ${CPACK_PACKAGE_NAME})
+set(CPACK_DEBIAN_PACKAGE_VERSION ${CPACK_PACKAGE_VERSION})
+set(CPACK_DEBIAN_PACKAGE_SECTION "libs")
+
+include(CPack)