From bbfbb2772be43eb34bc76f2dc856dce302172505 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Wed, 10 Apr 2024 10:54:02 -0700 Subject: [PATCH] Optimizing tests for less model upload calls Signed-off-by: Martin Gaievski --- .../processor/NormalizationProcessorIT.java | 52 ++-- .../TextImageEmbeddingProcessorIT.java | 18 +- .../rerank/MLOpenSearchRerankProcessorIT.java | 36 +-- .../neuralsearch/query/HybridQueryIT.java | 40 +-- .../neuralsearch/query/NeuralQueryIT.java | 227 ++++++------------ .../query/NeuralSparseQueryIT.java | 37 +-- .../neuralsearch/BaseNeuralSearchIT.java | 7 +- 7 files changed, 132 insertions(+), 285 deletions(-) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java index f6f59958f..4a7ccf8de 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java @@ -164,7 +164,7 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu } @SneakyThrows - public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful() { + public void testQueryMatches_whenMultipleShards_thenSuccessful() { String modelId = null; try { initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); @@ -223,55 +223,37 @@ public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful // verify that all ids are unique assertEquals(Set.copyOf(ids).size(), ids.size()); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE); - } - } - - @SneakyThrows - public void testResultProcessor_whenMultipleShardsAndNoMatches_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6)); - hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7)); + // very case when there are partial match + HybridQueryBuilder hybridQueryBuilderPartialMatch = new HybridQueryBuilder(); + hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); + hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4)); + hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7)); - Map searchResponseAsMap = search( + Map searchResponseAsMapPartialMatch = search( TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, + hybridQueryBuilderPartialMatch, null, 5, Map.of("search_pipeline", SEARCH_PIPELINE) ); - assertQueryResults(searchResponseAsMap, 0, true); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE); - } - } - - @SneakyThrows - public void testResultProcessor_whenMultipleShardsAndPartialMatches_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + assertQueryResults(searchResponseAsMapPartialMatch, 4, true, Range.between(0.33f, 1.0f)); - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); - hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4)); - hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7)); + // verify case when query doesn't have a match + HybridQueryBuilder hybridQueryBuilderNoMatches = new HybridQueryBuilder(); + hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6)); + hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7)); - Map searchResponseAsMap = search( + Map searchResponseAsMapNoMatches = search( TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, + hybridQueryBuilderNoMatches, null, 5, Map.of("search_pipeline", SEARCH_PIPELINE) ); - assertQueryResults(searchResponseAsMap, 4, true, Range.between(0.33f, 1.0f)); + assertQueryResults(searchResponseAsMapNoMatches, 0, true); } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE); + wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java index 43d629c71..cf1473293 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java @@ -34,29 +34,19 @@ public void setUp() throws Exception { updateClusterSettings(); } - public void testEmbeddingProcessor_whenIngestingDocumentWithSourceMatchingTextMapping_thenSuccessful() throws Exception { + public void testEmbeddingProcessor_whenIngestingDocumentWithOrWithoutSourceMatchingMapping_thenSuccessful() throws Exception { String modelId = null; try { modelId = uploadModel(); loadModel(modelId); createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); createTextImageEmbeddingIndex(); + // verify doc with mapping ingestDocumentWithTextMappedToEmbeddingField(); assertEquals(1, getDocCount(INDEX_NAME)); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } - } - - public void testEmbeddingProcessor_whenIngestingDocumentWithSourceWithoutMatchingInMapping_thenSuccessful() throws Exception { - String modelId = null; - try { - modelId = uploadModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); - createTextImageEmbeddingIndex(); + // verify doc without mapping ingestDocumentWithoutMappedFields(); - assertEquals(1, getDocCount(INDEX_NAME)); + assertEquals(2, getDocCount(INDEX_NAME)); } finally { wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java index 33620d149..f7dd7b647 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java @@ -12,8 +12,6 @@ import org.apache.hc.core5.http.HttpHeaders; import org.apache.hc.core5.http.io.entity.EntityUtils; import org.apache.hc.core5.http.message.BasicHeader; -import org.junit.After; -import org.junit.Before; import org.opensearch.client.Request; import org.opensearch.client.Response; import org.opensearch.common.xcontent.XContentHelper; @@ -36,33 +34,19 @@ public class MLOpenSearchRerankProcessorIT extends BaseNeuralSearchIT { private final static String TEXT_REP_1 = "Jacques loves fish. Fish make Jacques happy"; private final static String TEXT_REP_2 = "Fish like to eat plankton"; private final static String INDEX_CONFIG = "{\"mappings\": {\"properties\": {\"text_representation\": {\"type\": \"text\"}}}}"; - private String modelId; - - @After - @SneakyThrows - public void tearDown() { - super.tearDown(); - /* this is required to minimize chance of model not being deployed due to open memory CB, - * this happens in case we leave model from previous test case. We use new model for every test, and old model - * can be undeployed and deleted to free resources after each test case execution. - */ - deleteModel(modelId); - deleteSearchPipeline(PIPELINE_NAME); - deleteIndex(INDEX_NAME); - } - - @Before - @SneakyThrows - public void setup() { - modelId = uploadTextSimilarityModel(); - loadModel(modelId); - } @SneakyThrows public void testCrossEncoderRerankProcessor() { - createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json"); - setupIndex(); - runQueries(); + String modelId = null; + try { + modelId = uploadTextSimilarityModel(); + loadModel(modelId); + createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json"); + setupIndex(); + runQueries(); + } finally { + wipeOfTestResources(INDEX_NAME, null, modelId, PIPELINE_NAME); + } } private String uploadTextSimilarityModel() throws Exception { diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java index 12fa534dd..b198a51ee 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java @@ -566,28 +566,17 @@ public void testRequestCache_whenMultipleShardsQueryReturnResults_thenSuccessful @SneakyThrows public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNestedFields_thenSuccess() { - String modelId = null; String alias = "alias_with_filter"; try { initializeIndexIfNotExist(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME); - modelId = prepareModel(); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); // create alias for index QueryBuilder aliasFilter = QueryBuilders.boolQuery() .mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); createIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias, aliasFilter); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 5, - null, - null - ); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(neuralQueryBuilder); + hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); Map searchResponseAsMap = search( alias, @@ -608,34 +597,23 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNested assertEquals(RELATION_EQUAL_TO, total.get("relation")); } finally { deleteIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias); - wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, modelId, SEARCH_PIPELINE); + wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, null, SEARCH_PIPELINE); } } @SneakyThrows public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() { - String modelId = null; String alias = "alias_with_filter"; try { initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - modelId = prepareModel(); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); // create alias for index QueryBuilder aliasFilter = QueryBuilders.boolQuery() .mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); createIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias, aliasFilter); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 5, - null, - null - ); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(neuralQueryBuilder); + hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); Map searchResponseAsMap = search( alias, @@ -656,7 +634,7 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() { assertEquals(RELATION_EQUAL_TO, total.get("relation")); } finally { deleteIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias); - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, modelId, SEARCH_PIPELINE); + wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, SEARCH_PIPELINE); } } @@ -892,7 +870,15 @@ private void addDocsToIndex(final String testMultiDocIndexName) { Collections.singletonList(TEST_TEXT_FIELD_NAME_1), Collections.singletonList(TEST_DOC_TEXT2) ); - assertEquals(3, getDocCount(testMultiDocIndexName)); + addKnnDoc( + testMultiDocIndexName, + "4", + Collections.emptyList(), + Collections.emptyList(), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT3) + ); + assertEquals(4, getDocCount(testMultiDocIndexName)); } private List> getNestedHits(Map searchResponseAsMap) { diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java index 6f1e5f27e..9cc9dda71 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java @@ -27,7 +27,6 @@ public class NeuralQueryIT extends BaseNeuralSearchIT { private static final String TEST_BASIC_INDEX_NAME = "test-neural-basic-index"; private static final String TEST_MULTI_VECTOR_FIELD_INDEX_NAME = "test-neural-multi-vector-field-index"; - private static final String TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME = "test-neural-text-and-vector-field-index"; private static final String TEST_NESTED_INDEX_NAME = "test-neural-nested-index"; private static final String TEST_MULTI_DOC_INDEX_NAME = "test-neural-multi-doc-index"; private static final String TEST_QUERY_TEXT = "Hello world"; @@ -45,67 +44,40 @@ public void setUp() throws Exception { } /** - * Tests basic query: + * Tests basic query with boost parameter: * { * "query": { * "neural": { * "text_knn": { * "query_text": "Hello world", * "model_id": "dcsdcasd", - * "k": 1 + * "k": 1, + * "boost": 2.0 * } * } * } * } - */ - @SneakyThrows - public void testBasicQuery() { - String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - "", - modelId, - 1, - null, - null - ); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } - } - - /** - * Tests basic query with boost parameter: + * and query with image query part * { * "query": { * "neural": { * "text_knn": { * "query_text": "Hello world", + * "query_image": "base64_1234567890", * "model_id": "dcsdcasd", - * "k": 1, - * "boost": 2.0 + * "k": 1 * } * } * } * } */ @SneakyThrows - public void testBoostQuery() { + public void testQueryWithBoostAndImageQuery() { String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( + NeuralQueryBuilder neuralQueryBuilderTextQuery = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, TEST_QUERY_TEXT, "", @@ -116,13 +88,33 @@ public void testBoostQuery() { ); final float boost = 2.0f; - neuralQueryBuilder.boost(boost); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + neuralQueryBuilderTextQuery.boost(boost); + Map searchResponseAsMapTextQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderTextQuery, 1); + Map firstInnerHitTextQuery = getFirstInnerHit(searchResponseAsMapTextQuery); - assertEquals("1", firstInnerHit.get("_id")); + assertEquals("1", firstInnerHitTextQuery.get("_id")); float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION); + 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 + ); + Map searchResponseAsMapMultimodalQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderMultimodalQuery, 1); + Map firstInnerHitMultimodalQuery = getFirstInnerHit(searchResponseAsMapMultimodalQuery); + + assertEquals("1", firstInnerHitMultimodalQuery.get("_id")); + float expectedScoreMultimodalQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertEquals( + expectedScoreMultimodalQuery, + objectToFloat(firstInnerHitMultimodalQuery.get("_score")), + DELTA_FOR_SCORE_ASSERTION + ); } finally { wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); } @@ -200,6 +192,27 @@ public void testRescoreQuery() { * } * } * } + * and bool should with BM25 and neural query: + * { + * "query": { + * "bool" : { + * "should": [ + * "neural": { + * "field_1": { + * "query_text": "Hello world", + * "model_id": "dcsdcasd", + * "k": 1 + * }, + * }, + * "match": { + * "field_2": { + * "query": "Hello world" + * } + * } + * ] + * } + * } + * } */ @SneakyThrows public void testBooleanQuery_withMultipleNeuralQueries() { @@ -207,8 +220,8 @@ public void testBooleanQuery_withMultipleNeuralQueries() { try { initializeIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME); modelId = prepareModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - + // verify two neural queries wrapped into bool + BoolQueryBuilder boolQueryBuilderTwoNeuralQueries = new BoolQueryBuilder(); NeuralQueryBuilder neuralQueryBuilder1 = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, TEST_QUERY_TEXT, @@ -228,50 +241,21 @@ public void testBooleanQuery_withMultipleNeuralQueries() { null ); - boolQueryBuilder.should(neuralQueryBuilder1).should(neuralQueryBuilder2); + boolQueryBuilderTwoNeuralQueries.should(neuralQueryBuilder1).should(neuralQueryBuilder2); - Map searchResponseAsMap = search(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + Map searchResponseAsMapTwoNeuralQueries = search( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + boolQueryBuilderTwoNeuralQueries, + 1 + ); + Map firstInnerHitTwoNeuralQueries = getFirstInnerHit(searchResponseAsMapTwoNeuralQueries); - assertEquals("1", firstInnerHit.get("_id")); + assertEquals("1", firstInnerHitTwoNeuralQueries.get("_id")); float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, null, modelId, null); - } - } - - /** - * Tests bool should with BM25 and neural query: - * { - * "query": { - * "bool" : { - * "should": [ - * "neural": { - * "field_1": { - * "query_text": "Hello world", - * "model_id": "dcsdcasd", - * "k": 1 - * }, - * }, - * "match": { - * "field_2": { - * "query": "Hello world" - * } - * } - * ] - * } - * } - * } - */ - @SneakyThrows - public void testBooleanQuery_withNeuralAndBM25Queries() { - String modelId = null; - try { - initializeIndexIfNotExist(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME); - modelId = prepareModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + assertEquals(expectedScore, objectToFloat(firstInnerHitTwoNeuralQueries.get("_score")), DELTA_FOR_SCORE_ASSERTION); + // 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, @@ -284,16 +268,20 @@ public void testBooleanQuery_withNeuralAndBM25Queries() { MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); - boolQueryBuilder.should(neuralQueryBuilder).should(matchQueryBuilder); + boolQueryBuilderMixOfNeuralAndBM25.should(neuralQueryBuilder).should(matchQueryBuilder); - Map searchResponseAsMap = search(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + Map searchResponseAsMapMixOfNeuralAndBM25 = search( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + boolQueryBuilderMixOfNeuralAndBM25, + 1 + ); + Map firstInnerHitMixOfNeuralAndBM25 = getFirstInnerHit(searchResponseAsMapMixOfNeuralAndBM25); - assertEquals("1", firstInnerHit.get("_id")); + assertEquals("1", firstInnerHitMixOfNeuralAndBM25.get("_id")); float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); + assertTrue(minExpectedScore < objectToFloat(firstInnerHitMixOfNeuralAndBM25.get("_score"))); } finally { - wipeOfTestResources(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, null, modelId, null); } } @@ -389,47 +377,6 @@ public void testFilterQuery() { } } - /** - * Tests basic query for multimodal: - * { - * "query": { - * "neural": { - * "text_knn": { - * "query_text": "Hello world", - * "query_image": "base64_1234567890", - * "model_id": "dcsdcasd", - * "k": 1 - * } - * } - * } - * } - */ - @SneakyThrows - public void testMultimodalQuery() { - String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareModel(); - NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( - TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, - TEST_IMAGE_TEXT, - modelId, - 1, - null, - null - ); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } - } - @SneakyThrows private void initializeIndexIfNotExist(String indexName) { if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_INDEX_NAME)) { @@ -458,7 +405,9 @@ private void initializeIndexIfNotExist(String indexName) { TEST_MULTI_VECTOR_FIELD_INDEX_NAME, "1", List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), - List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray()) + List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_QUERY_TEXT) ); assertEquals(1, getDocCount(TEST_MULTI_VECTOR_FIELD_INDEX_NAME)); } @@ -477,22 +426,6 @@ private void initializeIndexIfNotExist(String indexName) { assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME)); } - if (TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)) { - prepareKnnIndex( - TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_QUERY_TEXT) - ); - assertEquals(1, getDocCount(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)); - } - if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) { prepareKnnIndex( TEST_MULTI_DOC_INDEX_NAME, diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java index 0caca4f43..d43d252b9 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java @@ -42,40 +42,7 @@ public void setUp() throws Exception { } /** - * Tests basic query: - * { - * "query": { - * "neural_sparse": { - * "text_sparse": { - * "query_text": "Hello world a b", - * "model_id": "dcsdcasd" - * } - * } - * } - * } - */ - @SneakyThrows - public void testBasicQueryUsingQueryText() { - String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } - } - - /** - * Tests basic query: + * Tests basic query with boost: * { * "query": { * "neural_sparse": { @@ -89,7 +56,7 @@ public void testBasicQueryUsingQueryText() { * } */ @SneakyThrows - public void testBoostQuery() { + public void testBasicQueryUsingQueryText() { String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index 32ba16696..1c21e5f5e 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -82,6 +82,7 @@ public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase { ProcessorType.TEXT_IMAGE_EMBEDDING, "processor/PipelineForTextImageEmbeddingProcessorConfiguration.json" ); + private static final Set SUCCESS_STATUSES = Set.of(RestStatus.CREATED, RestStatus.OK); protected final ClassLoader classLoader = this.getClass().getClassLoader(); @@ -114,6 +115,7 @@ protected void updateClusterSettings() { updateClusterSettings("plugins.ml_commons.only_run_on_ml_node", false); // default threshold for native circuit breaker is 90, it may be not enough on test runner machine updateClusterSettings("plugins.ml_commons.native_memory_threshold", 100); + updateClusterSettings("plugins.ml_commons.jvm_heap_memory_threshold", 95); updateClusterSettings("plugins.ml_commons.allow_registering_model_via_url", true); } @@ -633,7 +635,10 @@ protected void addKnnDoc( request.setJsonEntity(builder.toString()); Response response = client().performRequest(request); - assertEquals(request.getEndpoint() + ": failed", RestStatus.CREATED, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + assertTrue( + request.getEndpoint() + ": failed", + SUCCESS_STATUSES.contains(RestStatus.fromCode(response.getStatusLine().getStatusCode())) + ); } @SneakyThrows