Refactoring

Refactor, further delineate responsibilities of client vs. build request, add more abstraction

Signed-off-by: owenhalpert <ohalpert@gmail.com>

Author: owenhalpert

src/main/java/org/opensearch/knn/index/codec/nativeindex/remote/RemoteIndexBuildStrategy.java

@@ -13,10 +13,12 @@
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.NativeIndexBuildStrategy;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
+import org.opensearch.knn.index.remote.HTTPRemoteBuildRequest;
 import org.opensearch.knn.index.remote.RemoteBuildRequest;
+import org.opensearch.knn.index.remote.RemoteBuildRequestBuilder;
 import org.opensearch.knn.index.remote.RemoteBuildResponse;
 import org.opensearch.knn.index.remote.RemoteIndexClient;
-import org.opensearch.knn.index.remote.RemoteIndexHTTPClient;
+import org.opensearch.knn.index.remote.RemoteIndexClientFactory;
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.repositories.Repository;
 import org.opensearch.repositories.RepositoryMissingException;
@@ -27,7 +29,6 @@
 
 import static org.opensearch.knn.index.KNNSettings.KNN_INDEX_REMOTE_VECTOR_BUILD_SETTING;
 import static org.opensearch.knn.index.KNNSettings.KNN_INDEX_REMOTE_VECTOR_BUILD_THRESHOLD_SETTING;
-import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING;
 import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_VECTOR_REPO_SETTING;
 
 /**
@@ -129,15 +130,16 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Repository write took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
 
+            // TODO future implementations will set the following two params depending on some setting to denote the protocol
+            RemoteIndexClient client = RemoteIndexClientFactory.getRemoteIndexClient(RemoteIndexClientFactory.TYPE_HTTP);
+            RemoteBuildRequest request = RemoteBuildRequestBuilder.builder(HTTPRemoteBuildRequest.class)
+                .indexSettings(indexSettings)
+                .indexInfo(indexInfo)
+                .repositoryMetadata(getRepository().getMetadata())
+                .blobName(blobName)
+                .build();
             stopWatch = new StopWatch().start();
-            RemoteIndexClient client = getRemoteIndexClient();
-            RemoteBuildRequest remoteBuildRequest = client.constructBuildRequest(
-                indexSettings,
-                indexInfo,
-                getRepository().getMetadata(),
-                blobName
-            );
-            RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild(remoteBuildRequest);
+            RemoteBuildResponse remoteBuildResponse = client.submitVectorBuild(request);
             time_in_millis = stopWatch.stop().totalTime().millis();
             log.debug("Submit vector build took {} ms for vector field [{}]", time_in_millis, indexInfo.getFieldName());
 
@@ -174,16 +176,4 @@ private BlobStoreRepository getRepository() throws RepositoryMissingException {
         assert repository instanceof BlobStoreRepository : "Repository should be instance of BlobStoreRepository";
         return (BlobStoreRepository) repository;
     }
-
-    /**
-     * Determine which implementation of RemoteIndexClient to be used by the build strategy
-     * @return Concrete RemoteIndexClient implementation
-     */
-    private RemoteIndexClient getRemoteIndexClient() {
-        String endpoint = KNNSettings.state().getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey());
-        if (endpoint == null || endpoint.isEmpty()) {
-            throw new IllegalArgumentException("No endpoint set for RemoteIndexClient");
-        }
-        return RemoteIndexHTTPClient.getInstance();
-    }
 }

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

@@ -8,7 +8,6 @@
 import com.google.common.collect.ImmutableSet;
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.index.SpaceType;
-import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.engine.faiss.Faiss;
 import org.opensearch.knn.index.engine.lucene.Lucene;
 import org.opensearch.knn.index.engine.nmslib.Nmslib;
@@ -179,8 +178,8 @@ public KNNLibraryIndexingContext getKNNLibraryIndexingContext(
     }
 
     @Override
-    public Map<String, Object> getRemoteIndexingParameters(BuildIndexParams params) {
-        return knnLibrary.getRemoteIndexingParameters(params);
+    public Map<String, Object> getRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
+        return knnLibrary.getRemoteIndexingParameters(indexInfoParameters);
     }
 
     @Override

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

@@ -7,7 +7,6 @@
 
 import org.opensearch.common.ValidationException;
 import org.opensearch.knn.index.SpaceType;
-import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 
 import java.util.Collections;
 import java.util.List;
@@ -152,9 +151,9 @@ default boolean supportsRemoteIndexBuild() {
 
     /**
      * Get the remote build supported index parameter mapping to be sent to the remote build service.
-     * @param params to parse
+     * @param indexInfoParameters the index parameters from BuildIndexParams
      */
-    default Map<String, Object> getRemoteIndexingParameters(BuildIndexParams params) {
+    default Map<String, Object> getRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
         throw new UnsupportedOperationException("This method must be implemented by the implementing class");
     }
 }

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

@@ -8,7 +8,6 @@
 import com.google.common.collect.ImmutableMap;
 import org.opensearch.knn.common.KNNConstants;
 import org.opensearch.knn.index.SpaceType;
-import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.engine.KNNMethod;
 import org.opensearch.knn.index.engine.KNNMethodConfigContext;
 import org.opensearch.knn.index.engine.KNNMethodContext;
@@ -128,20 +127,22 @@ public Float scoreToRadialThreshold(Float score, SpaceType spaceType) {
         return spaceType.scoreToDistanceTranslation(score);
     }
 
+    // TODO refactor to make the index parameter fetching more intelligent and less cumbersome
     /**
      * Get the parameters that need to be passed to the remote build service for training
-     * @param indexInfo to parse
+     *
+     * @param indexInfoParameters result of indexInfo.getParameters() to parse
      * @return Map of parameters to be used as "index_parameters"
      */
     @Override
-    public Map<String, Object> getRemoteIndexingParameters(BuildIndexParams indexInfo) {
+    public Map<String, Object> getRemoteIndexingParameters(Map<String, Object> indexInfoParameters) {
         Map<String, Object> indexParameters = new HashMap<>();
-        String methodName = (String) indexInfo.getParameters().get(NAME);
+        String methodName = (String) indexInfoParameters.get(NAME);
         indexParameters.put(ALGORITHM, methodName);
-        indexParameters.put(METHOD_PARAMETER_SPACE_TYPE, indexInfo.getParameters().getOrDefault(SPACE_TYPE, INDEX_KNN_DEFAULT_SPACE_TYPE));
+        indexParameters.put(METHOD_PARAMETER_SPACE_TYPE, indexInfoParameters.getOrDefault(SPACE_TYPE, INDEX_KNN_DEFAULT_SPACE_TYPE));
 
-        assert (indexInfo.getParameters().containsKey(PARAMETERS));
-        Object innerParams = indexInfo.getParameters().get(PARAMETERS);
+        assert (indexInfoParameters.containsKey(PARAMETERS));
+        Object innerParams = indexInfoParameters.get(PARAMETERS);
         assert (innerParams instanceof Map);
         {
             Map<String, Object> algorithmParams = new HashMap<>();
@@ -156,7 +157,7 @@ public Map<String, Object> getRemoteIndexingParameters(BuildIndexParams indexInf
                         METHOD_PARAMETER_EF_SEARCH,
                         innerMap.getOrDefault(METHOD_PARAMETER_EF_SEARCH, INDEX_KNN_DEFAULT_ALGO_PARAM_EF_SEARCH)
                     );
-                    Object indexDescription = indexInfo.getParameters().get(INDEX_DESCRIPTION_PARAMETER);
+                    Object indexDescription = indexInfoParameters.get(INDEX_DESCRIPTION_PARAMETER);
                     assert indexDescription instanceof String;
                     algorithmParams.put(METHOD_PARAMETER_M, getMFromIndexDescription((String) indexDescription));
                 }

src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildRequest.java

@@ -10,39 +10,33 @@
 import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.index.IndexSettings;
-import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 
 import java.io.IOException;
-import java.net.URI;
 
 import static org.opensearch.knn.common.KNNConstants.CONTAINER_NAME;
 import static org.opensearch.knn.common.KNNConstants.DIMENSION;
 import static org.opensearch.knn.common.KNNConstants.DOC_COUNT;
-import static org.opensearch.knn.common.KNNConstants.VECTOR_PATH;
 import static org.opensearch.knn.common.KNNConstants.DOC_ID_PATH;
 import static org.opensearch.knn.common.KNNConstants.INDEX_PARAMETERS;
 import static org.opensearch.knn.common.KNNConstants.KNN_ENGINE;
 import static org.opensearch.knn.common.KNNConstants.REPOSITORY_TYPE;
 import static org.opensearch.knn.common.KNNConstants.TENANT_ID;
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
-import static org.opensearch.knn.index.KNNSettings.KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING;
+import static org.opensearch.knn.common.KNNConstants.VECTOR_PATH;
 
 /**
- * RemoteBuildRequest implementation for HTTP clients that sets the endpoint and offers a JSON conversion.
+ * HTTP-specific implementation of RemoteBuildRequest.
  */
 @Getter
 public class HTTPRemoteBuildRequest extends RemoteBuildRequest {
-    private final URI endpoint;
-
     public HTTPRemoteBuildRequest(
         IndexSettings indexSettings,
         BuildIndexParams indexInfo,
         RepositoryMetadata repositoryMetadata,
         String blobName
     ) throws IOException {
         super(indexSettings, indexInfo, repositoryMetadata, blobName);
-        this.endpoint = URI.create(KNNSettings.state().getSettingValue(KNN_REMOTE_BUILD_SERVICE_ENDPOINT_SETTING.getKey()));
     }
 
     public String toJson() throws IOException {
@@ -55,7 +49,7 @@ public String toJson() throws IOException {
             builder.field(TENANT_ID, tenantId);
             builder.field(DIMENSION, dimension);
             builder.field(DOC_COUNT, docCount);
-            builder.field(VECTOR_DATA_TYPE_FIELD, dataType);
+            builder.field(VECTOR_DATA_TYPE_FIELD, vectorDataType);
             builder.field(KNN_ENGINE, engine);
             builder.field(INDEX_PARAMETERS, indexParameters);
             builder.endObject();

src/main/java/org/opensearch/knn/index/remote/HTTPRemoteBuildResponse.java

@@ -5,9 +5,9 @@
 
 package org.opensearch.knn.index.remote;
 
+import lombok.Getter;
 import org.opensearch.cluster.ClusterName;
 import org.opensearch.cluster.metadata.RepositoryMetadata;
-import lombok.Getter;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
 import org.opensearch.knn.index.codec.nativeindex.remote.RemoteIndexBuildStrategy;
@@ -20,28 +20,22 @@
 import static org.opensearch.knn.common.KNNConstants.BUCKET;
 import static org.opensearch.knn.common.KNNConstants.S3;
 
+/**
+ * Abstract base class for Remote Build Requests.
+ */
 @Getter
 public abstract class RemoteBuildRequest {
-    public String repositoryType;
-    public String containerName;
-    public String vectorPath;
-    public String docIdPath;
-    public String tenantId;
-    public int dimension;
-    public int docCount;
-    public String dataType;
-    public String engine;
-    public Map<String, Object> indexParameters;
+    protected String repositoryType;
+    protected String containerName;
+    protected String vectorPath;
+    protected String docIdPath;
+    protected String tenantId;
+    protected int dimension;
+    protected int docCount;
+    protected String vectorDataType;
+    protected String engine;
+    protected Map<String, Object> indexParameters;
 
-    /**
-     * Construct the RemoteBuildRequest object for the index build request.
-     * Regardless of the Client implementation, the remote build service will need these parameters to build the index.
-     * This method being included in the interface ensures that the request parameters will be set for each Client implementation to use or transform for its needs.
-     * @param indexSettings Index settings
-     * @param indexInfo Index parameters
-     * @param repositoryMetadata Metadata of the repository containing the index
-     * @param blobName File name generated by the Build Strategy with a UUID
-     */
     public RemoteBuildRequest(
         IndexSettings indexSettings,
         BuildIndexParams indexInfo,
@@ -62,7 +56,7 @@ public RemoteBuildRequest(
         KNNCodecUtil.initializeVectorValues(vectorValues);
         assert (vectorValues.dimension() > 0);
 
-        Map<String, Object> indexParameters = indexInfo.getKnnEngine().getRemoteIndexingParameters(indexInfo);
+        Map<String, Object> indexParameters = indexInfo.getKnnEngine().getRemoteIndexingParameters(indexInfo.getParameters());
 
         this.repositoryType = repositoryType;
         this.containerName = containerName;
@@ -71,8 +65,13 @@ public RemoteBuildRequest(
         this.tenantId = indexSettings.getSettings().get(ClusterName.CLUSTER_NAME_SETTING.getKey());
         this.dimension = vectorValues.dimension();
         this.docCount = indexInfo.getTotalLiveDocs();
-        this.dataType = vectorDataType;
+        this.vectorDataType = vectorDataType;
         this.engine = indexInfo.getKnnEngine().getName();
         this.indexParameters = indexParameters;
     }
+
+    /**
+     * Convert the request to JSON format.
+     */
+    public abstract String toJson() throws IOException;
 }

src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequest.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.knn.index.remote;
+
+import org.opensearch.cluster.metadata.RepositoryMetadata;
+import org.opensearch.index.IndexSettings;
+import org.opensearch.knn.index.codec.nativeindex.model.BuildIndexParams;
+
+import java.io.IOException;
+
+/**
+ * Generic Builder for constructing RemoteBuildRequest objects per requestType (HTTP, etc).
+ */
+public class RemoteBuildRequestBuilder<T extends RemoteBuildRequest> {
+    private IndexSettings indexSettings;
+    private BuildIndexParams indexInfo;
+    private RepositoryMetadata repositoryMetadata;
+    private String blobName;
+    private final Class<T> requestType;
+
+    // Private constructor
+    private RemoteBuildRequestBuilder(Class<T> requestType) {
+        this.requestType = requestType;
+    }
+
+    /**
+     * Static factory method to create a builder instance.
+     * @param requestType Class type of the RemoteBuildRequest
+     * @return RemoteBuildRequestBuilder instance
+     */
+    public static <T extends RemoteBuildRequest> RemoteBuildRequestBuilder<T> builder(Class<T> requestType) {
+        return new RemoteBuildRequestBuilder<>(requestType);
+    }
+
+    public RemoteBuildRequestBuilder<T> indexSettings(IndexSettings indexSettings) {
+        this.indexSettings = indexSettings;
+        return this;
+    }
+
+    public RemoteBuildRequestBuilder<T> indexInfo(BuildIndexParams indexInfo) {
+        this.indexInfo = indexInfo;
+        return this;
+    }
+
+    public RemoteBuildRequestBuilder<T> repositoryMetadata(RepositoryMetadata repositoryMetadata) {
+        this.repositoryMetadata = repositoryMetadata;
+        return this;
+    }
+
+    public RemoteBuildRequestBuilder<T> blobName(String blobName) {
+        this.blobName = blobName;
+        return this;
+    }
+
+    public T build() throws IOException {
+        try {
+            return requestType.getConstructor(IndexSettings.class, BuildIndexParams.class, RepositoryMetadata.class, String.class)
+                .newInstance(indexSettings, indexInfo, repositoryMetadata, blobName);
+        } catch (Exception e) {
+            throw new IOException("Failed to instantiate RemoteBuildRequest", e);
+        }
+    }
+}

src/main/java/org/opensearch/knn/index/remote/RemoteBuildRequestBuilder.java

@@ -6,8 +6,8 @@
 package org.opensearch.knn.index.remote;
 
 /**
- * Generic remote build response interface
+ * Remote build response. Currently, this just contains the jobId from the server.
+ * In the future, this may be an interface if different clients expect different responses.
  */
-public interface RemoteBuildResponse {
-    String getJobId();
+public record RemoteBuildResponse(String jobId) {
 }