faiss interface refactoring to support multiple methods

jni/faiss/src/com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex.cpp

@@ -18,150 +18,222 @@
 
 #include <cmath>
 #include <cstdio>
-#include <cstdlib>
 #include <string>
 #include <vector>
-#include <sys/time.h>
-#include <omp.h>
 
 #include "faiss/index_factory.h"
 #include "faiss/MetaIndexes.h"
 #include "faiss/index_io.h"
 #include "faiss/IndexHNSW.h"
+#include "faiss/IndexIVFFlat.h"
 
 
 using std::string;
 using std::vector;
 
+// mapMetric is used to map a string from the plugin to a faiss metric. All translation should be done via this map
 std::unordered_map<string, faiss::MetricType> mapMetric = {
 	{"l2", faiss::METRIC_L2},
 	{"innerproduct", faiss::METRIC_INNER_PRODUCT}
 };
 
+void TrainIndex(faiss::Index * index, faiss::Index::idx_t n, const float* x) {
+    if (auto * indexIvf = dynamic_cast<faiss::IndexIVF*>(index)) {
+        if (indexIvf->quantizer_trains_alone == 2) {
+            TrainIndex(indexIvf->quantizer, n, x);
+        }
+        indexIvf->make_direct_map();
+    }
+
+    if (!index->is_trained) {
+        index->train(n, x);
+    }
+}
+
+void SetExtraParameters(JNIEnv *env, jobject parameterMap, faiss::Index * index) {
+    // Here, we parse parameterMap, which is a java Map<String, Object>. In order to implement this, I referred to
+    // https://stackoverflow.com/questions/4844022/jni-create-hashmap
+
+    // Load all of the class and methods to iterate over a map
+    jclass mapClass = knn_jni::FindClass(env, "java/util/Map");
+    jmethodID entrySet = knn_jni::FindMethod(env, mapClass, "entrySet", "()Ljava/util/Set;");
+
+    jobject parameterEntrySet = env->CallObjectMethod(parameterMap, entrySet);
+    knn_jni::HasExceptionInStack(env, "Unable to call \"entrySet\" method on \"java/util/Map\"");
+
+    jclass setClass = knn_jni::FindClass(env, "java/util/Set");
+
+    jmethodID iterator = knn_jni::FindMethod(env, setClass, "iterator", "()Ljava/util/Iterator;");
+
+    jclass iteratorClass = knn_jni::FindClass(env, "java/util/Iterator");
+
+    jobject iter = env->CallObjectMethod(parameterEntrySet, iterator);
+    knn_jni::HasExceptionInStack(env, "Call to \"iterator\" method failed");
+
+    jmethodID hasNext = knn_jni::FindMethod(env, iteratorClass, "hasNext", "()Z");
+    jmethodID next = knn_jni::FindMethod(env, iteratorClass, "next", "()Ljava/lang/Object;");
+
+    jclass entryClass = knn_jni::FindClass(env, "java/util/Map$Entry");
+
+    jmethodID getKey = knn_jni::FindMethod(env, entryClass, "getKey", "()Ljava/lang/Object;");
+
+    jmethodID getValue = knn_jni::FindMethod(env, entryClass, "getValue", "()Ljava/lang/Object;");
+
+    jclass integerClass = knn_jni::FindClass(env, "java/lang/Integer");
+
+    jmethodID intValue = knn_jni::FindMethod(env, integerClass, "intValue", "()I");
+
+    // Iterate over the entry Set
+    jobject entry;
+    std::string key;
+    jobject value;
+    while (env->CallBooleanMethod(iter, hasNext)) {
+        entry = env->CallObjectMethod(iter, next);
+        knn_jni::HasExceptionInStack(env, "Could not call \"next\" method");
+
+        key = knn_jni::GetStringJenv(env, (jstring) env->CallObjectMethod(entry, getKey));
+        knn_jni::HasExceptionInStack(env, "Could not call \"getKey\" method");
+
+        value = env->CallObjectMethod(entry, getValue);
+        knn_jni::HasExceptionInStack(env, "Could not call \"getValue\" method");
+
+        if (auto * indexIvf = dynamic_cast<faiss::IndexIVF*>(index)) {
+            if (key == "nprobes") {
+                if (env->IsInstanceOf(value, integerClass)) {
+                    indexIvf->nprobe = env->CallIntMethod(value, intValue);
+                    knn_jni::HasExceptionInStack(env, "Could not call \"intValue\" method on Integer");
+                } else {
+                    throw std::runtime_error("Cannot call IntMethod on non-integer class");
+                }
+            } else if (key == "coarse_quantizer" && indexIvf->quantizer != nullptr) {
+                SetExtraParameters(env, value, indexIvf->quantizer);
+            }
+            env->DeleteLocalRef(value);
+        }
+
+        if (auto * indexHnsw = dynamic_cast<faiss::IndexHNSW*>(index)) {
+            if (key == "ef_construction") {
+                if (env->IsInstanceOf(value, integerClass)) {
+                    indexHnsw->hnsw.efConstruction = env->CallIntMethod(value, intValue);
+                    knn_jni::HasExceptionInStack(env, "Could not call \"intValue\" method on Integer");
+                } else {
+                    throw std::runtime_error("Cannot call IntMethod on non-integer class");
+                }
+            } else if (key == "ef_search") {
+                if (env->IsInstanceOf(value, integerClass)) {
+                    indexHnsw->hnsw.efSearch = env->CallIntMethod(value, intValue);
+                    knn_jni::HasExceptionInStack(env, "Could not call \"intValue\" method on Integer");
+                } else {
+                    throw std::runtime_error("Cannot call IntMethod on non-integer class");
+                }
+            }
+            env->DeleteLocalRef(value);
+        }
+        env->DeleteLocalRef(entry);
+    }
+    env->DeleteLocalRef(parameterEntrySet);
+    knn_jni::HasExceptionInStack(env, "Could not call \"hasNext\" method");
+}
 
 /**
  * Method: saveIndex
  *
  */
-	JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_saveIndex
-(JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jobjectArray algoParams, jstring spaceType)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_save
+        (JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jobject parameterMap,
+         jstring spaceType, jstring indexDescription, jint trainingDatasetSizeLimit, jint minimumDatapoints)
 {
-	vector<int64_t> idVector;
-	vector<float>   dataset;
-	vector<string> paramsList;
-	//TODO we can support other FAISS index in the future, may be paramsList can add index=xxxx
-	string indexDescription = "HNSW32";
-	faiss::MetricType metric = faiss::METRIC_L2;
-	std::unique_ptr<faiss::Index> indexWriter;
-	int dim = 0;
-	try {
-		//---- ids
-		int* object_ids = NULL;
-		object_ids = env->GetIntArrayElements(ids, 0);
-		for(int i = 0; i < env->GetArrayLength(ids); ++i) {
-			idVector.push_back(object_ids[i]);
-		}
-		env->ReleaseIntArrayElements(ids, object_ids, 0);
-		knn_jni::HasExceptionInStack(env);
-
-		//---- vectors
-		for (int i = 0; i < env->GetArrayLength(vectors); ++i) {
-			jfloatArray vectorArray = (jfloatArray)env->GetObjectArrayElement(vectors, i);
-			float* vector = env->GetFloatArrayElements(vectorArray, 0);
-			dim = env->GetArrayLength(vectorArray);
-			for(int j = 0; j < dim; ++j) {
-				dataset.push_back(vector[j]);
-			}
-			env->ReleaseFloatArrayElements(vectorArray, vector, 0);
-		}
-        knn_jni::HasExceptionInStack(env);
+    vector<int64_t> idVector;
+    vector<float>   dataset;
+    vector<string> paramsList;
 
-		//---- indexPath
-		const char *indexString = env->GetStringUTFChars(indexPath, 0);
-		string indexPathString(indexString);
-		env->ReleaseStringUTFChars(indexPath, indexString);
+    faiss::MetricType metric = faiss::METRIC_L2;
+    std::unique_ptr<faiss::Index> indexWriter;
+    int dim = 0;
+    try {
+        //---- ids
+        int* object_ids = nullptr;
+        object_ids = env->GetIntArrayElements(ids, 0);
+        for(int i = 0; i < env->GetArrayLength(ids); ++i) {
+            idVector.push_back(object_ids[i]);
+        }
+        env->ReleaseIntArrayElements(ids, object_ids, 0);
         knn_jni::HasExceptionInStack(env);
 
-		//---- algoParams
-		int paramsCount = env->GetArrayLength(algoParams);
-		for (int i=0; i<paramsCount; i++) {
-			jstring param = (jstring) (env->GetObjectArrayElement(algoParams, i));
-			const char *rawString = env->GetStringUTFChars(param, 0);
-			paramsList.push_back(rawString);
+        //---- vectors
+        for (int i = 0; i < env->GetArrayLength(vectors); ++i) {
+            auto vectorArray = (jfloatArray)env->GetObjectArrayElement(vectors, i);
+            float* vector = env->GetFloatArrayElements(vectorArray, 0);
+            dim = env->GetArrayLength(vectorArray);
+            for(int j = 0; j < dim; ++j) {
+                dataset.push_back(vector[j]);
+            }
+            env->ReleaseFloatArrayElements(vectorArray, vector, 0);
+        }
+        knn_jni::HasExceptionInStack(env);
 
-			int M = 32;
-			if (sscanf(rawString, "M=%d", &M) == 1) {
-				indexDescription="HNSW"+std::to_string(M);
-			}
-			env->ReleaseStringUTFChars(param, rawString);
+        //---- indexPath
+        const char *indexString = env->GetStringUTFChars(indexPath, 0);
+        string indexPathString(indexString);
+        env->ReleaseStringUTFChars(indexPath, indexString);
+        knn_jni::HasExceptionInStack(env);
 
-		}
+        //---- space
+        const char *spaceTypeCStr = env->GetStringUTFChars(spaceType, 0);
+        string spaceTypeString(spaceTypeCStr);
+        env->ReleaseStringUTFChars(spaceType, spaceTypeCStr);
         knn_jni::HasExceptionInStack(env);
+        // space mapping faiss::MetricType
+        if(mapMetric.find(spaceTypeString) != mapMetric.end()) {
+            metric = mapMetric[spaceTypeString];
+        }
 
+        //---- Create IndexWriter from faiss index_factory
+        // If data is less than a certain amount, just create a flat index
+        if (idVector.size() < (int) minimumDatapoints) {
+            indexWriter.reset(faiss::index_factory(dim, "Flat", metric));
+        } else {
+            std::string description = knn_jni::GetStringJenv(env, indexDescription);
+            indexWriter.reset(faiss::index_factory(dim, description.c_str(), metric));
+        }
 
-		//---- space
-		const char *spaceTypeCStr = env->GetStringUTFChars(spaceType, 0);
-		string spaceTypeString(spaceTypeCStr);
-		env->ReleaseStringUTFChars(spaceType, spaceTypeCStr);
-        knn_jni::HasExceptionInStack(env);
-		// space mapping faiss::MetricType
-		if(mapMetric.find(spaceTypeString) != mapMetric.end()) {
-			metric = mapMetric[spaceTypeString];
-		}
+        // Add extra parameters that cant be configured with the index factory
+        SetExtraParameters(env, parameterMap, indexWriter.get());
+        env->DeleteLocalRef(parameterMap);
 
-		//---- Create IndexWriter from faiss index_factory
-		indexWriter.reset(faiss::index_factory(dim, indexDescription.data(), metric));
-
-		//Preparation And TODO Verify IndexWriter
-		//Some Param Can not Create from IndexFactory, Like HNSW efSearch and efCOnstruction
-		//----FOR HNSW 1st PARAM: M(HNSW32->M=32), efConstruction, efSearch
-		if(indexDescription.find("HNSW") != std::string::npos) {
-			for(int i = 0; i < paramsCount; ++i) {
-				const string& param = paramsList[i];
-				int efConstruction = 40; //default
-				int efSearch = 16;//default
-				if(param.find("efConstruction") != std::string::npos &&
-						sscanf(param.data(), "efConstruction=%d", &efConstruction) == 1) {
-					faiss::IndexHNSW* ihp = reinterpret_cast<faiss::IndexHNSW*>(indexWriter.get());
-					ihp->hnsw.efConstruction = efConstruction;
-				} else if (param.find("efSearch") != std::string::npos &&
-						sscanf(param.data(), "efSearch=%d", &efSearch) == 1){
-					faiss::IndexHNSW* ihp = reinterpret_cast<faiss::IndexHNSW*>(indexWriter.get());
-					ihp->hnsw.efSearch = efSearch;
-				}
-			}
-		}
+        //---- Do Index
+        if(!indexWriter->is_trained) {
+            if (idVector.size() <= (int) trainingDatasetSizeLimit) {
+                TrainIndex(indexWriter.get(), idVector.size(), dataset.data());
+            } else {
+                vector<float>::const_iterator first = dataset.begin();
+                vector<float>::const_iterator last = dataset.begin() + ((int) trainingDatasetSizeLimit)*dim;
+                vector<float> subDataVector(first, last);
+                TrainIndex(indexWriter.get(), (int) trainingDatasetSizeLimit, subDataVector.data());
+            }
+        }
 
-		//---- Do Index
-		//----- 1. Train
-		if(!indexWriter->is_trained) {
-			//TODO if we use like PQ, we have to train dataset
-			// but when a lucene segment only one document, it
-			// can not train the data.
-		}
-		//----- 2. Add IDMap
-		// default all use self defined IndexIDMap cause some class no add_with_ids
-		faiss::IndexIDMap idMap =  faiss::IndexIDMap(indexWriter.get());
-		idMap.add_with_ids(idVector.size(), dataset.data(), idVector.data());
+        //----- 2. Add IDMap
+        faiss::IndexIDMap idMap =  faiss::IndexIDMap(indexWriter.get());
+        idMap.add_with_ids(idVector.size(), dataset.data(), idVector.data());
 
 		//----- 3. WriteIndex
 		faiss::write_index(&idMap, indexPathString.c_str());
 
 		//Explicit delete object
 		faiss::Index* indexPointer = indexWriter.release();
-		if(indexPointer) delete indexPointer;
-
+		delete indexPointer;
 	}
 	catch(...) {
 		faiss::Index* indexPointer = indexWriter.release();
-		if(indexPointer) delete indexPointer;
+		delete indexPointer;
         knn_jni::CatchCppExceptionAndThrowJava(env);
 	}
 }
 
 
-	JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_queryIndex
-(JNIEnv* env, jclass cls, jlong indexPointer, jfloatArray queryVector, jint k)
+JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_query
+        (JNIEnv* env, jclass cls, jlong indexPointer, jfloatArray queryVector, jint k)
 {
 	faiss::Index *indexReader = nullptr;
 	try {
@@ -203,8 +275,8 @@ std::unordered_map<string, faiss::MetricType> mapMetric = {
 	return NULL;
 }
 
-	JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_init
-(JNIEnv* env, jclass cls,  jstring indexPath, jobjectArray algoParams, jstring spaceType)
+JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_init
+        (JNIEnv* env, jclass cls,  jstring indexPath)
 {
 
 	faiss::Index* indexReader = nullptr;
@@ -228,8 +300,8 @@ std::unordered_map<string, faiss::MetricType> mapMetric = {
  * When autoclose class do close, then delete the pointer
  * Method GC pointer
  */
-	JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_gc
-(JNIEnv* env, jclass cls,  jlong indexPointer)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_gc
+        (JNIEnv* env, jclass cls,  jlong indexPointer)
 {
 	try {
 		faiss::Index *indexWrapper = reinterpret_cast<faiss::Index*>(indexPointer);
@@ -246,7 +318,8 @@ std::unordered_map<string, faiss::MetricType> mapMetric = {
  * Method: Global Init
  *
  */
-JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_initLibrary(JNIEnv *, jclass)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_faiss_v165_KNNFaissIndex_initLibrary
+        (JNIEnv *, jclass)
 {
 	//set thread 1 cause ES has Search thread
 	//TODO make it different at search and write