From 2ecd32c5cc6db70aeb5832578afa41998b01ba40 Mon Sep 17 00:00:00 2001 From: Will Hwang <22586574+will-hwang@users.noreply.github.com> Date: Mon, 6 Jan 2025 18:49:34 -0800 Subject: [PATCH] add support for builder constructor in neural query builder (#1047) * add support for builder constructor in neural query builder Signed-off-by: will-hwang * create custom builder class to enforce valid neural query builder instantiation Signed-off-by: will-hwang * refactor code to remove duplicate Signed-off-by: will-hwang * include new constructor in qa packages Signed-off-by: will-hwang * refactor code to remove unnecessary code Signed-off-by: will-hwang * fix bug in neural query builder instantiation Signed-off-by: will-hwang --------- Signed-off-by: will-hwang --- CHANGELOG.md | 1 + .../neuralsearch/bwc/HybridSearchIT.java | 13 +- .../neuralsearch/bwc/KnnRadialSearchIT.java | 44 ++-- .../neuralsearch/bwc/MultiModalSearchIT.java | 21 +- .../bwc/NeuralQueryEnricherProcessorIT.java | 10 +- .../neuralsearch/bwc/SemanticSearchIT.java | 11 +- .../neuralsearch/bwc/HybridSearchIT.java | 11 +- .../neuralsearch/bwc/KnnRadialSearchIT.java | 44 ++-- .../neuralsearch/bwc/MultiModalSearchIT.java | 22 +- .../bwc/NeuralQueryEnricherProcessorIT.java | 10 +- .../neuralsearch/bwc/SemanticSearchIT.java | 11 +- .../query/NeuralQueryBuilder.java | 149 +++++++++++- .../NeuralQueryEnricherProcessorIT.java | 18 +- .../NeuralQueryEnricherProcessorTests.java | 2 +- .../processor/NormalizationProcessorIT.java | 63 ++--- .../processor/ScoreCombinationIT.java | 63 +---- .../processor/ScoreNormalizationIT.java | 90 +------- .../processor/TextEmbeddingProcessorIT.java | 21 +- .../MLOpenSearchRerankProcessorTests.java | 8 +- .../query/HybridQueryBuilderTests.java | 78 +++++-- .../query/NeuralQueryBuilderTests.java | 217 ++++++++++++++---- .../neuralsearch/query/NeuralQueryIT.java | 207 ++++++----------- .../NeuralSearchQueryVisitorTests.java | 18 +- 23 files changed, 592 insertions(+), 540 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eb31fa4d8..035d00fd4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Implement pruning for neural sparse ingestion pipeline and two phase search processor ([#988](https://github.com/opensearch-project/neural-search/pull/988)) - Support empty string for fields in text embedding processor ([#1041](https://github.com/opensearch-project/neural-search/pull/1041)) - Optimize ML inference connection retry logic ([#1054](https://github.com/opensearch-project/neural-search/pull/1054)) +- Support for builder constructor in Neural Query Builder ([#1047](https://github.com/opensearch-project/neural-search/pull/1047)) ### Bug Fixes - Address inconsistent scoring in hybrid query results ([#998](https://github.com/opensearch-project/neural-search/pull/998)) - Fix bug where ingested document has list of nested objects ([#1040](https://github.com/opensearch-project/neural-search/pull/1040)) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java index d4ae88a3f..56ffec24a 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java @@ -13,8 +13,6 @@ import org.opensearch.index.query.MatchQueryBuilder; -import static org.opensearch.knn.index.query.KNNQueryBuilder.EXPAND_NESTED_FIELD; -import static org.opensearch.neuralsearch.common.MinClusterVersionUtil.isClusterOnOrAfterMinReqVersion; import static org.opensearch.neuralsearch.util.TestUtils.getModelId; import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.PARAM_NAME_WEIGHTS; @@ -124,11 +122,12 @@ private HybridQueryBuilder getQueryBuilder( final Map methodParameters, final RescoreContext rescoreContext ) { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName("passage_embedding"); - neuralQueryBuilder.modelId(modelId); - neuralQueryBuilder.queryText(QUERY); - neuralQueryBuilder.k(5); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .modelId(modelId) + .queryText(QUERY) + .k(5) + .build(); if (expandNestedDocs != null) { neuralQueryBuilder.expandNested(expandNestedDocs); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java index 5411caa0f..c3f461871 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java @@ -51,37 +51,25 @@ public void testKnnRadialSearch_E2EFlow() throws Exception { } private void validateIndexQuery(final String modelId) { - NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder( - "passage_embedding", - TEXT, - TEST_IMAGE_TEXT, - modelId, - null, - null, - 0.01f, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(TEXT) + .queryImage(TEST_IMAGE_TEXT) + .modelId(modelId) + .minScore(0.01f) + .build(); + Map responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScoreQuery); - NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder( - "passage_embedding", - TEXT, - TEST_IMAGE_TEXT, - modelId, - null, - 100000f, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(TEXT) + .queryImage(TEST_IMAGE_TEXT) + .modelId(modelId) + .maxDistance(100000f) + .build(); + Map responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxDistanceQuery); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java index bbca26cf6..7c098ef42 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java @@ -53,20 +53,13 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { private void validateTestIndex(final String modelId) throws Exception { int docCount = getDocCount(getIndexNameForTest()); assertEquals(2, docCount); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - "passage_embedding", - TEXT, - TEST_IMAGE_TEXT, - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(TEXT) + .queryImage(TEST_IMAGE_TEXT) + .modelId(modelId) + .k(1) + .build(); Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java index 02edb486c..e9005f0d6 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java @@ -76,8 +76,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - NeuralQueryBuilder neuralQueryBuilderWithoutModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); - NeuralQueryBuilder neuralQueryBuilderWithModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); + NeuralQueryBuilder neuralQueryBuilderWithoutModelId = NeuralQueryBuilder.builder() + .fieldName(TEST_ENCODING_FIELD) + .queryText(TEXT_1) + .build(); + NeuralQueryBuilder neuralQueryBuilderWithModelId = NeuralQueryBuilder.builder() + .fieldName(TEST_ENCODING_FIELD) + .queryText(TEXT_1) + .build(); if (isRunningAgainstOldCluster()) { String modelId = uploadTextEmbeddingModel(); diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java index 27ca7f42d..dcbf58ed8 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java @@ -51,11 +51,12 @@ public void testTextEmbeddingProcessor_E2EFlow() throws Exception { private void validateTestIndex(final String modelId) throws Exception { int docCount = getDocCount(getIndexNameForTest()); assertEquals(2, docCount); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName("passage_embedding"); - neuralQueryBuilder.modelId(modelId); - neuralQueryBuilder.queryText(TEXT); - neuralQueryBuilder.k(1); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(TEXT) + .modelId(modelId) + .k(1) + .build(); Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java index 78caca0e0..f64ddd455 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java @@ -128,11 +128,12 @@ private HybridQueryBuilder getQueryBuilder( final Map methodParameters, final RescoreContext rescoreContextForNeuralQuery ) { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName(VECTOR_EMBEDDING_FIELD); - neuralQueryBuilder.modelId(modelId); - neuralQueryBuilder.queryText(QUERY); - neuralQueryBuilder.k(5); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_EMBEDDING_FIELD) + .modelId(modelId) + .queryText(QUERY) + .k(5) + .build(); if (expandNestedDocs != null) { neuralQueryBuilder.expandNested(expandNestedDocs); } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java index 9dabc8d37..88af8b757 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java @@ -77,37 +77,25 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo assertEquals(numberOfDocs, docCount); loadModel(modelId); - NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder( - "passage_embedding", - text, - imageText, - modelId, - null, - null, - 0.01f, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(text) + .queryImage(imageText) + .modelId(modelId) + .minScore(0.01f) + .build(); + Map responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScore); - NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder( - "passage_embedding", - text, - imageText, - modelId, - null, - 100000f, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(text) + .queryImage(imageText) + .modelId(modelId) + .maxDistance(100000f) + .build(); + Map responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxScore); } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java index 377397d0a..e2df88d6d 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java @@ -76,20 +76,14 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); - NeuralQueryBuilder neuralQueryBuilderWithKQuery = new NeuralQueryBuilder( - "passage_embedding", - text, - imageText, - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderWithKQuery = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .queryText(text) + .queryImage(imageText) + .modelId(modelId) + .k(1) + .build(); + Map responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1); assertNotNull(responseWithKQuery); } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java index c9897447e..58e362470 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java @@ -94,8 +94,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws // the feature is introduced from 2.11 public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - NeuralQueryBuilder neuralQueryBuilderWithoutModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); - NeuralQueryBuilder neuralQueryBuilderWithModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); + NeuralQueryBuilder neuralQueryBuilderWithoutModelId = NeuralQueryBuilder.builder() + .fieldName(TEST_ENCODING_FIELD) + .queryText(TEXT_1) + .build(); + NeuralQueryBuilder neuralQueryBuilderWithModelId = NeuralQueryBuilder.builder() + .fieldName(TEST_ENCODING_FIELD) + .queryText(TEXT_1) + .build(); switch (getClusterType()) { case OLD: diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java index 218ef6a7a..e30411709 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java @@ -71,11 +71,12 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName("passage_embedding"); - neuralQueryBuilder.modelId(modelId); - neuralQueryBuilder.queryText(text); - neuralQueryBuilder.k(1); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_embedding") + .modelId(modelId) + .queryText(text) + .k(1) + .build(); Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index ce22507cb..40af1be4f 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -67,8 +67,8 @@ @Getter @Setter @Accessors(chain = true, fluent = true) -@NoArgsConstructor -@AllArgsConstructor +@NoArgsConstructor(access = AccessLevel.PRIVATE) +@AllArgsConstructor(access = AccessLevel.PRIVATE) public class NeuralQueryBuilder extends AbstractQueryBuilder implements ModelInferenceQueryBuilder { public static final String NAME = "neural"; @@ -108,6 +108,125 @@ public static void initialize(MLCommonsClientAccessor mlClient) { private Map methodParameters; private RescoreContext rescoreContext; + /** + * A custom builder class to enforce valid Neural Query Builder instantiation + */ + public static class Builder { + private String fieldName; + private String queryText; + private String queryImage; + private String modelId; + private Integer k = null; + private Float maxDistance = null; + private Float minScore = null; + private Boolean expandNested; + private Supplier vectorSupplier; + private QueryBuilder filter; + private Map methodParameters; + private RescoreContext rescoreContext; + private String queryName; + private float boost = DEFAULT_BOOST; + + public Builder() {} + + public Builder fieldName(String fieldName) { + this.fieldName = fieldName; + return this; + } + + public Builder queryText(String queryText) { + this.queryText = queryText; + return this; + } + + public Builder queryImage(String queryImage) { + this.queryImage = queryImage; + return this; + } + + public Builder modelId(String modelId) { + this.modelId = modelId; + return this; + } + + public Builder k(Integer k) { + this.k = k; + return this; + } + + public Builder maxDistance(Float maxDistance) { + this.maxDistance = maxDistance; + return this; + } + + public Builder minScore(Float minScore) { + this.minScore = minScore; + return this; + } + + public Builder expandNested(Boolean expandNested) { + this.expandNested = expandNested; + return this; + } + + public Builder vectorSupplier(Supplier vectorSupplier) { + this.vectorSupplier = vectorSupplier; + return this; + } + + public Builder filter(QueryBuilder filter) { + this.filter = filter; + return this; + } + + public Builder methodParameters(Map methodParameters) { + this.methodParameters = methodParameters; + return this; + } + + public Builder queryName(String queryName) { + this.queryName = queryName; + return this; + } + + public Builder boost(float boost) { + this.boost = boost; + return this; + } + + public Builder rescoreContext(RescoreContext rescoreContext) { + this.rescoreContext = rescoreContext; + return this; + } + + public NeuralQueryBuilder build() { + validateQueryParameters(fieldName, queryText, queryImage); + boolean queryTypeIsProvided = validateKNNQueryType(k, maxDistance, minScore); + if (queryTypeIsProvided == false) { + k = DEFAULT_K; + } + return new NeuralQueryBuilder( + fieldName, + queryText, + queryImage, + modelId, + k, + maxDistance, + minScore, + expandNested, + vectorSupplier, + filter, + methodParameters, + rescoreContext + ).boost(boost).queryName(queryName); + } + + } + + public static NeuralQueryBuilder.Builder builder() { + return new NeuralQueryBuilder.Builder(); + } + /** * Constructor from stream input * @@ -246,15 +365,16 @@ public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOEx + "]" ); } - if (StringUtils.isBlank(neuralQueryBuilder.queryText()) && StringUtils.isBlank(neuralQueryBuilder.queryImage())) { - throw new IllegalArgumentException("Either query text or image text must be provided for neural query"); - } - requireValue(neuralQueryBuilder.fieldName(), "Field name must be provided for neural query"); + validateQueryParameters(neuralQueryBuilder.fieldName(), neuralQueryBuilder.queryText(), neuralQueryBuilder.queryImage()); if (!isClusterOnOrAfterMinReqVersionForDefaultModelIdSupport()) { requireValue(neuralQueryBuilder.modelId(), "Model ID must be provided for neural query"); } - boolean queryTypeIsProvided = validateKNNQueryType(neuralQueryBuilder); + boolean queryTypeIsProvided = validateKNNQueryType( + neuralQueryBuilder.k(), + neuralQueryBuilder.maxDistance(), + neuralQueryBuilder.minScore() + ); if (queryTypeIsProvided == false) { neuralQueryBuilder.k(DEFAULT_K); } @@ -397,15 +517,22 @@ public String getWriteableName() { return NAME; } - private static boolean validateKNNQueryType(NeuralQueryBuilder neuralQueryBuilder) { + private static void validateQueryParameters(String fieldName, String queryText, String queryImage) { + if (StringUtils.isBlank(queryText) && StringUtils.isBlank(queryImage)) { + throw new IllegalArgumentException("Either query text or image text must be provided for neural query"); + } + requireValue(fieldName, "Field name must be provided for neural query"); + } + + private static boolean validateKNNQueryType(Integer k, Float maxDistance, Float minScore) { int queryCount = 0; - if (neuralQueryBuilder.k() != null) { + if (k != null) { queryCount++; } - if (neuralQueryBuilder.maxDistance() != null) { + if (maxDistance != null) { queryCount++; } - if (neuralQueryBuilder.minScore() != null) { + if (minScore != null) { queryCount++; } if (queryCount > 1) { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index baee337ce..e352b7a88 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -54,10 +54,11 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() { createSearchRequestProcessor(modelId, search_pipeline); createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); 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); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText("Hello World") + .k(1) + .build(); Map response = search(index, neuralQueryBuilder, 2); assertFalse(response.isEmpty()); } finally { @@ -112,10 +113,11 @@ public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdP createSearchRequestProcessor(modelId, search_pipeline); createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); 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); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText("Hello World") + .k(1) + .build(); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); hybridQueryBuilder.add(neuralQueryBuilder); Map response = search(index, hybridQueryBuilder, 2); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java index 05d7298ee..f0e556d03 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorTests.java @@ -38,7 +38,7 @@ public void testFactory_whenModelIdIsNotString_thenFail() { public void testProcessRequest_whenVisitingQueryBuilder_thenSuccess() throws Exception { NeuralQueryEnricherProcessor.Factory factory = new NeuralQueryEnricherProcessor.Factory(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder().fieldName("field_name").queryText("query_text").build(); SearchRequest searchRequest = new SearchRequest(); searchRequest.source(new SearchSourceBuilder().query(neuralQueryBuilder)); NeuralQueryEnricherProcessor processor = createTestProcessor(factory); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java index 385ea9428..3c5fc08ef 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java @@ -87,20 +87,13 @@ public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { modelId = prepareModel(); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_DOC_TEXT1) + .modelId(modelId) + .k(5) + .build(); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); @@ -140,20 +133,13 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu modelId = prepareModel(); createSearchPipelineWithDefaultResultsPostProcessor(SEARCH_PIPELINE); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_DOC_TEXT1) + .modelId(modelId) + .k(5) + .build(); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); @@ -182,20 +168,13 @@ public void testQueryMatches_whenMultipleShards_thenSuccessful() { createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); int totalExpectedDocQty = 6; - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 6, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_DOC_TEXT1) + .modelId(modelId) + .k(6) + .build(); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java index 934fa94e0..feb914e30 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java @@ -224,20 +224,7 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); hybridQueryBuilderDefaultNorm.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderDefaultNorm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -262,20 +249,8 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); hybridQueryBuilderL2Norm.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() + ); hybridQueryBuilderL2Norm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -325,20 +300,8 @@ public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccess HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); hybridQueryBuilderDefaultNorm.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() + ); hybridQueryBuilderDefaultNorm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -363,20 +326,8 @@ public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccess HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); hybridQueryBuilderL2Norm.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() + ); hybridQueryBuilderL2Norm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java index 4ec0b49fa..87f9942af 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java @@ -85,20 +85,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); hybridQueryBuilderArithmeticMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderArithmeticMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -123,20 +110,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); hybridQueryBuilderHarmonicMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderHarmonicMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -161,20 +135,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); hybridQueryBuilderGeometricMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderGeometricMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -224,20 +185,7 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); hybridQueryBuilderArithmeticMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderArithmeticMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -262,20 +210,7 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); hybridQueryBuilderHarmonicMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderHarmonicMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -300,20 +235,7 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); hybridQueryBuilderGeometricMean.add( - new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_DOC_TEXT1, - "", - modelId, - 5, - null, - null, - null, - null, - null, - null, - null - ) + NeuralQueryBuilder.builder().fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).queryText(TEST_DOC_TEXT1).modelId(modelId).k(5).build() ); hybridQueryBuilderGeometricMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java index 280b5e812..1e09b48c7 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java @@ -114,20 +114,13 @@ public void testNestedFieldMapping_whenDocumentsIngested_thenSuccessful() throws ); assertDoc((Map) getDocById(INDEX_NAME, "4").get("_source"), TEXT_FIELD_VALUE_2, Optional.empty()); - NeuralQueryBuilder neuralQueryBuilderQuery = new NeuralQueryBuilder( - LEVEL_1_FIELD + "." + LEVEL_2_FIELD + "." + LEVEL_3_FIELD_CONTAINER + "." + LEVEL_3_FIELD_EMBEDDING, - QUERY_TEXT, - "", - modelId, - 10, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderQuery = NeuralQueryBuilder.builder() + .fieldName(LEVEL_1_FIELD + "." + LEVEL_2_FIELD + "." + LEVEL_3_FIELD_CONTAINER + "." + LEVEL_3_FIELD_EMBEDDING) + .queryText(QUERY_TEXT) + .modelId(modelId) + .k(10) + .build(); + QueryBuilder queryNestedLowerLevel = QueryBuilders.nestedQuery( LEVEL_1_FIELD + "." + LEVEL_2_FIELD, neuralQueryBuilderQuery, diff --git a/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorTests.java index dbd1c2bd6..c9f2d6b4e 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorTests.java @@ -102,8 +102,12 @@ public void setup() { private void setupParams(Map params) { SearchSourceBuilder ssb = new SearchSourceBuilder(); - NeuralQueryBuilder nqb = new NeuralQueryBuilder(); - nqb.fieldName("embedding").k(3).modelId("embedding_id").queryText("Question about dolphins"); + NeuralQueryBuilder nqb = NeuralQueryBuilder.builder() + .fieldName("embedding") + .k(3) + .modelId("embedding_id") + .queryText("Question about dolphins") + .build(); ssb.query(nqb); List exts = List.of( new RerankSearchExtBuilder(new HashMap<>(Map.of(QueryContextSourceFetcher.NAME, new HashMap<>(params)))) diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index f40d4bf59..3a453693f 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -131,11 +131,13 @@ public void testDoToQuery_whenOneSubquery_thenBuildSuccessfully() { when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) - .vectorSupplier(TEST_VECTOR_SUPPLIER); + .vectorSupplier(TEST_VECTOR_SUPPLIER) + .build(); queryBuilder.add(neuralQueryBuilder); Query queryOnlyNeural = queryBuilder.doToQuery(mockQueryShardContext); @@ -164,11 +166,13 @@ public void testDoToQuery_whenMultipleSubqueries_thenBuildSuccessfully() { when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) - .vectorSupplier(TEST_VECTOR_SUPPLIER); + .vectorSupplier(TEST_VECTOR_SUPPLIER) + .build(); queryBuilder.add(neuralQueryBuilder); @@ -415,12 +419,14 @@ public void testToXContent_whenIncomingJsonIsCorrect_thenSuccessful() { when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) - .filter(TEST_FILTER); + .filter(TEST_FILTER) + .build(); queryBuilder.add(neuralQueryBuilder); @@ -466,11 +472,13 @@ public void testToXContent_whenIncomingJsonIsCorrect_thenSuccessful() { public void testStreams_whenWrittingToStream_thenSuccessful() { setUpClusterService(); HybridQueryBuilder original = new HybridQueryBuilder(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) - .vectorSupplier(TEST_VECTOR_SUPPLIER); + .vectorSupplier(TEST_VECTOR_SUPPLIER) + .build(); original.add(neuralQueryBuilder); @@ -498,23 +506,27 @@ public void testStreams_whenWrittingToStream_thenSuccessful() { public void testHashAndEquals_whenSameOrIdenticalObject_thenReturnEqual() { HybridQueryBuilder hybridQueryBuilderBaseline = new HybridQueryBuilder(); hybridQueryBuilderBaseline.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderBaseline.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); HybridQueryBuilder hybridQueryBuilderBaselineCopy = new HybridQueryBuilder(); hybridQueryBuilderBaselineCopy.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderBaselineCopy.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); @@ -533,66 +545,78 @@ public void testHashAndEquals_whenSubQueriesDifferent_thenReturnNotEqual() { HybridQueryBuilder hybridQueryBuilderBaseline = new HybridQueryBuilder(); hybridQueryBuilderBaseline.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderBaseline.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); HybridQueryBuilder hybridQueryBuilderOnlyOneSubQuery = new HybridQueryBuilder(); hybridQueryBuilderOnlyOneSubQuery.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); HybridQueryBuilder hybridQueryBuilderOnlyDifferentModelId = new HybridQueryBuilder(); hybridQueryBuilderOnlyDifferentModelId.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(modelId) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderBaseline.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); HybridQueryBuilder hybridQueryBuilderOnlyDifferentFieldName = new HybridQueryBuilder(); hybridQueryBuilderOnlyDifferentFieldName.add( - new NeuralQueryBuilder().fieldName(fieldName) + NeuralQueryBuilder.builder() + .fieldName(fieldName) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderOnlyDifferentFieldName.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); HybridQueryBuilder hybridQueryBuilderOnlyDifferentQuery = new HybridQueryBuilder(); hybridQueryBuilderOnlyDifferentQuery.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(queryText) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderOnlyDifferentQuery.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, TERM_QUERY_TEXT)); HybridQueryBuilder hybridQueryBuilderOnlyDifferentTermValue = new HybridQueryBuilder(); hybridQueryBuilderOnlyDifferentTermValue.add( - new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) .filter(TEST_FILTER) + .build() ); hybridQueryBuilderOnlyDifferentTermValue.add(QueryBuilders.termQuery(TEXT_FIELD_NAME, termText)); @@ -615,11 +639,13 @@ public void testHashAndEquals_whenSubQueriesDifferent_thenReturnNotEqual() { @SneakyThrows public void testRewrite_whenMultipleSubQueries_thenReturnBuilderForEachSubQuery() { HybridQueryBuilder queryBuilder = new HybridQueryBuilder(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) - .vectorSupplier(TEST_VECTOR_SUPPLIER); + .vectorSupplier(TEST_VECTOR_SUPPLIER) + .build(); queryBuilder.add(neuralQueryBuilder); @@ -789,9 +815,17 @@ public void testBuild_whenValidParameters_thenCreateQuery() { @SneakyThrows public void testDoEquals_whenSameParameters_thenEqual() { // Create neural queries - NeuralQueryBuilder neuralQueryBuilder1 = new NeuralQueryBuilder().queryText("test").modelId("test_model"); + NeuralQueryBuilder neuralQueryBuilder1 = NeuralQueryBuilder.builder() + .fieldName("test") + .queryText("test") + .modelId("test_model") + .build(); - NeuralQueryBuilder neuralQueryBuilder2 = new NeuralQueryBuilder().queryText("test").modelId("test_model"); + NeuralQueryBuilder neuralQueryBuilder2 = NeuralQueryBuilder.builder() + .fieldName("test") + .queryText("test") + .modelId("test_model") + .build(); // Create neural sparse queries with queryTokensSupplier NeuralSparseQueryBuilder neuralSparseQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName("test_field") @@ -822,7 +856,9 @@ public void testValidate_whenInvalidParameters_thenThrowException() { } public void testVisit() { - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new NeuralQueryBuilder()).add(new NeuralSparseQueryBuilder()); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add( + NeuralQueryBuilder.builder().fieldName("test").queryText("test").build() + ).add(new NeuralSparseQueryBuilder()); List visitedQueries = new ArrayList<>(); hybridQueryBuilder.visit(createTestVisitor(visitedQueries)); assertEquals(3, visitedQueries.size()); diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java index 1d4a398cb..5791d57d2 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java @@ -425,15 +425,106 @@ public void testFromXContent_whenBuiltWithInvalidFilter_thenFail() { expectThrows(ParsingException.class, () -> NeuralQueryBuilder.fromXContent(contentParser)); } + @SneakyThrows + public void testBuilderInstantiation_whenBuiltWithRequiredFields_thenBuildSuccessfully() { + /* + { + "VECTOR_FIELD": { + "query_text": "string", + "query_image": "string", + "model_id": "string", + "k": int + } + } + */ + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder().fieldName(FIELD_NAME).queryText(QUERY_TEXT).build(); + + assertEquals(FIELD_NAME, neuralQueryBuilder.fieldName()); + assertEquals(QUERY_TEXT, neuralQueryBuilder.queryText()); + } + + @SneakyThrows + public void testBuilderInstantiation_whenBuiltWithOptionalFields_thenBuildSuccessfully() { + /* + { + "VECTOR_FIELD": { + "query_text": "string", + "model_id": "string", + "k": int, + "boost": 10.0, + "_name": "something", + "expandNestedDocs": true + } + } + */ + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) + .queryText(QUERY_TEXT) + .queryImage(IMAGE_TEXT) + .modelId(MODEL_ID) + .k(K) + .expandNested(Boolean.TRUE) + .boost(BOOST) + .queryName(QUERY_NAME) + .build(); + + assertEquals(FIELD_NAME, neuralQueryBuilder.fieldName()); + assertEquals(QUERY_TEXT, neuralQueryBuilder.queryText()); + assertEquals(IMAGE_TEXT, neuralQueryBuilder.queryImage()); + assertEquals(MODEL_ID, neuralQueryBuilder.modelId()); + assertEquals(K, neuralQueryBuilder.k()); + assertEquals(BOOST, neuralQueryBuilder.boost(), 0.0); + assertEquals(QUERY_NAME, neuralQueryBuilder.queryName()); + assertEquals(Boolean.TRUE, neuralQueryBuilder.expandNested()); + } + + @SneakyThrows + public void testBuilderInstantiation_whenMissingFieldName_thenFail() { + /* + { + null : { + "query_text": "string", + "model_id": "string", + "k": int + } + } + */ + + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> NeuralQueryBuilder.builder().queryText(QUERY_TEXT).modelId(MODEL_ID).k(K).build() + ); + assertEquals("Field name must be provided for neural query", exception.getMessage()); + } + + @SneakyThrows + public void testBuilderInstantiation_whenMissingBothQueryTextAndQueryImage_thenFail() { + /* + { + "VECTOR_FIELD": { + "model_id": "string", + "k": int + } + } + */ + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> NeuralQueryBuilder.builder().fieldName(FIELD_NAME).modelId(MODEL_ID).k(K).build() + ); + assertEquals("Either query text or image text must be provided for neural query", exception.getMessage()); + } + @SuppressWarnings("unchecked") @SneakyThrows public void testToXContent() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .modelId(MODEL_ID) .queryText(QUERY_TEXT) .k(K) .expandNested(Boolean.TRUE) - .filter(TEST_FILTER); + .filter(TEST_FILTER) + .build(); XContentBuilder builder = XContentFactory.jsonBuilder(); builder = neuralQueryBuilder.toXContent(builder, ToXContent.EMPTY_PARAMS); @@ -477,15 +568,16 @@ public void testStreams_whenClusterServiceWithDifferentVersions() { @SneakyThrows private void testStreams() { - NeuralQueryBuilder original = new NeuralQueryBuilder(); - original.fieldName(FIELD_NAME); - original.queryText(QUERY_TEXT); - original.queryImage(IMAGE_TEXT); - original.modelId(MODEL_ID); - original.k(K); - original.boost(BOOST); - original.queryName(QUERY_NAME); - original.filter(TEST_FILTER); + NeuralQueryBuilder original = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) + .queryText(QUERY_TEXT) + .queryImage(IMAGE_TEXT) + .modelId(MODEL_ID) + .k(K) + .boost(BOOST) + .queryName(QUERY_NAME) + .filter(TEST_FILTER) + .build(); BytesStreamOutput streamOutput = new BytesStreamOutput(); original.writeTo(streamOutput); @@ -519,101 +611,123 @@ public void testHashAndEquals() { QueryBuilder filter1 = new MatchAllQueryBuilder(); QueryBuilder filter2 = new MatchNoneQueryBuilder(); - NeuralQueryBuilder neuralQueryBuilder_baseline = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_baseline = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .queryImage(imageText1) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline - NeuralQueryBuilder neuralQueryBuilder_baselineCopy = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_baselineCopy = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except default boost and query name - NeuralQueryBuilder neuralQueryBuilder_defaultBoostAndQueryName = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_defaultBoostAndQueryName = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff field name - NeuralQueryBuilder neuralQueryBuilder_diffFieldName = new NeuralQueryBuilder().fieldName(fieldName2) + NeuralQueryBuilder neuralQueryBuilder_diffFieldName = NeuralQueryBuilder.builder() + .fieldName(fieldName2) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff query text - NeuralQueryBuilder neuralQueryBuilder_diffQueryText = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffQueryText = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText2) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff model ID - NeuralQueryBuilder neuralQueryBuilder_diffModelId = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffModelId = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId2) .k(k1) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff k - NeuralQueryBuilder neuralQueryBuilder_diffK = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffK = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k2) .boost(boost1) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff boost - NeuralQueryBuilder neuralQueryBuilder_diffBoost = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffBoost = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost2) .queryName(queryName1) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except diff query name - NeuralQueryBuilder neuralQueryBuilder_diffQueryName = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffQueryName = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName2) - .filter(filter1); + .filter(filter1) + .build(); // Identical to neuralQueryBuilder_baseline except no filter - NeuralQueryBuilder neuralQueryBuilder_noFilter = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_noFilter = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost1) - .queryName(queryName2); + .queryName(queryName2) + .build(); // Identical to neuralQueryBuilder_baseline except no filter - NeuralQueryBuilder neuralQueryBuilder_diffFilter = new NeuralQueryBuilder().fieldName(fieldName1) + NeuralQueryBuilder neuralQueryBuilder_diffFilter = NeuralQueryBuilder.builder() + .fieldName(fieldName1) .queryText(queryText1) .modelId(modelId1) .k(k1) .boost(boost1) .queryName(queryName2) - .filter(filter2); + .filter(filter2) + .build(); assertEquals(neuralQueryBuilder_baseline, neuralQueryBuilder_baseline); assertEquals(neuralQueryBuilder_baseline.hashCode(), neuralQueryBuilder_baseline.hashCode()); @@ -651,7 +765,12 @@ public void testHashAndEquals() { @SneakyThrows public void testRewrite_whenVectorSupplierNull_thenSetVectorSupplier() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME).queryText(QUERY_TEXT).modelId(MODEL_ID).k(K); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) + .queryText(QUERY_TEXT) + .modelId(MODEL_ID) + .k(K) + .build(); List expectedVector = Arrays.asList(1.0f, 2.0f, 3.0f, 4.0f, 5.0f); MLCommonsClientAccessor mlCommonsClientAccessor = mock(MLCommonsClientAccessor.class); doAnswer(invocation -> { @@ -683,11 +802,13 @@ public void testRewrite_whenVectorSupplierNull_thenSetVectorSupplier() { @SneakyThrows public void testRewrite_whenVectorSupplierNullAndQueryTextAndImageTextSet_thenSetVectorSupplier() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .queryText(QUERY_TEXT) .queryImage(IMAGE_TEXT) .modelId(MODEL_ID) - .k(K); + .k(K) + .build(); List expectedVector = Arrays.asList(1.0f, 2.0f, 3.0f, 4.0f, 5.0f); MLCommonsClientAccessor mlCommonsClientAccessor = mock(MLCommonsClientAccessor.class); doAnswer(invocation -> { @@ -719,19 +840,22 @@ public void testRewrite_whenVectorSupplierNullAndQueryTextAndImageTextSet_thenSe public void testRewrite_whenVectorNull_thenReturnCopy() { Supplier nullSupplier = () -> null; - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .queryText(QUERY_TEXT) .queryImage(IMAGE_TEXT) .modelId(MODEL_ID) .k(K) .expandNested(Boolean.TRUE) - .vectorSupplier(nullSupplier); + .vectorSupplier(nullSupplier) + .build(); QueryBuilder queryBuilder = neuralQueryBuilder.doRewrite(null); assertEquals(neuralQueryBuilder, queryBuilder); } public void testRewrite_whenVectorSupplierAndVectorSet_thenReturnKNNQueryBuilder() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .queryText(QUERY_TEXT) .queryImage(IMAGE_TEXT) .modelId(MODEL_ID) @@ -739,7 +863,8 @@ public void testRewrite_whenVectorSupplierAndVectorSet_thenReturnKNNQueryBuilder .expandNested(Boolean.TRUE) .methodParameters(Map.of("ef_search", 100)) .rescoreContext(RescoreContext.getDefault()) - .vectorSupplier(TEST_VECTOR_SUPPLIER); + .vectorSupplier(TEST_VECTOR_SUPPLIER) + .build(); KNNQueryBuilder expected = KNNQueryBuilder.builder() .k(K) @@ -755,12 +880,14 @@ public void testRewrite_whenVectorSupplierAndVectorSet_thenReturnKNNQueryBuilder } public void testRewrite_whenFilterSet_thenKNNQueryBuilderFilterSet() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) - .filter(TEST_FILTER); + .filter(TEST_FILTER) + .build(); QueryBuilder queryBuilder = neuralQueryBuilder.doRewrite(null); assertTrue(queryBuilder instanceof KNNQueryBuilder); KNNQueryBuilder knnQueryBuilder = (KNNQueryBuilder) queryBuilder; @@ -768,12 +895,14 @@ public void testRewrite_whenFilterSet_thenKNNQueryBuilderFilterSet() { } public void testQueryCreation_whenCreateQueryWithDoToQuery_thenFail() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(FIELD_NAME) + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(FIELD_NAME) .queryText(QUERY_TEXT) .modelId(MODEL_ID) .k(K) .vectorSupplier(TEST_VECTOR_SUPPLIER) - .filter(TEST_FILTER); + .filter(TEST_FILTER) + .build(); QueryShardContext queryShardContext = mock(QueryShardContext.class); UnsupportedOperationException exception = expectThrows( UnsupportedOperationException.class, diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java index a7d349f07..37d1320d4 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java @@ -102,20 +102,12 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilderTextQuery = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilderTextQuery = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); final float boost = 2.0f; neuralQueryBuilderTextQuery.boost(boost); @@ -126,20 +118,16 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); assertEquals(expectedScore, objectToFloat(firstInnerHitTextQuery.get("_score")), DELTA_FOR_SCORE_ASSERTION); - NeuralQueryBuilder neuralQueryBuilderMultimodalQuery = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - TEST_IMAGE_TEXT, - modelId, - 1, - null, - null, - null, - null, - null, - Map.of("ef_search", 10), - RescoreContext.getDefault() - ); + NeuralQueryBuilder neuralQueryBuilderMultimodalQuery = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .queryImage(TEST_IMAGE_TEXT) + .modelId(modelId) + .k(1) + .methodParameters(Map.of("ef_search", 10)) + .rescoreContext(RescoreContext.getDefault()) + .build(); + Map searchResponseAsMapMultimodalQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderMultimodalQuery, 1); Map firstInnerHitMultimodalQuery = getFirstInnerHit(searchResponseAsMapMultimodalQuery); @@ -155,20 +143,12 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { // Context: https://github.com/opensearch-project/neural-search/pull/697#discussion_r1571549776 // Test radial search max distance query - NeuralQueryBuilder neuralQueryWithMaxDistanceBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - null, - 100.0f, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryWithMaxDistanceBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .maxDistance(100.0f) + .build(); Map searchResponseAsMapWithMaxDistanceQuery = search( TEST_BASIC_INDEX_NAME, @@ -186,20 +166,12 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { ); // Test radial search min score query - NeuralQueryBuilder neuralQueryWithMinScoreBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - null, - null, - 0.01f, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryWithMinScoreBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .minScore(0.01f) + .build(); Map searchResponseAsMapWithMinScoreQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryWithMinScoreBuilder, 1); Map firstInnerHitWithMinScoreQuery = getFirstInnerHit(searchResponseAsMapWithMinScoreQuery); @@ -243,20 +215,12 @@ public void testRescoreQuery() { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); modelId = prepareModel(); MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); - NeuralQueryBuilder rescoreNeuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder rescoreNeuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, matchAllQueryBuilder, rescoreNeuralQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); @@ -323,34 +287,19 @@ public void testBooleanQuery_withMultipleNeuralQueries() { modelId = prepareModel(); // verify two neural queries wrapped into bool BoolQueryBuilder boolQueryBuilderTwoNeuralQueries = new BoolQueryBuilder(); - NeuralQueryBuilder neuralQueryBuilder1 = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); - NeuralQueryBuilder neuralQueryBuilder2 = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_2, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder1 = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); + + NeuralQueryBuilder neuralQueryBuilder2 = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_2) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); boolQueryBuilderTwoNeuralQueries.should(neuralQueryBuilder1).should(neuralQueryBuilder2); @@ -367,20 +316,12 @@ public void testBooleanQuery_withMultipleNeuralQueries() { // verify bool with one neural and one bm25 query BoolQueryBuilder boolQueryBuilderMixOfNeuralAndBM25 = new BoolQueryBuilder(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); @@ -425,20 +366,12 @@ public void testNestedQuery() { try { initializeIndexIfNotExist(TEST_NESTED_INDEX_NAME); modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_NESTED, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - null, - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_NESTED) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .build(); Map searchResponseAsMap = search(TEST_NESTED_INDEX_NAME, neuralQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); @@ -478,20 +411,14 @@ public void testFilterQuery() { try { initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null, - null, - null, - new MatchQueryBuilder("_id", "3"), - null, - null - ); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .k(1) + .filter(new MatchQueryBuilder("_id", "3")) + .build(); + Map searchResponseAsMap = search(TEST_MULTI_DOC_INDEX_NAME, neuralQueryBuilder, 3); assertEquals(1, getHitCount(searchResponseAsMap)); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); diff --git a/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java b/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java index 26dbb289b..c783449e5 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitorTests.java @@ -16,9 +16,11 @@ public class NeuralSearchQueryVisitorTests extends OpenSearchTestCase { public void testAccept_whenNeuralQueryBuilderWithoutModelId_thenSetModelId() { String modelId = "bdcvjkcdjvkddcjxdjsc"; - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName("passage_text"); - neuralQueryBuilder.k(768); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_text") + .queryText("query_text") + .k(768) + .build(); NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(modelId, null); neuralSearchQueryVisitor.accept(neuralQueryBuilder); @@ -29,9 +31,11 @@ public void testAccept_whenNeuralQueryBuilderWithoutModelId_thenSetModelId() { 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); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() + .fieldName("passage_text") + .queryText("query_text") + .k(768) + .build(); NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(null, neuralInfoMap); neuralSearchQueryVisitor.accept(neuralQueryBuilder); @@ -76,7 +80,7 @@ public void testAccept_whenNeuralSparseQueryBuilderWithoutFieldModelId_thenSetFi } public void testAccept_whenNullValuesInVisitor_thenFail() { - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder().fieldName("passage_text").queryText("query_text").build(); NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); NeuralSearchQueryVisitor neuralSearchQueryVisitor = new NeuralSearchQueryVisitor(null, null);