Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit bbfbb27

Browse files
martin-gaievskimartin-gaievski
committedApr 11, 2024·
Optimizing tests for less model upload calls
Signed-off-by: Martin Gaievski <gaievski@amazon.com>
1 parent cb83e20 commit bbfbb27

File tree

7 files changed

+132
-285
lines changed

7 files changed

+132
-285
lines changed
 

‎src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java

+17-35
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu
164164
}
165165

166166
@SneakyThrows
167-
public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful() {
167+
public void testQueryMatches_whenMultipleShards_thenSuccessful() {
168168
String modelId = null;
169169
try {
170170
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
@@ -223,55 +223,37 @@ public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful
223223

224224
// verify that all ids are unique
225225
assertEquals(Set.copyOf(ids).size(), ids.size());
226-
} finally {
227-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE);
228-
}
229-
}
230-
231-
@SneakyThrows
232-
public void testResultProcessor_whenMultipleShardsAndNoMatches_thenSuccessful() {
233-
try {
234-
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
235-
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
236226

237-
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
238-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6));
239-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
227+
// very case when there are partial match
228+
HybridQueryBuilder hybridQueryBuilderPartialMatch = new HybridQueryBuilder();
229+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
230+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4));
231+
hybridQueryBuilderPartialMatch.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
240232

241-
Map<String, Object> searchResponseAsMap = search(
233+
Map<String, Object> searchResponseAsMapPartialMatch = search(
242234
TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME,
243-
hybridQueryBuilder,
235+
hybridQueryBuilderPartialMatch,
244236
null,
245237
5,
246238
Map.of("search_pipeline", SEARCH_PIPELINE)
247239
);
248-
assertQueryResults(searchResponseAsMap, 0, true);
249-
} finally {
250-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE);
251-
}
252-
}
253-
254-
@SneakyThrows
255-
public void testResultProcessor_whenMultipleShardsAndPartialMatches_thenSuccessful() {
256-
try {
257-
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME);
258-
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
240+
assertQueryResults(searchResponseAsMapPartialMatch, 4, true, Range.between(0.33f, 1.0f));
259241

260-
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
261-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
262-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4));
263-
hybridQueryBuilder.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
242+
// verify case when query doesn't have a match
243+
HybridQueryBuilder hybridQueryBuilderNoMatches = new HybridQueryBuilder();
244+
hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT6));
245+
hybridQueryBuilderNoMatches.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT7));
264246

265-
Map<String, Object> searchResponseAsMap = search(
247+
Map<String, Object> searchResponseAsMapNoMatches = search(
266248
TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME,
267-
hybridQueryBuilder,
249+
hybridQueryBuilderNoMatches,
268250
null,
269251
5,
270252
Map.of("search_pipeline", SEARCH_PIPELINE)
271253
);
272-
assertQueryResults(searchResponseAsMap, 4, true, Range.between(0.33f, 1.0f));
254+
assertQueryResults(searchResponseAsMapNoMatches, 0, true);
273255
} finally {
274-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE);
256+
wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE);
275257
}
276258
}
277259

‎src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java

+4-14
Original file line numberDiff line numberDiff line change
@@ -34,29 +34,19 @@ public void setUp() throws Exception {
3434
updateClusterSettings();
3535
}
3636

37-
public void testEmbeddingProcessor_whenIngestingDocumentWithSourceMatchingTextMapping_thenSuccessful() throws Exception {
37+
public void testEmbeddingProcessor_whenIngestingDocumentWithOrWithoutSourceMatchingMapping_thenSuccessful() throws Exception {
3838
String modelId = null;
3939
try {
4040
modelId = uploadModel();
4141
loadModel(modelId);
4242
createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING);
4343
createTextImageEmbeddingIndex();
44+
// verify doc with mapping
4445
ingestDocumentWithTextMappedToEmbeddingField();
4546
assertEquals(1, getDocCount(INDEX_NAME));
46-
} finally {
47-
wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null);
48-
}
49-
}
50-
51-
public void testEmbeddingProcessor_whenIngestingDocumentWithSourceWithoutMatchingInMapping_thenSuccessful() throws Exception {
52-
String modelId = null;
53-
try {
54-
modelId = uploadModel();
55-
loadModel(modelId);
56-
createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING);
57-
createTextImageEmbeddingIndex();
47+
// verify doc without mapping
5848
ingestDocumentWithoutMappedFields();
59-
assertEquals(1, getDocCount(INDEX_NAME));
49+
assertEquals(2, getDocCount(INDEX_NAME));
6050
} finally {
6151
wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null);
6252
}

‎src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java

+10-26
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@
1212
import org.apache.hc.core5.http.HttpHeaders;
1313
import org.apache.hc.core5.http.io.entity.EntityUtils;
1414
import org.apache.hc.core5.http.message.BasicHeader;
15-
import org.junit.After;
16-
import org.junit.Before;
1715
import org.opensearch.client.Request;
1816
import org.opensearch.client.Response;
1917
import org.opensearch.common.xcontent.XContentHelper;
@@ -36,33 +34,19 @@ public class MLOpenSearchRerankProcessorIT extends BaseNeuralSearchIT {
3634
private final static String TEXT_REP_1 = "Jacques loves fish. Fish make Jacques happy";
3735
private final static String TEXT_REP_2 = "Fish like to eat plankton";
3836
private final static String INDEX_CONFIG = "{\"mappings\": {\"properties\": {\"text_representation\": {\"type\": \"text\"}}}}";
39-
private String modelId;
40-
41-
@After
42-
@SneakyThrows
43-
public void tearDown() {
44-
super.tearDown();
45-
/* this is required to minimize chance of model not being deployed due to open memory CB,
46-
* this happens in case we leave model from previous test case. We use new model for every test, and old model
47-
* can be undeployed and deleted to free resources after each test case execution.
48-
*/
49-
deleteModel(modelId);
50-
deleteSearchPipeline(PIPELINE_NAME);
51-
deleteIndex(INDEX_NAME);
52-
}
53-
54-
@Before
55-
@SneakyThrows
56-
public void setup() {
57-
modelId = uploadTextSimilarityModel();
58-
loadModel(modelId);
59-
}
6037

6138
@SneakyThrows
6239
public void testCrossEncoderRerankProcessor() {
63-
createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json");
64-
setupIndex();
65-
runQueries();
40+
String modelId = null;
41+
try {
42+
modelId = uploadTextSimilarityModel();
43+
loadModel(modelId);
44+
createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json");
45+
setupIndex();
46+
runQueries();
47+
} finally {
48+
wipeOfTestResources(INDEX_NAME, null, modelId, PIPELINE_NAME);
49+
}
6650
}
6751

6852
private String uploadTextSimilarityModel() throws Exception {

‎src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java

+13-27
Original file line numberDiff line numberDiff line change
@@ -566,28 +566,17 @@ public void testRequestCache_whenMultipleShardsQueryReturnResults_thenSuccessful
566566

567567
@SneakyThrows
568568
public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNestedFields_thenSuccess() {
569-
String modelId = null;
570569
String alias = "alias_with_filter";
571570
try {
572571
initializeIndexIfNotExist(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME);
573-
modelId = prepareModel();
574572
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
575573
// create alias for index
576574
QueryBuilder aliasFilter = QueryBuilders.boolQuery()
577575
.mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
578576
createIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias, aliasFilter);
579577

580-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
581-
TEST_KNN_VECTOR_FIELD_NAME_1,
582-
TEST_QUERY_TEXT,
583-
"",
584-
modelId,
585-
5,
586-
null,
587-
null
588-
);
589578
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
590-
hybridQueryBuilder.add(neuralQueryBuilder);
579+
hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1));
591580

592581
Map<String, Object> searchResponseAsMap = search(
593582
alias,
@@ -608,34 +597,23 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilterAndIndexWithNested
608597
assertEquals(RELATION_EQUAL_TO, total.get("relation"));
609598
} finally {
610599
deleteIndexAlias(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, alias);
611-
wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, modelId, SEARCH_PIPELINE);
600+
wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, null, SEARCH_PIPELINE);
612601
}
613602
}
614603

615604
@SneakyThrows
616605
public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() {
617-
String modelId = null;
618606
String alias = "alias_with_filter";
619607
try {
620608
initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME);
621-
modelId = prepareModel();
622609
createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE);
623610
// create alias for index
624611
QueryBuilder aliasFilter = QueryBuilders.boolQuery()
625612
.mustNot(QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3));
626613
createIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias, aliasFilter);
627614

628-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
629-
TEST_KNN_VECTOR_FIELD_NAME_1,
630-
TEST_QUERY_TEXT,
631-
"",
632-
modelId,
633-
5,
634-
null,
635-
null
636-
);
637615
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
638-
hybridQueryBuilder.add(neuralQueryBuilder);
616+
hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1));
639617

640618
Map<String, Object> searchResponseAsMap = search(
641619
alias,
@@ -656,7 +634,7 @@ public void testWrappedQueryWithFilter_whenIndexAliasHasFilters_thenSuccess() {
656634
assertEquals(RELATION_EQUAL_TO, total.get("relation"));
657635
} finally {
658636
deleteIndexAlias(TEST_MULTI_DOC_INDEX_NAME, alias);
659-
wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, modelId, SEARCH_PIPELINE);
637+
wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, SEARCH_PIPELINE);
660638
}
661639
}
662640

@@ -892,7 +870,15 @@ private void addDocsToIndex(final String testMultiDocIndexName) {
892870
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
893871
Collections.singletonList(TEST_DOC_TEXT2)
894872
);
895-
assertEquals(3, getDocCount(testMultiDocIndexName));
873+
addKnnDoc(
874+
testMultiDocIndexName,
875+
"4",
876+
Collections.emptyList(),
877+
Collections.emptyList(),
878+
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
879+
Collections.singletonList(TEST_DOC_TEXT3)
880+
);
881+
assertEquals(4, getDocCount(testMultiDocIndexName));
896882
}
897883

898884
private List<Map<String, Object>> getNestedHits(Map<String, Object> searchResponseAsMap) {

‎src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java

+80-147
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
public class NeuralQueryIT extends BaseNeuralSearchIT {
2828
private static final String TEST_BASIC_INDEX_NAME = "test-neural-basic-index";
2929
private static final String TEST_MULTI_VECTOR_FIELD_INDEX_NAME = "test-neural-multi-vector-field-index";
30-
private static final String TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME = "test-neural-text-and-vector-field-index";
3130
private static final String TEST_NESTED_INDEX_NAME = "test-neural-nested-index";
3231
private static final String TEST_MULTI_DOC_INDEX_NAME = "test-neural-multi-doc-index";
3332
private static final String TEST_QUERY_TEXT = "Hello world";
@@ -45,67 +44,40 @@ public void setUp() throws Exception {
4544
}
4645

4746
/**
48-
* Tests basic query:
47+
* Tests basic query with boost parameter:
4948
* {
5049
* "query": {
5150
* "neural": {
5251
* "text_knn": {
5352
* "query_text": "Hello world",
5453
* "model_id": "dcsdcasd",
55-
* "k": 1
54+
* "k": 1,
55+
* "boost": 2.0
5656
* }
5757
* }
5858
* }
5959
* }
60-
*/
61-
@SneakyThrows
62-
public void testBasicQuery() {
63-
String modelId = null;
64-
try {
65-
initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME);
66-
modelId = prepareModel();
67-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
68-
TEST_KNN_VECTOR_FIELD_NAME_1,
69-
TEST_QUERY_TEXT,
70-
"",
71-
modelId,
72-
1,
73-
null,
74-
null
75-
);
76-
Map<String, Object> searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1);
77-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
78-
79-
assertEquals("1", firstInnerHit.get("_id"));
80-
float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
81-
assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION);
82-
} finally {
83-
wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null);
84-
}
85-
}
86-
87-
/**
88-
* Tests basic query with boost parameter:
60+
* and query with image query part
8961
* {
9062
* "query": {
9163
* "neural": {
9264
* "text_knn": {
9365
* "query_text": "Hello world",
66+
* "query_image": "base64_1234567890",
9467
* "model_id": "dcsdcasd",
95-
* "k": 1,
96-
* "boost": 2.0
68+
* "k": 1
9769
* }
9870
* }
9971
* }
10072
* }
10173
*/
10274
@SneakyThrows
103-
public void testBoostQuery() {
75+
public void testQueryWithBoostAndImageQuery() {
10476
String modelId = null;
10577
try {
10678
initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME);
10779
modelId = prepareModel();
108-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
80+
NeuralQueryBuilder neuralQueryBuilderTextQuery = new NeuralQueryBuilder(
10981
TEST_KNN_VECTOR_FIELD_NAME_1,
11082
TEST_QUERY_TEXT,
11183
"",
@@ -116,13 +88,33 @@ public void testBoostQuery() {
11688
);
11789

11890
final float boost = 2.0f;
119-
neuralQueryBuilder.boost(boost);
120-
Map<String, Object> searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1);
121-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
91+
neuralQueryBuilderTextQuery.boost(boost);
92+
Map<String, Object> searchResponseAsMapTextQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderTextQuery, 1);
93+
Map<String, Object> firstInnerHitTextQuery = getFirstInnerHit(searchResponseAsMapTextQuery);
12294

123-
assertEquals("1", firstInnerHit.get("_id"));
95+
assertEquals("1", firstInnerHitTextQuery.get("_id"));
12496
float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
125-
assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION);
97+
assertEquals(expectedScore, objectToFloat(firstInnerHitTextQuery.get("_score")), DELTA_FOR_SCORE_ASSERTION);
98+
99+
NeuralQueryBuilder neuralQueryBuilderMultimodalQuery = new NeuralQueryBuilder(
100+
TEST_KNN_VECTOR_FIELD_NAME_1,
101+
TEST_QUERY_TEXT,
102+
TEST_IMAGE_TEXT,
103+
modelId,
104+
1,
105+
null,
106+
null
107+
);
108+
Map<String, Object> searchResponseAsMapMultimodalQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderMultimodalQuery, 1);
109+
Map<String, Object> firstInnerHitMultimodalQuery = getFirstInnerHit(searchResponseAsMapMultimodalQuery);
110+
111+
assertEquals("1", firstInnerHitMultimodalQuery.get("_id"));
112+
float expectedScoreMultimodalQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
113+
assertEquals(
114+
expectedScoreMultimodalQuery,
115+
objectToFloat(firstInnerHitMultimodalQuery.get("_score")),
116+
DELTA_FOR_SCORE_ASSERTION
117+
);
126118
} finally {
127119
wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null);
128120
}
@@ -200,15 +192,36 @@ public void testRescoreQuery() {
200192
* }
201193
* }
202194
* }
195+
* and bool should with BM25 and neural query:
196+
* {
197+
* "query": {
198+
* "bool" : {
199+
* "should": [
200+
* "neural": {
201+
* "field_1": {
202+
* "query_text": "Hello world",
203+
* "model_id": "dcsdcasd",
204+
* "k": 1
205+
* },
206+
* },
207+
* "match": {
208+
* "field_2": {
209+
* "query": "Hello world"
210+
* }
211+
* }
212+
* ]
213+
* }
214+
* }
215+
* }
203216
*/
204217
@SneakyThrows
205218
public void testBooleanQuery_withMultipleNeuralQueries() {
206219
String modelId = null;
207220
try {
208221
initializeIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME);
209222
modelId = prepareModel();
210-
BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
211-
223+
// verify two neural queries wrapped into bool
224+
BoolQueryBuilder boolQueryBuilderTwoNeuralQueries = new BoolQueryBuilder();
212225
NeuralQueryBuilder neuralQueryBuilder1 = new NeuralQueryBuilder(
213226
TEST_KNN_VECTOR_FIELD_NAME_1,
214227
TEST_QUERY_TEXT,
@@ -228,50 +241,21 @@ public void testBooleanQuery_withMultipleNeuralQueries() {
228241
null
229242
);
230243

231-
boolQueryBuilder.should(neuralQueryBuilder1).should(neuralQueryBuilder2);
244+
boolQueryBuilderTwoNeuralQueries.should(neuralQueryBuilder1).should(neuralQueryBuilder2);
232245

233-
Map<String, Object> searchResponseAsMap = search(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, boolQueryBuilder, 1);
234-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
246+
Map<String, Object> searchResponseAsMapTwoNeuralQueries = search(
247+
TEST_MULTI_VECTOR_FIELD_INDEX_NAME,
248+
boolQueryBuilderTwoNeuralQueries,
249+
1
250+
);
251+
Map<String, Object> firstInnerHitTwoNeuralQueries = getFirstInnerHit(searchResponseAsMapTwoNeuralQueries);
235252

236-
assertEquals("1", firstInnerHit.get("_id"));
253+
assertEquals("1", firstInnerHitTwoNeuralQueries.get("_id"));
237254
float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
238-
assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION);
239-
} finally {
240-
wipeOfTestResources(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, null, modelId, null);
241-
}
242-
}
243-
244-
/**
245-
* Tests bool should with BM25 and neural query:
246-
* {
247-
* "query": {
248-
* "bool" : {
249-
* "should": [
250-
* "neural": {
251-
* "field_1": {
252-
* "query_text": "Hello world",
253-
* "model_id": "dcsdcasd",
254-
* "k": 1
255-
* },
256-
* },
257-
* "match": {
258-
* "field_2": {
259-
* "query": "Hello world"
260-
* }
261-
* }
262-
* ]
263-
* }
264-
* }
265-
* }
266-
*/
267-
@SneakyThrows
268-
public void testBooleanQuery_withNeuralAndBM25Queries() {
269-
String modelId = null;
270-
try {
271-
initializeIndexIfNotExist(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME);
272-
modelId = prepareModel();
273-
BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
255+
assertEquals(expectedScore, objectToFloat(firstInnerHitTwoNeuralQueries.get("_score")), DELTA_FOR_SCORE_ASSERTION);
274256

257+
// verify bool with one neural and one bm25 query
258+
BoolQueryBuilder boolQueryBuilderMixOfNeuralAndBM25 = new BoolQueryBuilder();
275259
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
276260
TEST_KNN_VECTOR_FIELD_NAME_1,
277261
TEST_QUERY_TEXT,
@@ -284,16 +268,20 @@ public void testBooleanQuery_withNeuralAndBM25Queries() {
284268

285269
MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT);
286270

287-
boolQueryBuilder.should(neuralQueryBuilder).should(matchQueryBuilder);
271+
boolQueryBuilderMixOfNeuralAndBM25.should(neuralQueryBuilder).should(matchQueryBuilder);
288272

289-
Map<String, Object> searchResponseAsMap = search(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, boolQueryBuilder, 1);
290-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
273+
Map<String, Object> searchResponseAsMapMixOfNeuralAndBM25 = search(
274+
TEST_MULTI_VECTOR_FIELD_INDEX_NAME,
275+
boolQueryBuilderMixOfNeuralAndBM25,
276+
1
277+
);
278+
Map<String, Object> firstInnerHitMixOfNeuralAndBM25 = getFirstInnerHit(searchResponseAsMapMixOfNeuralAndBM25);
291279

292-
assertEquals("1", firstInnerHit.get("_id"));
280+
assertEquals("1", firstInnerHitMixOfNeuralAndBM25.get("_id"));
293281
float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
294-
assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score")));
282+
assertTrue(minExpectedScore < objectToFloat(firstInnerHitMixOfNeuralAndBM25.get("_score")));
295283
} finally {
296-
wipeOfTestResources(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, null, modelId, null);
284+
wipeOfTestResources(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, null, modelId, null);
297285
}
298286
}
299287

@@ -389,47 +377,6 @@ public void testFilterQuery() {
389377
}
390378
}
391379

392-
/**
393-
* Tests basic query for multimodal:
394-
* {
395-
* "query": {
396-
* "neural": {
397-
* "text_knn": {
398-
* "query_text": "Hello world",
399-
* "query_image": "base64_1234567890",
400-
* "model_id": "dcsdcasd",
401-
* "k": 1
402-
* }
403-
* }
404-
* }
405-
* }
406-
*/
407-
@SneakyThrows
408-
public void testMultimodalQuery() {
409-
String modelId = null;
410-
try {
411-
initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME);
412-
modelId = prepareModel();
413-
NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
414-
TEST_KNN_VECTOR_FIELD_NAME_1,
415-
TEST_QUERY_TEXT,
416-
TEST_IMAGE_TEXT,
417-
modelId,
418-
1,
419-
null,
420-
null
421-
);
422-
Map<String, Object> searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1);
423-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
424-
425-
assertEquals("1", firstInnerHit.get("_id"));
426-
float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT);
427-
assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION);
428-
} finally {
429-
wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null);
430-
}
431-
}
432-
433380
@SneakyThrows
434381
private void initializeIndexIfNotExist(String indexName) {
435382
if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_INDEX_NAME)) {
@@ -458,7 +405,9 @@ private void initializeIndexIfNotExist(String indexName) {
458405
TEST_MULTI_VECTOR_FIELD_INDEX_NAME,
459406
"1",
460407
List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2),
461-
List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray())
408+
List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray()),
409+
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
410+
Collections.singletonList(TEST_QUERY_TEXT)
462411
);
463412
assertEquals(1, getDocCount(TEST_MULTI_VECTOR_FIELD_INDEX_NAME));
464413
}
@@ -477,22 +426,6 @@ private void initializeIndexIfNotExist(String indexName) {
477426
assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME));
478427
}
479428

480-
if (TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)) {
481-
prepareKnnIndex(
482-
TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME,
483-
Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE))
484-
);
485-
addKnnDoc(
486-
TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME,
487-
"1",
488-
Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
489-
Collections.singletonList(Floats.asList(testVector).toArray()),
490-
Collections.singletonList(TEST_TEXT_FIELD_NAME_1),
491-
Collections.singletonList(TEST_QUERY_TEXT)
492-
);
493-
assertEquals(1, getDocCount(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME));
494-
}
495-
496429
if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) {
497430
prepareKnnIndex(
498431
TEST_MULTI_DOC_INDEX_NAME,

‎src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java

+2-35
Original file line numberDiff line numberDiff line change
@@ -42,40 +42,7 @@ public void setUp() throws Exception {
4242
}
4343

4444
/**
45-
* Tests basic query:
46-
* {
47-
* "query": {
48-
* "neural_sparse": {
49-
* "text_sparse": {
50-
* "query_text": "Hello world a b",
51-
* "model_id": "dcsdcasd"
52-
* }
53-
* }
54-
* }
55-
* }
56-
*/
57-
@SneakyThrows
58-
public void testBasicQueryUsingQueryText() {
59-
String modelId = null;
60-
try {
61-
initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME);
62-
modelId = prepareSparseEncodingModel();
63-
NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1)
64-
.queryText(TEST_QUERY_TEXT)
65-
.modelId(modelId);
66-
Map<String, Object> searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1);
67-
Map<String, Object> firstInnerHit = getFirstInnerHit(searchResponseAsMap);
68-
69-
assertEquals("1", firstInnerHit.get("_id"));
70-
float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT);
71-
assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA);
72-
} finally {
73-
wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null);
74-
}
75-
}
76-
77-
/**
78-
* Tests basic query:
45+
* Tests basic query with boost:
7946
* {
8047
* "query": {
8148
* "neural_sparse": {
@@ -89,7 +56,7 @@ public void testBasicQueryUsingQueryText() {
8956
* }
9057
*/
9158
@SneakyThrows
92-
public void testBoostQuery() {
59+
public void testBasicQueryUsingQueryText() {
9360
String modelId = null;
9461
try {
9562
initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME);

‎src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java

+6-1
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase {
8282
ProcessorType.TEXT_IMAGE_EMBEDDING,
8383
"processor/PipelineForTextImageEmbeddingProcessorConfiguration.json"
8484
);
85+
private static final Set<RestStatus> SUCCESS_STATUSES = Set.of(RestStatus.CREATED, RestStatus.OK);
8586

8687
protected final ClassLoader classLoader = this.getClass().getClassLoader();
8788

@@ -114,6 +115,7 @@ protected void updateClusterSettings() {
114115
updateClusterSettings("plugins.ml_commons.only_run_on_ml_node", false);
115116
// default threshold for native circuit breaker is 90, it may be not enough on test runner machine
116117
updateClusterSettings("plugins.ml_commons.native_memory_threshold", 100);
118+
updateClusterSettings("plugins.ml_commons.jvm_heap_memory_threshold", 95);
117119
updateClusterSettings("plugins.ml_commons.allow_registering_model_via_url", true);
118120
}
119121

@@ -633,7 +635,10 @@ protected void addKnnDoc(
633635

634636
request.setJsonEntity(builder.toString());
635637
Response response = client().performRequest(request);
636-
assertEquals(request.getEndpoint() + ": failed", RestStatus.CREATED, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
638+
assertTrue(
639+
request.getEndpoint() + ": failed",
640+
SUCCESS_STATUSES.contains(RestStatus.fromCode(response.getStatusLine().getStatusCode()))
641+
);
637642
}
638643

639644
@SneakyThrows

0 commit comments

Comments
 (0)
Please sign in to comment.