FEAT: support cosine similarity

jni/src/v1736/com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex.cpp

@@ -41,9 +41,9 @@ using similarity::KNNQueue;
 extern "C"
 
 struct IndexWrapper {
-  IndexWrapper() {
-    space.reset(SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams()));
-    index.reset(MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", "l2", *space, data));
+  IndexWrapper(string spaceType) {
+    space.reset(SpaceFactoryRegistry<float>::Instance().CreateSpace(spaceType, AnyParams()));
+    index.reset(MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", spaceType, *space, data));
   }
   std::unique_ptr<Space<float>> space;
   std::unique_ptr<Index<float>> index;
@@ -85,15 +85,17 @@ void catch_cpp_exception_and_throw_java(JNIEnv* env)
     }
 }
 
-JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_saveIndex(JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jobjectArray algoParams)
+JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_saveIndex(JNIEnv* env, jclass cls, jintArray ids, jobjectArray vectors, jstring indexPath, jstring spaceType, jobjectArray algoParams)
 {
     Space<float>* space = NULL;
     ObjectVector dataset;
     Index<float>* index = NULL;
     int* object_ids = NULL;
 
     try {
-        space = SpaceFactoryRegistry<float>::Instance().CreateSpace("l2", AnyParams());
+        const char *spaceTypeCStr = env->GetStringUTFChars(spaceType, 0);
+        string spaceTypeString(spaceTypeCStr);
+        space = SpaceFactoryRegistry<float>::Instance().CreateSpace(spaceTypeString, AnyParams());
         object_ids = env->GetIntArrayElements(ids, 0);
         for (int i = 0; i < env->GetArrayLength(vectors); i++) {
             jfloatArray vectorArray = (jfloatArray)env->GetObjectArrayElement(vectors, i);
@@ -103,7 +105,7 @@ JNIEXPORT void JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v173
         }
         // free up memory
         env->ReleaseIntArrayElements(ids, object_ids, 0);
-        index = MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", "l2", *space, dataset);
+        index = MethodFactoryRegistry<float>::Instance().CreateMethod(false, "hnsw", spaceTypeString, *space, dataset);
 
         int paramsCount = env->GetArrayLength(algoParams);
         vector<string> paramsList;
@@ -171,17 +173,19 @@ JNIEXPORT jobjectArray JNICALL Java_com_amazon_opendistroforelasticsearch_knn_in
     return NULL;
 }
 
-JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init(JNIEnv* env, jclass cls,  jstring indexPath, jobjectArray algoParams)
+JNIEXPORT jlong JNICALL Java_com_amazon_opendistroforelasticsearch_knn_index_v1736_KNNIndex_init(JNIEnv* env, jclass cls,  jstring indexPath, jstring spaceType, jobjectArray algoParams)
 {
     IndexWrapper *indexWrapper = NULL;
     try {
         const char *indexPathCStr = env->GetStringUTFChars(indexPath, 0);
         string indexPathString(indexPathCStr);
+        const char *spaceTypeCStr = env->GetStringUTFChars(spaceType, 0);
+        string spaceTypeString(spaceTypeCStr);
         env->ReleaseStringUTFChars(indexPath, indexPathCStr);
         has_exception_in_stack(env);
 
         // Load index from file (may throw)
-        IndexWrapper *indexWrapper = new IndexWrapper();
+        IndexWrapper *indexWrapper = new IndexWrapper(spaceTypeString);
         indexWrapper->index->LoadIndex(indexPathString);
 
         // Parse and set query params

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/KNNIndexCache.java

@@ -196,7 +196,7 @@ public KNNIndexCacheEntry loadIndex(String indexPathUrl, String indexName) throw
         // the entry
         fileWatcher.init();
 
-        final KNNIndex knnIndex = KNNIndex.loadIndex(indexPathUrl, getQueryParams(indexName));
+        final KNNIndex knnIndex = KNNIndex.loadIndex(indexPathUrl, KNNSettings.getSpaceType(indexName), getQueryParams(indexName));
 
         // TODO verify that this is safe - ideally we'd explicitly ensure that the FileWatcher is only checked
         // after the guava cache has finished loading the key to avoid a race condition where the watcher

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/KNNSettings.java

@@ -31,11 +31,14 @@
 import org.elasticsearch.monitor.jvm.JvmInfo;
 import org.elasticsearch.monitor.os.OsProbe;
 
+import java.security.InvalidParameterException;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -63,6 +66,7 @@ public class KNNSettings {
     /**
      * Settings name
      */
+    public static final String KNN_SPACE_TYPE = "index.knn.space_type";
     public static final String KNN_ALGO_PARAM_M = "index.knn.algo_param.m";
     public static final String KNN_ALGO_PARAM_EF_CONSTRUCTION = "index.knn.algo_param.ef_construction";
     public static final String KNN_ALGO_PARAM_EF_SEARCH = "index.knn.algo_param.ef_search";
@@ -80,6 +84,11 @@ public class KNNSettings {
      * Settings Definition
      */
 
+    public static final Setting<String> INDEX_KNN_SPACE_TYPE = Setting.simpleString(KNN_SPACE_TYPE,
+        "l2",
+        new SpaceTypeValidator(),
+        IndexScope);
+
     /**
      * M - the number of bi-directional links created for every new element during construction.
      * Reasonable range for M is 2-100. Higher M work better on datasets with high intrinsic
@@ -252,7 +261,8 @@ public Setting<?> getSetting(String key) {
     }
 
     public List<Setting<?>> getSettings() {
-        List<Setting<?>> settings =  Arrays.asList(INDEX_KNN_ALGO_PARAM_M_SETTING,
+        List<Setting<?>> settings =  Arrays.asList(INDEX_KNN_SPACE_TYPE,
+                INDEX_KNN_ALGO_PARAM_M_SETTING,
                 INDEX_KNN_ALGO_PARAM_EF_CONSTRUCTION_SETTING,
                 INDEX_KNN_ALGO_PARAM_EF_SEARCH_SETTING,
                 KNN_ALGO_PARAM_INDEX_THREAD_QTY_SETTING,
@@ -357,6 +367,16 @@ public static int getEfSearchParam(String index) {
         return getIndexSettingValue(index, KNN_ALGO_PARAM_EF_SEARCH, 512);
     }
 
+    /**
+     *
+     * @param index Name of the index
+     * @return spaceType value
+     */
+    public static String getSpaceType(String index) {
+        return KNNSettings.state().clusterService.state().getMetaData()
+            .index(index).getSettings().get(KNN_SPACE_TYPE, SpaceTypes.l2.getValue());
+    }
+
     public static int getIndexSettingValue(String index, String settingName, int defaultValue) {
         return KNNSettings.state().clusterService.state().getMetaData()
                                                  .index(index).getSettings()
@@ -367,4 +387,48 @@ public void setClusterService(ClusterService clusterService) {
         this.clusterService = clusterService;
     }
 
+    static class SpaceTypeValidator implements Setting.Validator<String> {
+
+        private Set<String> types = SpaceTypes.getValues();
+
+        @Override public void validate(String value) {
+            if (!types.contains(value.toLowerCase())){
+                throw new InvalidParameterException(String.format("Unsupported space type: %s", value));
+            }
+        }
+    }
+
+    /**
+     * Enum contains space types for k-NN similarity search
+     */
+    public enum SpaceTypes {
+        l2("l2"),
+        cosinesimil("cosinesimil");
+
+        private String value;
+
+        SpaceTypes(String value) { this.value = value; }
+
+        /**
+         * Get space type
+         *
+         * @return name
+         */
+        public String getValue() { return value; }
+
+        /**
+         * Get all space types
+         *
+         * @return set of all stat names
+         */
+        public static Set<String> getValues() {
+            Set<String> values = new HashSet<>();
+
+            for (SpaceTypes spaceType : SpaceTypes.values()) {
+                values.add(spaceType.getValue());
+            }
+            return values;
+        }
+    }
+
 }

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/KNNVectorFieldMapper.java

@@ -94,6 +94,11 @@ public Builder ignoreMalformed(boolean ignoreMalformed) {
             return builder;
         }
 
+        public Builder spaceTypeParam(String key, String paramValue) {
+            Defaults.FIELD_TYPE.putAttribute(key, paramValue.toLowerCase());
+            return builder;
+        }
+
         public Builder algoParams(String key, int paramValue) {
             Defaults.FIELD_TYPE.putAttribute(key, String.valueOf(paramValue));
             return builder;
@@ -135,6 +140,8 @@ public static class TypeParser implements Mapper.TypeParser {
         public Mapper.Builder<?, ?> parse(String name, Map<String, Object> node, ParserContext parserContext)
                 throws MapperParsingException {
             Builder builder = new KNNVectorFieldMapper.Builder(name);
+            builder.spaceTypeParam(KNNConstants.SPACE_TYPE, parserContext.mapperService().getIndexSettings().getValue(
+                KNNSettings.INDEX_KNN_SPACE_TYPE));
             builder.algoParams(KNNConstants.HNSW_ALGO_M, parserContext.mapperService().getIndexSettings().getValue(
                     KNNSettings.INDEX_KNN_ALGO_PARAM_M_SETTING));
             builder.algoParams(KNNConstants.HNSW_ALGO_EF_CONSTRUCTION, parserContext.mapperService().getIndexSettings()

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java

@@ -93,11 +93,13 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer)
 
             // Pass the path for the nms library to save the file
             String tempIndexPath = indexPath + TEMP_SUFFIX;
-            String[] algoParams = getKNNIndexParams(field.attributes());
+            Map<String, String> fieldAttributes = field.attributes();
+            String spaceType = fieldAttributes.getOrDefault(KNNConstants.SPACE_TYPE, KNNSettings.SpaceTypes.l2.getValue());
+            String[] algoParams = getKNNIndexParams(fieldAttributes);
             AccessController.doPrivileged(
                     new PrivilegedAction<Void>() {
                         public Void run() {
-                            KNNIndex.saveIndex(pair.docs, pair.vectors, tempIndexPath, algoParams);
+                            KNNIndex.saveIndex(pair.docs, pair.vectors, tempIndexPath, spaceType, algoParams);
                             return null;
                         }
                     }

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/util/KNNConstants.java

@@ -16,6 +16,7 @@
 package com.amazon.opendistroforelasticsearch.knn.index.util;
 
 public class KNNConstants {
+    public static final String SPACE_TYPE = "spaceType";
     public static final String HNSW_ALGO_M = "M";
     public static final String HNSW_ALGO_EF_CONSTRUCTION = "efConstruction";
     public static final String HNSW_ALGO_EF_SEARCH = "efSearch";

src/main/java/com/amazon/opendistroforelasticsearch/knn/index/v1736/KNNIndex.java

@@ -110,12 +110,13 @@ public void close() {
      * Loads the knn index to memory for querying the neighbours
      *
      * @param indexPath path where the hnsw index is stored
+     * @param spaceType space type of the index
      * @param algoParams hnsw algorithm parameters
      * @return knn index that can be queried for k nearest neighbours
      */
-    public static KNNIndex loadIndex(String indexPath, final String[] algoParams) {
+    public static KNNIndex loadIndex(String indexPath, String spaceType, final String[] algoParams) {
         long fileSize = computeFileSize(indexPath);
-        long indexPointer = init(indexPath, algoParams);
+        long indexPointer = init(indexPath, spaceType, algoParams);
         return new KNNIndex(indexPointer, fileSize);
     }
 
@@ -137,13 +138,13 @@ private static long computeFileSize(String indexPath) {
     }
 
     // Builds index and writes to disk (no index pointer escapes).
-    public static native void saveIndex(int[] ids, float[][] data, String indexPath, String[] algoParams);
+    public static native void saveIndex(int[] ids, float[][] data, String indexPath, String spaceType, String[] algoParams);
 
     // Queries index (thread safe with other readers, blocked by write lock)
     private static native KNNQueryResult[] queryIndex(long indexPointer, float[] query, int k);
 
     // Loads index and returns pointer to index
-    private static native long init(String indexPath, String[] algoParams);
+    private static native long init(String indexPath, String spaceType, String[] algoParams);
 
     // Deletes memory pointed to by index pointer (needs write lock)
     private static native void gc(long indexPointer);

src/test/java/com/amazon/opendistroforelasticsearch/knn/index/KNNJNITests.java → src/test/java/com/amazon/opendistroforelasticsearch/knn/index/KNNJNITestsIT.java

@@ -31,8 +31,8 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
-public class KNNJNITests extends ESTestCase {
-    private static final Logger logger = LogManager.getLogger(KNNJNITests.class);
+public class KNNJNITestsIT extends ESTestCase {
+    private static final Logger logger = LogManager.getLogger(KNNJNITestsIT.class);
 
     public void testCreateHnswIndex() throws Exception {
         int[] docs = {0, 1, 2};
@@ -52,7 +52,7 @@ public void testCreateHnswIndex() throws Exception {
                 AccessController.doPrivileged(
                 new PrivilegedAction<Void>() {
                     public Void run() {
-                        KNNIndex.saveIndex(docs, vectors, indexPath, algoParams);
+                        KNNIndex.saveIndex(docs, vectors, indexPath, "l2", algoParams);
                         return null;
                     }
                 }
@@ -80,7 +80,7 @@ public void testQueryHnswIndex() throws Exception {
         AccessController.doPrivileged(
                 new PrivilegedAction<Void>() {
                     public Void run() {
-                        KNNIndex.saveIndex(docs, vectors, indexPath, algoParams);
+                        KNNIndex.saveIndex(docs, vectors, indexPath, "l2", algoParams);
                         return null;
                     }
                 }
@@ -91,7 +91,7 @@ public Void run() {
         float[] queryVector = {1.0f, 1.0f, 1.0f, 1.0f};
         String[] algoQueryParams = {"efSearch=20"};
 
-        final KNNIndex knnIndex = KNNIndex.loadIndex(indexPath, algoQueryParams);
+        final KNNIndex knnIndex = KNNIndex.loadIndex(indexPath, "l2", algoQueryParams);
         final KNNQueryResult[] results = knnIndex.queryIndex(queryVector, 30);
 
         Map<Integer, Float> scores = Arrays.stream(results).collect(
@@ -126,7 +126,7 @@ public void testAssertExceptionFromJni() throws Exception {
         AccessController.doPrivileged(
                 new PrivilegedAction<Void>() {
                     public Void run() {
-                        KNNIndex index = KNNIndex.loadIndex(indexPath.toString(), new String[] {});
+                        KNNIndex index = KNNIndex.loadIndex(indexPath.toString(), "l2", new String[] {});
                         return null;
                     }
                 }
@@ -153,7 +153,7 @@ public void testQueryHnswIndexWithValidAlgoParams() throws Exception {
         AccessController.doPrivileged(
                 new PrivilegedAction<Void>() {
                     public Void run() {
-                        KNNIndex.saveIndex(docs, vectors, indexPath, algoIndexParams);
+                        KNNIndex.saveIndex(docs, vectors, indexPath, "l2", algoIndexParams);
                         return null;
                     }
                 }
@@ -165,7 +165,7 @@ public Void run() {
         float[] queryVector = {1.0f, 1.0f, 1.0f, 1.0f};
         String[] algoQueryParams = {"efSearch=200"};
 
-        final KNNIndex index = KNNIndex.loadIndex(indexPath, algoQueryParams);
+        final KNNIndex index = KNNIndex.loadIndex(indexPath, "l2", algoQueryParams);
         final KNNQueryResult[] results = index.queryIndex(queryVector, 30);
 
         Map<Integer, Float> scores = Arrays.stream(results).collect(