From dac50127733be45b8fedec9df6ddd5cf2bea817c Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Wed, 23 Aug 2023 13:29:13 -0700 Subject: [PATCH 01/17] Support for default Model Id Signed-off-by: Varun Jain --- .../neuralsearch/plugin/NeuralSearch.java | 10 ++ .../processor/DefaultValueProcessor.java | 100 ++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index b46d2bc6d..3de2a3133 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -5,6 +5,7 @@ package org.opensearch.neuralsearch.plugin; +import org.opensearch.neuralsearch.processor.DefaultValueProcessor; import static org.opensearch.neuralsearch.settings.NeuralSearchSettings.NEURAL_SEARCH_HYBRID_SEARCH_ENABLED; import java.util.Arrays; @@ -50,6 +51,7 @@ import org.opensearch.repositories.RepositoriesService; import org.opensearch.script.ScriptService; import org.opensearch.search.pipeline.SearchPhaseResultsProcessor; +import org.opensearch.search.pipeline.SearchRequestProcessor; import org.opensearch.search.query.QueryPhaseSearcher; import org.opensearch.threadpool.ThreadPool; import org.opensearch.watcher.ResourceWatcherService; @@ -125,4 +127,12 @@ public Map> getSettings() { return List.of(NEURAL_SEARCH_HYBRID_SEARCH_ENABLED); } + + @Override + public Map> getRequestProcessors(Parameters parameters) { + return Map.of( + DefaultValueProcessor.TYPE, + new DefaultValueProcessor.Factory(parameters.namedXContentRegistry) + ); + } } diff --git a/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java new file mode 100644 index 000000000..c3f3e6b7a --- /dev/null +++ b/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java @@ -0,0 +1,100 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.processor; + +import java.io.InputStream; +import java.util.Map; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.common.xcontent.LoggingDeprecationHandler; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.core.xcontent.XContentParser; +import static org.opensearch.index.query.AbstractQueryBuilder.parseInnerQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.ingest.ConfigurationUtils; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; +import org.opensearch.search.pipeline.AbstractProcessor; +import org.opensearch.search.pipeline.Processor; +import org.opensearch.search.pipeline.SearchRequestProcessor; + +public class DefaultValueProcessor extends AbstractProcessor implements SearchRequestProcessor { + + /** + * Key to reference this processor type from a search pipeline. + */ + public static final String TYPE = "default_query"; + + public static final String MODEL_ID_FIELD = "model_id"; + public static final String NEURAL_FIELD_MAP_FIELD = "neural_field_map"; + + final QueryBuilder queryBuilder; + + /** + * Returns the type of the processor. + * + * @return The processor type. + */ + @Override + public String getType() { + return TYPE; + } + + + protected DefaultValueProcessor(String tag, String description, boolean ignoreFailure, QueryBuilder neuralQueryBuilder) { + super(tag, description, ignoreFailure); + this.queryBuilder = neuralQueryBuilder; + } + + @Override + public SearchRequest processRequest(SearchRequest searchRequest) throws Exception { + QueryBuilder originalQuery = null; + if (searchRequest.source() != null) { + originalQuery = searchRequest.source().query(); + } + + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(originalQuery, MODEL_ID_FIELD, NEURAL_FIELD_MAP_FIELD); + + searchRequest.source().query(neuralQueryBuilder); + + return searchRequest; + } + + public static class Factory implements Processor.Factory { + private static final String QUERY_KEY = "query"; + private final NamedXContentRegistry namedXContentRegistry; + + public Factory(NamedXContentRegistry namedXContentRegistry) { + this.namedXContentRegistry = namedXContentRegistry; + } + + @Override + public DefaultValueProcessor create( + Map> processorFactories, + String tag, + String description, + boolean ignoreFailure, + Map config, + PipelineContext pipelineContext + ) throws Exception { + Map query = ConfigurationUtils.readOptionalMap(TYPE, tag, config, QUERY_KEY); + if (query == null) { + throw new IllegalArgumentException("Did not specify the " + QUERY_KEY + " property in processor of type " + TYPE); + } + try ( + XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent).map(query); + InputStream stream = BytesReference.bytes(builder).streamInput(); + XContentParser parser = MediaTypeRegistry.JSON.xContent() + .createParser(namedXContentRegistry, LoggingDeprecationHandler.INSTANCE, stream) + ) { + return new DefaultValueProcessor(tag, description, ignoreFailure, parseInnerQueryBuilder(parser)); + } + } + } + +} From 6903f9415be28b4b0a4532468272cc8667523d5f Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Mon, 25 Sep 2023 12:56:13 -0700 Subject: [PATCH 02/17] Support for Default Model id Signed-off-by: Varun Jain --- .../neuralsearch/plugin/NeuralSearch.java | 8 +- .../processor/DefaultValueProcessor.java | 100 ------------------ .../processor/NeuralQueryProcessor.java | 81 ++++++++++++++ .../query/NeuralQueryBuilder.java | 63 ++++++----- .../visitor/NeuralSearchQueryVisitor.java | 48 +++++++++ .../util/NeuralSearchClusterUtil.java | 57 ++++++++++ 6 files changed, 229 insertions(+), 128 deletions(-) delete mode 100644 src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java create mode 100644 src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java create mode 100644 src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java create mode 100644 src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index 3de2a3133..233acadcb 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -5,7 +5,7 @@ package org.opensearch.neuralsearch.plugin; -import org.opensearch.neuralsearch.processor.DefaultValueProcessor; +import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; import static org.opensearch.neuralsearch.settings.NeuralSearchSettings.NEURAL_SEARCH_HYBRID_SEARCH_ENABLED; import java.util.Arrays; @@ -42,6 +42,7 @@ import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; import org.opensearch.neuralsearch.search.query.HybridQueryPhaseSearcher; +import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.ExtensiblePlugin; import org.opensearch.plugins.IngestPlugin; @@ -80,6 +81,7 @@ public Collection createComponents( final IndexNameExpressionResolver indexNameExpressionResolver, final Supplier repositoriesServiceSupplier ) { + NeuralSearchClusterUtil.instance().initialize(clusterService); NeuralQueryBuilder.initialize(clientAccessor); normalizationProcessorWorkflow = new NormalizationProcessorWorkflow(new ScoreNormalizer(), new ScoreCombiner()); return List.of(clientAccessor); @@ -131,8 +133,8 @@ public List> getSettings() { @Override public Map> getRequestProcessors(Parameters parameters) { return Map.of( - DefaultValueProcessor.TYPE, - new DefaultValueProcessor.Factory(parameters.namedXContentRegistry) + NeuralQueryProcessor.TYPE, + new NeuralQueryProcessor.Factory() ); } } diff --git a/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java deleted file mode 100644 index c3f3e6b7a..000000000 --- a/src/main/java/org/opensearch/neuralsearch/processor/DefaultValueProcessor.java +++ /dev/null @@ -1,100 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.neuralsearch.processor; - -import java.io.InputStream; -import java.util.Map; -import org.opensearch.action.search.SearchRequest; -import org.opensearch.common.xcontent.LoggingDeprecationHandler; -import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.core.common.bytes.BytesReference; -import org.opensearch.core.xcontent.MediaTypeRegistry; -import org.opensearch.core.xcontent.NamedXContentRegistry; -import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.core.xcontent.XContentParser; -import static org.opensearch.index.query.AbstractQueryBuilder.parseInnerQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.ingest.ConfigurationUtils; -import org.opensearch.neuralsearch.query.NeuralQueryBuilder; -import org.opensearch.search.pipeline.AbstractProcessor; -import org.opensearch.search.pipeline.Processor; -import org.opensearch.search.pipeline.SearchRequestProcessor; - -public class DefaultValueProcessor extends AbstractProcessor implements SearchRequestProcessor { - - /** - * Key to reference this processor type from a search pipeline. - */ - public static final String TYPE = "default_query"; - - public static final String MODEL_ID_FIELD = "model_id"; - public static final String NEURAL_FIELD_MAP_FIELD = "neural_field_map"; - - final QueryBuilder queryBuilder; - - /** - * Returns the type of the processor. - * - * @return The processor type. - */ - @Override - public String getType() { - return TYPE; - } - - - protected DefaultValueProcessor(String tag, String description, boolean ignoreFailure, QueryBuilder neuralQueryBuilder) { - super(tag, description, ignoreFailure); - this.queryBuilder = neuralQueryBuilder; - } - - @Override - public SearchRequest processRequest(SearchRequest searchRequest) throws Exception { - QueryBuilder originalQuery = null; - if (searchRequest.source() != null) { - originalQuery = searchRequest.source().query(); - } - - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(originalQuery, MODEL_ID_FIELD, NEURAL_FIELD_MAP_FIELD); - - searchRequest.source().query(neuralQueryBuilder); - - return searchRequest; - } - - public static class Factory implements Processor.Factory { - private static final String QUERY_KEY = "query"; - private final NamedXContentRegistry namedXContentRegistry; - - public Factory(NamedXContentRegistry namedXContentRegistry) { - this.namedXContentRegistry = namedXContentRegistry; - } - - @Override - public DefaultValueProcessor create( - Map> processorFactories, - String tag, - String description, - boolean ignoreFailure, - Map config, - PipelineContext pipelineContext - ) throws Exception { - Map query = ConfigurationUtils.readOptionalMap(TYPE, tag, config, QUERY_KEY); - if (query == null) { - throw new IllegalArgumentException("Did not specify the " + QUERY_KEY + " property in processor of type " + TYPE); - } - try ( - XContentBuilder builder = XContentBuilder.builder(JsonXContent.jsonXContent).map(query); - InputStream stream = BytesReference.bytes(builder).streamInput(); - XContentParser parser = MediaTypeRegistry.JSON.xContent() - .createParser(namedXContentRegistry, LoggingDeprecationHandler.INSTANCE, stream) - ) { - return new DefaultValueProcessor(tag, description, ignoreFailure, parseInnerQueryBuilder(parser)); - } - } - } - -} diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java new file mode 100644 index 000000000..2b9906e84 --- /dev/null +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.processor; + +import java.util.Map; + +import org.opensearch.action.search.SearchRequest; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.ingest.ConfigurationUtils; +import org.opensearch.neuralsearch.query.visitor.NeuralSearchQueryVisitor; +import org.opensearch.search.pipeline.AbstractProcessor; +import org.opensearch.search.pipeline.Processor; +import org.opensearch.search.pipeline.SearchRequestProcessor; + +public class NeuralQueryProcessor extends AbstractProcessor implements SearchRequestProcessor { + + /** + * Key to reference this processor type from a search pipeline. + */ + public static final String TYPE = "default_query"; + + private final String modelId; + + private final Map fieldInfoMap; + + /** + * Returns the type of the processor. + * + * @return The processor type. + */ + @Override + public String getType() { + return TYPE; + } + + protected NeuralQueryProcessor( + String tag, + String description, + boolean ignoreFailure, + String modelId, + Map fieldInfoMap + ) { + super(tag, description, ignoreFailure); + this.modelId = modelId; + this.fieldInfoMap = fieldInfoMap; + } + + @Override + public SearchRequest processRequest(SearchRequest searchRequest) throws Exception { + QueryBuilder queryBuilder = searchRequest.source().query(); + queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, fieldInfoMap)); + return searchRequest; + } + + public static class Factory implements Processor.Factory { + private static final String DEFAULT_MODEL_ID = "default_model_id"; + private static final String NEURAL_FIELD_MAP = "neural_field_map"; + + @Override + public NeuralQueryProcessor create( + Map> processorFactories, + String tag, + String description, + boolean ignoreFailure, + Map config, + PipelineContext pipelineContext + ) throws Exception { + String modelId = (String) config.remove(DEFAULT_MODEL_ID); + Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_MAP); + + if (modelId == null && neuralInfoMap == null) { + throw new IllegalArgumentException("model Id or neural info map either of them should be provided"); + } + + return new NeuralQueryProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); + } + } +} diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index ebcd9a88b..c59a1173c 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -22,6 +22,7 @@ import org.apache.commons.lang.builder.EqualsBuilder; import org.apache.commons.lang.builder.HashCodeBuilder; import org.apache.lucene.search.Query; +import org.opensearch.Version; import org.opensearch.common.SetOnce; import org.opensearch.core.ParseField; import org.opensearch.core.action.ActionListener; @@ -31,12 +32,10 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.index.query.AbstractQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryRewriteContext; -import org.opensearch.index.query.QueryShardContext; +import org.opensearch.index.query.*; import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; +import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import com.google.common.annotations.VisibleForTesting; @@ -82,6 +81,7 @@ public static void initialize(MLCommonsClientAccessor mlClient) { @Setter(AccessLevel.PACKAGE) private Supplier vectorSupplier; private QueryBuilder filter; + private static final Version MINIMAL_SUPPORTED_VERSION_DEFAULT_MODEL_ID = Version.V_2_11_0; /** * Constructor from stream input @@ -93,7 +93,11 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { super(in); this.fieldName = in.readString(); this.queryText = in.readString(); - this.modelId = in.readString(); + if (isClusterOnOrAfterMinRequiredVersion()) { + this.modelId = in.readOptionalString(); + } else { + this.modelId = in.readString(); + } this.k = in.readVInt(); this.filter = in.readOptionalNamedWriteable(QueryBuilder.class); } @@ -102,7 +106,11 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { protected void doWriteTo(StreamOutput out) throws IOException { out.writeString(this.fieldName); out.writeString(this.queryText); - out.writeString(this.modelId); + if (isClusterOnOrAfterMinRequiredVersion()) { + out.writeOptionalString(this.modelId); + } else { + out.writeString(this.modelId); + } out.writeVInt(this.k); out.writeOptionalNamedWriteable(this.filter); } @@ -152,20 +160,21 @@ public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOEx parseQueryParams(parser, neuralQueryBuilder); if (parser.nextToken() != XContentParser.Token.END_OBJECT) { throw new ParsingException( - parser.getTokenLocation(), - "[" - + NAME - + "] query doesn't support multiple fields, found [" - + neuralQueryBuilder.fieldName() - + "] and [" - + parser.currentName() - + "]" + parser.getTokenLocation(), + "[" + + NAME + + "] query doesn't support multiple fields, found [" + + neuralQueryBuilder.fieldName() + + "] and [" + + parser.currentName() + + "]" ); } requireValue(neuralQueryBuilder.queryText(), "Query text must be provided for neural query"); requireValue(neuralQueryBuilder.fieldName(), "Field name must be provided for neural query"); - requireValue(neuralQueryBuilder.modelId(), "Model ID must be provided for neural query"); - + if (!isClusterOnOrAfterMinRequiredVersion()) { + requireValue(neuralQueryBuilder.modelId(), "Model ID must be provided for neural query"); + } return neuralQueryBuilder; } @@ -188,8 +197,8 @@ private static void parseQueryParams(XContentParser parser, NeuralQueryBuilder n neuralQueryBuilder.boost(parser.floatValue()); } else { throw new ParsingException( - parser.getTokenLocation(), - "[" + NAME + "] query does not support [" + currentFieldName + "]" + parser.getTokenLocation(), + "[" + NAME + "] query does not support [" + currentFieldName + "]" ); } } else if (token == XContentParser.Token.START_OBJECT) { @@ -198,8 +207,8 @@ private static void parseQueryParams(XContentParser parser, NeuralQueryBuilder n } } else { throw new ParsingException( - parser.getTokenLocation(), - "[" + NAME + "] unknown token [" + token + "] after [" + currentFieldName + "]" + parser.getTokenLocation(), + "[" + NAME + "] unknown token [" + token + "] after [" + currentFieldName + "]" ); } } @@ -222,10 +231,10 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { SetOnce vectorSetOnce = new SetOnce<>(); queryRewriteContext.registerAsyncAction( - ((client, actionListener) -> ML_CLIENT.inferenceSentence(modelId(), queryText(), ActionListener.wrap(floatList -> { - vectorSetOnce.set(vectorAsListToArray(floatList)); - actionListener.onResponse(null); - }, actionListener::onFailure))) + ((client, actionListener) -> ML_CLIENT.inferenceSentence(modelId(), queryText(), ActionListener.wrap(floatList -> { + vectorSetOnce.set(vectorAsListToArray(floatList)); + actionListener.onResponse(null); + }, actionListener::onFailure))) ); return new NeuralQueryBuilder(fieldName(), queryText(), modelId(), k(), vectorSetOnce::get, filter()); } @@ -258,4 +267,8 @@ protected int doHashCode() { public String getWriteableName() { return NAME; } -} + + private static boolean isClusterOnOrAfterMinRequiredVersion() { + return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_DEFAULT_MODEL_ID); + } +} \ No newline at end of file diff --git a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java new file mode 100644 index 000000000..c1f66ee9e --- /dev/null +++ b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java @@ -0,0 +1,48 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.query.visitor; + +import java.util.Map; + +import org.apache.lucene.search.BooleanClause; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; + +public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { + + private String modelId; + private Map neuralFieldMap; + + public NeuralSearchQueryVisitor(String modelId, Map neuralFieldMap) { + this.modelId = modelId; + this.neuralFieldMap = neuralFieldMap; + } + + @Override + public void accept(QueryBuilder queryBuilder) { + if (queryBuilder instanceof NeuralQueryBuilder) { + NeuralQueryBuilder neuralQueryBuilder = (NeuralQueryBuilder) queryBuilder; + if (neuralFieldMap != null + && neuralQueryBuilder.fieldName() != null + && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { + String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); + neuralQueryBuilder.modelId(fieldDefaultModelId); + } else if (modelId != null) { + neuralQueryBuilder.modelId(modelId); + } else { + throw new IllegalArgumentException( + "model id must be provided in neural query or a default model id must be set in search request processor" + ); + } + } + } + + @Override + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + return this; + } +} \ No newline at end of file diff --git a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java new file mode 100644 index 000000000..1e0b22094 --- /dev/null +++ b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.util; + +import lombok.AccessLevel; +import lombok.NoArgsConstructor; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.Version; +import org.opensearch.cluster.service.ClusterService; + +@NoArgsConstructor(access = AccessLevel.PRIVATE) +@Log4j2 +public class NeuralSearchClusterUtil { + private ClusterService clusterService; + + private static NeuralSearchClusterUtil instance; + + /** + * Return instance of the cluster context, must be initialized first for proper usage + * @return instance of cluster context + */ + public static synchronized NeuralSearchClusterUtil instance() { + if (instance == null) { + instance = new NeuralSearchClusterUtil(); + } + return instance; + } + + /** + * Initializes instance of cluster context by injecting dependencies + * @param clusterService + */ + public void initialize(final ClusterService clusterService) { + this.clusterService = clusterService; + } + + /** + * Return minimal OpenSearch version based on all nodes currently discoverable in the cluster + * @return minimal installed OpenSearch version, default to Version.CURRENT which is typically the latest version + */ + public Version getClusterMinVersion() { + try { + return this.clusterService.state().getNodes().getMinNodeVersion(); + } catch (Exception exception) { + log.error( + String.format("Failed to get cluster minimum node version, returning current node version %s instead.", Version.CURRENT), + exception + ); + return Version.CURRENT; + } + } + +} \ No newline at end of file From c06d07a0fa271b0cc02885bf92370fd348c85b38 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Mon, 25 Sep 2023 17:54:06 -0700 Subject: [PATCH 03/17] Support for default model Id Signed-off-by: Varun Jain --- .../neuralsearch/plugin/NeuralSearch.java | 12 ++--- .../processor/NeuralQueryProcessor.java | 36 ++++++------- .../query/NeuralQueryBuilder.java | 34 ++++++------- .../visitor/NeuralSearchQueryVisitor.java | 8 +-- .../util/NeuralSearchClusterUtil.java | 12 +++-- .../common/BaseNeuralSearchIT.java | 23 +++++++++ .../plugin/NeuralSearchTests.java | 12 +++++ .../processor/NeuralQueryProcessorTests.java | 50 ++++++++++++++++++ .../util/NeuralSearchClusterTestUtils.java | 32 ++++++++++++ .../util/NeuralSearchClusterUtilTests.java | 51 +++++++++++++++++++ .../SearchRequestPipelineConfiguration.json | 11 ++++ 11 files changed, 233 insertions(+), 48 deletions(-) create mode 100644 src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java create mode 100644 src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterTestUtils.java create mode 100644 src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java create mode 100644 src/test/resources/processor/SearchRequestPipelineConfiguration.json diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index b01ef703d..57ac4feb5 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -5,8 +5,8 @@ package org.opensearch.neuralsearch.plugin; -import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; import static org.opensearch.neuralsearch.settings.NeuralSearchSettings.NEURAL_SEARCH_HYBRID_SEARCH_DISABLED; + import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -29,6 +29,7 @@ import org.opensearch.ingest.Processor; import org.opensearch.ml.client.MachineLearningNodeClient; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; +import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; import org.opensearch.neuralsearch.processor.NormalizationProcessor; import org.opensearch.neuralsearch.processor.NormalizationProcessorWorkflow; import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; @@ -132,10 +133,9 @@ public List> getSettings() { } @Override - public Map> getRequestProcessors(Parameters parameters) { - return Map.of( - NeuralQueryProcessor.TYPE, - new NeuralQueryProcessor.Factory() - ); + public Map> getRequestProcessors( + Parameters parameters + ) { + return Map.of(NeuralQueryProcessor.TYPE, new NeuralQueryProcessor.Factory()); } } diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java index 2b9906e84..59a8688ca 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -20,11 +20,11 @@ public class NeuralQueryProcessor extends AbstractProcessor implements SearchReq /** * Key to reference this processor type from a search pipeline. */ - public static final String TYPE = "default_query"; + public static final String TYPE = "neural_query"; - private final String modelId; + final String modelId; - private final Map fieldInfoMap; + final Map neuralFieldMap; /** * Returns the type of the processor. @@ -37,39 +37,39 @@ public String getType() { } protected NeuralQueryProcessor( - String tag, - String description, - boolean ignoreFailure, - String modelId, - Map fieldInfoMap + String tag, + String description, + boolean ignoreFailure, + String modelId, + Map fieldInfoMap ) { super(tag, description, ignoreFailure); this.modelId = modelId; - this.fieldInfoMap = fieldInfoMap; + this.neuralFieldMap = fieldInfoMap; } @Override public SearchRequest processRequest(SearchRequest searchRequest) throws Exception { QueryBuilder queryBuilder = searchRequest.source().query(); - queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, fieldInfoMap)); + queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, neuralFieldMap)); return searchRequest; } public static class Factory implements Processor.Factory { private static final String DEFAULT_MODEL_ID = "default_model_id"; - private static final String NEURAL_FIELD_MAP = "neural_field_map"; + private static final String NEURAL_FIELD_DEFAULT_ID = "neural_field_default_id"; @Override public NeuralQueryProcessor create( - Map> processorFactories, - String tag, - String description, - boolean ignoreFailure, - Map config, - PipelineContext pipelineContext + Map> processorFactories, + String tag, + String description, + boolean ignoreFailure, + Map config, + PipelineContext pipelineContext ) throws Exception { String modelId = (String) config.remove(DEFAULT_MODEL_ID); - Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_MAP); + Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_DEFAULT_ID); if (modelId == null && neuralInfoMap == null) { throw new IllegalArgumentException("model Id or neural info map either of them should be provided"); diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index c59a1173c..799cc22de 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -160,14 +160,14 @@ public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOEx parseQueryParams(parser, neuralQueryBuilder); if (parser.nextToken() != XContentParser.Token.END_OBJECT) { throw new ParsingException( - parser.getTokenLocation(), - "[" - + NAME - + "] query doesn't support multiple fields, found [" - + neuralQueryBuilder.fieldName() - + "] and [" - + parser.currentName() - + "]" + parser.getTokenLocation(), + "[" + + NAME + + "] query doesn't support multiple fields, found [" + + neuralQueryBuilder.fieldName() + + "] and [" + + parser.currentName() + + "]" ); } requireValue(neuralQueryBuilder.queryText(), "Query text must be provided for neural query"); @@ -197,8 +197,8 @@ private static void parseQueryParams(XContentParser parser, NeuralQueryBuilder n neuralQueryBuilder.boost(parser.floatValue()); } else { throw new ParsingException( - parser.getTokenLocation(), - "[" + NAME + "] query does not support [" + currentFieldName + "]" + parser.getTokenLocation(), + "[" + NAME + "] query does not support [" + currentFieldName + "]" ); } } else if (token == XContentParser.Token.START_OBJECT) { @@ -207,8 +207,8 @@ private static void parseQueryParams(XContentParser parser, NeuralQueryBuilder n } } else { throw new ParsingException( - parser.getTokenLocation(), - "[" + NAME + "] unknown token [" + token + "] after [" + currentFieldName + "]" + parser.getTokenLocation(), + "[" + NAME + "] unknown token [" + token + "] after [" + currentFieldName + "]" ); } } @@ -231,10 +231,10 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { SetOnce vectorSetOnce = new SetOnce<>(); queryRewriteContext.registerAsyncAction( - ((client, actionListener) -> ML_CLIENT.inferenceSentence(modelId(), queryText(), ActionListener.wrap(floatList -> { - vectorSetOnce.set(vectorAsListToArray(floatList)); - actionListener.onResponse(null); - }, actionListener::onFailure))) + ((client, actionListener) -> ML_CLIENT.inferenceSentence(modelId(), queryText(), ActionListener.wrap(floatList -> { + vectorSetOnce.set(vectorAsListToArray(floatList)); + actionListener.onResponse(null); + }, actionListener::onFailure))) ); return new NeuralQueryBuilder(fieldName(), queryText(), modelId(), k(), vectorSetOnce::get, filter()); } @@ -271,4 +271,4 @@ public String getWriteableName() { private static boolean isClusterOnOrAfterMinRequiredVersion() { return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_DEFAULT_MODEL_ID); } -} \ No newline at end of file +} diff --git a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java index c1f66ee9e..2d0b53eaa 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java +++ b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java @@ -27,15 +27,15 @@ public void accept(QueryBuilder queryBuilder) { if (queryBuilder instanceof NeuralQueryBuilder) { NeuralQueryBuilder neuralQueryBuilder = (NeuralQueryBuilder) queryBuilder; if (neuralFieldMap != null - && neuralQueryBuilder.fieldName() != null - && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { + && neuralQueryBuilder.fieldName() != null + && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); neuralQueryBuilder.modelId(fieldDefaultModelId); } else if (modelId != null) { neuralQueryBuilder.modelId(modelId); } else { throw new IllegalArgumentException( - "model id must be provided in neural query or a default model id must be set in search request processor" + "model id must be provided in neural query or a default model id must be set in search request processor" ); } } @@ -45,4 +45,4 @@ public void accept(QueryBuilder queryBuilder) { public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { return this; } -} \ No newline at end of file +} diff --git a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java index 1e0b22094..f9e044c53 100644 --- a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java +++ b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java @@ -5,6 +5,8 @@ package org.opensearch.neuralsearch.util; +import java.util.Locale; + import lombok.AccessLevel; import lombok.NoArgsConstructor; import lombok.extern.log4j.Log4j2; @@ -47,11 +49,15 @@ public Version getClusterMinVersion() { return this.clusterService.state().getNodes().getMinNodeVersion(); } catch (Exception exception) { log.error( - String.format("Failed to get cluster minimum node version, returning current node version %s instead.", Version.CURRENT), - exception + String.format( + Locale.ROOT, + "Failed to get cluster minimum node version, returning current node version %s instead.", + Version.CURRENT + ), + exception ); return Version.CURRENT; } } -} \ No newline at end of file +} diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java index b144ade6c..fac32b538 100644 --- a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java +++ b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java @@ -253,6 +253,29 @@ protected void createPipelineProcessor(String modelId, String pipelineName) thro assertEquals("true", node.get("acknowledged").toString()); } + protected void createSearchRequestProcessor(String modelId, String pipelineName) throws Exception { + Response pipelineCreateResponse = makeRequest( + client(), + "PUT", + "/_search/pipeline/" + pipelineName, + null, + toHttpEntity( + String.format( + LOCALE, + Files.readString(Path.of(classLoader.getResource("processor/SearchRequestPipelineConfiguration.json").toURI())), + modelId + ) + ), + ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT)) + ); + Map node = XContentHelper.convertToMap( + XContentType.JSON.xContent(), + EntityUtils.toString(pipelineCreateResponse.getEntity()), + false + ); + assertEquals("true", node.get("acknowledged").toString()); + } + /** * Get the number of documents in a particular index * diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java index 8918e174c..9c8a23d03 100644 --- a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java +++ b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java @@ -12,6 +12,7 @@ import java.util.Optional; import org.opensearch.ingest.Processor; +import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; import org.opensearch.neuralsearch.processor.NormalizationProcessor; import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory; @@ -22,6 +23,7 @@ import org.opensearch.plugins.SearchPipelinePlugin; import org.opensearch.plugins.SearchPlugin; import org.opensearch.search.pipeline.SearchPhaseResultsProcessor; +import org.opensearch.search.pipeline.SearchRequestProcessor; import org.opensearch.search.query.QueryPhaseSearcher; public class NeuralSearchTests extends OpenSearchQueryTestCase { @@ -73,4 +75,14 @@ public void testSearchPhaseResultsProcessors() { ); assertTrue(scoringProcessor instanceof NormalizationProcessorFactory); } + + public void testRequestProcessors() { + NeuralSearch plugin = new NeuralSearch(); + SearchPipelinePlugin.Parameters parameters = mock(SearchPipelinePlugin.Parameters.class); + Map> processors = plugin.getRequestProcessors( + parameters + ); + assertNotNull(processors); + assertNotNull(processors.get(NeuralQueryProcessor.TYPE)); + } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java new file mode 100644 index 000000000..0cd9bea3a --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java @@ -0,0 +1,50 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.processor; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import org.opensearch.action.search.SearchRequest; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; +import org.opensearch.search.builder.SearchSourceBuilder; +import org.opensearch.test.OpenSearchTestCase; + +public class NeuralQueryProcessorTests extends OpenSearchTestCase { + + public void testFactory() throws Exception { + NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + NeuralQueryProcessor processor = createTestProcessor(factory); + assertEquals("vasdcvkcjkbldbjkd", processor.modelId); + assertEquals("bahbkcdkacb", processor.neuralFieldMap.get("fieldName").toString()); + + // Missing "query" parameter: + expectThrows( + IllegalArgumentException.class, + () -> factory.create(Collections.emptyMap(), null, null, false, Collections.emptyMap(), null) + ); + } + + public void testProcessRequest() throws Exception { + NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + SearchRequest searchRequest = new SearchRequest(); + searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder)); + NeuralQueryProcessor processor = createTestProcessor(factory); + SearchRequest processSearchRequest = processor.processRequest(searchRequest); + assertEquals(processSearchRequest, searchRequest); + } + + public NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory factory) throws Exception { + Map configMap = new HashMap<>(); + configMap.put("default_model_id", "vasdcvkcjkbldbjkd"); + configMap.put("neural_field_default_id", Map.of("fieldName", "bahbkcdkacb")); + NeuralQueryProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); + return processor; + } + +} diff --git a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterTestUtils.java b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterTestUtils.java new file mode 100644 index 000000000..30399cfea --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterTestUtils.java @@ -0,0 +1,32 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.util; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import org.opensearch.Version; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.service.ClusterService; + +public class NeuralSearchClusterTestUtils { + + /** + * Create new mock for ClusterService + * @param version min version for cluster nodes + * @return + */ + public static ClusterService mockClusterService(final Version version) { + ClusterService clusterService = mock(ClusterService.class); + ClusterState clusterState = mock(ClusterState.class); + when(clusterService.state()).thenReturn(clusterState); + DiscoveryNodes discoveryNodes = mock(DiscoveryNodes.class); + when(clusterState.getNodes()).thenReturn(discoveryNodes); + when(discoveryNodes.getMinNodeVersion()).thenReturn(version); + return clusterService; + } +} diff --git a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java new file mode 100644 index 000000000..9f619d4c2 --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java @@ -0,0 +1,51 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.util; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils.mockClusterService; + +import org.opensearch.Version; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.test.OpenSearchTestCase; + +public class NeuralSearchClusterUtilTests extends OpenSearchTestCase { + + public void testSingleNodeCluster() { + ClusterService clusterService = mockClusterService(Version.V_2_4_0); + + final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); + neuralSearchClusterUtil.initialize(clusterService); + + final Version minVersion = neuralSearchClusterUtil.getClusterMinVersion(); + + assertTrue(Version.V_2_4_0.equals(minVersion)); + } + + public void testMultipleNodesCluster() { + ClusterService clusterService = mockClusterService(Version.V_2_3_0); + + final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); + neuralSearchClusterUtil.initialize(clusterService); + + final Version minVersion = neuralSearchClusterUtil.getClusterMinVersion(); + + assertTrue(Version.V_2_3_0.equals(minVersion)); + } + + public void testWhenErrorOnClusterStateDiscover() { + ClusterService clusterService = mock(ClusterService.class); + when(clusterService.state()).thenThrow(new RuntimeException("Cluster state is not ready")); + + final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); + neuralSearchClusterUtil.initialize(clusterService); + + final Version minVersion = neuralSearchClusterUtil.getClusterMinVersion(); + + assertTrue(Version.CURRENT.equals(minVersion)); + } +} diff --git a/src/test/resources/processor/SearchRequestPipelineConfiguration.json b/src/test/resources/processor/SearchRequestPipelineConfiguration.json new file mode 100644 index 000000000..3f3675ce3 --- /dev/null +++ b/src/test/resources/processor/SearchRequestPipelineConfiguration.json @@ -0,0 +1,11 @@ +{ + "request_processors": [ + { + "default_query": { + "tag": "tag1", + "description": "This processor is going to restrict to publicly visible documents", + "default_model_id": "%s" + } + } + ] +} \ No newline at end of file From 7c3a3026f1ecf8c96823d321452f8fd0994db677 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Mon, 25 Sep 2023 17:57:55 -0700 Subject: [PATCH 04/17] Removing wildcard Imports Signed-off-by: Varun Jain --- .../opensearch/neuralsearch/query/NeuralQueryBuilder.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 799cc22de..51263a068 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -32,7 +32,10 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.mapper.NumberFieldMapper; -import org.opensearch.index.query.*; +import org.opensearch.index.query.AbstractQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryRewriteContext; +import org.opensearch.index.query.QueryShardContext; import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; From dee67cabf405f80722cfc65d758f2ec2fc5478ca Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Mon, 25 Sep 2023 18:01:12 -0700 Subject: [PATCH 05/17] Typo fix Signed-off-by: Varun Jain --- .../neuralsearch/processor/NeuralQueryProcessor.java | 10 +++++----- .../processor/NeuralQueryProcessorTests.java | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java index 59a8688ca..c6476da42 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -24,7 +24,7 @@ public class NeuralQueryProcessor extends AbstractProcessor implements SearchReq final String modelId; - final Map neuralFieldMap; + final Map neuralFieldDefaultIdMap; /** * Returns the type of the processor. @@ -41,17 +41,17 @@ protected NeuralQueryProcessor( String description, boolean ignoreFailure, String modelId, - Map fieldInfoMap + Map neuralFieldDefaultIdMap ) { super(tag, description, ignoreFailure); this.modelId = modelId; - this.neuralFieldMap = fieldInfoMap; + this.neuralFieldDefaultIdMap = neuralFieldDefaultIdMap; } @Override - public SearchRequest processRequest(SearchRequest searchRequest) throws Exception { + public SearchRequest processRequest(SearchRequest searchRequest) { QueryBuilder queryBuilder = searchRequest.source().query(); - queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, neuralFieldMap)); + queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, neuralFieldDefaultIdMap)); return searchRequest; } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java index 0cd9bea3a..176ffbad3 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java @@ -20,7 +20,7 @@ public void testFactory() throws Exception { NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); NeuralQueryProcessor processor = createTestProcessor(factory); assertEquals("vasdcvkcjkbldbjkd", processor.modelId); - assertEquals("bahbkcdkacb", processor.neuralFieldMap.get("fieldName").toString()); + assertEquals("bahbkcdkacb", processor.neuralFieldDefaultIdMap.get("fieldName").toString()); // Missing "query" parameter: expectThrows( From b13eb809770befd918af905a4900b8ce1b8db49b Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 26 Sep 2023 13:18:32 -0700 Subject: [PATCH 06/17] Integ test cases Signed-off-by: Varun Jain --- .../query/NeuralQueryBuilder.java | 4 +- .../common/BaseNeuralSearchIT.java | 26 ++++++ .../processor/NeuralQueryProcessorIT.java | 83 +++++++++++++++++++ .../resources/processor/IndexMappings.json | 7 ++ .../SearchRequestPipelineConfiguration.json | 2 +- 5 files changed, 120 insertions(+), 2 deletions(-) create mode 100644 src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 51263a068..728966428 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -123,7 +123,9 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws xContentBuilder.startObject(NAME); xContentBuilder.startObject(fieldName); xContentBuilder.field(QUERY_TEXT_FIELD.getPreferredName(), queryText); - xContentBuilder.field(MODEL_ID_FIELD.getPreferredName(), modelId); + if (!isClusterOnOrAfterMinRequiredVersion() || (isClusterOnOrAfterMinRequiredVersion() && modelId != null)) { + xContentBuilder.field(MODEL_ID_FIELD.getPreferredName(), modelId); + } xContentBuilder.field(K_FIELD.getPreferredName(), k); if (filter != null) { xContentBuilder.field(FILTER_FIELD.getPreferredName(), filter); diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java index fac32b538..faafa6d59 100644 --- a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java +++ b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java @@ -39,6 +39,8 @@ import org.opensearch.client.Response; import org.opensearch.client.RestClient; import org.opensearch.client.WarningsHandler; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; @@ -48,10 +50,16 @@ import org.opensearch.index.query.QueryBuilder; import org.opensearch.knn.index.SpaceType; import org.opensearch.neuralsearch.OpenSearchSecureRestTestCase; +import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; import com.google.common.collect.ImmutableList; +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase { private static final Locale LOCALE = Locale.ROOT; @@ -66,11 +74,29 @@ public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase { protected final ClassLoader classLoader = this.getClass().getClassLoader(); + protected ThreadPool threadPool; + protected ClusterService clusterService; + @Before public void setupSettings() { + threadPool = setUpThreadPool(); + clusterService = createClusterService(threadPool); if (isUpdateClusterSettings()) { updateClusterSettings(); } + NeuralSearchClusterUtil.instance().initialize(clusterService); + } + + protected ThreadPool setUpThreadPool() { + return new TestThreadPool(getClass().getName(), threadPoolSettings()); + } + + public Settings threadPoolSettings() { + return Settings.EMPTY; + } + + public static ClusterService createClusterService(ThreadPool threadPool) { + return ClusterServiceUtils.createClusterService(threadPool); } protected void updateClusterSettings() { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java new file mode 100644 index 000000000..9f845720e --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.processor; + +import static org.opensearch.neuralsearch.TestUtils.createRandomVector; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; + +import lombok.SneakyThrows; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.common.settings.Settings; +import org.opensearch.knn.index.SpaceType; +import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; + +import com.google.common.primitives.Floats; + +public class NeuralQueryProcessorIT extends BaseNeuralSearchIT { + + public static final String index = "my-nlp-index"; + public static final String search_pipeline = "search-pipeline"; + public static final String ingest_pipeline = "nlp-pipeline"; + private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; + private static final int TEST_DIMENSION = 768; + private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; + private final float[] testVector = createRandomVector(TEST_DIMENSION); + + @Before + public void setUp() throws Exception { + super.setUp(); + updateClusterSettings(); + prepareModel(); + } + + @After + @SneakyThrows + public void tearDown() { + super.tearDown(); + deleteSearchPipeline(search_pipeline); + findDeployedModels().forEach(this::deleteModel); + } + + public void testNeuralQueryProcessor() throws Exception { + initializeIndexIfNotExist(index); + String modelId = getDeployedModelId(); + createSearchRequestProcessor(modelId, search_pipeline); + createPipelineProcessor(modelId, ingest_pipeline); + updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + neuralQueryBuilder.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1); + neuralQueryBuilder.queryText("Hello World"); + neuralQueryBuilder.k(1); + Map response = search(index, neuralQueryBuilder, 2); + + assertFalse(response.isEmpty()); + + assertEquals(modelId, neuralQueryBuilder.modelId()); + + } + + private void initializeIndexIfNotExist(String indexName) throws IOException { + if (index.equals(indexName) && !indexExists(index)) { + prepareKnnIndex( + index, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + index, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + assertEquals(1, getDocCount(index)); + } + } +} diff --git a/src/test/resources/processor/IndexMappings.json b/src/test/resources/processor/IndexMappings.json index 6cc3fabbe..5464a9311 100644 --- a/src/test/resources/processor/IndexMappings.json +++ b/src/test/resources/processor/IndexMappings.json @@ -67,6 +67,13 @@ } } } + }, + "passage_embedding": { + "type": "knn_vector", + "dimension": 768 + }, + "passage_text": { + "type": "text" } } } diff --git a/src/test/resources/processor/SearchRequestPipelineConfiguration.json b/src/test/resources/processor/SearchRequestPipelineConfiguration.json index 3f3675ce3..ad5a0ba52 100644 --- a/src/test/resources/processor/SearchRequestPipelineConfiguration.json +++ b/src/test/resources/processor/SearchRequestPipelineConfiguration.json @@ -1,7 +1,7 @@ { "request_processors": [ { - "default_query": { + "neural_query": { "tag": "tag1", "description": "This processor is going to restrict to publicly visible documents", "default_model_id": "%s" From badced1aa09e344e52a9b69ae1dd8ea7f6c43870 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 26 Sep 2023 13:27:24 -0700 Subject: [PATCH 07/17] Fixing Integ Test case Signed-off-by: Varun Jain --- .../neuralsearch/processor/NeuralQueryProcessorIT.java | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java index 9f845720e..5c7a58346 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java @@ -48,7 +48,7 @@ public void tearDown() { } public void testNeuralQueryProcessor() throws Exception { - initializeIndexIfNotExist(index); + initializeIndexIfNotExist(); String modelId = getDeployedModelId(); createSearchRequestProcessor(modelId, search_pipeline); createPipelineProcessor(modelId, ingest_pipeline); @@ -61,12 +61,10 @@ public void testNeuralQueryProcessor() throws Exception { assertFalse(response.isEmpty()); - assertEquals(modelId, neuralQueryBuilder.modelId()); - } - private void initializeIndexIfNotExist(String indexName) throws IOException { - if (index.equals(indexName) && !indexExists(index)) { + private void initializeIndexIfNotExist() throws IOException { + if (index.equals(NeuralQueryProcessorIT.index) && !indexExists(index)) { prepareKnnIndex( index, Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) From 9c010e747e151b1fe8968c0fa2e9045ccdf2453a Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 26 Sep 2023 16:18:06 -0700 Subject: [PATCH 08/17] Addressing Comments Signed-off-by: Varun Jain --- CHANGELOG.md | 1 + .../neuralsearch/plugin/NeuralSearch.java | 2 +- .../processor/NeuralQueryProcessor.java | 13 ++++++ .../visitor/NeuralSearchQueryVisitor.java | 44 ++++++++++++------- .../util/NeuralSearchClusterUtil.java | 3 ++ 5 files changed, 46 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b3ea389e..e3f371181 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 3.0](https://github.com/opensearch-project/neural-search/compare/2.x...HEAD) ### Features +- Enabled support for applying default modelId in neural search query ([#337](https://github.com/opensearch-project/neural-search/pull/337) ### Enhancements ### Bug Fixes ### Infrastructure diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index 57ac4feb5..b44e99338 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -65,7 +65,7 @@ public class NeuralSearch extends Plugin implements ActionPlugin, SearchPlugin, private MLCommonsClientAccessor clientAccessor; private NormalizationProcessorWorkflow normalizationProcessorWorkflow; private final ScoreNormalizationFactory scoreNormalizationFactory = new ScoreNormalizationFactory(); - private final ScoreCombinationFactory scoreCombinationFactory = new ScoreCombinationFactory();; + private final ScoreCombinationFactory scoreCombinationFactory = new ScoreCombinationFactory(); @Override public Collection createComponents( diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java index c6476da42..8c52385e2 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -15,6 +15,9 @@ import org.opensearch.search.pipeline.Processor; import org.opensearch.search.pipeline.SearchRequestProcessor; +/** + * Neural Search Query Request Processor + */ public class NeuralQueryProcessor extends AbstractProcessor implements SearchRequestProcessor { /** @@ -48,6 +51,11 @@ protected NeuralQueryProcessor( this.neuralFieldDefaultIdMap = neuralFieldDefaultIdMap; } + /** + * Processes the Search Request. + * + * @return The Search Request. + */ @Override public SearchRequest processRequest(SearchRequest searchRequest) { QueryBuilder queryBuilder = searchRequest.source().query(); @@ -59,6 +67,11 @@ public static class Factory implements Processor.Factory private static final String DEFAULT_MODEL_ID = "default_model_id"; private static final String NEURAL_FIELD_DEFAULT_ID = "neural_field_default_id"; + /** + * Create the processor object. + * + * @return NeuralQueryProcessor + */ @Override public NeuralQueryProcessor create( Map> processorFactories, diff --git a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java index 2d0b53eaa..06148a5b5 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java +++ b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java @@ -7,40 +7,52 @@ import java.util.Map; +import lombok.AllArgsConstructor; + import org.apache.lucene.search.BooleanClause; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilderVisitor; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; +/** + * Neural Search Query Visitor. It visits the each and every component of query buikder tree. + */ +@AllArgsConstructor public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { private String modelId; private Map neuralFieldMap; - public NeuralSearchQueryVisitor(String modelId, Map neuralFieldMap) { - this.modelId = modelId; - this.neuralFieldMap = neuralFieldMap; - } - + /** + * Accept method accepts every query builder from the search request, + * and processes it if the required conditions in accept method are satisfied. + */ @Override public void accept(QueryBuilder queryBuilder) { if (queryBuilder instanceof NeuralQueryBuilder) { NeuralQueryBuilder neuralQueryBuilder = (NeuralQueryBuilder) queryBuilder; - if (neuralFieldMap != null - && neuralQueryBuilder.fieldName() != null - && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { - String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); - neuralQueryBuilder.modelId(fieldDefaultModelId); - } else if (modelId != null) { - neuralQueryBuilder.modelId(modelId); - } else { - throw new IllegalArgumentException( - "model id must be provided in neural query or a default model id must be set in search request processor" - ); + if (neuralQueryBuilder.modelId() == null) { + if (neuralFieldMap != null + && neuralQueryBuilder.fieldName() != null + && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { + String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); + neuralQueryBuilder.modelId(fieldDefaultModelId); + } else if (modelId != null) { + neuralQueryBuilder.modelId(modelId); + } else { + throw new IllegalArgumentException( + "model id must be provided in neural query or a default model id must be set in search request processor" + ); + } } } } + /** + * Retrieves the child visitor from the Visitor object. + * + * @return The sub Query Visitor + */ @Override public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { return this; diff --git a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java index f9e044c53..c91cb0fae 100644 --- a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java +++ b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java @@ -14,6 +14,9 @@ import org.opensearch.Version; import org.opensearch.cluster.service.ClusterService; +/** + * Class abstracts information related to underlying OpenSearch cluster + */ @NoArgsConstructor(access = AccessLevel.PRIVATE) @Log4j2 public class NeuralSearchClusterUtil { From c9ada839e3b8fc009bfaf1ad7aadaaf7677c4db2 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Wed, 27 Sep 2023 15:38:08 -0700 Subject: [PATCH 09/17] Added Visitor test cases and addressed comments Signed-off-by: Varun Jain --- .../processor/NeuralQueryProcessorIT.java | 7 ++- .../NeuralSearchQueryVisitorTests.java | 57 +++++++++++++++++++ .../util/NeuralSearchClusterUtilTests.java | 6 +- 3 files changed, 64 insertions(+), 6 deletions(-) create mode 100644 src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java index 5c7a58346..584388957 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java @@ -24,9 +24,9 @@ public class NeuralQueryProcessorIT extends BaseNeuralSearchIT { - public static final String index = "my-nlp-index"; - public static final String search_pipeline = "search-pipeline"; - public static final String ingest_pipeline = "nlp-pipeline"; + private static final String index = "my-nlp-index"; + private static final String search_pipeline = "search-pipeline"; + private static final String ingest_pipeline = "nlp-pipeline"; private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; private static final int TEST_DIMENSION = 768; private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; @@ -47,6 +47,7 @@ public void tearDown() { findDeployedModels().forEach(this::deleteModel); } + @SneakyThrows public void testNeuralQueryProcessor() throws Exception { initializeIndexIfNotExist(); String modelId = getDeployedModelId(); diff --git a/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java b/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java new file mode 100644 index 000000000..7570ece54 --- /dev/null +++ b/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.query.visitor; + +import java.util.HashMap; +import java.util.Map; + +import org.apache.lucene.search.BooleanClause; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; +import org.opensearch.test.OpenSearchTestCase; + +public class NeuralSearchQueryVisitorTests extends OpenSearchTestCase { + + public void testAccept_whenNeuralQueryBuilderWithoutModelId_thenSetModelId() { + String modelId = "bdcvjkcdjvkddcjxdjsc"; + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + neuralQueryBuilder.fieldName("passage_text"); + neuralQueryBuilder.k(768); + + NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(modelId, null); + neuralSearchQueryVisitor.accept(neuralQueryBuilder); + + assertEquals(modelId, neuralQueryBuilder.modelId()); + } + + public void testAccept_whenNeuralQueryBuilderWithoutFieldModelId_thenSetFieldModelId() { + Map neuralInfoMap = new HashMap<>(); + neuralInfoMap.put("passage_text", "bdcvjkcdjvkddcjxdjsc"); + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + neuralQueryBuilder.fieldName("passage_text"); + neuralQueryBuilder.k(768); + + NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(null, neuralInfoMap); + neuralSearchQueryVisitor.accept(neuralQueryBuilder); + + assertEquals("bdcvjkcdjvkddcjxdjsc", neuralQueryBuilder.modelId()); + } + + public void testAccept_whenNullValuesInVisitor_thenFail() { + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(null, null); + + expectThrows(IllegalArgumentException.class, () -> neuralSearchQueryVisitor.accept(neuralQueryBuilder)); + } + + public void testGetChildVisitor() { + NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(null, null); + + NeuralSearchQueryVisitor subVisitor = (NeuralSearchQueryVisitor) neuralSearchQueryVisitor.getChildVisitor(BooleanClause.Occur.MUST); + + assertEquals(subVisitor, neuralSearchQueryVisitor); + + } +} diff --git a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java index 9f619d4c2..548f651e8 100644 --- a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java +++ b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java @@ -15,7 +15,7 @@ public class NeuralSearchClusterUtilTests extends OpenSearchTestCase { - public void testSingleNodeCluster() { + public void testMinNodeVersion_whenSingleNodeCluster_thenSuccess() { ClusterService clusterService = mockClusterService(Version.V_2_4_0); final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); @@ -26,7 +26,7 @@ public void testSingleNodeCluster() { assertTrue(Version.V_2_4_0.equals(minVersion)); } - public void testMultipleNodesCluster() { + public void testMinNodeVersion_whenMultipleNodesCluster_thenSuccess() { ClusterService clusterService = mockClusterService(Version.V_2_3_0); final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); @@ -37,7 +37,7 @@ public void testMultipleNodesCluster() { assertTrue(Version.V_2_3_0.equals(minVersion)); } - public void testWhenErrorOnClusterStateDiscover() { + public void testMinNodeVersion_WhenErrorOnClusterState_thenMatchCurrentVersion() { ClusterService clusterService = mock(ClusterService.class); when(clusterService.state()).thenThrow(new RuntimeException("Cluster state is not ready")); From 4be198c97bb0f0bd566f47c71599d4388f2b733a Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Wed, 27 Sep 2023 17:08:51 -0700 Subject: [PATCH 10/17] Comments Addressed of Jack Signed-off-by: Varun Jain --- .../neuralsearch/processor/NeuralQueryProcessor.java | 2 +- .../neuralsearch/query/NeuralQueryBuilder.java | 10 +++++----- .../query/visitor/NeuralSearchQueryVisitor.java | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java index 8c52385e2..9c479b48e 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -80,7 +80,7 @@ public NeuralQueryProcessor create( boolean ignoreFailure, Map config, PipelineContext pipelineContext - ) throws Exception { + ) throws IllegalArgumentException { String modelId = (String) config.remove(DEFAULT_MODEL_ID); Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_DEFAULT_ID); diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 728966428..6701b5835 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -96,7 +96,7 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { super(in); this.fieldName = in.readString(); this.queryText = in.readString(); - if (isClusterOnOrAfterMinRequiredVersion()) { + if (isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { this.modelId = in.readOptionalString(); } else { this.modelId = in.readString(); @@ -109,7 +109,7 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { protected void doWriteTo(StreamOutput out) throws IOException { out.writeString(this.fieldName); out.writeString(this.queryText); - if (isClusterOnOrAfterMinRequiredVersion()) { + if (isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { out.writeOptionalString(this.modelId); } else { out.writeString(this.modelId); @@ -123,7 +123,7 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws xContentBuilder.startObject(NAME); xContentBuilder.startObject(fieldName); xContentBuilder.field(QUERY_TEXT_FIELD.getPreferredName(), queryText); - if (!isClusterOnOrAfterMinRequiredVersion() || (isClusterOnOrAfterMinRequiredVersion() && modelId != null)) { + if (modelId != null) { xContentBuilder.field(MODEL_ID_FIELD.getPreferredName(), modelId); } xContentBuilder.field(K_FIELD.getPreferredName(), k); @@ -177,7 +177,7 @@ public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOEx } requireValue(neuralQueryBuilder.queryText(), "Query text must be provided for neural query"); requireValue(neuralQueryBuilder.fieldName(), "Field name must be provided for neural query"); - if (!isClusterOnOrAfterMinRequiredVersion()) { + if (!isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { requireValue(neuralQueryBuilder.modelId(), "Model ID must be provided for neural query"); } return neuralQueryBuilder; @@ -273,7 +273,7 @@ public String getWriteableName() { return NAME; } - private static boolean isClusterOnOrAfterMinRequiredVersion() { + private static boolean isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport() { return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_DEFAULT_MODEL_ID); } } diff --git a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java index 06148a5b5..9d746190c 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java +++ b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java @@ -15,7 +15,7 @@ import org.opensearch.neuralsearch.query.NeuralQueryBuilder; /** - * Neural Search Query Visitor. It visits the each and every component of query buikder tree. + * Neural Search Query Visitor. It visits each and every component of query buikder tree. */ @AllArgsConstructor public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { From e9db5525d188a2833dac04b6a9e16040c827a7ee Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Thu, 28 Sep 2023 12:47:12 -0700 Subject: [PATCH 11/17] Addressed changes requested by Martin Signed-off-by: Varun Jain --- .../processor/NeuralQueryProcessor.java | 28 +++++++++++++------ .../query/NeuralQueryBuilder.java | 2 ++ .../visitor/NeuralSearchQueryVisitor.java | 4 +-- .../util/NeuralSearchClusterUtil.java | 16 +---------- .../processor/NeuralQueryProcessorTests.java | 6 ++-- .../util/NeuralSearchClusterUtilTests.java | 14 ---------- 6 files changed, 28 insertions(+), 42 deletions(-) diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java index 9c479b48e..e38d96941 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java @@ -7,7 +7,11 @@ import java.util.Map; +import lombok.Getter; +import lombok.Setter; + import org.opensearch.action.search.SearchRequest; +import org.opensearch.common.Nullable; import org.opensearch.index.query.QueryBuilder; import org.opensearch.ingest.ConfigurationUtils; import org.opensearch.neuralsearch.query.visitor.NeuralSearchQueryVisitor; @@ -16,18 +20,21 @@ import org.opensearch.search.pipeline.SearchRequestProcessor; /** - * Neural Search Query Request Processor + * Neural Search Query Request Processor, It modifies the search request with neural query clause + * and adds model Id if not present in the search query. */ +@Setter +@Getter public class NeuralQueryProcessor extends AbstractProcessor implements SearchRequestProcessor { /** * Key to reference this processor type from a search pipeline. */ - public static final String TYPE = "neural_query"; + public static final String TYPE = "enriching_query_defaults"; - final String modelId; + private final String modelId; - final Map neuralFieldDefaultIdMap; + private final Map neuralFieldDefaultIdMap; /** * Returns the type of the processor. @@ -39,12 +46,12 @@ public String getType() { return TYPE; } - protected NeuralQueryProcessor( + private NeuralQueryProcessor( String tag, String description, boolean ignoreFailure, - String modelId, - Map neuralFieldDefaultIdMap + @Nullable String modelId, + @Nullable Map neuralFieldDefaultIdMap ) { super(tag, description, ignoreFailure); this.modelId = modelId; @@ -81,7 +88,12 @@ public NeuralQueryProcessor create( Map config, PipelineContext pipelineContext ) throws IllegalArgumentException { - String modelId = (String) config.remove(DEFAULT_MODEL_ID); + String modelId; + try { + modelId = (String) config.remove(DEFAULT_MODEL_ID); + } catch (ClassCastException e) { + throw new IllegalArgumentException("model Id must of String type"); + } Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_DEFAULT_ID); if (modelId == null && neuralInfoMap == null) { diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 6701b5835..7b78be269 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -96,6 +96,7 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { super(in); this.fieldName = in.readString(); this.queryText = in.readString(); + // If cluster version is on or after 2.11 then default model Id support is enabled if (isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { this.modelId = in.readOptionalString(); } else { @@ -109,6 +110,7 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { protected void doWriteTo(StreamOutput out) throws IOException { out.writeString(this.fieldName); out.writeString(this.queryText); + // If cluster version is on or after 2.11 then default model Id support is enabled if (isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { out.writeOptionalString(this.modelId); } else { diff --git a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java index 9d746190c..febb35294 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java +++ b/src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java @@ -20,8 +20,8 @@ @AllArgsConstructor public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { - private String modelId; - private Map neuralFieldMap; + private final String modelId; + private final Map neuralFieldMap; /** * Accept method accepts every query builder from the search request, diff --git a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java index c91cb0fae..5a97120e0 100644 --- a/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java +++ b/src/main/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtil.java @@ -5,8 +5,6 @@ package org.opensearch.neuralsearch.util; -import java.util.Locale; - import lombok.AccessLevel; import lombok.NoArgsConstructor; import lombok.extern.log4j.Log4j2; @@ -48,19 +46,7 @@ public void initialize(final ClusterService clusterService) { * @return minimal installed OpenSearch version, default to Version.CURRENT which is typically the latest version */ public Version getClusterMinVersion() { - try { - return this.clusterService.state().getNodes().getMinNodeVersion(); - } catch (Exception exception) { - log.error( - String.format( - Locale.ROOT, - "Failed to get cluster minimum node version, returning current node version %s instead.", - Version.CURRENT - ), - exception - ); - return Version.CURRENT; - } + return this.clusterService.state().getNodes().getMinNodeVersion(); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java index 176ffbad3..0a90d5c70 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java @@ -19,8 +19,8 @@ public class NeuralQueryProcessorTests extends OpenSearchTestCase { public void testFactory() throws Exception { NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); NeuralQueryProcessor processor = createTestProcessor(factory); - assertEquals("vasdcvkcjkbldbjkd", processor.modelId); - assertEquals("bahbkcdkacb", processor.neuralFieldDefaultIdMap.get("fieldName").toString()); + assertEquals("vasdcvkcjkbldbjkd", processor.getModelId()); + assertEquals("bahbkcdkacb", processor.getNeuralFieldDefaultIdMap().get("fieldName").toString()); // Missing "query" parameter: expectThrows( @@ -39,7 +39,7 @@ public void testProcessRequest() throws Exception { assertEquals(processSearchRequest, searchRequest); } - public NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory factory) throws Exception { + private NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory factory) throws Exception { Map configMap = new HashMap<>(); configMap.put("default_model_id", "vasdcvkcjkbldbjkd"); configMap.put("neural_field_default_id", Map.of("fieldName", "bahbkcdkacb")); diff --git a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java index 548f651e8..f85be25d5 100644 --- a/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java +++ b/src/test/java/org/opensearch/neuralsearch/util/NeuralSearchClusterUtilTests.java @@ -5,8 +5,6 @@ package org.opensearch.neuralsearch.util; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; import static org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils.mockClusterService; import org.opensearch.Version; @@ -36,16 +34,4 @@ public void testMinNodeVersion_whenMultipleNodesCluster_thenSuccess() { assertTrue(Version.V_2_3_0.equals(minVersion)); } - - public void testMinNodeVersion_WhenErrorOnClusterState_thenMatchCurrentVersion() { - ClusterService clusterService = mock(ClusterService.class); - when(clusterService.state()).thenThrow(new RuntimeException("Cluster state is not ready")); - - final NeuralSearchClusterUtil neuralSearchClusterUtil = NeuralSearchClusterUtil.instance(); - neuralSearchClusterUtil.initialize(clusterService); - - final Version minVersion = neuralSearchClusterUtil.getClusterMinVersion(); - - assertTrue(Version.CURRENT.equals(minVersion)); - } } From e9efd7249222701f7f8f234e1a9d4a1284890a52 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Thu, 28 Sep 2023 13:01:05 -0700 Subject: [PATCH 12/17] Addressed changes requested by Martin Signed-off-by: Varun Jain --- .../resources/processor/SearchRequestPipelineConfiguration.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/test/resources/processor/SearchRequestPipelineConfiguration.json b/src/test/resources/processor/SearchRequestPipelineConfiguration.json index ad5a0ba52..517b89491 100644 --- a/src/test/resources/processor/SearchRequestPipelineConfiguration.json +++ b/src/test/resources/processor/SearchRequestPipelineConfiguration.json @@ -1,7 +1,7 @@ { "request_processors": [ { - "neural_query": { + "enriching_query_defaults": { "tag": "tag1", "description": "This processor is going to restrict to publicly visible documents", "default_model_id": "%s" From ff10aba483ca4d730eac83f089dd6e346146ae4c Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Thu, 28 Sep 2023 13:52:07 -0700 Subject: [PATCH 13/17] Fixing test cases Signed-off-by: Varun Jain --- .../neuralsearch/query/HybridQueryBuilderTests.java | 11 +++++++++++ .../neuralsearch/query/NeuralQueryBuilderTests.java | 13 +++++++++++++ 2 files changed, 24 insertions(+) diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index 49d1ba974..0544feff8 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -28,6 +28,8 @@ import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TermQuery; +import org.opensearch.Version; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.ParseField; @@ -51,6 +53,8 @@ import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; import org.opensearch.knn.index.query.KNNQuery; import org.opensearch.knn.index.query.KNNQueryBuilder; +import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils; +import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import com.carrotsearch.randomizedtesting.RandomizedTest; @@ -235,6 +239,7 @@ public void testDoToQuery_whenTooManySubqueries_thenFail() { */ @SneakyThrows public void testFromXContent_whenMultipleSubQueries_thenBuildSuccessfully() { + setUpClusterService(); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startArray("queries") @@ -412,6 +417,7 @@ public void testToXContent_whenIncomingJsonIsCorrect_thenSuccessful() { @SneakyThrows public void testStreams_whenWrittingToStream_thenSuccessful() { + setUpClusterService(); HybridQueryBuilder original = new HybridQueryBuilder(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) @@ -716,4 +722,9 @@ private Map getInnerMap(Object innerObject, String queryName, St Map vectorFieldInnerMap = (Map) neuralInnerMap.get(fieldName); return vectorFieldInnerMap; } + + private void setUpClusterService() { + ClusterService clusterService = NeuralSearchClusterTestUtils.mockClusterService(Version.CURRENT); + NeuralSearchClusterUtil.instance().initialize(clusterService); + } } diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java index f389dfd22..d883fe616 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java @@ -29,7 +29,9 @@ import lombok.SneakyThrows; +import org.opensearch.Version; import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.ParseField; @@ -50,6 +52,8 @@ import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.neuralsearch.common.VectorUtil; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; +import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils; +import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import org.opensearch.test.OpenSearchTestCase; public class NeuralQueryBuilderTests extends OpenSearchTestCase { @@ -75,6 +79,7 @@ public void testFromXContent_whenBuiltWithDefaults_thenBuildSuccessfully() { } } */ + setUpClusterService(); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -107,6 +112,7 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() { } } */ + setUpClusterService(); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -146,6 +152,7 @@ public void testFromXContent_whenBuiltWithFilter_thenBuildSuccessfully() { } } */ + setUpClusterService(); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -334,6 +341,7 @@ public void testToXContent() { @SneakyThrows public void testStreams() { + setUpClusterService(); NeuralQueryBuilder original = new NeuralQueryBuilder(); original.fieldName(FIELD_NAME); original.queryText(QUERY_TEXT); @@ -572,4 +580,9 @@ public void testRewrite_whenFilterSet_thenKNNQueryBuilderFilterSet() { KNNQueryBuilder knnQueryBuilder = (KNNQueryBuilder) queryBuilder; assertEquals(neuralQueryBuilder.filter(), knnQueryBuilder.getFilter()); } + + private void setUpClusterService() { + ClusterService clusterService = NeuralSearchClusterTestUtils.mockClusterService(Version.CURRENT); + NeuralSearchClusterUtil.instance().initialize(clusterService); + } } From c3d06499a39c074be9a192427ae852695c035d51 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Thu, 28 Sep 2023 14:29:43 -0700 Subject: [PATCH 14/17] Increasing test coverage Signed-off-by: Varun Jain --- .../processor/NeuralQueryProcessorTests.java | 14 ++++++++++++- .../query/NeuralQueryBuilderTests.java | 21 ++++++++++++------- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java index 0a90d5c70..d1540524f 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java @@ -29,6 +29,13 @@ public void testFactory() throws Exception { ); } + public void testFactory_whenModelIdIsNotString_thenFail() { + NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + Map configMap = new HashMap<>(); + configMap.put("default_model_id", 55555L); + expectThrows(IllegalArgumentException.class, () -> factory.create(Collections.emptyMap(), null, null, false, configMap, null)); + } + public void testProcessRequest() throws Exception { NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); @@ -39,6 +46,12 @@ public void testProcessRequest() throws Exception { assertEquals(processSearchRequest, searchRequest); } + public void testType() throws Exception { + NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + NeuralQueryProcessor processor = createTestProcessor(factory); + assertEquals(NeuralQueryProcessor.TYPE, processor.getType()); + } + private NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory factory) throws Exception { Map configMap = new HashMap<>(); configMap.put("default_model_id", "vasdcvkcjkbldbjkd"); @@ -46,5 +59,4 @@ private NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory fa NeuralQueryProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); return processor; } - } diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java index d883fe616..681c1247d 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java @@ -79,7 +79,7 @@ public void testFromXContent_whenBuiltWithDefaults_thenBuildSuccessfully() { } } */ - setUpClusterService(); + setUpClusterService(Version.V_2_10_0); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -112,7 +112,7 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() { } } */ - setUpClusterService(); + setUpClusterService(Version.CURRENT); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -152,7 +152,7 @@ public void testFromXContent_whenBuiltWithFilter_thenBuildSuccessfully() { } } */ - setUpClusterService(); + setUpClusterService(Version.CURRENT); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() .startObject(FIELD_NAME) @@ -340,8 +340,15 @@ public void testToXContent() { } @SneakyThrows - public void testStreams() { - setUpClusterService(); + public void testStreams_whenClusterServiceWithDifferentVersions() { + setUpClusterService(Version.V_2_10_0); + testStreams(); + setUpClusterService(Version.CURRENT); + testStreams(); + } + + @SneakyThrows + private void testStreams() { NeuralQueryBuilder original = new NeuralQueryBuilder(); original.fieldName(FIELD_NAME); original.queryText(QUERY_TEXT); @@ -581,8 +588,8 @@ public void testRewrite_whenFilterSet_thenKNNQueryBuilderFilterSet() { assertEquals(neuralQueryBuilder.filter(), knnQueryBuilder.getFilter()); } - private void setUpClusterService() { - ClusterService clusterService = NeuralSearchClusterTestUtils.mockClusterService(Version.CURRENT); + private void setUpClusterService(Version version) { + ClusterService clusterService = NeuralSearchClusterTestUtils.mockClusterService(version); NeuralSearchClusterUtil.instance().initialize(clusterService); } } From ef5f9394bb550797299742c2f7b4c8384ed64a17 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Fri, 29 Sep 2023 09:53:25 -0700 Subject: [PATCH 15/17] Renaming and addressing comments of Martin Signed-off-by: Varun Jain --- .../neuralsearch/plugin/NeuralSearch.java | 9 +++---- ...va => EnrichingQueryDefaultProcessor.java} | 10 +++---- .../plugin/NeuralSearchTests.java | 4 +-- ... => EnrichingQueryDefaultProcessorIT.java} | 10 +++---- ... EnrichingQueryDefaultProcessorTests.java} | 26 +++++++++---------- 5 files changed, 28 insertions(+), 31 deletions(-) rename src/main/java/org/opensearch/neuralsearch/processor/{NeuralQueryProcessor.java => EnrichingQueryDefaultProcessor.java} (89%) rename src/test/java/org/opensearch/neuralsearch/processor/{NeuralQueryProcessorIT.java => EnrichingQueryDefaultProcessorIT.java} (89%) rename src/test/java/org/opensearch/neuralsearch/processor/{NeuralQueryProcessorTests.java => EnrichingQueryDefaultProcessorTests.java} (59%) diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index 9be60fed1..47882f25f 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -28,11 +28,8 @@ import org.opensearch.ingest.Processor; import org.opensearch.ml.client.MachineLearningNodeClient; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; -import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; -import org.opensearch.neuralsearch.processor.NormalizationProcessor; -import org.opensearch.neuralsearch.processor.NormalizationProcessorWorkflow; -import org.opensearch.neuralsearch.processor.SparseEncodingProcessor; -import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; +import org.opensearch.neuralsearch.processor.*; +import org.opensearch.neuralsearch.processor.EnrichingQueryDefaultProcessor; import org.opensearch.neuralsearch.processor.combination.ScoreCombinationFactory; import org.opensearch.neuralsearch.processor.combination.ScoreCombiner; import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory; @@ -145,6 +142,6 @@ public List> getSettings() { public Map> getRequestProcessors( Parameters parameters ) { - return Map.of(NeuralQueryProcessor.TYPE, new NeuralQueryProcessor.Factory()); + return Map.of(EnrichingQueryDefaultProcessor.TYPE, new EnrichingQueryDefaultProcessor.Factory()); } } diff --git a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java similarity index 89% rename from src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java rename to src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java index e38d96941..4b01a8f88 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java @@ -25,7 +25,7 @@ */ @Setter @Getter -public class NeuralQueryProcessor extends AbstractProcessor implements SearchRequestProcessor { +public class EnrichingQueryDefaultProcessor extends AbstractProcessor implements SearchRequestProcessor { /** * Key to reference this processor type from a search pipeline. @@ -46,7 +46,7 @@ public String getType() { return TYPE; } - private NeuralQueryProcessor( + private EnrichingQueryDefaultProcessor( String tag, String description, boolean ignoreFailure, @@ -77,10 +77,10 @@ public static class Factory implements Processor.Factory /** * Create the processor object. * - * @return NeuralQueryProcessor + * @return EnrichingQueryDefaultProcessor */ @Override - public NeuralQueryProcessor create( + public EnrichingQueryDefaultProcessor create( Map> processorFactories, String tag, String description, @@ -100,7 +100,7 @@ public NeuralQueryProcessor create( throw new IllegalArgumentException("model Id or neural info map either of them should be provided"); } - return new NeuralQueryProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); + return new EnrichingQueryDefaultProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); } } } diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java index 9c8a23d03..e3fe2bf0f 100644 --- a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java +++ b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java @@ -12,7 +12,7 @@ import java.util.Optional; import org.opensearch.ingest.Processor; -import org.opensearch.neuralsearch.processor.NeuralQueryProcessor; +import org.opensearch.neuralsearch.processor.EnrichingQueryDefaultProcessor; import org.opensearch.neuralsearch.processor.NormalizationProcessor; import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory; @@ -83,6 +83,6 @@ public void testRequestProcessors() { parameters ); assertNotNull(processors); - assertNotNull(processors.get(NeuralQueryProcessor.TYPE)); + assertNotNull(processors.get(EnrichingQueryDefaultProcessor.TYPE)); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java similarity index 89% rename from src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java rename to src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java index 584388957..436a12d7e 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java @@ -7,7 +7,6 @@ import static org.opensearch.neuralsearch.TestUtils.createRandomVector; -import java.io.IOException; import java.util.Collections; import java.util.Map; @@ -22,7 +21,7 @@ import com.google.common.primitives.Floats; -public class NeuralQueryProcessorIT extends BaseNeuralSearchIT { +public class EnrichingQueryDefaultProcessorIT extends BaseNeuralSearchIT { private static final String index = "my-nlp-index"; private static final String search_pipeline = "search-pipeline"; @@ -48,7 +47,7 @@ public void tearDown() { } @SneakyThrows - public void testNeuralQueryProcessor() throws Exception { + public void testEnrichingQueryProcessor_whenNoModelIdPassed_thenSuccess() { initializeIndexIfNotExist(); String modelId = getDeployedModelId(); createSearchRequestProcessor(modelId, search_pipeline); @@ -64,8 +63,9 @@ public void testNeuralQueryProcessor() throws Exception { } - private void initializeIndexIfNotExist() throws IOException { - if (index.equals(NeuralQueryProcessorIT.index) && !indexExists(index)) { + @SneakyThrows + private void initializeIndexIfNotExist() { + if (index.equals(EnrichingQueryDefaultProcessorIT.index) && !indexExists(index)) { prepareKnnIndex( index, Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java similarity index 59% rename from src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java rename to src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java index d1540524f..897a2b8e0 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java @@ -14,11 +14,11 @@ import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.test.OpenSearchTestCase; -public class NeuralQueryProcessorTests extends OpenSearchTestCase { +public class EnrichingQueryDefaultProcessorTests extends OpenSearchTestCase { - public void testFactory() throws Exception { - NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); - NeuralQueryProcessor processor = createTestProcessor(factory); + public void testFactory_whenMissingQueryParam_thenThrowException() throws Exception { + EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); + EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); assertEquals("vasdcvkcjkbldbjkd", processor.getModelId()); assertEquals("bahbkcdkacb", processor.getNeuralFieldDefaultIdMap().get("fieldName").toString()); @@ -30,33 +30,33 @@ public void testFactory() throws Exception { } public void testFactory_whenModelIdIsNotString_thenFail() { - NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); Map configMap = new HashMap<>(); configMap.put("default_model_id", 55555L); expectThrows(IllegalArgumentException.class, () -> factory.create(Collections.emptyMap(), null, null, false, configMap, null)); } - public void testProcessRequest() throws Exception { - NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); + public void testProcessRequest_whenVisitingQueryBuilder_thenSuccess() throws Exception { + EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); SearchRequest searchRequest = new SearchRequest(); searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder)); - NeuralQueryProcessor processor = createTestProcessor(factory); + EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); SearchRequest processSearchRequest = processor.processRequest(searchRequest); assertEquals(processSearchRequest, searchRequest); } public void testType() throws Exception { - NeuralQueryProcessor.Factory factory = new NeuralQueryProcessor.Factory(); - NeuralQueryProcessor processor = createTestProcessor(factory); - assertEquals(NeuralQueryProcessor.TYPE, processor.getType()); + EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); + EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); + assertEquals(EnrichingQueryDefaultProcessor.TYPE, processor.getType()); } - private NeuralQueryProcessor createTestProcessor(NeuralQueryProcessor.Factory factory) throws Exception { + private EnrichingQueryDefaultProcessor createTestProcessor(EnrichingQueryDefaultProcessor.Factory factory) throws Exception { Map configMap = new HashMap<>(); configMap.put("default_model_id", "vasdcvkcjkbldbjkd"); configMap.put("neural_field_default_id", Map.of("fieldName", "bahbkcdkacb")); - NeuralQueryProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); + EnrichingQueryDefaultProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); return processor; } } From 29017baf479d27561d1fdf24a0539b20f32fbc84 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Fri, 29 Sep 2023 11:21:12 -0700 Subject: [PATCH 16/17] Addressing Comments of Navneet Signed-off-by: Varun Jain --- .../neuralsearch/plugin/NeuralSearch.java | 9 ++++--- ...java => NeuralQueryEnricherProcessor.java} | 22 ++++++++-------- .../plugin/NeuralSearchTests.java | 4 +-- ...va => NeuralQueryEnricherProcessorIT.java} | 6 ++--- ...=> NeuralQueryEnricherProcessorTests.java} | 25 ++++++++++--------- .../SearchRequestPipelineConfiguration.json | 2 +- 6 files changed, 35 insertions(+), 33 deletions(-) rename src/main/java/org/opensearch/neuralsearch/processor/{EnrichingQueryDefaultProcessor.java => NeuralQueryEnricherProcessor.java} (80%) rename src/test/java/org/opensearch/neuralsearch/processor/{EnrichingQueryDefaultProcessorIT.java => NeuralQueryEnricherProcessorIT.java} (91%) rename src/test/java/org/opensearch/neuralsearch/processor/{EnrichingQueryDefaultProcessorTests.java => NeuralQueryEnricherProcessorTests.java} (62%) diff --git a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java index 47882f25f..dd7dfee49 100644 --- a/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java +++ b/src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java @@ -28,8 +28,11 @@ import org.opensearch.ingest.Processor; import org.opensearch.ml.client.MachineLearningNodeClient; import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor; -import org.opensearch.neuralsearch.processor.*; -import org.opensearch.neuralsearch.processor.EnrichingQueryDefaultProcessor; +import org.opensearch.neuralsearch.processor.NeuralQueryEnricherProcessor; +import org.opensearch.neuralsearch.processor.NormalizationProcessor; +import org.opensearch.neuralsearch.processor.NormalizationProcessorWorkflow; +import org.opensearch.neuralsearch.processor.SparseEncodingProcessor; +import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; import org.opensearch.neuralsearch.processor.combination.ScoreCombinationFactory; import org.opensearch.neuralsearch.processor.combination.ScoreCombiner; import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory; @@ -142,6 +145,6 @@ public List> getSettings() { public Map> getRequestProcessors( Parameters parameters ) { - return Map.of(EnrichingQueryDefaultProcessor.TYPE, new EnrichingQueryDefaultProcessor.Factory()); + return Map.of(NeuralQueryEnricherProcessor.TYPE, new NeuralQueryEnricherProcessor.Factory()); } } diff --git a/src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessor.java similarity index 80% rename from src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java rename to src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessor.java index 4b01a8f88..379c6a8cc 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessor.java @@ -5,6 +5,9 @@ package org.opensearch.neuralsearch.processor; +import static org.opensearch.ingest.ConfigurationUtils.*; +import static org.opensearch.neuralsearch.processor.TextEmbeddingProcessor.TYPE; + import java.util.Map; import lombok.Getter; @@ -25,12 +28,12 @@ */ @Setter @Getter -public class EnrichingQueryDefaultProcessor extends AbstractProcessor implements SearchRequestProcessor { +public class NeuralQueryEnricherProcessor extends AbstractProcessor implements SearchRequestProcessor { /** * Key to reference this processor type from a search pipeline. */ - public static final String TYPE = "enriching_query_defaults"; + public static final String TYPE = "neural_query_enricher"; private final String modelId; @@ -46,7 +49,7 @@ public String getType() { return TYPE; } - private EnrichingQueryDefaultProcessor( + private NeuralQueryEnricherProcessor( String tag, String description, boolean ignoreFailure, @@ -77,10 +80,10 @@ public static class Factory implements Processor.Factory /** * Create the processor object. * - * @return EnrichingQueryDefaultProcessor + * @return NeuralQueryEnricherProcessor */ @Override - public EnrichingQueryDefaultProcessor create( + public NeuralQueryEnricherProcessor create( Map> processorFactories, String tag, String description, @@ -88,19 +91,14 @@ public EnrichingQueryDefaultProcessor create( Map config, PipelineContext pipelineContext ) throws IllegalArgumentException { - String modelId; - try { - modelId = (String) config.remove(DEFAULT_MODEL_ID); - } catch (ClassCastException e) { - throw new IllegalArgumentException("model Id must of String type"); - } + String modelId = readOptionalStringProperty(TYPE, tag, config, DEFAULT_MODEL_ID); Map neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_DEFAULT_ID); if (modelId == null && neuralInfoMap == null) { throw new IllegalArgumentException("model Id or neural info map either of them should be provided"); } - return new EnrichingQueryDefaultProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); + return new NeuralQueryEnricherProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); } } } diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java index e3fe2bf0f..8cae15678 100644 --- a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java +++ b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java @@ -12,7 +12,7 @@ import java.util.Optional; import org.opensearch.ingest.Processor; -import org.opensearch.neuralsearch.processor.EnrichingQueryDefaultProcessor; +import org.opensearch.neuralsearch.processor.NeuralQueryEnricherProcessor; import org.opensearch.neuralsearch.processor.NormalizationProcessor; import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor; import org.opensearch.neuralsearch.processor.factory.NormalizationProcessorFactory; @@ -83,6 +83,6 @@ public void testRequestProcessors() { parameters ); assertNotNull(processors); - assertNotNull(processors.get(EnrichingQueryDefaultProcessor.TYPE)); + assertNotNull(processors.get(NeuralQueryEnricherProcessor.TYPE)); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java similarity index 91% rename from src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java rename to src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index 436a12d7e..b85f4acb6 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -21,7 +21,7 @@ import com.google.common.primitives.Floats; -public class EnrichingQueryDefaultProcessorIT extends BaseNeuralSearchIT { +public class NeuralQueryEnricherProcessorIT extends BaseNeuralSearchIT { private static final String index = "my-nlp-index"; private static final String search_pipeline = "search-pipeline"; @@ -47,7 +47,7 @@ public void tearDown() { } @SneakyThrows - public void testEnrichingQueryProcessor_whenNoModelIdPassed_thenSuccess() { + public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() { initializeIndexIfNotExist(); String modelId = getDeployedModelId(); createSearchRequestProcessor(modelId, search_pipeline); @@ -65,7 +65,7 @@ public void testEnrichingQueryProcessor_whenNoModelIdPassed_thenSuccess() { @SneakyThrows private void initializeIndexIfNotExist() { - if (index.equals(EnrichingQueryDefaultProcessorIT.index) && !indexExists(index)) { + if (index.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(index)) { prepareKnnIndex( index, Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java similarity index 62% rename from src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java rename to src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java index 897a2b8e0..f6de3e58d 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/EnrichingQueryDefaultProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java @@ -9,16 +9,17 @@ import java.util.HashMap; import java.util.Map; +import org.opensearch.OpenSearchParseException; import org.opensearch.action.search.SearchRequest; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.test.OpenSearchTestCase; -public class EnrichingQueryDefaultProcessorTests extends OpenSearchTestCase { +public class NeuralQueryEnricherProcessorTests extends OpenSearchTestCase { public void testFactory_whenMissingQueryParam_thenThrowException() throws Exception { - EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); - EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); + NeuralQueryEnricherProcessor.Factory factory = new NeuralQueryEnricherProcessor.Factory(); + NeuralQueryEnricherProcessor processor = createTestProcessor(factory); assertEquals("vasdcvkcjkbldbjkd", processor.getModelId()); assertEquals("bahbkcdkacb", processor.getNeuralFieldDefaultIdMap().get("fieldName").toString()); @@ -30,33 +31,33 @@ public void testFactory_whenMissingQueryParam_thenThrowException() throws Except } public void testFactory_whenModelIdIsNotString_thenFail() { - EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); + NeuralQueryEnricherProcessor.Factory factory = new NeuralQueryEnricherProcessor.Factory(); Map configMap = new HashMap<>(); configMap.put("default_model_id", 55555L); - expectThrows(IllegalArgumentException.class, () -> factory.create(Collections.emptyMap(), null, null, false, configMap, null)); + expectThrows(OpenSearchParseException.class, () -> factory.create(Collections.emptyMap(), null, null, false, configMap, null)); } public void testProcessRequest_whenVisitingQueryBuilder_thenSuccess() throws Exception { - EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); + NeuralQueryEnricherProcessor.Factory factory = new NeuralQueryEnricherProcessor.Factory(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); SearchRequest searchRequest = new SearchRequest(); searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder)); - EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); + NeuralQueryEnricherProcessor processor = createTestProcessor(factory); SearchRequest processSearchRequest = processor.processRequest(searchRequest); assertEquals(processSearchRequest, searchRequest); } public void testType() throws Exception { - EnrichingQueryDefaultProcessor.Factory factory = new EnrichingQueryDefaultProcessor.Factory(); - EnrichingQueryDefaultProcessor processor = createTestProcessor(factory); - assertEquals(EnrichingQueryDefaultProcessor.TYPE, processor.getType()); + NeuralQueryEnricherProcessor.Factory factory = new NeuralQueryEnricherProcessor.Factory(); + NeuralQueryEnricherProcessor processor = createTestProcessor(factory); + assertEquals(NeuralQueryEnricherProcessor.TYPE, processor.getType()); } - private EnrichingQueryDefaultProcessor createTestProcessor(EnrichingQueryDefaultProcessor.Factory factory) throws Exception { + private NeuralQueryEnricherProcessor createTestProcessor(NeuralQueryEnricherProcessor.Factory factory) throws Exception { Map configMap = new HashMap<>(); configMap.put("default_model_id", "vasdcvkcjkbldbjkd"); configMap.put("neural_field_default_id", Map.of("fieldName", "bahbkcdkacb")); - EnrichingQueryDefaultProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); + NeuralQueryEnricherProcessor processor = factory.create(Collections.emptyMap(), null, null, false, configMap, null); return processor; } } diff --git a/src/test/resources/processor/SearchRequestPipelineConfiguration.json b/src/test/resources/processor/SearchRequestPipelineConfiguration.json index 517b89491..44d3b3ef0 100644 --- a/src/test/resources/processor/SearchRequestPipelineConfiguration.json +++ b/src/test/resources/processor/SearchRequestPipelineConfiguration.json @@ -1,7 +1,7 @@ { "request_processors": [ { - "enriching_query_defaults": { + "neural_query_enricher": { "tag": "tag1", "description": "This processor is going to restrict to publicly visible documents", "default_model_id": "%s" From ce28954ea1045d151fed779ada055ab2c3c50aa1 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Fri, 29 Sep 2023 11:23:24 -0700 Subject: [PATCH 17/17] Updating tests Signed-off-by: Varun Jain --- .../neuralsearch/processor/NeuralQueryEnricherProcessorIT.java | 1 + 1 file changed, 1 insertion(+) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index b85f4acb6..7e7660457 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -44,6 +44,7 @@ public void tearDown() { super.tearDown(); deleteSearchPipeline(search_pipeline); findDeployedModels().forEach(this::deleteModel); + deleteIndex(index); } @SneakyThrows