Skip to content

Commit

Permalink
Optimizing tests for less model upload calls
Browse files Browse the repository at this point in the history 
Signed-off-by: Martin Gaievski <[email protected]>
  • Loading branch information
@martin-gaievski
martin-gaievski committed Apr 11, 2024
1 parent cb83e20 commit bbfbb27
Show file tree
Hide file tree
Showing 7 changed files with 132 additions and 285 deletions.
52 changes: 17 additions & 35 deletions src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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<String, Object> searchResponseAsMap = search(
Map<String, Object> 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<String, Object> searchResponseAsMap = search(
Map<String, Object> 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);
}
}

Expand Down
18 changes: 4 additions & 14 deletions src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
36 changes: 10 additions & 26 deletions ...test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 {
Expand Down
40 changes: 13 additions & 27 deletions src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, Object> searchResponseAsMap = search(
alias,
Expand All @@ -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<String, Object> searchResponseAsMap = search(
alias,
Expand All @@ -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);
}
}

Expand Down Expand Up @@ -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<Map<String, Object>> getNestedHits(Map<String, Object> searchResponseAsMap) {
Expand Down
Loading

0 comments on commit bbfbb27

Please sign in to comment.