Generalize lib interface to return context objects (#1927)

Generalizes the KNNLibrary to return an object for both search and
indexing so that the plugin can search/index against them. This will
help properly pass information that does not need to be sent to the JNI
for search and index builds.

Signed-off-by: John Mazanec <[email protected]>
(cherry picked from commit a16b8aa)

CHANGELOG.md

@@ -26,3 +26,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Clean up parsing for query [#1824](https://github.com/opensearch-project/k-NN/pull/1824)
 * Refactor engine package structure [#1913](https://github.com/opensearch-project/k-NN/pull/1913)
 * Refactor method structure and definitions [#1920](https://github.com/opensearch-project/k-NN/pull/1920)
+* Generalize lib interface to return context objects [#1925](https://github.com/opensearch-project/k-NN/pull/1925)

src/main/java/org/opensearch/knn/index/engine/AbstractKNNLibrary.java

@@ -24,10 +24,18 @@ public abstract class AbstractKNNLibrary implements KNNLibrary {
     protected final String version;
 
     @Override
-    public EngineSpecificMethodContext getMethodContext(String methodName) {
+    public KNNLibrarySearchContext getKNNLibrarySearchContext(String methodName) {
         validateMethodExists(methodName);
         KNNMethod method = methods.get(methodName);
-        return method.getEngineSpecificMethodContext();
+        return method.getKNNLibrarySearchContext();
+    }
+
+    @Override
+    public KNNLibraryIndexingContext getKNNLibraryIndexingContext(KNNMethodContext knnMethodContext) {
+        String method = knnMethodContext.getMethodComponentContext().getName();
+        validateMethodExists(method);
+        KNNMethod knnMethod = methods.get(method);
+        return knnMethod.getKNNLibraryIndexingContext(knnMethodContext);
     }
 
     @Override
@@ -51,14 +59,6 @@ public boolean isTrainingRequired(KNNMethodContext knnMethodContext) {
         return methods.get(methodName).isTrainingRequired(knnMethodContext);
     }
 
-    @Override
-    public Map<String, Object> getMethodAsMap(KNNMethodContext knnMethodContext) {
-        String method = knnMethodContext.getMethodComponentContext().getName();
-        validateMethodExists(method);
-        KNNMethod knnMethod = methods.get(method);
-        return knnMethod.getAsMap(knnMethodContext);
-    }
-
     private void validateMethodExists(String methodName) {
         KNNMethod method = methods.get(methodName);
         if (method == null) {

src/main/java/org/opensearch/knn/index/engine/AbstractKNNMethod.java

@@ -27,7 +27,7 @@ public abstract class AbstractKNNMethod implements KNNMethod {
 
     protected final MethodComponent methodComponent;
     protected final Set<SpaceType> spaces;
-    protected final EngineSpecificMethodContext engineSpecificMethodContext;
+    protected final KNNLibrarySearchContext knnLibrarySearchContext;
 
     @Override
     public boolean isSpaceTypeSupported(SpaceType space) {
@@ -106,14 +106,14 @@ public int estimateOverheadInKB(KNNMethodContext knnMethodContext, int dimension
     }
 
     @Override
-    public Map<String, Object> getAsMap(KNNMethodContext knnMethodContext) {
+    public KNNLibraryIndexingContext getKNNLibraryIndexingContext(KNNMethodContext knnMethodContext) {
         Map<String, Object> parameterMap = new HashMap<>(methodComponent.getAsMap(knnMethodContext.getMethodComponentContext()));
         parameterMap.put(KNNConstants.SPACE_TYPE, knnMethodContext.getSpaceType().getValue());
-        return parameterMap;
+        return KNNLibraryIndexingContextImpl.builder().parameters(parameterMap).build();
     }
 
     @Override
-    public EngineSpecificMethodContext getEngineSpecificMethodContext() {
-        return engineSpecificMethodContext;
+    public KNNLibrarySearchContext getKNNLibrarySearchContext() {
+        return knnLibrarySearchContext;
     }
 }

src/main/java/org/opensearch/knn/index/engine/DefaultHnswContext.java → src/main/java/org/opensearch/knn/index/engine/DefaultHnswSearchContext.java

@@ -14,7 +14,7 @@
 /**
  * Default HNSW context for all engines. Have a different implementation if engine context differs.
  */
-public final class DefaultHnswContext implements EngineSpecificMethodContext {
+public final class DefaultHnswSearchContext implements KNNLibrarySearchContext {
 
     private final Map<String, Parameter<?>> supportedMethodParameters = ImmutableMap.<String, Parameter<?>>builder()
         .put(MethodParameter.EF_SEARCH.getName(), new Parameter.IntegerParameter(MethodParameter.EF_SEARCH.getName(), null, value -> true))

src/main/java/org/opensearch/knn/index/engine/DefaultIVFContext.java → src/main/java/org/opensearch/knn/index/engine/DefaultIVFSearchContext.java

@@ -11,7 +11,7 @@
 
 import java.util.Map;
 
-public final class DefaultIVFContext implements EngineSpecificMethodContext {
+public final class DefaultIVFSearchContext implements KNNLibrarySearchContext {
 
     private final Map<String, Parameter<?>> supportedMethodParameters = ImmutableMap.<String, Parameter<?>>builder()
         .put(MethodParameter.NPROBE.getName(), new Parameter.IntegerParameter(MethodParameter.NPROBE.getName(), null, value -> true))

src/main/java/org/opensearch/knn/index/engine/KNNEngine.java

@@ -145,11 +145,6 @@ public String getCompoundExtension() {
         return knnLibrary.getCompoundExtension();
     }
 
-    @Override
-    public EngineSpecificMethodContext getMethodContext(String methodName) {
-        return knnLibrary.getMethodContext(methodName);
-    }
-
     @Override
     public float score(float rawScore, SpaceType spaceType) {
         return knnLibrary.score(rawScore, spaceType);
@@ -181,8 +176,13 @@ public boolean isTrainingRequired(KNNMethodContext knnMethodContext) {
     }
 
     @Override
-    public Map<String, Object> getMethodAsMap(KNNMethodContext knnMethodContext) {
-        return knnLibrary.getMethodAsMap(knnMethodContext);
+    public KNNLibraryIndexingContext getKNNLibraryIndexingContext(KNNMethodContext knnMethodContext) {
+        return knnLibrary.getKNNLibraryIndexingContext(knnMethodContext);
+    }
+
+    @Override
+    public KNNLibrarySearchContext getKNNLibrarySearchContext(String methodName) {
+        return knnLibrary.getKNNLibrarySearchContext(methodName);
     }
 
     @Override

src/main/java/org/opensearch/knn/index/engine/KNNLibrary.java

@@ -11,7 +11,6 @@
 
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 
 /**
  * KNNLibrary is an interface that helps the plugin communicate with k-NN libraries
@@ -41,13 +40,6 @@ public interface KNNLibrary {
      */
     String getCompoundExtension();
 
-    /**
-     * Gets metadata related to methods supported by the library
-     * @param methodName
-     * @return
-     */
-    EngineSpecificMethodContext getMethodContext(String methodName);
-
     /**
      * Generate the Lucene score from the rawScore returned by the library. With k-NN, often times the library
      * will return a score where the lower the score, the better the result. This is the opposite of how Lucene scores
@@ -116,12 +108,20 @@ public interface KNNLibrary {
     int estimateOverheadInKB(KNNMethodContext knnMethodContext, int dimension);
 
     /**
-     * Generate method as map that can be used to configure the knn index from the jni
+     * Get the context from the library needed to build the index.
      *
-     * @param knnMethodContext to generate parameter map from
+     * @param knnMethodContext to get build context for
      * @return parameter map
      */
-    Map<String, Object> getMethodAsMap(KNNMethodContext knnMethodContext);
+    KNNLibraryIndexingContext getKNNLibraryIndexingContext(KNNMethodContext knnMethodContext);
+
+    /**
+     * Gets metadata related to methods supported by the library
+     *
+     * @param methodName name of method
+     * @return KNNLibrarySearchContext
+     */
+    KNNLibrarySearchContext getKNNLibrarySearchContext(String methodName);
 
     /**
      * Getter for initialized

src/main/java/org/opensearch/knn/index/engine/KNNLibraryIndexingContext.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.engine;
+
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Context a library gives to build one of its indices
+ */
+public interface KNNLibraryIndexingContext {
+    /**
+     * Get map of parameters that get passed to the library to build the index
+     *
+     * @return Map of parameters
+     */
+    Map<String, Object> getLibraryParameters();
+
+    KNNLibraryIndexingContext EMPTY = Collections::emptyMap;
+}

src/main/java/org/opensearch/knn/index/engine/KNNLibraryIndexingContextImpl.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.engine;
+
+import lombok.Builder;
+
+import java.util.Map;
+
+/**
+ * Simple implementation of {@link KNNLibraryIndexingContext}
+ */
+@Builder
+public class KNNLibraryIndexingContextImpl implements KNNLibraryIndexingContext {
+
+    private Map<String, Object> parameters;
+
+    @Override
+    public Map<String, Object> getLibraryParameters() {
+        return parameters;
+    }
+}

src/main/java/org/opensearch/knn/index/engine/KNNLibrarySearchContext.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.engine;
+
+import org.opensearch.knn.index.engine.model.QueryContext;
+
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Holds the context needed to search a knn library.
+ */
+public interface KNNLibrarySearchContext {
+
+    /**
+     * Returns supported parameters for the library.
+     *
+     * @param ctx QueryContext
+     * @return parameters supported by the library
+     */
+    Map<String, Parameter<?>> supportedMethodParameters(QueryContext ctx);
+
+    KNNLibrarySearchContext EMPTY = ctx -> Collections.emptyMap();
+}

src/main/java/org/opensearch/knn/index/engine/KNNMethod.java

@@ -9,8 +9,6 @@
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.training.VectorSpaceInfo;
 
-import java.util.Map;
-
 /**
  * KNNMethod defines the structure of a method supported by a particular k-NN library. It is used to validate
  * the KNNMethodContext passed in by the user, where the KNNMethodContext provides the configuration that the user may
@@ -61,17 +59,17 @@ public interface KNNMethod {
     int estimateOverheadInKB(KNNMethodContext knnMethodContext, int dimension);
 
     /**
-     * Parse knnMethodContext into a map that the library can use to configure the index
+     * Parse knnMethodContext into context that the library can use to build the index
      *
-     * @param knnMethodContext from which to generate map
-     * @return KNNMethod as a map
+     * @param knnMethodContext to generate the context for
+     * @return KNNLibraryIndexingContext
      */
-    Map<String, Object> getAsMap(KNNMethodContext knnMethodContext);
+    KNNLibraryIndexingContext getKNNLibraryIndexingContext(KNNMethodContext knnMethodContext);
 
     /**
-     * Get the method context for a particular method
+     * Get the search context for a particular method
      *
-     * @return EngineSpecificMethodContext for the method
+     * @return KNNLibrarySearchContext
      */
-    EngineSpecificMethodContext getEngineSpecificMethodContext();
+    KNNLibrarySearchContext getKNNLibrarySearchContext();
 }

src/main/java/org/opensearch/knn/index/engine/faiss/FaissHNSWMethod.java

@@ -9,7 +9,7 @@
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.engine.AbstractKNNMethod;
-import org.opensearch.knn.index.engine.DefaultHnswContext;
+import org.opensearch.knn.index.engine.DefaultHnswSearchContext;
 import org.opensearch.knn.index.engine.Encoder;
 import org.opensearch.knn.index.engine.MethodComponent;
 import org.opensearch.knn.index.engine.MethodComponentContext;
@@ -51,7 +51,7 @@ public class FaissHNSWMethod extends AbstractKNNMethod {
      * @see AbstractKNNMethod
      */
     public FaissHNSWMethod() {
-        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultHnswContext());
+        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultHnswSearchContext());
     }
 
     private static MethodComponent initMethodComponent() {

src/main/java/org/opensearch/knn/index/engine/faiss/FaissIVFMethod.java

@@ -8,7 +8,7 @@
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.engine.AbstractKNNMethod;
-import org.opensearch.knn.index.engine.DefaultIVFContext;
+import org.opensearch.knn.index.engine.DefaultIVFSearchContext;
 import org.opensearch.knn.index.engine.Encoder;
 import org.opensearch.knn.index.engine.MethodComponent;
 import org.opensearch.knn.index.engine.MethodComponentContext;
@@ -55,7 +55,7 @@ public class FaissIVFMethod extends AbstractKNNMethod {
      * @see AbstractKNNMethod
      */
     public FaissIVFMethod() {
-        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultIVFContext());
+        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultIVFSearchContext());
     }
 
     private static MethodComponent initMethodComponent() {

src/main/java/org/opensearch/knn/index/engine/lucene/LuceneHNSWMethod.java

@@ -49,7 +49,7 @@ public class LuceneHNSWMethod extends AbstractKNNMethod {
      * @see AbstractKNNMethod
      */
     public LuceneHNSWMethod() {
-        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new LuceneHNSWContext());
+        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new LuceneHNSWSearchContext());
     }
 
     private static MethodComponent initMethodComponent() {

src/main/java/org/opensearch/knn/index/engine/lucene/LuceneHNSWContext.java → src/main/java/org/opensearch/knn/index/engine/lucene/LuceneHNSWSearchContext.java

@@ -6,15 +6,15 @@
 package org.opensearch.knn.index.engine.lucene;
 
 import com.google.common.collect.ImmutableMap;
-import org.opensearch.knn.index.engine.EngineSpecificMethodContext;
+import org.opensearch.knn.index.engine.KNNLibrarySearchContext;
 import org.opensearch.knn.index.engine.Parameter;
 import org.opensearch.knn.index.engine.model.QueryContext;
 import org.opensearch.knn.index.query.request.MethodParameter;
 
 import java.util.Collections;
 import java.util.Map;
 
-public class LuceneHNSWContext implements EngineSpecificMethodContext {
+public class LuceneHNSWSearchContext implements KNNLibrarySearchContext {
 
     private final Map<String, Parameter<?>> supportedMethodParameters = ImmutableMap.<String, Parameter<?>>builder()
         .put(MethodParameter.EF_SEARCH.getName(), new Parameter.IntegerParameter(MethodParameter.EF_SEARCH.getName(), null, value -> true))

src/main/java/org/opensearch/knn/index/engine/nmslib/NmslibHNSWMethod.java

@@ -8,7 +8,7 @@
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.SpaceType;
 import org.opensearch.knn.index.engine.AbstractKNNMethod;
-import org.opensearch.knn.index.engine.DefaultHnswContext;
+import org.opensearch.knn.index.engine.DefaultHnswSearchContext;
 import org.opensearch.knn.index.engine.MethodComponent;
 import org.opensearch.knn.index.engine.Parameter;
 
@@ -39,7 +39,7 @@ public class NmslibHNSWMethod extends AbstractKNNMethod {
      * @see AbstractKNNMethod
      */
     public NmslibHNSWMethod() {
-        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultHnswContext());
+        super(initMethodComponent(), Set.copyOf(SUPPORTED_SPACES), new DefaultHnswSearchContext());
     }
 
     private static MethodComponent initMethodComponent() {