From e83ad5d6c4c7a5a26c2f58873ba0a691de55d4f9 Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Thu, 30 Jan 2025 13:23:23 -0800 Subject: [PATCH 1/5] Handle exceptions gracefully when delete non-existent resources during integ test resource clean up Signed-off-by: Weijia Zhao --- .../neuralsearch/BaseNeuralSearchIT.java | 42 ++++++++++++------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index 2182d8d79..54edfd260 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -4,6 +4,7 @@ */ package org.opensearch.neuralsearch; +import org.opensearch.client.ResponseException; import org.opensearch.ml.common.model.MLModelState; import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray; @@ -1504,23 +1505,32 @@ protected void wipeOfTestResources( final String modelId, final String searchPipeline ) throws IOException { - if (ingestPipeline != null) { - deletePipeline(ingestPipeline); - } - if (searchPipeline != null) { - deleteSearchPipeline(searchPipeline); - } - if (modelId != null) { - try { - deleteModel(modelId); - } catch (AssertionError e) { - // sometimes we have flaky test that the model state doesn't change after call undeploy api - // for this case we can call undeploy api one more time - deleteModel(modelId); + try { + if (ingestPipeline != null) { + deletePipeline(ingestPipeline); + } + if (searchPipeline != null) { + deleteSearchPipeline(searchPipeline); + } + if (modelId != null) { + try { + deleteModel(modelId); + } catch (AssertionError e) { + // sometimes we have flaky test that the model state doesn't change after call undeploy api + // for this case we can call undeploy api one more time + deleteModel(modelId); + } + } + if (indexName != null) { + deleteIndex(indexName); + } + } catch (ResponseException e) { + // It's possible that test fails during resources (index, model, pipeline, etc.) creation, when clean up + // these resources after test run, the delete methods will throw ResponseException with 404 Not Found code + // In this case, we just need to ignore this exception, for other exceptions, continue throwing + if (RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode()) != RestStatus.NOT_FOUND) { + throw e; } - } - if (indexName != null) { - deleteIndex(indexName); } } From b9dd581f98bc316196eb5e4e284e43ba15a1aca2 Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Tue, 4 Feb 2025 14:00:17 -0800 Subject: [PATCH 2/5] Clean up test resource dynamically after each test case Signed-off-by: Weijia Zhao --- .../bwc/restart/BatchIngestionIT.java | 10 +- .../bwc/restart/HybridSearchIT.java | 19 +- .../restart/HybridSearchWithRescoreIT.java | 21 +- .../bwc/restart/KnnRadialSearchIT.java | 13 +- .../bwc/restart/MultiModalSearchIT.java | 13 +- .../NeuralQueryEnricherProcessorIT.java | 38 +- .../bwc/restart/NeuralSparseSearchIT.java | 27 +- .../NeuralSparseTwoPhaseProcessorIT.java | 20 +- .../bwc/restart/SemanticSearchIT.java | 13 +- .../bwc/restart/TextChunkingProcessorIT.java | 8 +- .../bwc/rolling/BatchIngestionIT.java | 14 +- .../bwc/rolling/HybridSearchIT.java | 20 +- .../rolling/HybridSearchWithRescoreIT.java | 22 +- .../bwc/rolling/KnnRadialSearchIT.java | 14 +- .../bwc/rolling/MultiModalSearchIT.java | 14 +- .../NeuralQueryEnricherProcessorIT.java | 36 +- .../bwc/rolling/NeuralSparseSearchIT.java | 28 +- .../NeuralSparseTwoPhaseProcessorIT.java | 17 +- .../bwc/rolling/SemanticSearchIT.java | 14 +- .../bwc/rolling/TextChunkingProcessorIT.java | 10 +- .../NeuralQueryEnricherProcessorIT.java | 110 +- .../NeuralSparseTwoPhaseProcessorIT.java | 476 +++---- .../processor/NormalizationProcessorIT.java | 272 ++-- .../processor/RRFProcessorIT.java | 37 +- .../processor/ScoreCombinationIT.java | 397 +++--- .../processor/ScoreNormalizationIT.java | 260 ++-- .../processor/SparseEncodingProcessIT.java | 74 +- .../processor/TextChunkingProcessorIT.java | 172 +-- .../processor/TextEmbeddingProcessorIT.java | 299 ++-- .../TextImageEmbeddingProcessorIT.java | 44 +- .../rerank/ByFieldRerankProcessorIT.java | 12 +- .../rerank/MLOpenSearchRerankProcessorIT.java | 15 +- .../query/HybridQueryAggregationsIT.java | 612 ++++---- .../query/HybridQueryExplainIT.java | 1226 ++++++++--------- .../neuralsearch/query/HybridQueryIT.java | 798 +++++------ .../query/HybridQueryPostFilterIT.java | 56 +- .../neuralsearch/query/HybridQuerySortIT.java | 587 ++++---- .../neuralsearch/query/NeuralQueryIT.java | 380 +++-- .../query/NeuralSparseQueryIT.java | 159 +-- .../BucketAggregationsWithHybridQueryIT.java | 744 +++++----- .../MetricAggregationsWithHybridQueryIT.java | 434 +++--- ...PipelineAggregationsWithHybridQueryIT.java | 447 +++--- .../neuralsearch/BaseNeuralSearchIT.java | 244 +++- .../neuralsearch/util/TestUtils.java | 3 + 44 files changed, 3846 insertions(+), 4383 deletions(-) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java index 7fa0b1301..4ad0bdc6d 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java @@ -40,13 +40,9 @@ public void testBatchIngestionWithNeuralSparseProcessor_E2EFlow() throws Excepti String modelId = null; modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); loadModel(modelId); - try { - List> docs = prepareDataForBulkIngestion(5, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); - validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); - } finally { - wipeOfTestResources(indexName, PIPELINE_NAME, modelId, null); - } + List> docs = prepareDataForBulkIngestion(5, 5); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); + validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java index d08d208da..d42a29c59 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java @@ -67,18 +67,13 @@ private void validateNormalizationProcessor(final String fileName, final String addDocuments(getIndexNameForTest(), true); createSearchPipeline(searchPipelineName); } else { - String modelId = null; - try { - modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocuments(getIndexNameForTest(), false); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); - hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); - } finally { - wipeOfTestResources(getIndexNameForTest(), pipelineName, modelId, searchPipelineName); - } + String modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocuments(getIndexNameForTest(), false); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java index 9329a934b..2ca84745b 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java @@ -59,19 +59,14 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f })) ); } else { - String modelId = null; - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java index 2e3ac34ab..02e52ce9d 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java @@ -38,15 +38,10 @@ public void testKnnRadialSearch_E2EFlow() throws Exception { ); addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - String modelId = null; - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateIndexQuery(modelId); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateIndexQuery(modelId); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java index df3f8e94e..4b89eb901 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java @@ -38,15 +38,10 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { ); addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - String modelId = null; - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateTestIndex(modelId); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateTestIndex(modelId); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java index 97182bfa0..eb29e909e 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java @@ -59,18 +59,13 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); } else { - String modelId = null; - try { - modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); - sparseEncodingQueryBuilderWithModelId.modelId(modelId); - assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") - ); - } finally { - wipeOfTestResources(getIndexNameForTest(), SPARSE_INGEST_PIPELINE_NAME, modelId, SPARSE_SEARCH_PIPELINE_NAME); - } + String modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + sparseEncodingQueryBuilderWithModelId.modelId(modelId); + assertEquals( + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + ); } } @@ -105,19 +100,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") ); } else { - String modelId = null; - try { - modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - neuralQueryBuilderWithModelId.modelId(modelId); + String modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + neuralQueryBuilderWithModelId.modelId(modelId); - assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") - ); - } finally { - wipeOfTestResources(getIndexNameForTest(), DENSE_INGEST_PIPELINE_NAME, modelId, DENSE_SEARCH_PIPELINE_NAME); - } + assertEquals( + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + ); } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java index 5978f71ab..893da4c87 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java @@ -51,22 +51,17 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEXT_1) ); } else { - String modelId = null; - try { - modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); - addSparseEncodingDoc( - getIndexNameForTest(), - "1", - List.of(TEST_SPARSE_ENCODING_FIELD), - List.of(testRankFeaturesDoc2), - List.of(TEST_TEXT_FIELD), - List.of(TEXT_2) - ); - validateTestIndex(modelId); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + String modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + addSparseEncodingDoc( + getIndexNameForTest(), + "1", + List.of(TEST_SPARSE_ENCODING_FIELD), + List.of(testRankFeaturesDoc2), + List.of(TEST_TEXT_FIELD), + List.of(TEXT_2) + ); + validateTestIndex(modelId); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java index 6a6994809..3ee154d8d 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java @@ -45,21 +45,11 @@ public void testNeuralSparseQueryTwoPhaseProcessor_NeuralSearch_E2EFlow() throws Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); assertNotNull(resultWith2PhasePipeline); } else { - String modelId = null; - try { - modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); - neuralSparseQueryBuilder.modelId(modelId); - Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); - assertNotNull(resultWith2PhasePipeline); - } finally { - wipeOfTestResources( - getIndexNameForTest(), - NEURAL_SPARSE_INGEST_PIPELINE_NAME, - modelId, - NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME - ); - } + String modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + neuralSparseQueryBuilder.modelId(modelId); + Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); + assertNotNull(resultWith2PhasePipeline); } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java index 9dc35c02d..d378ab12f 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java @@ -36,15 +36,10 @@ public void testTextEmbeddingProcessor_E2EFlow() throws Exception { ); addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); } else { - String modelId = null; - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null); - validateTestIndex(modelId); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null); + validateTestIndex(modelId); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java index d68903478..254d49041 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java @@ -40,12 +40,8 @@ public void testTextChunkingProcessor_E2EFlow() throws Exception { addDocument(indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null); validateTestIndex(indexName, OUTPUT_FIELD, 1, expectedPassages); } else { - try { - addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages); - } finally { - wipeOfTestResources(indexName, PIPELINE_NAME, null, null); - } + addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java index 01fce83f0..6fc64fa96 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java @@ -46,15 +46,11 @@ public void testBatchIngestion_SparseEncodingProcessor_E2EFlow() throws Exceptio validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); break; case UPGRADED: - try { - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - List> docsForUpgraded = prepareDataForBulkIngestion(10, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); - validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); - } finally { - wipeOfTestResources(indexName, SPARSE_PIPELINE, sparseModelId, null); - } + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + List> docsForUpgraded = prepareDataForBulkIngestion(10, 5); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); + validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java index de7ddef55..c4beb3845 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java @@ -76,18 +76,14 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); - hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java index 3cea2400d..afed627e8 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java @@ -78,19 +78,15 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java index e9fb1a4a7..e59f6e911 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java @@ -56,15 +56,11 @@ public void testKnnRadialSearch_E2EFlow() throws Exception { } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateIndexQueryOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateIndexQueryOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java index 6ced20d95..11be782a4 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java @@ -56,15 +56,11 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java index e22e65e4a..6f66d0f72 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java @@ -73,17 +73,13 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws ); break; case UPGRADED: - try { - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); - assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") - ); - } finally { - wipeOfTestResources(getIndexNameForTest(), SPARSE_INGEST_PIPELINE_NAME, sparseModelId, SPARSE_SEARCH_PIPELINE_NAME); - } + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); + assertEquals( + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + ); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -138,18 +134,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep ); break; case UPGRADED: - try { - denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(denseModelId); - neuralQueryBuilderWithModelId.modelId(denseModelId); + denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(denseModelId); + neuralQueryBuilderWithModelId.modelId(denseModelId); - assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") - ); - } finally { - wipeOfTestResources(getIndexNameForTest(), DENSE_INGEST_PIPELINE_NAME, denseModelId, DENSE_SEARCH_PIPELINE_NAME); - } + assertEquals( + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + ); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java index 8cebae4a2..b91211859 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java @@ -78,22 +78,18 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addSparseEncodingDoc( - getIndexNameForTest(), - "2", - List.of(TEST_SPARSE_ENCODING_FIELD), - List.of(testRankFeaturesDoc3), - List.of(TEST_TEXT_FIELD), - List.of(TEXT_UPGRADED) - ); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addSparseEncodingDoc( + getIndexNameForTest(), + "2", + List.of(TEST_SPARSE_ENCODING_FIELD), + List.of(testRankFeaturesDoc3), + List.of(TEST_TEXT_FIELD), + List.of(TEXT_UPGRADED) + ); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java index 5338b777f..60dd9eac7 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java @@ -57,19 +57,10 @@ public void testNeuralSparseTwoPhaseProcessorIT_NeuralSparseSearch_E2EFlow() thr assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); break; case UPGRADED: - try { - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - neuralSparseQueryBuilder.modelId(sparseModelId); - assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); - } finally { - wipeOfTestResources( - getIndexNameForTest(), - SPARSE_INGEST_PIPELINE_NAME, - sparseModelId, - SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME - ); - } + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + neuralSparseQueryBuilder.modelId(sparseModelId); + assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java index 989d53897..99107e79c 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java @@ -51,15 +51,11 @@ public void testSemanticSearch_E2EFlow() throws Exception { } break; case UPGRADED: - try { - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED); - } finally { - wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); - } + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java index 29ec1fed5..227b64f85 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java @@ -54,13 +54,9 @@ public void testTextChunkingProcessor_E2EFlow() throws Exception { } break; case UPGRADED: - try { - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - addDocument(indexName, "2", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(indexName, OUTPUT_FIELD, totalDocsCountUpgraded, expectedPassages); - } finally { - wipeOfTestResources(indexName, PIPELINE_NAME, null, null); - } + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + addDocument(indexName, "2", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(indexName, OUTPUT_FIELD, totalDocsCountUpgraded, expectedPassages); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index e352b7a88..d21bd1526 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -48,84 +48,68 @@ public void setUp() throws Exception { @SneakyThrows public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() { String modelId = null; - try { - initializeIndexIfNotExist(index); - modelId = prepareModel(); - createSearchRequestProcessor(modelId, search_pipeline); - createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); - updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); - 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 { - wipeOfTestResources(index, ingest_pipeline, modelId, search_pipeline); - } + initializeIndexIfNotExist(index); + modelId = prepareModel(); + createSearchRequestProcessor(modelId, search_pipeline); + createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); + updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + 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()); } @SneakyThrows public void testNeuralQueryEnricherProcessor_whenNoModelIdPassedInNeuralSparseQuery_thenSuccess() { String modelId = null; - try { - initializeIndexIfNotExist(sparseIndex); - modelId = prepareSparseEncodingModel(); - createSearchRequestProcessor(modelId, search_pipeline); - createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.SPARSE_ENCODING); - updateIndexSettings(sparseIndex, Settings.builder().put("index.search.default_pipeline", search_pipeline)); - NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); - neuralSparseQueryBuilder.fieldName(TEST_RANK_FEATURES_FIELD_NAME_1); - neuralSparseQueryBuilder.queryText("hello"); - Map response = search(sparseIndex, neuralSparseQueryBuilder, 2); - assertFalse(response.isEmpty()); - } finally { - wipeOfTestResources(sparseIndex, ingest_pipeline, modelId, search_pipeline); - } + initializeIndexIfNotExist(sparseIndex); + modelId = prepareSparseEncodingModel(); + createSearchRequestProcessor(modelId, search_pipeline); + createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.SPARSE_ENCODING); + updateIndexSettings(sparseIndex, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); + neuralSparseQueryBuilder.fieldName(TEST_RANK_FEATURES_FIELD_NAME_1); + neuralSparseQueryBuilder.queryText("hello"); + Map response = search(sparseIndex, neuralSparseQueryBuilder, 2); + assertFalse(response.isEmpty()); } @SneakyThrows public void testNeuralQueryEnricherProcessor_whenGetEmptyQueryBody_thenSuccess() { - try { - initializeIndexIfNotExist(index); - createSearchRequestProcessor(null, search_pipeline); - createPipelineProcessor(null, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); - updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); - Request request = new Request("POST", "/" + index + "/_search"); - Response response = client().performRequest(request); - assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); - String responseBody = EntityUtils.toString(response.getEntity()); - Map responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody, false); - assertFalse(responseInMap.isEmpty()); - assertEquals(3, ((Map) responseInMap.get("hits")).size()); - } finally { - wipeOfTestResources(index, ingest_pipeline, null, search_pipeline); - } + initializeIndexIfNotExist(index); + createSearchRequestProcessor(null, search_pipeline); + createPipelineProcessor(null, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); + updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + Request request = new Request("POST", "/" + index + "/_search"); + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + String responseBody = EntityUtils.toString(response.getEntity()); + Map responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody, false); + assertFalse(responseInMap.isEmpty()); + assertEquals(3, ((Map) responseInMap.get("hits")).size()); } @SneakyThrows public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdPassed_thenSuccess() { String modelId = null; - try { - initializeIndexIfNotExist(index); - modelId = prepareModel(); - createSearchRequestProcessor(modelId, search_pipeline); - createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); - updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); - 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); + initializeIndexIfNotExist(index); + modelId = prepareModel(); + createSearchRequestProcessor(modelId, search_pipeline); + createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING); + updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + 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); - assertFalse(response.isEmpty()); - } finally { - wipeOfTestResources(index, ingest_pipeline, modelId, search_pipeline); - } + assertFalse(response.isEmpty()); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessorIT.java index bc61f7c29..5c2bdf72d 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralSparseTwoPhaseProcessorIT.java @@ -65,28 +65,24 @@ public void testCreateOutOfRangePipeline_thenThrowsException() { @SneakyThrows public void testBooleanQuery_withMultipleSparseEncodingQueries_whenTwoPhaseEnabled() { - try { - initializeTwoPhaseProcessor(); - initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); - setDefaultSearchPipelineForIndex(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - Map randomTokenWeight = createRandomTokenWeightMap(TWO_PHASE_TEST_TOKEN); - Supplier> randomTokenWeightSupplier = () -> randomTokenWeight; - NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(randomTokenWeightSupplier); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) - .queryTokensSupplier(randomTokenWeightSupplier); - boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); - - Map searchResponseAsMap = search(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, randomTokenWeightSupplier.get()); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, null, null, search_pipeline); - } + initializeTwoPhaseProcessor(); + initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); + setDefaultSearchPipelineForIndex(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + Map randomTokenWeight = createRandomTokenWeightMap(TWO_PHASE_TEST_TOKEN); + Supplier> randomTokenWeightSupplier = () -> randomTokenWeight; + NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(randomTokenWeightSupplier); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) + .queryTokensSupplier(randomTokenWeightSupplier); + boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); + + Map searchResponseAsMap = search(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, randomTokenWeightSupplier.get()); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows @@ -120,50 +116,42 @@ private void setDefaultSearchPipelineForIndex(String indexName) { */ @SneakyThrows public void testBasicQueryUsingQueryTokens_whenTwoPhaseEnabled_thenGetExpectedScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - initializeTwoPhaseProcessor(); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeTwoPhaseProcessor(); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testBasicQueryUsingQueryTokens_whenTwoPhaseEnabledAndDisabled_thenGetSameScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - initializeTwoPhaseProcessor(); - - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float scoreWithoutTwoPhase = objectToFloat(firstInnerHit.get("_score")); - - sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float scoreWithTwoPhase = objectToFloat(firstInnerHit.get("_score")); - assertEquals(scoreWithTwoPhase, scoreWithoutTwoPhase, DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeTwoPhaseProcessor(); + + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float scoreWithoutTwoPhase = objectToFloat(firstInnerHit.get("_score")); + + sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float scoreWithTwoPhase = objectToFloat(firstInnerHit.get("_score")); + assertEquals(scoreWithTwoPhase, scoreWithoutTwoPhase, DELTA); } /** @@ -192,23 +180,19 @@ public void testBasicQueryUsingQueryTokens_whenTwoPhaseEnabledAndDisabled_thenGe */ @SneakyThrows public void testNeuralSparseQueryAsRescoreQuery_whenTwoPhase_thenGetExpectedScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - initializeTwoPhaseProcessor(); - - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - QueryBuilder queryBuilder = new MatchAllQueryBuilder(); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, queryBuilder, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeTwoPhaseProcessor(); + + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + QueryBuilder queryBuilder = new MatchAllQueryBuilder(); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, queryBuilder, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -240,24 +224,20 @@ public void testNeuralSparseQueryAsRescoreQuery_whenTwoPhase_thenGetExpectedScor */ @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhase_thenGetExpectedScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - initializeTwoPhaseProcessor(); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - boolQueryBuilder.should(sparseEncodingQueryBuilder); - boolQueryBuilder.should(sparseEncodingQueryBuilder); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 4 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeTwoPhaseProcessor(); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + boolQueryBuilder.should(sparseEncodingQueryBuilder); + boolQueryBuilder.should(sparseEncodingQueryBuilder); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 4 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -268,179 +248,151 @@ public void testMultiNeuralSparseQuery_whenTwoPhase_thenGetExpectedScore() { */ @SneakyThrows public void testNeuralSparseQuery_whenDifferentTwoPhaseParameter_thenGetDifferentResult() { - try { - initializeIndexIfNotExist(TEST_TWO_PHASE_BASIC_INDEX_NAME); - Map queryToken = new HashMap<>(); - for (int i = 1; i < 6; i++) { - queryToken.put(String.valueOf(i + 10), (float) i); - } - for (int i = 1; i < 8; i++) { - queryToken.put(String.valueOf(i), (float) i); - } - Supplier> queryTokenSupplier = () -> queryToken; - - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(queryTokenSupplier); - assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.3f, 1f, 10000); - setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); - assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.7f, 1f, 10000); - setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); - assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 61); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.7f, 30f, 10000); - setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); - assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); - } finally { - wipeOfTestResources(TEST_TWO_PHASE_BASIC_INDEX_NAME, null, null, search_pipeline); + initializeIndexIfNotExist(TEST_TWO_PHASE_BASIC_INDEX_NAME); + Map queryToken = new HashMap<>(); + for (int i = 1; i < 6; i++) { + queryToken.put(String.valueOf(i + 10), (float) i); + } + for (int i = 1; i < 8; i++) { + queryToken.put(String.valueOf(i), (float) i); } + Supplier> queryTokenSupplier = () -> queryToken; + + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(queryTokenSupplier); + assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.3f, 1f, 10000); + setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); + assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.7f, 1f, 10000); + setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); + assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 61); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.7f, 30f, 10000); + setDefaultSearchPipelineForIndex(TEST_TWO_PHASE_BASIC_INDEX_NAME); + assertSearchScore(sparseEncodingQueryBuilder, TEST_TWO_PHASE_BASIC_INDEX_NAME, 110); } @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhaseAndFilter_thenGetExpectedScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .8f, 5f, 1000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(2.0f); - boolQueryBuilder.should(sparseEncodingQueryBuilder); - boolQueryBuilder.filter(sparseEncodingQueryBuilder); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .8f, 5f, 1000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(2.0f); + boolQueryBuilder.should(sparseEncodingQueryBuilder); + boolQueryBuilder.filter(sparseEncodingQueryBuilder); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhaseAndMultiBoolean_thenGetExpectedScore() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .6f, 5f, 1000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId) - .boost(1.0f); - boolQueryBuilder.should(sparseEncodingQueryBuilder1); - boolQueryBuilder.should(sparseEncodingQueryBuilder1); - BoolQueryBuilder subBoolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId) - .boost(2.0f); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder3 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId) - .boost(3.0f); - subBoolQueryBuilder.should(sparseEncodingQueryBuilder2); - subBoolQueryBuilder.should(sparseEncodingQueryBuilder3); - subBoolQueryBuilder.boost(2.0f); - boolQueryBuilder.should(subBoolQueryBuilder); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 12 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .6f, 5f, 1000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .boost(1.0f); + boolQueryBuilder.should(sparseEncodingQueryBuilder1); + boolQueryBuilder.should(sparseEncodingQueryBuilder1); + BoolQueryBuilder subBoolQueryBuilder = new BoolQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .boost(2.0f); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder3 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .boost(3.0f); + subBoolQueryBuilder.should(sparseEncodingQueryBuilder2); + subBoolQueryBuilder.should(sparseEncodingQueryBuilder3); + subBoolQueryBuilder.boost(2.0f); + boolQueryBuilder.should(subBoolQueryBuilder); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 12 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhaseAndNoLowScoreToken_thenGetExpectedScore() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .6f, 5f, 1000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - Map queryTokens = new HashMap<>(); - queryTokens.put("hello", 10.0f); - queryTokens.put("world", 10.0f); - queryTokens.put("a", 10.0f); - queryTokens.put("b", 10.0f); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(() -> queryTokens) - .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, sparseEncodingQueryBuilder.queryTokensSupplier().get()); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, .6f, 5f, 1000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + Map queryTokens = new HashMap<>(); + queryTokens.put("hello", 10.0f); + queryTokens.put("world", 10.0f); + queryTokens.put("a", 10.0f); + queryTokens.put("b", 10.0f); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(() -> queryTokens) + .boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, sparseEncodingQueryBuilder.queryTokensSupplier().get()); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testNeuralSParseQuery_whenTwoPhaseAndNestedInConstantScoreQuery_thenSuccess() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(1.0f); - ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(sparseEncodingQueryBuilder); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, constantScoreQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - assertEquals(1.0f, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(1.0f); + ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(sparseEncodingQueryBuilder); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, constantScoreQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + assertEquals(1.0f, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testNeuralSParseQuery_whenTwoPhaseAndNestedInDisjunctionMaxQuery_thenSuccess() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(5.0f); - DisMaxQueryBuilder disMaxQueryBuilder = new DisMaxQueryBuilder(); - disMaxQueryBuilder.add(sparseEncodingQueryBuilder); - disMaxQueryBuilder.add(new MatchAllQueryBuilder()); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, disMaxQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 5f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(5.0f); + DisMaxQueryBuilder disMaxQueryBuilder = new DisMaxQueryBuilder(); + disMaxQueryBuilder.add(sparseEncodingQueryBuilder); + disMaxQueryBuilder.add(new MatchAllQueryBuilder()); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, disMaxQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 5f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows public void testNeuralSparseQuery_whenTwoPhaseAndNestedInFunctionScoreQuery_thenSuccess() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(testFixedQueryTokenSupplier) - .boost(5.0f); - FunctionScoreQueryBuilder functionScoreQueryBuilder = new FunctionScoreQueryBuilder(sparseEncodingQueryBuilder); - functionScoreQueryBuilder.boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, functionScoreQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 10f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + createNeuralSparseTwoPhaseSearchProcessor(search_pipeline, 0.6f, 5f, 10000); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(testFixedQueryTokenSupplier) + .boost(5.0f); + FunctionScoreQueryBuilder functionScoreQueryBuilder = new FunctionScoreQueryBuilder(sparseEncodingQueryBuilder); + functionScoreQueryBuilder.boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, functionScoreQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 10f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -473,26 +425,22 @@ public void testNeuralSparseQuery_whenTwoPhaseAndNestedInFunctionScoreQuery_then @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhaseAndModelInference_thenGetExpectedScore() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - initializeTwoPhaseProcessor(); - setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId) - .boost(3.0f); - boolQueryBuilder.should(sparseEncodingQueryBuilder); - boolQueryBuilder.should(sparseEncodingQueryBuilder); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 6 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, search_pipeline); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeTwoPhaseProcessor(); + setDefaultSearchPipelineForIndex(TEST_BASIC_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .boost(3.0f); + boolQueryBuilder.should(sparseEncodingQueryBuilder); + boolQueryBuilder.should(sparseEncodingQueryBuilder); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 6 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java index 3c5fc08ef..e083475dc 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java @@ -82,35 +82,31 @@ public void setUp() throws Exception { @SneakyThrows public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - - 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(); - hybridQueryBuilder.add(neuralQueryBuilder); - hybridQueryBuilder.add(termQueryBuilder); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertQueryResults(searchResponseAsMap, 5, false); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + + 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(); + hybridQueryBuilder.add(neuralQueryBuilder); + hybridQueryBuilder.add(termQueryBuilder); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilder, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertQueryResults(searchResponseAsMap, 5, false); } /** @@ -128,127 +124,119 @@ public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { @SneakyThrows public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipelineWithDefaultResultsPostProcessor(SEARCH_PIPELINE); - - 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(); - hybridQueryBuilder.add(neuralQueryBuilder); - hybridQueryBuilder.add(termQueryBuilder); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertQueryResults(searchResponseAsMap, 5, false); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipelineWithDefaultResultsPostProcessor(SEARCH_PIPELINE); + + 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(); + hybridQueryBuilder.add(neuralQueryBuilder); + hybridQueryBuilder.add(termQueryBuilder); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilder, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertQueryResults(searchResponseAsMap, 5, false); } @SneakyThrows public void testQueryMatches_whenMultipleShards_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); - modelId = prepareModel(); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - int totalExpectedDocQty = 6; - - 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(); - hybridQueryBuilder.add(neuralQueryBuilder); - hybridQueryBuilder.add(termQueryBuilder); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, - null, - 6, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + modelId = prepareModel(); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + int totalExpectedDocQty = 6; + + 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(); + hybridQueryBuilder.add(neuralQueryBuilder); + hybridQueryBuilder.add(termQueryBuilder); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + hybridQueryBuilder, + null, + 6, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertNotNull(searchResponseAsMap); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(totalExpectedDocQty, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - assertTrue(Range.between(.5f, 1.0f).contains(getMaxScore(searchResponseAsMap).get())); - List> hitsNestedList = getNestedHits(searchResponseAsMap); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } - // verify scores order - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - - // verify the scores are normalized. we need special assert logic because combined score may vary as neural search query - // based on random vectors and return results for every doc. In some cases that may affect 1.0 score from term query and make it - // lower. - float highestScore = scores.stream().max(Double::compare).get().floatValue(); - assertTrue(Range.between(.5f, 1.0f).contains(highestScore)); - float lowestScore = scores.stream().min(Double::compare).get().floatValue(); - assertTrue(Range.between(.0f, .5f).contains(lowestScore)); - - // verify that all ids are unique - assertEquals(Set.copyOf(ids).size(), ids.size()); - - // verify 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 searchResponseAsMapPartialMatch = search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilderPartialMatch, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertQueryResults(searchResponseAsMapPartialMatch, 4, true, Range.between(0.33f, 1.0f)); + assertNotNull(searchResponseAsMap); + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(totalExpectedDocQty, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + assertTrue(Range.between(.5f, 1.0f).contains(getMaxScore(searchResponseAsMap).get())); + List> hitsNestedList = getNestedHits(searchResponseAsMap); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); + } + // verify scores order + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - // 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)); + // verify the scores are normalized. we need special assert logic because combined score may vary as neural search query + // based on random vectors and return results for every doc. In some cases that may affect 1.0 score from term query and make it + // lower. + float highestScore = scores.stream().max(Double::compare).get().floatValue(); + assertTrue(Range.between(.5f, 1.0f).contains(highestScore)); + float lowestScore = scores.stream().min(Double::compare).get().floatValue(); + assertTrue(Range.between(.0f, .5f).contains(lowestScore)); - Map searchResponseAsMapNoMatches = search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilderNoMatches, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertQueryResults(searchResponseAsMapNoMatches, 0, true); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, modelId, SEARCH_PIPELINE); - } + // verify that all ids are unique + assertEquals(Set.copyOf(ids).size(), ids.size()); + + // verify 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 searchResponseAsMapPartialMatch = search( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + hybridQueryBuilderPartialMatch, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertQueryResults(searchResponseAsMapPartialMatch, 4, true, Range.between(0.33f, 1.0f)); + + // 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 searchResponseAsMapNoMatches = search( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + hybridQueryBuilderNoMatches, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertQueryResults(searchResponseAsMapNoMatches, 0, true); } private void initializeIndexIfNotExist(String indexName) throws IOException { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/RRFProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/RRFProcessorIT.java index fccabab5c..25613e227 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/RRFProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/RRFProcessorIT.java @@ -29,33 +29,20 @@ public class RRFProcessorIT extends BaseNeuralSearchIT { @SneakyThrows public void testRRF_whenValidInput_thenSucceed() { - try { - createPipelineProcessor(null, RRF_INGEST_PIPELINE, ProcessorType.TEXT_EMBEDDING); - prepareKnnIndex( - RRF_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig("passage_embedding", RRF_DIMENSION, TEST_SPACE_TYPE)) - ); - addDocuments(); - createDefaultRRFSearchPipeline(); + createPipelineProcessor(null, RRF_INGEST_PIPELINE, ProcessorType.TEXT_EMBEDDING); + prepareKnnIndex(RRF_INDEX_NAME, Collections.singletonList(new KNNFieldConfig("passage_embedding", RRF_DIMENSION, TEST_SPACE_TYPE))); + addDocuments(); + createDefaultRRFSearchPipeline(); - HybridQueryBuilder hybridQueryBuilder = getHybridQueryBuilder(); + HybridQueryBuilder hybridQueryBuilder = getHybridQueryBuilder(); - Map results = search( - RRF_INDEX_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", RRF_SEARCH_PIPELINE) - ); - Map hits = (Map) results.get("hits"); - ArrayList> hitsList = (ArrayList>) hits.get("hits"); - assertEquals(3, hitsList.size()); - assertEquals(0.016393442, (Double) hitsList.getFirst().get("_score"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(0.016129032, (Double) hitsList.get(1).get("_score"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(0.015873017, (Double) hitsList.getLast().get("_score"), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(RRF_INDEX_NAME, RRF_INGEST_PIPELINE, null, RRF_SEARCH_PIPELINE); - } + Map results = search(RRF_INDEX_NAME, hybridQueryBuilder, null, 5, Map.of("search_pipeline", RRF_SEARCH_PIPELINE)); + Map hits = (Map) results.get("hits"); + ArrayList> hitsList = (ArrayList>) hits.get("hits"); + assertEquals(3, hitsList.size()); + assertEquals(0.016393442, (Double) hitsList.getFirst().get("_score"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(0.016129032, (Double) hitsList.get(1).get("_score"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(0.015873017, (Double) hitsList.getLast().get("_score"), DELTA_FOR_SCORE_ASSERTION); } private HybridQueryBuilder getHybridQueryBuilder() { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java index feb914e30..d76ee2c5b 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java @@ -85,110 +85,94 @@ public void setUp() throws Exception { */ @SneakyThrows public void testArithmeticWeightedMean_whenWeightsPassed_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); - // check case when number of weights and sub-queries are same - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.4f, 0.3f, 0.3f })) - ); - - 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)); - - Map searchResponseWithWeights1AsMap = search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - - assertWeightedScores(searchResponseWithWeights1AsMap, 0.4, 0.3, 0.001); - - // delete existing pipeline and create a new one with another set of weights - deleteSearchPipeline(SEARCH_PIPELINE); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.233f, 0.666f, 0.1f })) - ); - - Map searchResponseWithWeights2AsMap = search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - - assertWeightedScores(searchResponseWithWeights2AsMap, 0.6666, 0.2332, 0.001); - - // check case when number of weights is less than number of sub-queries - // delete existing pipeline and create a new one with another set of weights - deleteSearchPipeline(SEARCH_PIPELINE); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 1.0f })) - ); - - ResponseException exception1 = expectThrows( - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ) - ); - org.hamcrest.MatcherAssert.assertThat( - exception1.getMessage(), - allOf( - containsString("number of weights"), - containsString("must match number of sub-queries"), - containsString("in hybrid query") - ) - ); - - // check case when number of weights is more than number of sub-queries - // delete existing pipeline and create a new one with another set of weights - deleteSearchPipeline(SEARCH_PIPELINE); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.25f, 0.25f, 0.2f })) - ); - - ResponseException exception2 = expectThrows( - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - hybridQueryBuilder, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ) - ); - org.hamcrest.MatcherAssert.assertThat( - exception2.getMessage(), - allOf( - containsString("number of weights"), - containsString("must match number of sub-queries"), - containsString("in hybrid query") - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, null, null, SEARCH_PIPELINE); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + // check case when number of weights and sub-queries are same + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.4f, 0.3f, 0.3f })) + ); + + 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)); + + Map searchResponseWithWeights1AsMap = search( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + hybridQueryBuilder, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + + assertWeightedScores(searchResponseWithWeights1AsMap, 0.4, 0.3, 0.001); + + // delete existing pipeline and create a new one with another set of weights + deleteSearchPipeline(SEARCH_PIPELINE); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.233f, 0.666f, 0.1f })) + ); + + Map searchResponseWithWeights2AsMap = search( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + hybridQueryBuilder, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + + assertWeightedScores(searchResponseWithWeights2AsMap, 0.6666, 0.2332, 0.001); + + // check case when number of weights is less than number of sub-queries + // delete existing pipeline and create a new one with another set of weights + deleteSearchPipeline(SEARCH_PIPELINE); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 1.0f })) + ); + + ResponseException exception1 = expectThrows( + ResponseException.class, + () -> search(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, hybridQueryBuilder, null, 5, Map.of("search_pipeline", SEARCH_PIPELINE)) + ); + org.hamcrest.MatcherAssert.assertThat( + exception1.getMessage(), + allOf( + containsString("number of weights"), + containsString("must match number of sub-queries"), + containsString("in hybrid query") + ) + ); + + // check case when number of weights is more than number of sub-queries + // delete existing pipeline and create a new one with another set of weights + deleteSearchPipeline(SEARCH_PIPELINE); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.25f, 0.25f, 0.2f })) + ); + + ResponseException exception2 = expectThrows( + ResponseException.class, + () -> search(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, hybridQueryBuilder, null, 5, Map.of("search_pipeline", SEARCH_PIPELINE)) + ); + org.hamcrest.MatcherAssert.assertThat( + exception2.getMessage(), + allOf( + containsString("number of weights"), + containsString("must match number of sub-queries"), + containsString("in hybrid query") + ) + ); } /** @@ -212,59 +196,56 @@ public void testArithmeticWeightedMean_whenWeightsPassed_thenSuccessful() { @SneakyThrows public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - HARMONIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); - - HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); - hybridQueryBuilderDefaultNorm.add( - 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)); - - Map searchResponseAsMapDefaultNorm = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderDefaultNorm, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + HARMONIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); + + HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); + hybridQueryBuilderDefaultNorm.add( + 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)); + + Map searchResponseAsMapDefaultNorm = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderDefaultNorm, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + + assertHybridSearchResults(searchResponseAsMapDefaultNorm, 5, new float[] { 0.5f, 1.0f }); + + deleteSearchPipeline(SEARCH_PIPELINE); + + createSearchPipeline( + SEARCH_PIPELINE, + L2_NORMALIZATION_METHOD, + HARMONIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); + + HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); + hybridQueryBuilderL2Norm.add( + 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)); + + Map searchResponseAsMapL2Norm = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderL2Norm, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapL2Norm, 5, new float[] { 0.5f, 1.0f }); - assertHybridSearchResults(searchResponseAsMapDefaultNorm, 5, new float[] { 0.5f, 1.0f }); - - deleteSearchPipeline(SEARCH_PIPELINE); - - createSearchPipeline( - SEARCH_PIPELINE, - L2_NORMALIZATION_METHOD, - HARMONIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); - - HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); - hybridQueryBuilderL2Norm.add( - 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)); - - Map searchResponseAsMapL2Norm = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderL2Norm, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapL2Norm, 5, new float[] { 0.5f, 1.0f }); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } } /** @@ -288,60 +269,56 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf @SneakyThrows public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - GEOMETRIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); - - HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); - hybridQueryBuilderDefaultNorm.add( - 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)); - - Map searchResponseAsMapDefaultNorm = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderDefaultNorm, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - - assertHybridSearchResults(searchResponseAsMapDefaultNorm, 5, new float[] { 0.5f, 1.0f }); - - deleteSearchPipeline(SEARCH_PIPELINE); - - createSearchPipeline( - SEARCH_PIPELINE, - L2_NORMALIZATION_METHOD, - GEOMETRIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); - - HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); - hybridQueryBuilderL2Norm.add( - 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)); - - Map searchResponseAsMapL2Norm = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderL2Norm, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapL2Norm, 5, new float[] { 0.5f, 1.0f }); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + GEOMETRIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); + + HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); + hybridQueryBuilderDefaultNorm.add( + 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)); + + Map searchResponseAsMapDefaultNorm = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderDefaultNorm, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + + assertHybridSearchResults(searchResponseAsMapDefaultNorm, 5, new float[] { 0.5f, 1.0f }); + + deleteSearchPipeline(SEARCH_PIPELINE); + + createSearchPipeline( + SEARCH_PIPELINE, + L2_NORMALIZATION_METHOD, + GEOMETRIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); + + HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); + hybridQueryBuilderL2Norm.add( + 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)); + + Map searchResponseAsMapL2Norm = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderL2Norm, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapL2Norm, 5, new float[] { 0.5f, 1.0f }); } private void initializeIndexIfNotExist(String indexName) throws IOException { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java index 87f9942af..038d8a19b 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java @@ -73,83 +73,79 @@ public void setUp() throws Exception { @SneakyThrows public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipeline( - SEARCH_PIPELINE, - L2_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipeline( + SEARCH_PIPELINE, + L2_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); - hybridQueryBuilderArithmeticMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); + hybridQueryBuilderArithmeticMean.add( + 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)); - Map searchResponseAsMapArithmeticMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderArithmeticMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapArithmeticMean, 5, new float[] { 0.6f, 1.0f }); + Map searchResponseAsMapArithmeticMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderArithmeticMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapArithmeticMean, 5, new float[] { 0.6f, 1.0f }); - deleteSearchPipeline(SEARCH_PIPELINE); + deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - createSearchPipeline( - SEARCH_PIPELINE, - L2_NORMALIZATION_METHOD, - HARMONIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + createSearchPipeline( + SEARCH_PIPELINE, + L2_NORMALIZATION_METHOD, + HARMONIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); - hybridQueryBuilderHarmonicMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); + hybridQueryBuilderHarmonicMean.add( + 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)); - Map searchResponseAsMapHarmonicMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderHarmonicMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapHarmonicMean, 5, new float[] { 0.5f, 1.0f }); + Map searchResponseAsMapHarmonicMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderHarmonicMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapHarmonicMean, 5, new float[] { 0.5f, 1.0f }); - deleteSearchPipeline(SEARCH_PIPELINE); + deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - createSearchPipeline( - SEARCH_PIPELINE, - L2_NORMALIZATION_METHOD, - GEOMETRIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + createSearchPipeline( + SEARCH_PIPELINE, + L2_NORMALIZATION_METHOD, + GEOMETRIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); - hybridQueryBuilderGeometricMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); + hybridQueryBuilderGeometricMean.add( + 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)); - Map searchResponseAsMapGeometricMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderGeometricMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapGeometricMean, 5, new float[] { 0.5f, 1.0f }); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } + Map searchResponseAsMapGeometricMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderGeometricMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapGeometricMean, 5, new float[] { 0.5f, 1.0f }); } /** @@ -173,83 +169,79 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { @SneakyThrows public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - modelId = prepareModel(); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + modelId = prepareModel(); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); - hybridQueryBuilderArithmeticMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); + hybridQueryBuilderArithmeticMean.add( + 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)); - Map searchResponseAsMapArithmeticMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderArithmeticMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapArithmeticMean, 5, new float[] { 0.5f, 1.0f }); + Map searchResponseAsMapArithmeticMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderArithmeticMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapArithmeticMean, 5, new float[] { 0.5f, 1.0f }); - deleteSearchPipeline(SEARCH_PIPELINE); + deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - HARMONIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + HARMONIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); - hybridQueryBuilderHarmonicMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); + hybridQueryBuilderHarmonicMean.add( + 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)); - Map searchResponseAsMapHarmonicMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderHarmonicMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapHarmonicMean, 5, new float[] { 0.6f, 1.0f }); + Map searchResponseAsMapHarmonicMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderHarmonicMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapHarmonicMean, 5, new float[] { 0.6f, 1.0f }); - deleteSearchPipeline(SEARCH_PIPELINE); + deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); - createSearchPipeline( - SEARCH_PIPELINE, - DEFAULT_NORMALIZATION_METHOD, - GEOMETRIC_MEAN_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + createSearchPipeline( + SEARCH_PIPELINE, + DEFAULT_NORMALIZATION_METHOD, + GEOMETRIC_MEAN_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.533f, 0.466f })) + ); - HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); - hybridQueryBuilderGeometricMean.add( - 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)); + HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); + hybridQueryBuilderGeometricMean.add( + 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)); - Map searchResponseAsMapGeometricMean = search( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - hybridQueryBuilderGeometricMean, - null, - 5, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); - assertHybridSearchResults(searchResponseAsMapGeometricMean, 5, new float[] { 0.6f, 1.0f }); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, null, modelId, SEARCH_PIPELINE); - } + Map searchResponseAsMapGeometricMean = search( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + hybridQueryBuilderGeometricMean, + null, + 5, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); + assertHybridSearchResults(searchResponseAsMapGeometricMean, 5, new float[] { 0.6f, 1.0f }); } private void initializeIndexIfNotExist(String indexName) throws IOException { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java b/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java index b5e14a11f..1de1d7bcc 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java @@ -39,44 +39,35 @@ public void setUp() throws Exception { public void testSparseEncodingProcessor() throws Exception { String modelId = null; - try { - modelId = prepareSparseEncodingModel(); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING); - createIndexWithPipeline(INDEX_NAME, "SparseEncodingIndexMappings.json", PIPELINE_NAME); - ingestDocument(INDEX_NAME, INGEST_DOCUMENT); - assertEquals(1, getDocCount(INDEX_NAME)); + modelId = prepareSparseEncodingModel(); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING); + createIndexWithPipeline(INDEX_NAME, "SparseEncodingIndexMappings.json", PIPELINE_NAME); + ingestDocument(INDEX_NAME, INGEST_DOCUMENT); + assertEquals(1, getDocCount(INDEX_NAME)); - NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); - neuralSparseQueryBuilder.fieldName("title_sparse"); - neuralSparseQueryBuilder.queryTokensSupplier(() -> Map.of("good", 1.0f, "a", 2.0f)); - Map searchResponse = search(INDEX_NAME, neuralSparseQueryBuilder, 2); - assertFalse(searchResponse.isEmpty()); - double maxScore = (Double) ((Map) searchResponse.get("hits")).get("max_score"); - assertEquals(4.4433594, maxScore, 1e-3); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); + neuralSparseQueryBuilder.fieldName("title_sparse"); + neuralSparseQueryBuilder.queryTokensSupplier(() -> Map.of("good", 1.0f, "a", 2.0f)); + Map searchResponse = search(INDEX_NAME, neuralSparseQueryBuilder, 2); + assertFalse(searchResponse.isEmpty()); + double maxScore = (Double) ((Map) searchResponse.get("hits")).get("max_score"); + assertEquals(4.4433594, maxScore, 1e-3); } public void testSparseEncodingProcessorWithPrune() throws Exception { - String modelId = null; - try { - modelId = prepareSparseEncodingModel(); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING_PRUNE); - createIndexWithPipeline(INDEX_NAME, "SparseEncodingIndexMappings.json", PIPELINE_NAME); - ingestDocument(INDEX_NAME, INGEST_DOCUMENT); - assertEquals(1, getDocCount(INDEX_NAME)); + String modelId = prepareSparseEncodingModel(); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING_PRUNE); + createIndexWithPipeline(INDEX_NAME, "SparseEncodingIndexMappings.json", PIPELINE_NAME); + ingestDocument(INDEX_NAME, INGEST_DOCUMENT); + assertEquals(1, getDocCount(INDEX_NAME)); - NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); - neuralSparseQueryBuilder.fieldName("title_sparse"); - neuralSparseQueryBuilder.queryTokensSupplier(() -> Map.of("good", 1.0f, "a", 2.0f)); - Map searchResponse = search(INDEX_NAME, neuralSparseQueryBuilder, 2); - assertFalse(searchResponse.isEmpty()); - double maxScore = (Double) ((Map) searchResponse.get("hits")).get("max_score"); - assertEquals(3.640625, maxScore, 1e-3); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder(); + neuralSparseQueryBuilder.fieldName("title_sparse"); + neuralSparseQueryBuilder.queryTokensSupplier(() -> Map.of("good", 1.0f, "a", 2.0f)); + Map searchResponse = search(INDEX_NAME, neuralSparseQueryBuilder, 2); + assertFalse(searchResponse.isEmpty()); + double maxScore = (Double) ((Map) searchResponse.get("hits")).get("max_score"); + assertEquals(3.640625, maxScore, 1e-3); } public void testSparseEncodingProcessorWithReindex() throws Exception { @@ -85,16 +76,11 @@ public void testSparseEncodingProcessorWithReindex() throws Exception { createIndexWithConfiguration(fromIndexName, "{ \"settings\": { \"number_of_shards\": 1, \"number_of_replicas\": 0 } }", null); ingestDocument(fromIndexName, "{ \"text\": \"hello world\" }"); // create text embedding index for reindex - String modelId = null; - try { - modelId = prepareSparseEncodingModel(); - String toIndexName = "test-reindex-to"; - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING); - createIndexWithPipeline(toIndexName, "SparseEncodingIndexMappings.json", PIPELINE_NAME); - reindex(fromIndexName, toIndexName); - assertEquals(1, getDocCount(toIndexName)); - } finally { - wipeOfTestResources(fromIndexName, PIPELINE_NAME, modelId, null); - } + String modelId = prepareSparseEncodingModel(); + String toIndexName = "test-reindex-to"; + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.SPARSE_ENCODING); + createIndexWithPipeline(toIndexName, "SparseEncodingIndexMappings.json", PIPELINE_NAME); + reindex(fromIndexName, toIndexName); + assertEquals(1, getDocCount(toIndexName)); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextChunkingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextChunkingProcessorIT.java index e3b79b94e..6e7cc3d46 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextChunkingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextChunkingProcessorIT.java @@ -63,140 +63,108 @@ public void setUp() throws Exception { @SneakyThrows public void testTextChunkingProcessor_withFixedTokenLengthAlgorithmStandardTokenizer_thenSucceed() { - try { - createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(INDEX_NAME, document); + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(INDEX_NAME, document); - List expectedPassages = new ArrayList<>(); - expectedPassages.add("This is an example document to be chunked. The document "); - expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by "); - expectedPassages.add("standard tokenizer in OpenSearch."); - validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); - } finally { - wipeOfTestResources(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME, null, null); - } + List expectedPassages = new ArrayList<>(); + expectedPassages.add("This is an example document to be chunked. The document "); + expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by "); + expectedPassages.add("standard tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); } @SneakyThrows public void testTextChunkingProcessor_withFixedTokenLengthAlgorithmLetterTokenizer_thenSucceed() { - try { - createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_LETTER_TOKENIZER_NAME); - createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LETTER_TOKENIZER_NAME); + createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_LETTER_TOKENIZER_NAME); + createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LETTER_TOKENIZER_NAME); - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(INDEX_NAME, document); + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(INDEX_NAME, document); - List expectedPassages = new ArrayList<>(); - expectedPassages.add("This is an example document to be chunked. The document "); - expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by standard "); - expectedPassages.add("tokenizer in OpenSearch."); - validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); - } finally { - wipeOfTestResources(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LETTER_TOKENIZER_NAME, null, null); - } + List expectedPassages = new ArrayList<>(); + expectedPassages.add("This is an example document to be chunked. The document "); + expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by standard "); + expectedPassages.add("tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); } @SneakyThrows public void testTextChunkingProcessor_withFixedTokenLengthAlgorithmLowercaseTokenizer_thenSucceed() { - try { - createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_LOWERCASE_TOKENIZER_NAME); - createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LOWERCASE_TOKENIZER_NAME); + createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_LOWERCASE_TOKENIZER_NAME); + createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LOWERCASE_TOKENIZER_NAME); - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(INDEX_NAME, document); + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(INDEX_NAME, document); - List expectedPassages = new ArrayList<>(); - expectedPassages.add("This is an example document to be chunked. The document "); - expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by standard "); - expectedPassages.add("tokenizer in OpenSearch."); - validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); - } finally { - wipeOfTestResources(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_LOWERCASE_TOKENIZER_NAME, null, null); - } + List expectedPassages = new ArrayList<>(); + expectedPassages.add("This is an example document to be chunked. The document "); + expectedPassages.add("contains a single paragraph, two sentences and 24 tokens by standard "); + expectedPassages.add("tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); } @SneakyThrows public void testTextChunkingProcessor_withFixedTokenLengthAlgorithmStandardTokenizer_whenExceedMaxTokenCount_thenFail() { - try { - createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - Exception exception = assertThrows(Exception.class, () -> { - String document = getDocumentFromFilePath(TEST_LONG_DOCUMENT); - ingestDocument(INDEX_NAME, document); - }); - // max_token_count is 100 by index settings - assert (exception.getMessage() - .contains("The number of tokens produced by calling _analyze has exceeded the allowed maximum of [100].")); - assertEquals(0, getDocCount(INDEX_NAME)); - } finally { - wipeOfTestResources(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME, null, null); - } + createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + Exception exception = assertThrows(Exception.class, () -> { + String document = getDocumentFromFilePath(TEST_LONG_DOCUMENT); + ingestDocument(INDEX_NAME, document); + }); + // max_token_count is 100 by index settings + assert (exception.getMessage() + .contains("The number of tokens produced by calling _analyze has exceeded the allowed maximum of [100].")); + assertEquals(0, getDocCount(INDEX_NAME)); } @SneakyThrows public void testTextChunkingProcessor_withDelimiterAlgorithm_successful() { - try { - createPipelineProcessor(DELIMITER_PIPELINE_NAME); - createTextChunkingIndex(INDEX_NAME, DELIMITER_PIPELINE_NAME); + createPipelineProcessor(DELIMITER_PIPELINE_NAME); + createTextChunkingIndex(INDEX_NAME, DELIMITER_PIPELINE_NAME); - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(INDEX_NAME, document); + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(INDEX_NAME, document); - List expectedPassages = new ArrayList<>(); - expectedPassages.add("This is an example document to be chunked."); - expectedPassages.add( - " The document contains a single paragraph, two sentences and 24 tokens by standard tokenizer in OpenSearch." - ); - validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); - } finally { - wipeOfTestResources(INDEX_NAME, DELIMITER_PIPELINE_NAME, null, null); - } + List expectedPassages = new ArrayList<>(); + expectedPassages.add("This is an example document to be chunked."); + expectedPassages.add(" The document contains a single paragraph, two sentences and 24 tokens by standard tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); } @SneakyThrows public void testTextChunkingProcessor_withCascadePipeline_successful() { - try { - createPipelineProcessor(CASCADE_PIPELINE_NAME); - createTextChunkingIndex(INDEX_NAME, CASCADE_PIPELINE_NAME); - - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(INDEX_NAME, document); - - List expectedPassages = new ArrayList<>(); - expectedPassages.add("This is an example document to be chunked."); - expectedPassages.add(" The document contains a single paragraph, two sentences and 24 "); - expectedPassages.add("tokens by standard tokenizer in OpenSearch."); - validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); - - expectedPassages.clear(); - expectedPassages.add("This is an example document to be chunked."); - expectedPassages.add( - " The document contains a single paragraph, two sentences and 24 tokens by standard tokenizer in OpenSearch." - ); - validateIndexIngestResults(INDEX_NAME, INTERMEDIATE_FIELD, expectedPassages); - } finally { - wipeOfTestResources(INDEX_NAME, CASCADE_PIPELINE_NAME, null, null); - } + createPipelineProcessor(CASCADE_PIPELINE_NAME); + createTextChunkingIndex(INDEX_NAME, CASCADE_PIPELINE_NAME); + + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(INDEX_NAME, document); + + List expectedPassages = new ArrayList<>(); + expectedPassages.add("This is an example document to be chunked."); + expectedPassages.add(" The document contains a single paragraph, two sentences and 24 "); + expectedPassages.add("tokens by standard tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, OUTPUT_FIELD, expectedPassages); + + expectedPassages.clear(); + expectedPassages.add("This is an example document to be chunked."); + expectedPassages.add(" The document contains a single paragraph, two sentences and 24 tokens by standard tokenizer in OpenSearch."); + validateIndexIngestResults(INDEX_NAME, INTERMEDIATE_FIELD, expectedPassages); } public void testTextChunkingProcessor_withFixedTokenLengthAlgorithmStandardTokenizer_whenReindexingDocument_thenSuccessful() throws Exception { - try { - String fromIndexName = "test-reindex-from"; - createIndexWithConfiguration(fromIndexName, "{ \"settings\": { \"number_of_shards\": 1, \"number_of_replicas\": 0 } }", null); - String document = getDocumentFromFilePath(TEST_DOCUMENT); - ingestDocument(fromIndexName, document); - - createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); - reindex(fromIndexName, INDEX_NAME); - assertEquals(1, getDocCount(INDEX_NAME)); - } finally { - wipeOfTestResources(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME, null, null); - } + String fromIndexName = "test-reindex-from"; + createIndexWithConfiguration(fromIndexName, "{ \"settings\": { \"number_of_shards\": 1, \"number_of_replicas\": 0 } }", null); + String document = getDocumentFromFilePath(TEST_DOCUMENT); + ingestDocument(fromIndexName, document); + + createPipelineProcessor(FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + createTextChunkingIndex(INDEX_NAME, FIXED_TOKEN_LENGTH_PIPELINE_WITH_STANDARD_TOKENIZER_NAME); + reindex(fromIndexName, INDEX_NAME); + assertEquals(1, getDocCount(INDEX_NAME)); } private void validateIndexIngestResults(String indexName, String fieldName, Object expected) { diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java index 88d90ef35..156e25ad7 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java @@ -64,91 +64,72 @@ public void setUp() throws Exception { } public void testTextEmbeddingProcessor() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - ingestDocument(INDEX_NAME, INGEST_DOC1); - assertEquals(1, getDocCount(INDEX_NAME)); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + ingestDocument(INDEX_NAME, INGEST_DOC1); + assertEquals(1, getDocCount(INDEX_NAME)); } public void testTextEmbeddingProcessor_batch() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING, 2); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - ingestBatchDocumentWithBulk("batch_", 2, Collections.emptySet(), Collections.emptySet()); - assertEquals(2, getDocCount(INDEX_NAME)); - - ingestDocument(INDEX_NAME, String.format(LOCALE, INGEST_DOC1, "success"), "1"); - ingestDocument(INDEX_NAME, String.format(LOCALE, INGEST_DOC2, "success"), "2"); - - assertEquals(getDocById(INDEX_NAME, "1").get("_source"), getDocById(INDEX_NAME, "batch_1").get("_source")); - assertEquals(getDocById(INDEX_NAME, "2").get("_source"), getDocById(INDEX_NAME, "batch_2").get("_source")); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING, 2); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + ingestBatchDocumentWithBulk("batch_", 2, Collections.emptySet(), Collections.emptySet()); + assertEquals(2, getDocCount(INDEX_NAME)); + + ingestDocument(INDEX_NAME, String.format(LOCALE, INGEST_DOC1, "success"), "1"); + ingestDocument(INDEX_NAME, String.format(LOCALE, INGEST_DOC2, "success"), "2"); + + assertEquals(getDocById(INDEX_NAME, "1").get("_source"), getDocById(INDEX_NAME, "batch_1").get("_source")); + assertEquals(getDocById(INDEX_NAME, "2").get("_source"), getDocById(INDEX_NAME, "batch_2").get("_source")); } public void testNestedFieldMapping_whenDocumentsIngested_thenSuccessful() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_NESTED_FIELDS_MAPPING); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - ingestDocument(INDEX_NAME, INGEST_DOC3, "3"); - ingestDocument(INDEX_NAME, INGEST_DOC4, "4"); - - assertDoc( - (Map) getDocById(INDEX_NAME, "3").get("_source"), - TEXT_FIELD_VALUE_1, - Optional.of(TEXT_FIELD_VALUE_3) - ); - assertDoc((Map) getDocById(INDEX_NAME, "4").get("_source"), TEXT_FIELD_VALUE_2, Optional.empty()); - - 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, - ScoreMode.Total - ); - QueryBuilder queryNestedHighLevel = QueryBuilders.nestedQuery(LEVEL_1_FIELD, queryNestedLowerLevel, ScoreMode.Total); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_NESTED_FIELDS_MAPPING); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + ingestDocument(INDEX_NAME, INGEST_DOC3, "3"); + ingestDocument(INDEX_NAME, INGEST_DOC4, "4"); + + assertDoc((Map) getDocById(INDEX_NAME, "3").get("_source"), TEXT_FIELD_VALUE_1, Optional.of(TEXT_FIELD_VALUE_3)); + assertDoc((Map) getDocById(INDEX_NAME, "4").get("_source"), TEXT_FIELD_VALUE_2, Optional.empty()); + + 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, + ScoreMode.Total + ); + QueryBuilder queryNestedHighLevel = QueryBuilders.nestedQuery(LEVEL_1_FIELD, queryNestedLowerLevel, ScoreMode.Total); - Map searchResponseAsMap = search(INDEX_NAME, queryNestedHighLevel, 2); - assertNotNull(searchResponseAsMap); + Map searchResponseAsMap = search(INDEX_NAME, queryNestedHighLevel, 2); + assertNotNull(searchResponseAsMap); - Map hits = (Map) searchResponseAsMap.get("hits"); - assertNotNull(hits); + Map hits = (Map) searchResponseAsMap.get("hits"); + assertNotNull(hits); - assertEquals(1.0, hits.get("max_score")); - List> listOfHits = (List>) hits.get("hits"); - assertNotNull(listOfHits); - assertEquals(2, listOfHits.size()); + assertEquals(1.0, hits.get("max_score")); + List> listOfHits = (List>) hits.get("hits"); + assertNotNull(listOfHits); + assertEquals(2, listOfHits.size()); - Map innerHitDetails = listOfHits.get(0); - assertEquals("3", innerHitDetails.get("_id")); - assertEquals(1.0, innerHitDetails.get("_score")); + Map innerHitDetails = listOfHits.get(0); + assertEquals("3", innerHitDetails.get("_id")); + assertEquals(1.0, innerHitDetails.get("_score")); - innerHitDetails = listOfHits.get(1); - assertEquals("4", innerHitDetails.get("_id")); - assertTrue((double) innerHitDetails.get("_score") <= 1.0); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + innerHitDetails = listOfHits.get(1); + assertEquals("4", innerHitDetails.get("_id")); + assertTrue((double) innerHitDetails.get("_score") <= 1.0); } private void assertDoc(Map sourceMap, String textFieldValue, Optional level3ExpectedValue) { @@ -187,105 +168,90 @@ private void assertDocWithLevel2AsList(Map sourceMap) { } public void testTextEmbeddingProcessor_withBatchSizeInProcessor() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - URL pipelineURLPath = classLoader.getResource("processor/PipelineConfigurationWithBatchSize.json"); - Objects.requireNonNull(pipelineURLPath); - String requestBody = Files.readString(Path.of(pipelineURLPath.toURI())); - createPipelineProcessor(requestBody, PIPELINE_NAME, modelId, null); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - int docCount = 5; - ingestBatchDocumentWithBulk("batch_", docCount, Collections.emptySet(), Collections.emptySet()); - assertEquals(5, getDocCount(INDEX_NAME)); - - for (int i = 0; i < docCount; ++i) { - String template = List.of(INGEST_DOC1, INGEST_DOC2).get(i % 2); - String payload = String.format(LOCALE, template, "success"); - ingestDocument(INDEX_NAME, payload, String.valueOf(i + 1)); - } + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + URL pipelineURLPath = classLoader.getResource("processor/PipelineConfigurationWithBatchSize.json"); + Objects.requireNonNull(pipelineURLPath); + String requestBody = Files.readString(Path.of(pipelineURLPath.toURI())); + createPipelineProcessor(requestBody, PIPELINE_NAME, modelId, null); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + int docCount = 5; + ingestBatchDocumentWithBulk("batch_", docCount, Collections.emptySet(), Collections.emptySet()); + assertEquals(5, getDocCount(INDEX_NAME)); - for (int i = 0; i < docCount; ++i) { - assertEquals( - getDocById(INDEX_NAME, String.valueOf(i + 1)).get("_source"), - getDocById(INDEX_NAME, "batch_" + (i + 1)).get("_source") - ); + for (int i = 0; i < docCount; ++i) { + String template = List.of(INGEST_DOC1, INGEST_DOC2).get(i % 2); + String payload = String.format(LOCALE, template, "success"); + ingestDocument(INDEX_NAME, payload, String.valueOf(i + 1)); + } + + for (int i = 0; i < docCount; ++i) { + assertEquals( + getDocById(INDEX_NAME, String.valueOf(i + 1)).get("_source"), + getDocById(INDEX_NAME, "batch_" + (i + 1)).get("_source") + ); - } - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); } } public void testTextEmbeddingProcessor_withFailureAndSkip() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - URL pipelineURLPath = classLoader.getResource("processor/PipelineConfigurationWithBatchSize.json"); - Objects.requireNonNull(pipelineURLPath); - String requestBody = Files.readString(Path.of(pipelineURLPath.toURI())); - createPipelineProcessor(requestBody, PIPELINE_NAME, modelId, null); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - int docCount = 5; - ingestBatchDocumentWithBulk("batch_", docCount, Set.of(0), Set.of(1)); - assertEquals(3, getDocCount(INDEX_NAME)); - - for (int i = 2; i < docCount; ++i) { - String template = List.of(INGEST_DOC1, INGEST_DOC2).get(i % 2); - String payload = String.format(LOCALE, template, "success"); - ingestDocument(INDEX_NAME, payload, String.valueOf(i + 1)); - } + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + URL pipelineURLPath = classLoader.getResource("processor/PipelineConfigurationWithBatchSize.json"); + Objects.requireNonNull(pipelineURLPath); + String requestBody = Files.readString(Path.of(pipelineURLPath.toURI())); + createPipelineProcessor(requestBody, PIPELINE_NAME, modelId, null); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + int docCount = 5; + ingestBatchDocumentWithBulk("batch_", docCount, Set.of(0), Set.of(1)); + assertEquals(3, getDocCount(INDEX_NAME)); + + for (int i = 2; i < docCount; ++i) { + String template = List.of(INGEST_DOC1, INGEST_DOC2).get(i % 2); + String payload = String.format(LOCALE, template, "success"); + ingestDocument(INDEX_NAME, payload, String.valueOf(i + 1)); + } - for (int i = 2; i < docCount; ++i) { - assertEquals( - getDocById(INDEX_NAME, String.valueOf(i + 1)).get("_source"), - getDocById(INDEX_NAME, "batch_" + (i + 1)).get("_source") - ); + for (int i = 2; i < docCount; ++i) { + assertEquals( + getDocById(INDEX_NAME, String.valueOf(i + 1)).get("_source"), + getDocById(INDEX_NAME, "batch_" + (i + 1)).get("_source") + ); - } - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); } } @SuppressWarnings("unchecked") public void testNestedFieldMapping_whenDocumentInListIngested_thenSuccessful() throws Exception { - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_NESTED_FIELDS_MAPPING); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - ingestDocument(INDEX_NAME, INGEST_DOC5, "5"); - - assertDocWithLevel2AsList((Map) getDocById(INDEX_NAME, "5").get("_source")); - - 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, - ScoreMode.Total - ); - QueryBuilder queryNestedHighLevel = QueryBuilders.nestedQuery(LEVEL_1_FIELD, queryNestedLowerLevel, ScoreMode.Total); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING_WITH_NESTED_FIELDS_MAPPING); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + ingestDocument(INDEX_NAME, INGEST_DOC5, "5"); + + assertDocWithLevel2AsList((Map) getDocById(INDEX_NAME, "5").get("_source")); + + 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, + ScoreMode.Total + ); + QueryBuilder queryNestedHighLevel = QueryBuilders.nestedQuery(LEVEL_1_FIELD, queryNestedLowerLevel, ScoreMode.Total); - Map searchResponseAsMap = search(INDEX_NAME, queryNestedHighLevel, 2); - assertNotNull(searchResponseAsMap); + Map searchResponseAsMap = search(INDEX_NAME, queryNestedHighLevel, 2); + assertNotNull(searchResponseAsMap); - assertEquals(1, getHitCount(searchResponseAsMap)); + assertEquals(1, getHitCount(searchResponseAsMap)); - Map innerHitDetails = getFirstInnerHit(searchResponseAsMap); - assertEquals("5", innerHitDetails.get("_id")); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + Map innerHitDetails = getFirstInnerHit(searchResponseAsMap); + assertEquals("5", innerHitDetails.get("_id")); } private String uploadTextEmbeddingModel() throws Exception { @@ -345,17 +311,12 @@ public void testTextEmbeddingProcessorWithReindexOperation() throws Exception { createIndexWithConfiguration(fromIndexName, "{ \"settings\": { \"number_of_shards\": 1, \"number_of_replicas\": 0 } }", null); ingestDocument(fromIndexName, "{ \"text\": \"hello world\" }"); // create text embedding index for reindex - String modelId = null; - try { - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - String toIndexName = "test-reindex-to"; - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING); - createIndexWithPipeline(toIndexName, "IndexMappings.json", PIPELINE_NAME); - reindex(fromIndexName, toIndexName); - assertEquals(1, getDocCount(toIndexName)); - } finally { - wipeOfTestResources(fromIndexName, PIPELINE_NAME, modelId, null); - } + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + String toIndexName = "test-reindex-to"; + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_EMBEDDING); + createIndexWithPipeline(toIndexName, "IndexMappings.json", PIPELINE_NAME); + reindex(fromIndexName, toIndexName); + assertEquals(1, getDocCount(toIndexName)); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java index c2bf16959..667591789 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java @@ -41,21 +41,16 @@ public void setUp() 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); - createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); - // verify doc with mapping - ingestDocument(INDEX_NAME, INGEST_DOCUMENT); - assertEquals(1, getDocCount(INDEX_NAME)); - // verify doc without mapping - ingestDocument(INDEX_NAME, INGEST_DOCUMENT_UNMAPPED_FIELDS); - assertEquals(2, getDocCount(INDEX_NAME)); - } finally { - wipeOfTestResources(INDEX_NAME, PIPELINE_NAME, modelId, null); - } + String modelId = uploadModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); + createIndexWithPipeline(INDEX_NAME, "IndexMappings.json", PIPELINE_NAME); + // verify doc with mapping + ingestDocument(INDEX_NAME, INGEST_DOCUMENT); + assertEquals(1, getDocCount(INDEX_NAME)); + // verify doc without mapping + ingestDocument(INDEX_NAME, INGEST_DOCUMENT_UNMAPPED_FIELDS); + assertEquals(2, getDocCount(INDEX_NAME)); } private String uploadModel() throws Exception { @@ -68,17 +63,12 @@ public void testEmbeddingProcessor_whenReindexingDocument_thenSuccessful() throw String fromIndexName = "test-reindex-from"; createIndexWithConfiguration(fromIndexName, "{ \"settings\": { \"number_of_shards\": 1, \"number_of_replicas\": 0 } }", null); ingestDocument(fromIndexName, "{ \"text\": \"hello world\" }"); - String modelId = null; - try { - modelId = uploadModel(); - loadModel(modelId); - String toIndexName = "test-reindex-to"; - createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); - createIndexWithPipeline(toIndexName, "IndexMappings.json", PIPELINE_NAME); - reindex(fromIndexName, toIndexName); - assertEquals(1, getDocCount(toIndexName)); - } finally { - wipeOfTestResources(fromIndexName, PIPELINE_NAME, modelId, null); - } + String modelId = uploadModel(); + loadModel(modelId); + String toIndexName = "test-reindex-to"; + createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); + createIndexWithPipeline(toIndexName, "IndexMappings.json", PIPELINE_NAME); + reindex(fromIndexName, toIndexName); + assertEquals(1, getDocCount(toIndexName)); } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/rerank/ByFieldRerankProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/rerank/ByFieldRerankProcessorIT.java index 93cf182c8..652d1913f 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/rerank/ByFieldRerankProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/rerank/ByFieldRerankProcessorIT.java @@ -81,14 +81,10 @@ public class ByFieldRerankProcessorIT extends BaseNeuralSearchIT { */ @SneakyThrows public void testByFieldRerankProcessor() throws IOException { - try { - createAndPopulateIndex(); - createPipeline(); - applyPipeLine(); - testSearchResponse(); - } finally { - wipeOfTestResources(INDEX_NAME, null, null, PIPELINE_NAME); - } + createAndPopulateIndex(); + createPipeline(); + applyPipeLine(); + testSearchResponse(); } private void createAndPopulateIndex() throws Exception { 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 fcb946d84..db74929f5 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/rerank/MLOpenSearchRerankProcessorIT.java @@ -37,16 +37,11 @@ public class MLOpenSearchRerankProcessorIT extends BaseNeuralSearchIT { @SneakyThrows public void testCrossEncoderRerankProcessor() { - 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); - } + String modelId = uploadTextSimilarityModel(); + loadModel(modelId); + createSearchPipelineViaConfig(modelId, PIPELINE_NAME, "processor/RerankMLOpenSearchPipelineConfiguration.json"); + setupIndex(); + runQueries(); } private String uploadTextSimilarityModel() throws Exception { diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java index cf0a753dc..3607251bd 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java @@ -152,41 +152,37 @@ public void testBucketAndNestedAggs_whenConcurrentSearchEnabled_thenSuccessful() public void testAggregationNotSupportedConcurrentSearch_whenUseSamplerAgg_thenSuccessful() { updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.sampler(GENERIC_AGGREGATION_NAME) - .shardSize(2) - .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_1)); + AggregationBuilder aggsBuilder = AggregationBuilders.sampler(GENERIC_AGGREGATION_NAME) + .shardSize(2) + .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_1)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggsBuilder), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - 3 - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggsBuilder), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + 3 + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertEquals(2, aggValue.size()); - assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); - Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(nestedAggs); - assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); - List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertEquals(2, buckets.size()); - - Map firstBucket = buckets.get(0); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("likeable", firstBucket.get(KEY)); - - Map secondBucket = buckets.get(1); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("workable", secondBucket.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertEquals(2, aggValue.size()); + assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); + Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertNotNull(nestedAggs); + assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); + List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertEquals(2, buckets.size()); + + Map firstBucket = buckets.get(0); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("likeable", firstBucket.get(KEY)); + + Map secondBucket = buckets.get(1); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("workable", secondBucket.get(KEY)); } @SneakyThrows @@ -205,263 +201,251 @@ public void testPostFilterOnIndexWithMultipleShards_WhenConcurrentSearchEnabled_ @SneakyThrows private void testPostFilterWithSimpleHybridQuery(boolean isSingleShard, boolean hasPostFilterQuery) { - try { - if (isSingleShard) { - prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); - } else { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - } + if (isSingleShard) { + prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); + } else { + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + } - HybridQueryBuilder simpleHybridQueryBuilder = createHybridQueryBuilder(false); + HybridQueryBuilder simpleHybridQueryBuilder = createHybridQueryBuilder(false); - QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_1).gte(2000).lte(5000); + QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_1).gte(2000).lte(5000); - Map searchResponseAsMap; - - if (isSingleShard && hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, - simpleHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - rangeFilterQuery, - null, - false, - null, - 0 - ); - - assertHitResultsFromQuery(1, searchResponseAsMap); - } else if (isSingleShard && !hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, - simpleHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(2, searchResponseAsMap); - } else if (!isSingleShard && hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - simpleHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - rangeFilterQuery, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(2, searchResponseAsMap); - } else { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - simpleHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(3, searchResponseAsMap); - } + Map searchResponseAsMap; - // assert post-filter - List> hitsNestedList = getNestedHits(searchResponseAsMap); + if (isSingleShard && hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, + simpleHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + rangeFilterQuery, + null, + false, + null, + 0 + ); - List docIndexes = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - assertNotNull(oneHit.get("_source")); - Map source = (Map) oneHit.get("_source"); - int docIndex = (int) source.get(INTEGER_FIELD_1); - docIndexes.add(docIndex); - } - if (isSingleShard && hasPostFilterQuery) { - assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + assertHitResultsFromQuery(1, searchResponseAsMap); + } else if (isSingleShard && !hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, + simpleHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(2, searchResponseAsMap); + } else if (!isSingleShard && hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + simpleHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + rangeFilterQuery, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(2, searchResponseAsMap); + } else { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + simpleHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(3, searchResponseAsMap); + } - } else if (isSingleShard && !hasPostFilterQuery) { - assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + // assert post-filter + List> hitsNestedList = getNestedHits(searchResponseAsMap); - } else if (!isSingleShard && hasPostFilterQuery) { - assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); - } else { - assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); - } - } finally { - if (isSingleShard) { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } else { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + List docIndexes = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + assertNotNull(oneHit.get("_source")); + Map source = (Map) oneHit.get("_source"); + int docIndex = (int) source.get(INTEGER_FIELD_1); + docIndexes.add(docIndex); + } + if (isSingleShard && hasPostFilterQuery) { + assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + + } else if (isSingleShard && !hasPostFilterQuery) { + assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + + } else if (!isSingleShard && hasPostFilterQuery) { + assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + } else { + assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); } } @SneakyThrows private void testPostFilterWithComplexHybridQuery(boolean isSingleShard, boolean hasPostFilterQuery) { - try { - if (isSingleShard) { - prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); - } else { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - } + if (isSingleShard) { + prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); + } else { + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + } - HybridQueryBuilder complexHybridQueryBuilder = createHybridQueryBuilder(true); + HybridQueryBuilder complexHybridQueryBuilder = createHybridQueryBuilder(true); - QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_1).gte(2000).lte(5000); + QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_1).gte(2000).lte(5000); - Map searchResponseAsMap; - - if (isSingleShard && hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, - complexHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - rangeFilterQuery, - null, - false, - null, - 0 - ); - - assertHitResultsFromQuery(1, searchResponseAsMap); - } else if (isSingleShard && !hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, - complexHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(2, searchResponseAsMap); - } else if (!isSingleShard && hasPostFilterQuery) { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - complexHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - rangeFilterQuery, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(4, searchResponseAsMap); - } else { - searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - complexHybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 0 - ); - assertHitResultsFromQuery(3, searchResponseAsMap); - } + Map searchResponseAsMap; - // assert post-filter - List> hitsNestedList = getNestedHits(searchResponseAsMap); + if (isSingleShard && hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, + complexHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + rangeFilterQuery, + null, + false, + null, + 0 + ); - List docIndexes = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - assertNotNull(oneHit.get("_source")); - Map source = (Map) oneHit.get("_source"); - int docIndex = (int) source.get(INTEGER_FIELD_1); - docIndexes.add(docIndex); - } - if (isSingleShard && hasPostFilterQuery) { - assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + assertHitResultsFromQuery(1, searchResponseAsMap); + } else if (isSingleShard && !hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, + complexHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(2, searchResponseAsMap); + } else if (!isSingleShard && hasPostFilterQuery) { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + complexHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + rangeFilterQuery, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(4, searchResponseAsMap); + } else { + searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + complexHybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ); + assertHitResultsFromQuery(3, searchResponseAsMap); + } - } else if (isSingleShard && !hasPostFilterQuery) { - assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + // assert post-filter + List> hitsNestedList = getNestedHits(searchResponseAsMap); - } else if (!isSingleShard && hasPostFilterQuery) { - assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); - } else { - assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); - } - } finally { - if (isSingleShard) { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } else { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + List docIndexes = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + assertNotNull(oneHit.get("_source")); + Map source = (Map) oneHit.get("_source"); + int docIndex = (int) source.get(INTEGER_FIELD_1); + docIndexes.add(docIndex); + } + if (isSingleShard && hasPostFilterQuery) { + assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + + } else if (isSingleShard && !hasPostFilterQuery) { + assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + + } else if (!isSingleShard && hasPostFilterQuery) { + assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); + } else { + assertEquals(1, docIndexes.stream().filter(docIndex -> docIndex < 2000 || docIndex > 5000).count()); } } @SneakyThrows private void testAvgSumMinMaxAggs() { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD_1) - .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_1)); + AggregationBuilder aggsBuilder = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD_1) + .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_1)); - BucketMetricsPipelineAggregationBuilder aggAvgBucket = PipelineAggregatorBuilders - .avgBucket(BUCKETS_AGGREGATION_NAME_1, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggAvgBucket = PipelineAggregatorBuilders.avgBucket( + BUCKETS_AGGREGATION_NAME_1, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggSumBucket = PipelineAggregatorBuilders - .sumBucket(BUCKETS_AGGREGATION_NAME_2, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggSumBucket = PipelineAggregatorBuilders.sumBucket( + BUCKETS_AGGREGATION_NAME_2, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggMinBucket = PipelineAggregatorBuilders - .minBucket(BUCKETS_AGGREGATION_NAME_3, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggMinBucket = PipelineAggregatorBuilders.minBucket( + BUCKETS_AGGREGATION_NAME_3, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggMaxBucket = PipelineAggregatorBuilders - .maxBucket(BUCKETS_AGGREGATION_NAME_4, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggMaxBucket = PipelineAggregatorBuilders.maxBucket( + BUCKETS_AGGREGATION_NAME_4, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - Map searchResponseAsMapAnngsBoolQuery = executeQueryAndGetAggsResults( - List.of(aggsBuilder, aggAvgBucket, aggSumBucket, aggMinBucket, aggMaxBucket), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - 3 - ); + Map searchResponseAsMapAnngsBoolQuery = executeQueryAndGetAggsResults( + List.of(aggsBuilder, aggAvgBucket, aggSumBucket, aggMinBucket, aggMaxBucket), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + 3 + ); - assertResultsOfPipelineSumtoDateHistogramAggs(searchResponseAsMapAnngsBoolQuery); + assertResultsOfPipelineSumtoDateHistogramAggs(searchResponseAsMapAnngsBoolQuery); - // test only aggregation without query (handled as match_all query) - Map searchResponseAsMapAggsNoQuery = executeQueryAndGetAggsResults( - List.of(aggsBuilder, aggAvgBucket), - null, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - 6 - ); + // test only aggregation without query (handled as match_all query) + Map searchResponseAsMapAggsNoQuery = executeQueryAndGetAggsResults( + List.of(aggsBuilder, aggAvgBucket), + null, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + 6 + ); - assertResultsOfPipelineSumtoDateHistogramAggsForMatchAllQuery(searchResponseAsMapAggsNoQuery); + assertResultsOfPipelineSumtoDateHistogramAggsForMatchAllQuery(searchResponseAsMapAggsNoQuery); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } } @SneakyThrows @@ -481,23 +465,15 @@ public void testPostFilterOnIndexWithSingleShards_WhenConcurrentSearchEnabled_th @SneakyThrows public void testNestedAggs_whenMultipleShardsAndConcurrentSearchDisabled_thenSuccessful() { updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - try { - prepareResourcesForNestegAggregationsScenario(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); - assertNestedAggregations(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesForNestegAggregationsScenario(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); + assertNestedAggregations(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); } @SneakyThrows public void testNestedAggs_whenMultipleShardsAndConcurrentSearchEnabled_thenSuccessful() { updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - try { - prepareResourcesForNestegAggregationsScenario(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); - assertNestedAggregations(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesForNestegAggregationsScenario(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); + assertNestedAggregations(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS); } private void prepareResourcesForNestegAggregationsScenario(String index) throws Exception { @@ -659,72 +635,64 @@ private void assertNestedAggregations(String index) { } private void testMaxAggsOnSingleShardCluster() throws Exception { - try { - prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); + prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); - AggregationBuilder aggsBuilder = AggregationBuilders.max(MAX_AGGREGATION_NAME).field(INTEGER_FIELD_1); - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, - hybridQueryBuilderNeuralThenTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - List.of(aggsBuilder) - ); + AggregationBuilder aggsBuilder = AggregationBuilders.max(MAX_AGGREGATION_NAME).field(INTEGER_FIELD_1); + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, + hybridQueryBuilderNeuralThenTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + List.of(aggsBuilder) + ); - assertHitResultsFromQuery(2, searchResponseAsMap); + assertHitResultsFromQuery(2, searchResponseAsMap); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(MAX_AGGREGATION_NAME)); - double maxAggsValue = getAggregationValue(aggregations, MAX_AGGREGATION_NAME); - assertTrue(maxAggsValue >= 0); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(MAX_AGGREGATION_NAME)); + double maxAggsValue = getAggregationValue(aggregations, MAX_AGGREGATION_NAME); + assertTrue(maxAggsValue >= 0); } private void testDateRange() throws IOException { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.dateRange(DATE_AGGREGATION_NAME) - .field(DATE_FIELD_1) - .format("MM-yyyy") - .addRange("01-2014", "02-2024"); + AggregationBuilder aggsBuilder = AggregationBuilders.dateRange(DATE_AGGREGATION_NAME) + .field(DATE_FIELD_1) + .format("MM-yyyy") + .addRange("01-2014", "02-2024"); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggsBuilder), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - 3 - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggsBuilder), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + 3 + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); - assertNotNull(buckets); - assertEquals(1, buckets.size()); - - Map bucket = buckets.get(0); - - assertEquals(6, bucket.size()); - assertEquals("01-2014", bucket.get("from_as_string")); - assertEquals(2, bucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("02-2024", bucket.get("to_as_string")); - assertTrue(bucket.containsKey("from")); - assertTrue(bucket.containsKey("to")); - assertTrue(bucket.containsKey(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); + assertNotNull(buckets); + assertEquals(1, buckets.size()); + + Map bucket = buckets.get(0); + + assertEquals(6, bucket.size()); + assertEquals("01-2014", bucket.get("from_as_string")); + assertEquals(2, bucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("02-2024", bucket.get("to_as_string")); + assertTrue(bucket.containsKey("from")); + assertTrue(bucket.containsKey("to")); + assertTrue(bucket.containsKey(KEY)); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryExplainIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryExplainIT.java index 113bfee77..ed763f840 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryExplainIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryExplainIT.java @@ -76,659 +76,635 @@ protected boolean preserveClusterUponCompletion() { @SneakyThrows public void testExplain_whenMultipleSubqueriesAndOneShard_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); - // create search pipeline with both normalization processor and explain response processor - createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); - - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); - - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); - - Map searchResponseAsMap1 = search( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - hybridQueryBuilderNeuralThenTerm, - null, - 10, - Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); - // Assert - // search hits - assertEquals(3, getHitCount(searchResponseAsMap1)); - - List> hitsNestedList = getNestedHits(searchResponseAsMap1); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } - - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - assertEquals(Set.copyOf(ids).size(), ids.size()); - - Map total = getTotalHits(searchResponseAsMap1); - assertNotNull(total.get("value")); - assertEquals(3, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // explain - Map searchHit1 = hitsNestedList.get(0); - Map topLevelExplanationsHit1 = getValueByKey(searchHit1, "_explanation"); - assertNotNull(topLevelExplanationsHit1); - assertEquals((double) searchHit1.get("_score"), (double) topLevelExplanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - String expectedTopLevelDescription = "arithmetic_mean combination of:"; - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit1.get("description")); - List> normalizationExplanationHit1 = getListOfValues(topLevelExplanationsHit1, "details"); - assertEquals(1, normalizationExplanationHit1.size()); - Map hit1DetailsForHit1 = normalizationExplanationHit1.get(0); - assertEquals(1.0, hit1DetailsForHit1.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit1.get("description")); - assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); - - Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); - assertEquals("sum of:", explanationsHit1.get("description")); - assertEquals(0.754f, (double) explanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1, ((List) explanationsHit1.get("details")).size()); - - // search hit 2 - Map searchHit2 = hitsNestedList.get(1); - Map topLevelExplanationsHit2 = getValueByKey(searchHit2, "_explanation"); - assertNotNull(topLevelExplanationsHit2); - assertEquals((double) searchHit2.get("_score"), (double) topLevelExplanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit2.get("description")); - List> normalizationExplanationHit2 = getListOfValues(topLevelExplanationsHit2, "details"); - assertEquals(1, normalizationExplanationHit2.size()); - - Map hit1DetailsForHit2 = normalizationExplanationHit2.get(0); - assertEquals(1.0, hit1DetailsForHit2.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit2.get("description")); - assertEquals(1, getListOfValues(hit1DetailsForHit2, "details").size()); - - Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); - assertEquals(0.287f, (double) explanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", explanationsHit2.get("description")); - assertEquals(1, getListOfValues(explanationsHit2, "details").size()); - - Map explanationsHit2Details = getListOfValues(explanationsHit2, "details").get(0); - assertEquals(0.287f, (double) explanationsHit2Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit2Details.get("description")); - assertEquals(3, getListOfValues(explanationsHit2Details, "details").size()); - - // search hit 3 - Map searchHit3 = hitsNestedList.get(1); - Map topLevelExplanationsHit3 = getValueByKey(searchHit3, "_explanation"); - assertNotNull(topLevelExplanationsHit3); - assertEquals((double) searchHit2.get("_score"), (double) topLevelExplanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit3.get("description")); - List> normalizationExplanationHit3 = getListOfValues(topLevelExplanationsHit3, "details"); - assertEquals(1, normalizationExplanationHit3.size()); - - Map hit1DetailsForHit3 = normalizationExplanationHit3.get(0); - assertEquals(1.0, hit1DetailsForHit3.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit3.get("description")); - assertEquals(1, getListOfValues(hit1DetailsForHit3, "details").size()); - - Map explanationsHit3 = getListOfValues(hit1DetailsForHit3, "details").get(0); - assertEquals(0.287f, (double) explanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", explanationsHit3.get("description")); - assertEquals(1, getListOfValues(explanationsHit3, "details").size()); - - Map explanationsHit3Details = getListOfValues(explanationsHit3, "details").get(0); - assertEquals(0.287f, (double) explanationsHit3Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit3Details.get("description")); - assertEquals(3, getListOfValues(explanationsHit3Details, "details").size()); - } finally { - wipeOfTestResources(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, null, null, NORMALIZATION_SEARCH_PIPELINE); + initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); + // create search pipeline with both normalization processor and explain response processor + createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); + + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); + + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); + + Map searchResponseAsMap1 = search( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + hybridQueryBuilderNeuralThenTerm, + null, + 10, + Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); + // Assert + // search hits + assertEquals(3, getHitCount(searchResponseAsMap1)); + + List> hitsNestedList = getNestedHits(searchResponseAsMap1); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); } + + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + assertEquals(Set.copyOf(ids).size(), ids.size()); + + Map total = getTotalHits(searchResponseAsMap1); + assertNotNull(total.get("value")); + assertEquals(3, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // explain + Map searchHit1 = hitsNestedList.get(0); + Map topLevelExplanationsHit1 = getValueByKey(searchHit1, "_explanation"); + assertNotNull(topLevelExplanationsHit1); + assertEquals((double) searchHit1.get("_score"), (double) topLevelExplanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + String expectedTopLevelDescription = "arithmetic_mean combination of:"; + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit1.get("description")); + List> normalizationExplanationHit1 = getListOfValues(topLevelExplanationsHit1, "details"); + assertEquals(1, normalizationExplanationHit1.size()); + Map hit1DetailsForHit1 = normalizationExplanationHit1.get(0); + assertEquals(1.0, hit1DetailsForHit1.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit1.get("description")); + assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); + + Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); + assertEquals("sum of:", explanationsHit1.get("description")); + assertEquals(0.754f, (double) explanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1, ((List) explanationsHit1.get("details")).size()); + + // search hit 2 + Map searchHit2 = hitsNestedList.get(1); + Map topLevelExplanationsHit2 = getValueByKey(searchHit2, "_explanation"); + assertNotNull(topLevelExplanationsHit2); + assertEquals((double) searchHit2.get("_score"), (double) topLevelExplanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit2.get("description")); + List> normalizationExplanationHit2 = getListOfValues(topLevelExplanationsHit2, "details"); + assertEquals(1, normalizationExplanationHit2.size()); + + Map hit1DetailsForHit2 = normalizationExplanationHit2.get(0); + assertEquals(1.0, hit1DetailsForHit2.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit2.get("description")); + assertEquals(1, getListOfValues(hit1DetailsForHit2, "details").size()); + + Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); + assertEquals(0.287f, (double) explanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", explanationsHit2.get("description")); + assertEquals(1, getListOfValues(explanationsHit2, "details").size()); + + Map explanationsHit2Details = getListOfValues(explanationsHit2, "details").get(0); + assertEquals(0.287f, (double) explanationsHit2Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit2Details.get("description")); + assertEquals(3, getListOfValues(explanationsHit2Details, "details").size()); + + // search hit 3 + Map searchHit3 = hitsNestedList.get(1); + Map topLevelExplanationsHit3 = getValueByKey(searchHit3, "_explanation"); + assertNotNull(topLevelExplanationsHit3); + assertEquals((double) searchHit2.get("_score"), (double) topLevelExplanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit3.get("description")); + List> normalizationExplanationHit3 = getListOfValues(topLevelExplanationsHit3, "details"); + assertEquals(1, normalizationExplanationHit3.size()); + + Map hit1DetailsForHit3 = normalizationExplanationHit3.get(0); + assertEquals(1.0, hit1DetailsForHit3.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit3.get("description")); + assertEquals(1, getListOfValues(hit1DetailsForHit3, "details").size()); + + Map explanationsHit3 = getListOfValues(hit1DetailsForHit3, "details").get(0); + assertEquals(0.287f, (double) explanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", explanationsHit3.get("description")); + assertEquals(1, getListOfValues(explanationsHit3, "details").size()); + + Map explanationsHit3Details = getListOfValues(explanationsHit3, "details").get(0); + assertEquals(0.287f, (double) explanationsHit3Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit3Details.get("description")); + assertEquals(3, getListOfValues(explanationsHit3Details, "details").size()); } @SneakyThrows public void testExplain_whenMultipleSubqueriesAndMultipleShards_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - createSearchPipeline( - NORMALIZATION_SEARCH_PIPELINE, - NORMALIZATION_TECHNIQUE_L2, - DEFAULT_COMBINATION_METHOD, - Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f })), - true - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + createSearchPipeline( + NORMALIZATION_SEARCH_PIPELINE, + NORMALIZATION_TECHNIQUE_L2, + DEFAULT_COMBINATION_METHOD, + Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f })), + true + ); - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder() - .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) - .vector(createRandomVector(TEST_DIMENSION)) - .k(10) - .build(); - hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); - hybridQueryBuilder.add(knnQueryBuilder); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_NAME, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); - // Assert - // basic sanity check for search hits - assertEquals(4, getHitCount(searchResponseAsMap)); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - float actualMaxScore = getMaxScore(searchResponseAsMap).get(); - assertTrue(actualMaxScore > 0); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(4, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // explain, hit 1 - List> hitsNestedList = getNestedHits(searchResponseAsMap); - Map searchHit1 = hitsNestedList.get(0); - Map explanationForHit1 = getValueByKey(searchHit1, "_explanation"); - assertNotNull(explanationForHit1); - assertEquals((double) searchHit1.get("_score"), (double) explanationForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - String expectedTopLevelDescription = "arithmetic_mean, weights [0.3, 0.7] combination of:"; - assertEquals(expectedTopLevelDescription, explanationForHit1.get("description")); - List> hit1Details = getListOfValues(explanationForHit1, "details"); - assertEquals(2, hit1Details.size()); - // two sub-queries meaning we do have two detail objects with separate query level details - Map hit1DetailsForHit1 = hit1Details.get(0); - assertTrue((double) hit1DetailsForHit1.get("value") > 0.5f); - assertEquals("l2 normalization of:", hit1DetailsForHit1.get("description")); - assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); - - Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); - assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit1.get("description")); - assertTrue((double) explanationsHit1.get("value") > 0.5f); - assertEquals(0, ((List) explanationsHit1.get("details")).size()); - - Map hit1DetailsForHit2 = hit1Details.get(1); - assertTrue((double) hit1DetailsForHit2.get("value") > 0.0f); - assertEquals("l2 normalization of:", hit1DetailsForHit2.get("description")); - assertEquals(1, ((List) hit1DetailsForHit2.get("details")).size()); - - Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); - assertEquals("within top 10", explanationsHit2.get("description")); - assertTrue((double) explanationsHit2.get("value") > 0.0f); - assertEquals(0, ((List) explanationsHit2.get("details")).size()); - - // hit 2 - Map searchHit2 = hitsNestedList.get(1); - Map explanationForHit2 = getValueByKey(searchHit2, "_explanation"); - assertNotNull(explanationForHit2); - assertEquals((double) searchHit2.get("_score"), (double) explanationForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit2.get("description")); - List> hit2Details = getListOfValues(explanationForHit2, "details"); - assertEquals(2, hit2Details.size()); - - Map hit2DetailsForHit1 = hit2Details.get(0); - assertTrue((double) hit2DetailsForHit1.get("value") > 0.5f); - assertEquals("l2 normalization of:", hit2DetailsForHit1.get("description")); - assertEquals(1, ((List) hit2DetailsForHit1.get("details")).size()); - - Map hit2DetailsForHit2 = hit2Details.get(1); - assertTrue((double) hit2DetailsForHit2.get("value") > 0.0f); - assertEquals("l2 normalization of:", hit2DetailsForHit2.get("description")); - assertEquals(1, ((List) hit2DetailsForHit2.get("details")).size()); - - // hit 3 - Map searchHit3 = hitsNestedList.get(2); - Map explanationForHit3 = getValueByKey(searchHit3, "_explanation"); - assertNotNull(explanationForHit3); - assertEquals((double) searchHit3.get("_score"), (double) explanationForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit3.get("description")); - List> hit3Details = getListOfValues(explanationForHit3, "details"); - assertEquals(1, hit3Details.size()); - - Map hit3DetailsForHit1 = hit3Details.get(0); - assertTrue((double) hit3DetailsForHit1.get("value") > 0.5f); - assertEquals("l2 normalization of:", hit3DetailsForHit1.get("description")); - assertEquals(1, ((List) hit3DetailsForHit1.get("details")).size()); - - Map explanationsHit3 = getListOfValues(hit3DetailsForHit1, "details").get(0); - assertEquals("within top 10", explanationsHit3.get("description")); - assertEquals(0, getListOfValues(explanationsHit3, "details").size()); - assertTrue((double) explanationsHit3.get("value") > 0.0f); - - // hit 4 - Map searchHit4 = hitsNestedList.get(3); - Map explanationForHit4 = getValueByKey(searchHit4, "_explanation"); - assertNotNull(explanationForHit4); - assertEquals((double) searchHit4.get("_score"), (double) explanationForHit4.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit4.get("description")); - List> hit4Details = getListOfValues(explanationForHit4, "details"); - assertEquals(1, hit4Details.size()); - - Map hit4DetailsForHit1 = hit4Details.get(0); - assertTrue((double) hit4DetailsForHit1.get("value") > 0.5f); - assertEquals("l2 normalization of:", hit4DetailsForHit1.get("description")); - assertEquals(1, ((List) hit4DetailsForHit1.get("details")).size()); - - Map explanationsHit4 = getListOfValues(hit4DetailsForHit1, "details").get(0); - assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit4.get("description")); - assertEquals(0, getListOfValues(explanationsHit4, "details").size()); - assertTrue((double) explanationsHit4.get("value") > 0.0f); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, NORMALIZATION_SEARCH_PIPELINE); - } + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .vector(createRandomVector(TEST_DIMENSION)) + .k(10) + .build(); + hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); + hybridQueryBuilder.add(knnQueryBuilder); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_NAME, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); + // Assert + // basic sanity check for search hits + assertEquals(4, getHitCount(searchResponseAsMap)); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + float actualMaxScore = getMaxScore(searchResponseAsMap).get(); + assertTrue(actualMaxScore > 0); + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(4, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // explain, hit 1 + List> hitsNestedList = getNestedHits(searchResponseAsMap); + Map searchHit1 = hitsNestedList.get(0); + Map explanationForHit1 = getValueByKey(searchHit1, "_explanation"); + assertNotNull(explanationForHit1); + assertEquals((double) searchHit1.get("_score"), (double) explanationForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + String expectedTopLevelDescription = "arithmetic_mean, weights [0.3, 0.7] combination of:"; + assertEquals(expectedTopLevelDescription, explanationForHit1.get("description")); + List> hit1Details = getListOfValues(explanationForHit1, "details"); + assertEquals(2, hit1Details.size()); + // two sub-queries meaning we do have two detail objects with separate query level details + Map hit1DetailsForHit1 = hit1Details.get(0); + assertTrue((double) hit1DetailsForHit1.get("value") > 0.5f); + assertEquals("l2 normalization of:", hit1DetailsForHit1.get("description")); + assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); + + Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); + assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit1.get("description")); + assertTrue((double) explanationsHit1.get("value") > 0.5f); + assertEquals(0, ((List) explanationsHit1.get("details")).size()); + + Map hit1DetailsForHit2 = hit1Details.get(1); + assertTrue((double) hit1DetailsForHit2.get("value") > 0.0f); + assertEquals("l2 normalization of:", hit1DetailsForHit2.get("description")); + assertEquals(1, ((List) hit1DetailsForHit2.get("details")).size()); + + Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); + assertEquals("within top 10", explanationsHit2.get("description")); + assertTrue((double) explanationsHit2.get("value") > 0.0f); + assertEquals(0, ((List) explanationsHit2.get("details")).size()); + + // hit 2 + Map searchHit2 = hitsNestedList.get(1); + Map explanationForHit2 = getValueByKey(searchHit2, "_explanation"); + assertNotNull(explanationForHit2); + assertEquals((double) searchHit2.get("_score"), (double) explanationForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit2.get("description")); + List> hit2Details = getListOfValues(explanationForHit2, "details"); + assertEquals(2, hit2Details.size()); + + Map hit2DetailsForHit1 = hit2Details.get(0); + assertTrue((double) hit2DetailsForHit1.get("value") > 0.5f); + assertEquals("l2 normalization of:", hit2DetailsForHit1.get("description")); + assertEquals(1, ((List) hit2DetailsForHit1.get("details")).size()); + + Map hit2DetailsForHit2 = hit2Details.get(1); + assertTrue((double) hit2DetailsForHit2.get("value") > 0.0f); + assertEquals("l2 normalization of:", hit2DetailsForHit2.get("description")); + assertEquals(1, ((List) hit2DetailsForHit2.get("details")).size()); + + // hit 3 + Map searchHit3 = hitsNestedList.get(2); + Map explanationForHit3 = getValueByKey(searchHit3, "_explanation"); + assertNotNull(explanationForHit3); + assertEquals((double) searchHit3.get("_score"), (double) explanationForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit3.get("description")); + List> hit3Details = getListOfValues(explanationForHit3, "details"); + assertEquals(1, hit3Details.size()); + + Map hit3DetailsForHit1 = hit3Details.get(0); + assertTrue((double) hit3DetailsForHit1.get("value") > 0.5f); + assertEquals("l2 normalization of:", hit3DetailsForHit1.get("description")); + assertEquals(1, ((List) hit3DetailsForHit1.get("details")).size()); + + Map explanationsHit3 = getListOfValues(hit3DetailsForHit1, "details").get(0); + assertEquals("within top 10", explanationsHit3.get("description")); + assertEquals(0, getListOfValues(explanationsHit3, "details").size()); + assertTrue((double) explanationsHit3.get("value") > 0.0f); + + // hit 4 + Map searchHit4 = hitsNestedList.get(3); + Map explanationForHit4 = getValueByKey(searchHit4, "_explanation"); + assertNotNull(explanationForHit4); + assertEquals((double) searchHit4.get("_score"), (double) explanationForHit4.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit4.get("description")); + List> hit4Details = getListOfValues(explanationForHit4, "details"); + assertEquals(1, hit4Details.size()); + + Map hit4DetailsForHit1 = hit4Details.get(0); + assertTrue((double) hit4DetailsForHit1.get("value") > 0.5f); + assertEquals("l2 normalization of:", hit4DetailsForHit1.get("description")); + assertEquals(1, ((List) hit4DetailsForHit1.get("details")).size()); + + Map explanationsHit4 = getListOfValues(hit4DetailsForHit1, "details").get(0); + assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit4.get("description")); + assertEquals(0, getListOfValues(explanationsHit4, "details").size()); + assertTrue((double) explanationsHit4.get("value") > 0.0f); } @SneakyThrows public void testExplanationResponseProcessor_whenProcessorIsNotConfigured_thenResponseHasQueryExplanations() { - try { - initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); - // create search pipeline with normalization processor, no explanation response processor - createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), false); - - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); - - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); - - Map searchResponseAsMap1 = search( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - hybridQueryBuilderNeuralThenTerm, - null, - 10, - Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); - // Assert - // search hits - assertEquals(3, getHitCount(searchResponseAsMap1)); - - List> hitsNestedList = getNestedHits(searchResponseAsMap1); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } - - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - assertEquals(Set.copyOf(ids).size(), ids.size()); - - Map total = getTotalHits(searchResponseAsMap1); - assertNotNull(total.get("value")); - assertEquals(3, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // explain - Map searchHit1 = hitsNestedList.get(0); - Map topLevelExplanationsHit1 = getValueByKey(searchHit1, "_explanation"); - assertNotNull(topLevelExplanationsHit1); - assertEquals(0.754f, (double) topLevelExplanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - String expectedTopLevelDescription = "combined score of:"; - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit1.get("description")); - List> normalizationExplanationHit1 = getListOfValues(topLevelExplanationsHit1, "details"); - assertEquals(2, normalizationExplanationHit1.size()); - - Map noMatchDetailsForHit1 = normalizationExplanationHit1.get(0); - assertEquals(0.0f, (double) noMatchDetailsForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("no matching term", noMatchDetailsForHit1.get("description")); - assertEquals(0, ((List) noMatchDetailsForHit1.get("details")).size()); - - Map hit1DetailsForHit1 = normalizationExplanationHit1.get(1); - assertEquals(0.754f, (double) hit1DetailsForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("sum of:", hit1DetailsForHit1.get("description")); - assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); - - Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); - assertEquals("weight(test-text-field-1:place in 0) [PerFieldSimilarity], result of:", explanationsHit1.get("description")); - assertEquals(0.754f, (double) explanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1, ((List) explanationsHit1.get("details")).size()); - - Map explanationsHit1Details = getListOfValues(explanationsHit1, "details").get(0); - assertEquals(0.754f, (double) explanationsHit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit1Details.get("description")); - assertEquals(3, getListOfValues(explanationsHit1Details, "details").size()); - - Map explanationsDetails1Hit1Details = getListOfValues(explanationsHit1Details, "details").get(0); - assertEquals(2.2f, (double) explanationsDetails1Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("boost", explanationsDetails1Hit1Details.get("description")); - assertEquals(0, getListOfValues(explanationsDetails1Hit1Details, "details").size()); - - Map explanationsDetails2Hit1Details = getListOfValues(explanationsHit1Details, "details").get(1); - assertEquals(0.693f, (double) explanationsDetails2Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("idf, computed as log(1 + (N - n + 0.5) / (n + 0.5)) from:", explanationsDetails2Hit1Details.get("description")); - assertFalse(getListOfValues(explanationsDetails2Hit1Details, "details").isEmpty()); - - Map explanationsDetails3Hit1Details = getListOfValues(explanationsHit1Details, "details").get(2); - assertEquals(0.495f, (double) explanationsDetails3Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals( - "tf, computed as freq / (freq + k1 * (1 - b + b * dl / avgdl)) from:", - explanationsDetails3Hit1Details.get("description") - ); - assertFalse(getListOfValues(explanationsDetails3Hit1Details, "details").isEmpty()); - - // search hit 2 - Map searchHit2 = hitsNestedList.get(1); - Map topLevelExplanationsHit2 = getValueByKey(searchHit2, "_explanation"); - assertNotNull(topLevelExplanationsHit2); - assertEquals(0.287f, (double) topLevelExplanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit2.get("description")); - List> normalizationExplanationHit2 = getListOfValues(topLevelExplanationsHit2, "details"); - assertEquals(2, normalizationExplanationHit2.size()); - - Map hit1DetailsForHit2 = normalizationExplanationHit2.get(0); - assertEquals(0.287f, (double) hit1DetailsForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit2.get("description")); - assertEquals(1, getListOfValues(hit1DetailsForHit2, "details").size()); - - Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); - assertEquals(0.287f, (double) explanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit2.get("description")); - assertEquals(3, getListOfValues(explanationsHit2, "details").size()); - - Map explanationsHit2Details = getListOfValues(explanationsHit2, "details").get(0); - assertEquals(2.2f, (double) explanationsHit2Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("boost", explanationsHit2Details.get("description")); - assertEquals(0, getListOfValues(explanationsHit2Details, "details").size()); - - Map hit1DetailsForHit2NoMatch = normalizationExplanationHit2.get(1); - assertEquals(0.0f, (double) hit1DetailsForHit2NoMatch.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("No matching clauses", hit1DetailsForHit2NoMatch.get("description")); - assertEquals(0, ((List) hit1DetailsForHit2NoMatch.get("details")).size()); - - // search hit 3 - Map searchHit3 = hitsNestedList.get(1); - Map topLevelExplanationsHit3 = getValueByKey(searchHit3, "_explanation"); - assertNotNull(topLevelExplanationsHit3); - assertEquals(0.287f, (double) topLevelExplanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, topLevelExplanationsHit3.get("description")); - List> normalizationExplanationHit3 = getListOfValues(topLevelExplanationsHit3, "details"); - assertEquals(2, normalizationExplanationHit3.size()); - - Map hit1DetailsForHit3 = normalizationExplanationHit3.get(0); - assertEquals(0.287, (double) hit1DetailsForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit3.get("description")); - assertEquals(1, getListOfValues(hit1DetailsForHit3, "details").size()); - - Map explanationsHit3 = getListOfValues(hit1DetailsForHit3, "details").get(0); - assertEquals(0.287f, (double) explanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit3.get("description")); - assertEquals(3, getListOfValues(explanationsHit3, "details").size()); - - Map explanationsHit3Details = getListOfValues(explanationsHit3, "details").get(0); - assertEquals(2.2f, (double) explanationsHit3Details.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("boost", explanationsHit3Details.get("description")); - assertEquals(0, getListOfValues(explanationsHit3Details, "details").size()); - - Map hit1DetailsForHit3NoMatch = normalizationExplanationHit2.get(1); - assertEquals(0.0f, (double) hit1DetailsForHit3NoMatch.get("value"), DELTA_FOR_SCORE_ASSERTION); - assertEquals("No matching clauses", hit1DetailsForHit3NoMatch.get("description")); - assertEquals(0, ((List) hit1DetailsForHit3NoMatch.get("details")).size()); - } finally { - wipeOfTestResources(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, null, null, NORMALIZATION_SEARCH_PIPELINE); + initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); + // create search pipeline with normalization processor, no explanation response processor + createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), false); + + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); + + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); + + Map searchResponseAsMap1 = search( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + hybridQueryBuilderNeuralThenTerm, + null, + 10, + Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); + // Assert + // search hits + assertEquals(3, getHitCount(searchResponseAsMap1)); + + List> hitsNestedList = getNestedHits(searchResponseAsMap1); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); } + + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + assertEquals(Set.copyOf(ids).size(), ids.size()); + + Map total = getTotalHits(searchResponseAsMap1); + assertNotNull(total.get("value")); + assertEquals(3, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // explain + Map searchHit1 = hitsNestedList.get(0); + Map topLevelExplanationsHit1 = getValueByKey(searchHit1, "_explanation"); + assertNotNull(topLevelExplanationsHit1); + assertEquals(0.754f, (double) topLevelExplanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + String expectedTopLevelDescription = "combined score of:"; + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit1.get("description")); + List> normalizationExplanationHit1 = getListOfValues(topLevelExplanationsHit1, "details"); + assertEquals(2, normalizationExplanationHit1.size()); + + Map noMatchDetailsForHit1 = normalizationExplanationHit1.get(0); + assertEquals(0.0f, (double) noMatchDetailsForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("no matching term", noMatchDetailsForHit1.get("description")); + assertEquals(0, ((List) noMatchDetailsForHit1.get("details")).size()); + + Map hit1DetailsForHit1 = normalizationExplanationHit1.get(1); + assertEquals(0.754f, (double) hit1DetailsForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("sum of:", hit1DetailsForHit1.get("description")); + assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); + + Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); + assertEquals("weight(test-text-field-1:place in 0) [PerFieldSimilarity], result of:", explanationsHit1.get("description")); + assertEquals(0.754f, (double) explanationsHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1, ((List) explanationsHit1.get("details")).size()); + + Map explanationsHit1Details = getListOfValues(explanationsHit1, "details").get(0); + assertEquals(0.754f, (double) explanationsHit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit1Details.get("description")); + assertEquals(3, getListOfValues(explanationsHit1Details, "details").size()); + + Map explanationsDetails1Hit1Details = getListOfValues(explanationsHit1Details, "details").get(0); + assertEquals(2.2f, (double) explanationsDetails1Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("boost", explanationsDetails1Hit1Details.get("description")); + assertEquals(0, getListOfValues(explanationsDetails1Hit1Details, "details").size()); + + Map explanationsDetails2Hit1Details = getListOfValues(explanationsHit1Details, "details").get(1); + assertEquals(0.693f, (double) explanationsDetails2Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("idf, computed as log(1 + (N - n + 0.5) / (n + 0.5)) from:", explanationsDetails2Hit1Details.get("description")); + assertFalse(getListOfValues(explanationsDetails2Hit1Details, "details").isEmpty()); + + Map explanationsDetails3Hit1Details = getListOfValues(explanationsHit1Details, "details").get(2); + assertEquals(0.495f, (double) explanationsDetails3Hit1Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals( + "tf, computed as freq / (freq + k1 * (1 - b + b * dl / avgdl)) from:", + explanationsDetails3Hit1Details.get("description") + ); + assertFalse(getListOfValues(explanationsDetails3Hit1Details, "details").isEmpty()); + + // search hit 2 + Map searchHit2 = hitsNestedList.get(1); + Map topLevelExplanationsHit2 = getValueByKey(searchHit2, "_explanation"); + assertNotNull(topLevelExplanationsHit2); + assertEquals(0.287f, (double) topLevelExplanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit2.get("description")); + List> normalizationExplanationHit2 = getListOfValues(topLevelExplanationsHit2, "details"); + assertEquals(2, normalizationExplanationHit2.size()); + + Map hit1DetailsForHit2 = normalizationExplanationHit2.get(0); + assertEquals(0.287f, (double) hit1DetailsForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit2.get("description")); + assertEquals(1, getListOfValues(hit1DetailsForHit2, "details").size()); + + Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); + assertEquals(0.287f, (double) explanationsHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit2.get("description")); + assertEquals(3, getListOfValues(explanationsHit2, "details").size()); + + Map explanationsHit2Details = getListOfValues(explanationsHit2, "details").get(0); + assertEquals(2.2f, (double) explanationsHit2Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("boost", explanationsHit2Details.get("description")); + assertEquals(0, getListOfValues(explanationsHit2Details, "details").size()); + + Map hit1DetailsForHit2NoMatch = normalizationExplanationHit2.get(1); + assertEquals(0.0f, (double) hit1DetailsForHit2NoMatch.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("No matching clauses", hit1DetailsForHit2NoMatch.get("description")); + assertEquals(0, ((List) hit1DetailsForHit2NoMatch.get("details")).size()); + + // search hit 3 + Map searchHit3 = hitsNestedList.get(1); + Map topLevelExplanationsHit3 = getValueByKey(searchHit3, "_explanation"); + assertNotNull(topLevelExplanationsHit3); + assertEquals(0.287f, (double) topLevelExplanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, topLevelExplanationsHit3.get("description")); + List> normalizationExplanationHit3 = getListOfValues(topLevelExplanationsHit3, "details"); + assertEquals(2, normalizationExplanationHit3.size()); + + Map hit1DetailsForHit3 = normalizationExplanationHit3.get(0); + assertEquals(0.287, (double) hit1DetailsForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("weight(test-text-field-1:hello in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit3.get("description")); + assertEquals(1, getListOfValues(hit1DetailsForHit3, "details").size()); + + Map explanationsHit3 = getListOfValues(hit1DetailsForHit3, "details").get(0); + assertEquals(0.287f, (double) explanationsHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("score(freq=1.0), computed as boost * idf * tf from:", explanationsHit3.get("description")); + assertEquals(3, getListOfValues(explanationsHit3, "details").size()); + + Map explanationsHit3Details = getListOfValues(explanationsHit3, "details").get(0); + assertEquals(2.2f, (double) explanationsHit3Details.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("boost", explanationsHit3Details.get("description")); + assertEquals(0, getListOfValues(explanationsHit3Details, "details").size()); + + Map hit1DetailsForHit3NoMatch = normalizationExplanationHit2.get(1); + assertEquals(0.0f, (double) hit1DetailsForHit3NoMatch.get("value"), DELTA_FOR_SCORE_ASSERTION); + assertEquals("No matching clauses", hit1DetailsForHit3NoMatch.get("description")); + assertEquals(0, ((List) hit1DetailsForHit3NoMatch.get("details")).size()); } @SneakyThrows public void testExplain_whenLargeNumberOfDocuments_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_LARGE_DOCS_INDEX_NAME); - // create search pipeline with both normalization processor and explain response processor - createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); - - TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(termQueryBuilder); - - Map searchResponseAsMap = search( - TEST_LARGE_DOCS_INDEX_NAME, - hybridQueryBuilder, - null, - MAX_NUMBER_OF_DOCS_IN_LARGE_INDEX, - Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); + initializeIndexIfNotExist(TEST_LARGE_DOCS_INDEX_NAME); + // create search pipeline with both normalization processor and explain response processor + createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); + + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(termQueryBuilder); + + Map searchResponseAsMap = search( + TEST_LARGE_DOCS_INDEX_NAME, + hybridQueryBuilder, + null, + MAX_NUMBER_OF_DOCS_IN_LARGE_INDEX, + Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); - List> hitsNestedList = getNestedHits(searchResponseAsMap); - assertNotNull(hitsNestedList); - assertFalse(hitsNestedList.isEmpty()); - - // Verify total hits - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertTrue((int) total.get("value") > 0); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // Sanity checks for each hit's explanation - for (Map hit : hitsNestedList) { - // Verify score is positive - double score = (double) hit.get("_score"); - assertTrue("Score should be positive", score > 0.0); - - // Basic explanation structure checks - Map explanation = getValueByKey(hit, "_explanation"); - assertNotNull(explanation); - assertEquals("arithmetic_mean combination of:", explanation.get("description")); - Map hitDetailsForHit = getListOfValues(explanation, "details").get(0); - assertTrue((double) hitDetailsForHit.get("value") > 0.0f); - assertEquals("min_max normalization of:", hitDetailsForHit.get("description")); - Map subQueryDetailsForHit = getListOfValues(hitDetailsForHit, "details").get(0); - assertTrue((double) subQueryDetailsForHit.get("value") > 0.0f); - assertFalse(subQueryDetailsForHit.get("description").toString().isEmpty()); - assertEquals(1, getListOfValues(subQueryDetailsForHit, "details").size()); - } - // Verify scores are properly ordered - List scores = new ArrayList<>(); - for (Map hit : hitsNestedList) { - scores.add((Double) hit.get("_score")); - } - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(i -> scores.get(i) < scores.get(i + 1))); - } finally { - wipeOfTestResources(TEST_LARGE_DOCS_INDEX_NAME, null, null, NORMALIZATION_SEARCH_PIPELINE); + List> hitsNestedList = getNestedHits(searchResponseAsMap); + assertNotNull(hitsNestedList); + assertFalse(hitsNestedList.isEmpty()); + + // Verify total hits + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertTrue((int) total.get("value") > 0); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // Sanity checks for each hit's explanation + for (Map hit : hitsNestedList) { + // Verify score is positive + double score = (double) hit.get("_score"); + assertTrue("Score should be positive", score > 0.0); + + // Basic explanation structure checks + Map explanation = getValueByKey(hit, "_explanation"); + assertNotNull(explanation); + assertEquals("arithmetic_mean combination of:", explanation.get("description")); + Map hitDetailsForHit = getListOfValues(explanation, "details").get(0); + assertTrue((double) hitDetailsForHit.get("value") > 0.0f); + assertEquals("min_max normalization of:", hitDetailsForHit.get("description")); + Map subQueryDetailsForHit = getListOfValues(hitDetailsForHit, "details").get(0); + assertTrue((double) subQueryDetailsForHit.get("value") > 0.0f); + assertFalse(subQueryDetailsForHit.get("description").toString().isEmpty()); + assertEquals(1, getListOfValues(subQueryDetailsForHit, "details").size()); + } + // Verify scores are properly ordered + List scores = new ArrayList<>(); + for (Map hit : hitsNestedList) { + scores.add((Double) hit.get("_score")); } + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(i -> scores.get(i) < scores.get(i + 1))); } @SneakyThrows public void testSpecificQueryTypes_whenMultiMatchAndKnn_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_LARGE_DOCS_INDEX_NAME); - // create search pipeline with both normalization processor and explain response processor - createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); - - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(QueryBuilders.multiMatchQuery(TEST_QUERY_TEXT3, TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2)); - hybridQueryBuilder.add( - KNNQueryBuilder.builder().k(10).fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).vector(TEST_VECTOR_SUPPLIER.get()).build() - ); + initializeIndexIfNotExist(TEST_LARGE_DOCS_INDEX_NAME); + // create search pipeline with both normalization processor and explain response processor + createSearchPipeline(NORMALIZATION_SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); + + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(QueryBuilders.multiMatchQuery(TEST_QUERY_TEXT3, TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2)); + hybridQueryBuilder.add( + KNNQueryBuilder.builder().k(10).fieldName(TEST_KNN_VECTOR_FIELD_NAME_1).vector(TEST_VECTOR_SUPPLIER.get()).build() + ); - Map searchResponseAsMap = search( - TEST_LARGE_DOCS_INDEX_NAME, - hybridQueryBuilder, - null, - MAX_NUMBER_OF_DOCS_IN_LARGE_INDEX, - Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); + Map searchResponseAsMap = search( + TEST_LARGE_DOCS_INDEX_NAME, + hybridQueryBuilder, + null, + MAX_NUMBER_OF_DOCS_IN_LARGE_INDEX, + Map.of("search_pipeline", NORMALIZATION_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); - List> hitsNestedList = getNestedHits(searchResponseAsMap); - assertNotNull(hitsNestedList); - assertFalse(hitsNestedList.isEmpty()); - - // Verify total hits - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertTrue((int) total.get("value") > 0); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // Sanity checks for each hit's explanation - for (Map hit : hitsNestedList) { - // Verify score is positive - double score = (double) hit.get("_score"); - assertTrue("Score should be positive", score > 0.0); - - // Basic explanation structure checks - Map explanation = getValueByKey(hit, "_explanation"); - assertNotNull(explanation); - assertEquals("arithmetic_mean combination of:", explanation.get("description")); - Map hitDetailsForHit = getListOfValues(explanation, "details").get(0); - assertTrue((double) hitDetailsForHit.get("value") > 0.0f); - assertEquals("min_max normalization of:", hitDetailsForHit.get("description")); - Map subQueryDetailsForHit = getListOfValues(hitDetailsForHit, "details").get(0); - assertTrue((double) subQueryDetailsForHit.get("value") > 0.0f); - assertFalse(subQueryDetailsForHit.get("description").toString().isEmpty()); - assertNotNull(getListOfValues(subQueryDetailsForHit, "details")); - } - // Verify scores are properly ordered - List scores = new ArrayList<>(); - for (Map hit : hitsNestedList) { - scores.add((Double) hit.get("_score")); - } - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(i -> scores.get(i) < scores.get(i + 1))); - } finally { - wipeOfTestResources(TEST_LARGE_DOCS_INDEX_NAME, null, null, NORMALIZATION_SEARCH_PIPELINE); + List> hitsNestedList = getNestedHits(searchResponseAsMap); + assertNotNull(hitsNestedList); + assertFalse(hitsNestedList.isEmpty()); + + // Verify total hits + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertTrue((int) total.get("value") > 0); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // Sanity checks for each hit's explanation + for (Map hit : hitsNestedList) { + // Verify score is positive + double score = (double) hit.get("_score"); + assertTrue("Score should be positive", score > 0.0); + + // Basic explanation structure checks + Map explanation = getValueByKey(hit, "_explanation"); + assertNotNull(explanation); + assertEquals("arithmetic_mean combination of:", explanation.get("description")); + Map hitDetailsForHit = getListOfValues(explanation, "details").get(0); + assertTrue((double) hitDetailsForHit.get("value") > 0.0f); + assertEquals("min_max normalization of:", hitDetailsForHit.get("description")); + Map subQueryDetailsForHit = getListOfValues(hitDetailsForHit, "details").get(0); + assertTrue((double) subQueryDetailsForHit.get("value") > 0.0f); + assertFalse(subQueryDetailsForHit.get("description").toString().isEmpty()); + assertNotNull(getListOfValues(subQueryDetailsForHit, "details")); } + // Verify scores are properly ordered + List scores = new ArrayList<>(); + for (Map hit : hitsNestedList) { + scores.add((Double) hit.get("_score")); + } + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(i -> scores.get(i) < scores.get(i + 1))); } @SneakyThrows public void testExplain_whenRRFProcessor_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - createRRFSearchPipeline(RRF_SEARCH_PIPELINE, true); - - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder() - .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) - .vector(createRandomVector(TEST_DIMENSION)) - .k(10) - .build(); - hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); - hybridQueryBuilder.add(knnQueryBuilder); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_NAME, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", RRF_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) - ); - // Assert - // basic sanity check for search hits - assertEquals(4, getHitCount(searchResponseAsMap)); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - float actualMaxScore = getMaxScore(searchResponseAsMap).get(); - assertTrue(actualMaxScore > 0); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(4, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - - // explain, hit 1 - List> hitsNestedList = getNestedHits(searchResponseAsMap); - Map searchHit1 = hitsNestedList.get(0); - Map explanationForHit1 = getValueByKey(searchHit1, "_explanation"); - assertNotNull(explanationForHit1); - assertEquals((double) searchHit1.get("_score"), (double) explanationForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); - String expectedTopLevelDescription = "rrf combination of:"; - assertEquals(expectedTopLevelDescription, explanationForHit1.get("description")); - List> hit1Details = getListOfValues(explanationForHit1, "details"); - assertEquals(2, hit1Details.size()); - // two sub-queries meaning we do have two detail objects with separate query level details - Map hit1DetailsForHit1 = hit1Details.get(0); - assertTrue((double) hit1DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); - assertEquals("rrf, rank_constant [60] normalization of:", hit1DetailsForHit1.get("description")); - assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); - - Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); - assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit1.get("description")); - assertTrue((double) explanationsHit1.get("value") > 0.5f); - assertEquals(0, ((List) explanationsHit1.get("details")).size()); - - Map hit1DetailsForHit2 = hit1Details.get(1); - assertTrue((double) hit1DetailsForHit2.get("value") > 0.0f); - assertEquals("rrf, rank_constant [60] normalization of:", hit1DetailsForHit2.get("description")); - assertEquals(1, ((List) hit1DetailsForHit2.get("details")).size()); - - Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); - assertEquals("within top 10", explanationsHit2.get("description")); - assertTrue((double) explanationsHit2.get("value") > 0.0f); - assertEquals(0, ((List) explanationsHit2.get("details")).size()); - - // hit 2 - Map searchHit2 = hitsNestedList.get(1); - Map explanationForHit2 = getValueByKey(searchHit2, "_explanation"); - assertNotNull(explanationForHit2); - assertEquals((double) searchHit2.get("_score"), (double) explanationForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit2.get("description")); - List> hit2Details = getListOfValues(explanationForHit2, "details"); - assertEquals(2, hit2Details.size()); - - Map hit2DetailsForHit1 = hit2Details.get(0); - assertTrue((double) hit2DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); - assertEquals("rrf, rank_constant [60] normalization of:", hit2DetailsForHit1.get("description")); - assertEquals(1, ((List) hit2DetailsForHit1.get("details")).size()); - - Map hit2DetailsForHit2 = hit2Details.get(1); - assertTrue((double) hit2DetailsForHit2.get("value") > DELTA_FOR_SCORE_ASSERTION); - assertEquals("rrf, rank_constant [60] normalization of:", hit2DetailsForHit2.get("description")); - assertEquals(1, ((List) hit2DetailsForHit2.get("details")).size()); - - // hit 3 - Map searchHit3 = hitsNestedList.get(2); - Map explanationForHit3 = getValueByKey(searchHit3, "_explanation"); - assertNotNull(explanationForHit3); - assertEquals((double) searchHit3.get("_score"), (double) explanationForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit3.get("description")); - List> hit3Details = getListOfValues(explanationForHit3, "details"); - assertEquals(1, hit3Details.size()); - - Map hit3DetailsForHit1 = hit3Details.get(0); - assertTrue((double) hit3DetailsForHit1.get("value") > .0f); - assertEquals("rrf, rank_constant [60] normalization of:", hit3DetailsForHit1.get("description")); - assertEquals(1, ((List) hit3DetailsForHit1.get("details")).size()); - - Map explanationsHit3 = getListOfValues(hit3DetailsForHit1, "details").get(0); - assertEquals("within top 10", explanationsHit3.get("description")); - assertEquals(0, getListOfValues(explanationsHit3, "details").size()); - assertTrue((double) explanationsHit3.get("value") > 0.0f); - - // hit 4 - Map searchHit4 = hitsNestedList.get(3); - Map explanationForHit4 = getValueByKey(searchHit4, "_explanation"); - assertNotNull(explanationForHit4); - assertEquals((double) searchHit4.get("_score"), (double) explanationForHit4.get("value"), DELTA_FOR_SCORE_ASSERTION); - - assertEquals(expectedTopLevelDescription, explanationForHit4.get("description")); - List> hit4Details = getListOfValues(explanationForHit4, "details"); - assertEquals(1, hit4Details.size()); - - Map hit4DetailsForHit1 = hit4Details.get(0); - assertTrue((double) hit4DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); - assertEquals("rrf, rank_constant [60] normalization of:", hit4DetailsForHit1.get("description")); - assertEquals(1, ((List) hit4DetailsForHit1.get("details")).size()); - - Map explanationsHit4 = getListOfValues(hit4DetailsForHit1, "details").get(0); - assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit4.get("description")); - assertEquals(0, getListOfValues(explanationsHit4, "details").size()); - assertTrue((double) explanationsHit4.get("value") > 0.0f); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, RRF_SEARCH_PIPELINE); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + createRRFSearchPipeline(RRF_SEARCH_PIPELINE, true); + + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + KNNQueryBuilder knnQueryBuilder = KNNQueryBuilder.builder() + .fieldName(TEST_KNN_VECTOR_FIELD_NAME_1) + .vector(createRandomVector(TEST_DIMENSION)) + .k(10) + .build(); + hybridQueryBuilder.add(QueryBuilders.existsQuery(TEST_TEXT_FIELD_NAME_1)); + hybridQueryBuilder.add(knnQueryBuilder); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_NAME, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", RRF_SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()) + ); + // Assert + // basic sanity check for search hits + assertEquals(4, getHitCount(searchResponseAsMap)); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + float actualMaxScore = getMaxScore(searchResponseAsMap).get(); + assertTrue(actualMaxScore > 0); + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(4, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); + + // explain, hit 1 + List> hitsNestedList = getNestedHits(searchResponseAsMap); + Map searchHit1 = hitsNestedList.get(0); + Map explanationForHit1 = getValueByKey(searchHit1, "_explanation"); + assertNotNull(explanationForHit1); + assertEquals((double) searchHit1.get("_score"), (double) explanationForHit1.get("value"), DELTA_FOR_SCORE_ASSERTION); + String expectedTopLevelDescription = "rrf combination of:"; + assertEquals(expectedTopLevelDescription, explanationForHit1.get("description")); + List> hit1Details = getListOfValues(explanationForHit1, "details"); + assertEquals(2, hit1Details.size()); + // two sub-queries meaning we do have two detail objects with separate query level details + Map hit1DetailsForHit1 = hit1Details.get(0); + assertTrue((double) hit1DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); + assertEquals("rrf, rank_constant [60] normalization of:", hit1DetailsForHit1.get("description")); + assertEquals(1, ((List) hit1DetailsForHit1.get("details")).size()); + + Map explanationsHit1 = getListOfValues(hit1DetailsForHit1, "details").get(0); + assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit1.get("description")); + assertTrue((double) explanationsHit1.get("value") > 0.5f); + assertEquals(0, ((List) explanationsHit1.get("details")).size()); + + Map hit1DetailsForHit2 = hit1Details.get(1); + assertTrue((double) hit1DetailsForHit2.get("value") > 0.0f); + assertEquals("rrf, rank_constant [60] normalization of:", hit1DetailsForHit2.get("description")); + assertEquals(1, ((List) hit1DetailsForHit2.get("details")).size()); + + Map explanationsHit2 = getListOfValues(hit1DetailsForHit2, "details").get(0); + assertEquals("within top 10", explanationsHit2.get("description")); + assertTrue((double) explanationsHit2.get("value") > 0.0f); + assertEquals(0, ((List) explanationsHit2.get("details")).size()); + + // hit 2 + Map searchHit2 = hitsNestedList.get(1); + Map explanationForHit2 = getValueByKey(searchHit2, "_explanation"); + assertNotNull(explanationForHit2); + assertEquals((double) searchHit2.get("_score"), (double) explanationForHit2.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit2.get("description")); + List> hit2Details = getListOfValues(explanationForHit2, "details"); + assertEquals(2, hit2Details.size()); + + Map hit2DetailsForHit1 = hit2Details.get(0); + assertTrue((double) hit2DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); + assertEquals("rrf, rank_constant [60] normalization of:", hit2DetailsForHit1.get("description")); + assertEquals(1, ((List) hit2DetailsForHit1.get("details")).size()); + + Map hit2DetailsForHit2 = hit2Details.get(1); + assertTrue((double) hit2DetailsForHit2.get("value") > DELTA_FOR_SCORE_ASSERTION); + assertEquals("rrf, rank_constant [60] normalization of:", hit2DetailsForHit2.get("description")); + assertEquals(1, ((List) hit2DetailsForHit2.get("details")).size()); + + // hit 3 + Map searchHit3 = hitsNestedList.get(2); + Map explanationForHit3 = getValueByKey(searchHit3, "_explanation"); + assertNotNull(explanationForHit3); + assertEquals((double) searchHit3.get("_score"), (double) explanationForHit3.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit3.get("description")); + List> hit3Details = getListOfValues(explanationForHit3, "details"); + assertEquals(1, hit3Details.size()); + + Map hit3DetailsForHit1 = hit3Details.get(0); + assertTrue((double) hit3DetailsForHit1.get("value") > .0f); + assertEquals("rrf, rank_constant [60] normalization of:", hit3DetailsForHit1.get("description")); + assertEquals(1, ((List) hit3DetailsForHit1.get("details")).size()); + + Map explanationsHit3 = getListOfValues(hit3DetailsForHit1, "details").get(0); + assertEquals("within top 10", explanationsHit3.get("description")); + assertEquals(0, getListOfValues(explanationsHit3, "details").size()); + assertTrue((double) explanationsHit3.get("value") > 0.0f); + + // hit 4 + Map searchHit4 = hitsNestedList.get(3); + Map explanationForHit4 = getValueByKey(searchHit4, "_explanation"); + assertNotNull(explanationForHit4); + assertEquals((double) searchHit4.get("_score"), (double) explanationForHit4.get("value"), DELTA_FOR_SCORE_ASSERTION); + + assertEquals(expectedTopLevelDescription, explanationForHit4.get("description")); + List> hit4Details = getListOfValues(explanationForHit4, "details"); + assertEquals(1, hit4Details.size()); + + Map hit4DetailsForHit1 = hit4Details.get(0); + assertTrue((double) hit4DetailsForHit1.get("value") > DELTA_FOR_SCORE_ASSERTION); + assertEquals("rrf, rank_constant [60] normalization of:", hit4DetailsForHit1.get("description")); + assertEquals(1, ((List) hit4DetailsForHit1.get("details")).size()); + + Map explanationsHit4 = getListOfValues(hit4DetailsForHit1, "details").get(0); + assertEquals("ConstantScore(FieldExistsQuery [field=test-text-field-1])", explanationsHit4.get("description")); + assertEquals(0, getListOfValues(explanationsHit4, "details").size()); + assertTrue((double) explanationsHit4.get("value") > 0.0f); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java index 91c9abb20..fa01fe86f 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java @@ -133,50 +133,46 @@ protected boolean preserveClusterUponCompletion() { */ @SneakyThrows public void testComplexQuery_whenMultipleSubqueries_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); + initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + boolQueryBuilder.should(termQueryBuilder2).should(termQueryBuilder3); - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(boolQueryBuilder); - Map searchResponseAsMap1 = search( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - hybridQueryBuilderNeuralThenTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + Map searchResponseAsMap1 = search( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + hybridQueryBuilderNeuralThenTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertEquals(3, getHitCount(searchResponseAsMap1)); + assertEquals(3, getHitCount(searchResponseAsMap1)); - List> hits1NestedList = getNestedHits(searchResponseAsMap1); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hits1NestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } + List> hits1NestedList = getNestedHits(searchResponseAsMap1); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hits1NestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); + } - // verify that scores are in desc order - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - // verify that all ids are unique - assertEquals(Set.copyOf(ids).size(), ids.size()); + // verify that scores are in desc order + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + // verify that all ids are unique + assertEquals(Set.copyOf(ids).size(), ids.size()); - Map total = getTotalHits(searchResponseAsMap1); - assertNotNull(total.get("value")); - assertEquals(3, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - } finally { - wipeOfTestResources(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, null, null, SEARCH_PIPELINE); - } + Map total = getTotalHits(searchResponseAsMap1); + assertNotNull(total.get("value")); + assertEquals(3, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); } @SneakyThrows @@ -266,363 +262,333 @@ public void testMaxScoreCalculation_whenMaxScoreIsTrackedAtCollectorLevel_thenSu */ @SneakyThrows public void testComplexQuery_whenMultipleIdenticalSubQueries_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - HybridQueryBuilder hybridQueryBuilderThreeTerms = new HybridQueryBuilder(); - hybridQueryBuilderThreeTerms.add(termQueryBuilder1); - hybridQueryBuilderThreeTerms.add(termQueryBuilder2); - hybridQueryBuilderThreeTerms.add(termQueryBuilder3); + HybridQueryBuilder hybridQueryBuilderThreeTerms = new HybridQueryBuilder(); + hybridQueryBuilderThreeTerms.add(termQueryBuilder1); + hybridQueryBuilderThreeTerms.add(termQueryBuilder2); + hybridQueryBuilderThreeTerms.add(termQueryBuilder3); - Map searchResponseAsMap1 = search( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - hybridQueryBuilderThreeTerms, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + Map searchResponseAsMap1 = search( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + hybridQueryBuilderThreeTerms, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertEquals(2, getHitCount(searchResponseAsMap1)); + assertEquals(2, getHitCount(searchResponseAsMap1)); - List> hits1NestedList = getNestedHits(searchResponseAsMap1); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hits1NestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } + List> hits1NestedList = getNestedHits(searchResponseAsMap1); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hits1NestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); + } - // verify that scores are in desc order - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - // verify that all ids are unique - assertEquals(Set.copyOf(ids).size(), ids.size()); + // verify that scores are in desc order + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + // verify that all ids are unique + assertEquals(Set.copyOf(ids).size(), ids.size()); - Map total = getTotalHits(searchResponseAsMap1); - assertNotNull(total.get("value")); - assertEquals(2, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - } finally { - wipeOfTestResources(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, null, null, SEARCH_PIPELINE); - } + Map total = getTotalHits(searchResponseAsMap1); + assertNotNull(total.get("value")); + assertEquals(2, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); } @SneakyThrows public void testNoMatchResults_whenOnlyTermSubQueryWithoutMatch_thenEmptyResult() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); - TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); - HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); - hybridQueryBuilderOnlyTerm.add(termQueryBuilder); - hybridQueryBuilderOnlyTerm.add(termQuery2Builder); + initializeIndexIfNotExist(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); + HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); + hybridQueryBuilderOnlyTerm.add(termQueryBuilder); + hybridQueryBuilderOnlyTerm.add(termQuery2Builder); - Map searchResponseAsMap = search( - TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, - hybridQueryBuilderOnlyTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + Map searchResponseAsMap = search( + TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, + hybridQueryBuilderOnlyTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertEquals(0, getHitCount(searchResponseAsMap)); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - assertEquals(0.0f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); + assertEquals(0, getHitCount(searchResponseAsMap)); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + assertEquals(0.0f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(0, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_WITH_NESTED_FIELDS_INDEX_NAME, null, null, SEARCH_PIPELINE); - } + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(0, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); } @SneakyThrows public void testNestedQuery_whenHybridQueryIsWrappedIntoOtherQuery_thenFail() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - MatchQueryBuilder matchQuery2Builder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); - hybridQueryBuilderOnlyTerm.add(matchQueryBuilder); - hybridQueryBuilderOnlyTerm.add(matchQuery2Builder); - MatchQueryBuilder matchQuery3Builder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().should(hybridQueryBuilderOnlyTerm).should(matchQuery3Builder); - - ResponseException exceptionNoNestedTypes = expectThrows( - ResponseException.class, - () -> search(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, boolQueryBuilder, null, 10, Map.of("search_pipeline", SEARCH_PIPELINE)) - ); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + MatchQueryBuilder matchQuery2Builder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); + hybridQueryBuilderOnlyTerm.add(matchQueryBuilder); + hybridQueryBuilderOnlyTerm.add(matchQuery2Builder); + MatchQueryBuilder matchQuery3Builder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().should(hybridQueryBuilderOnlyTerm).should(matchQuery3Builder); + + ResponseException exceptionNoNestedTypes = expectThrows( + ResponseException.class, + () -> search(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, boolQueryBuilder, null, 10, Map.of("search_pipeline", SEARCH_PIPELINE)) + ); - org.hamcrest.MatcherAssert.assertThat( - exceptionNoNestedTypes.getMessage(), - allOf( - containsString("hybrid query must be a top level query and cannot be wrapped into other queries"), - containsString("illegal_argument_exception") - ) - ); + org.hamcrest.MatcherAssert.assertThat( + exceptionNoNestedTypes.getMessage(), + allOf( + containsString("hybrid query must be a top level query and cannot be wrapped into other queries"), + containsString("illegal_argument_exception") + ) + ); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); - ResponseException exceptionQWithNestedTypes = expectThrows( - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, - boolQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ) - ); + ResponseException exceptionQWithNestedTypes = expectThrows( + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, + boolQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ) + ); - org.hamcrest.MatcherAssert.assertThat( - exceptionQWithNestedTypes.getMessage(), - allOf( - containsString("hybrid query must be a top level query and cannot be wrapped into other queries"), - containsString("illegal_argument_exception") - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + org.hamcrest.MatcherAssert.assertThat( + exceptionQWithNestedTypes.getMessage(), + allOf( + containsString("hybrid query must be a top level query and cannot be wrapped into other queries"), + containsString("illegal_argument_exception") + ) + ); } @SneakyThrows public void testIndexWithNestedFields_whenHybridQuery_thenSuccess() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); - HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); - hybridQueryBuilderOnlyTerm.add(termQueryBuilder); - hybridQueryBuilderOnlyTerm.add(termQuery2Builder); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); + HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); + hybridQueryBuilderOnlyTerm.add(termQueryBuilder); + hybridQueryBuilderOnlyTerm.add(termQuery2Builder); - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, - hybridQueryBuilderOnlyTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, + hybridQueryBuilderOnlyTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertEquals(1, getHitCount(searchResponseAsMap)); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1, getHitCount(searchResponseAsMap)); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(1, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(1, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); } @SneakyThrows public void testIndexWithNestedFields_whenHybridQueryIncludesNested_thenSuccess() { - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); - NestedQueryBuilder nestedQueryBuilder = QueryBuilders.nestedQuery( - TEST_NESTED_TYPE_FIELD_NAME_1, - matchQuery(TEST_NESTED_TYPE_FIELD_NAME_1 + "." + NESTED_FIELD_1, NESTED_FIELD_1_VALUE), - ScoreMode.Total - ); - HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); - hybridQueryBuilderOnlyTerm.add(termQueryBuilder); - hybridQueryBuilderOnlyTerm.add(nestedQueryBuilder); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + NestedQueryBuilder nestedQueryBuilder = QueryBuilders.nestedQuery( + TEST_NESTED_TYPE_FIELD_NAME_1, + matchQuery(TEST_NESTED_TYPE_FIELD_NAME_1 + "." + NESTED_FIELD_1, NESTED_FIELD_1_VALUE), + ScoreMode.Total + ); + HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder(); + hybridQueryBuilderOnlyTerm.add(termQueryBuilder); + hybridQueryBuilderOnlyTerm.add(nestedQueryBuilder); - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, - hybridQueryBuilderOnlyTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE) - ); + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, + hybridQueryBuilderOnlyTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE) + ); - assertEquals(1, getHitCount(searchResponseAsMap)); - assertTrue(getMaxScore(searchResponseAsMap).isPresent()); - assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1, getHitCount(searchResponseAsMap)); + assertTrue(getMaxScore(searchResponseAsMap).isPresent()); + assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION); - Map total = getTotalHits(searchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(1, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + Map total = getTotalHits(searchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(1, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); } @SneakyThrows public void testRequestCache_whenOneShardAndQueryReturnResults_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_INDEX_WITH_KEYWORDS_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(KEYWORD_FIELD_1, KEYWORD_FIELD_2_VALUE); - RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(INTEGER_FIELD_PRICE).gte(10).lte(1000); - - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(matchQueryBuilder); - hybridQueryBuilder.add(rangeQueryBuilder); - - // first query with cache flag executed normally by reading documents from index - Map firstSearchResponseAsMap = search( - TEST_INDEX_WITH_KEYWORDS_ONE_SHARD, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) - ); + initializeIndexIfNotExist(TEST_INDEX_WITH_KEYWORDS_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(KEYWORD_FIELD_1, KEYWORD_FIELD_2_VALUE); + RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(INTEGER_FIELD_PRICE).gte(10).lte(1000); - int firstQueryHitCount = getHitCount(firstSearchResponseAsMap); - assertTrue(firstQueryHitCount > 0); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(matchQueryBuilder); + hybridQueryBuilder.add(rangeQueryBuilder); - List> hitsNestedList = getNestedHits(firstSearchResponseAsMap); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } + // first query with cache flag executed normally by reading documents from index + Map firstSearchResponseAsMap = search( + TEST_INDEX_WITH_KEYWORDS_ONE_SHARD, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) + ); - // verify that scores are in desc order - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - // verify that all ids are unique - assertEquals(Set.copyOf(ids).size(), ids.size()); + int firstQueryHitCount = getHitCount(firstSearchResponseAsMap); + assertTrue(firstQueryHitCount > 0); - Map total = getTotalHits(firstSearchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(firstQueryHitCount, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); + List> hitsNestedList = getNestedHits(firstSearchResponseAsMap); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); + } - // second query is served from the cache - Map secondSearchResponseAsMap = search( - TEST_INDEX_WITH_KEYWORDS_ONE_SHARD, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) - ); + // verify that scores are in desc order + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + // verify that all ids are unique + assertEquals(Set.copyOf(ids).size(), ids.size()); - assertEquals(firstQueryHitCount, getHitCount(secondSearchResponseAsMap)); + Map total = getTotalHits(firstSearchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(firstQueryHitCount, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); - List> hitsNestedListSecondQuery = getNestedHits(secondSearchResponseAsMap); - List idsSecondQuery = new ArrayList<>(); - List scoresSecondQuery = new ArrayList<>(); - for (Map oneHit : hitsNestedListSecondQuery) { - idsSecondQuery.add((String) oneHit.get("_id")); - scoresSecondQuery.add((Double) oneHit.get("_score")); - } + // second query is served from the cache + Map secondSearchResponseAsMap = search( + TEST_INDEX_WITH_KEYWORDS_ONE_SHARD, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) + ); - // verify that scores are in desc order - assertTrue( - IntStream.range(0, scoresSecondQuery.size() - 1) - .noneMatch(idx -> scoresSecondQuery.get(idx) < scoresSecondQuery.get(idx + 1)) - ); - // verify that all ids are unique - assertEquals(Set.copyOf(idsSecondQuery).size(), idsSecondQuery.size()); + assertEquals(firstQueryHitCount, getHitCount(secondSearchResponseAsMap)); - Map totalSecondQuery = getTotalHits(secondSearchResponseAsMap); - assertNotNull(totalSecondQuery.get("value")); - assertEquals(firstQueryHitCount, totalSecondQuery.get("value")); - assertNotNull(totalSecondQuery.get("relation")); - assertEquals(RELATION_EQUAL_TO, totalSecondQuery.get("relation")); - } finally { - wipeOfTestResources(TEST_INDEX_WITH_KEYWORDS_ONE_SHARD, null, null, SEARCH_PIPELINE); + List> hitsNestedListSecondQuery = getNestedHits(secondSearchResponseAsMap); + List idsSecondQuery = new ArrayList<>(); + List scoresSecondQuery = new ArrayList<>(); + for (Map oneHit : hitsNestedListSecondQuery) { + idsSecondQuery.add((String) oneHit.get("_id")); + scoresSecondQuery.add((Double) oneHit.get("_score")); } + + // verify that scores are in desc order + assertTrue( + IntStream.range(0, scoresSecondQuery.size() - 1).noneMatch(idx -> scoresSecondQuery.get(idx) < scoresSecondQuery.get(idx + 1)) + ); + // verify that all ids are unique + assertEquals(Set.copyOf(idsSecondQuery).size(), idsSecondQuery.size()); + + Map totalSecondQuery = getTotalHits(secondSearchResponseAsMap); + assertNotNull(totalSecondQuery.get("value")); + assertEquals(firstQueryHitCount, totalSecondQuery.get("value")); + assertNotNull(totalSecondQuery.get("relation")); + assertEquals(RELATION_EQUAL_TO, totalSecondQuery.get("relation")); } @SneakyThrows public void testRequestCache_whenMultipleShardsQueryReturnResults_thenSuccessful() { - try { - initializeIndexIfNotExist(TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(KEYWORD_FIELD_1, KEYWORD_FIELD_2_VALUE); - RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(INTEGER_FIELD_PRICE).gte(10).lte(1000); - - HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); - hybridQueryBuilder.add(matchQueryBuilder); - hybridQueryBuilder.add(rangeQueryBuilder); - - // first query with cache flag executed normally by reading documents from index - Map firstSearchResponseAsMap = search( - TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) - ); + initializeIndexIfNotExist(TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(KEYWORD_FIELD_1, KEYWORD_FIELD_2_VALUE); + RangeQueryBuilder rangeQueryBuilder = QueryBuilders.rangeQuery(INTEGER_FIELD_PRICE).gte(10).lte(1000); - int firstQueryHitCount = getHitCount(firstSearchResponseAsMap); - assertTrue(firstQueryHitCount > 0); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(matchQueryBuilder); + hybridQueryBuilder.add(rangeQueryBuilder); - List> hitsNestedList = getNestedHits(firstSearchResponseAsMap); - List ids = new ArrayList<>(); - List scores = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - ids.add((String) oneHit.get("_id")); - scores.add((Double) oneHit.get("_score")); - } + // first query with cache flag executed normally by reading documents from index + Map firstSearchResponseAsMap = search( + TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) + ); - // verify that scores are in desc order - assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); - // verify that all ids are unique - assertEquals(Set.copyOf(ids).size(), ids.size()); + int firstQueryHitCount = getHitCount(firstSearchResponseAsMap); + assertTrue(firstQueryHitCount > 0); - Map total = getTotalHits(firstSearchResponseAsMap); - assertNotNull(total.get("value")); - assertEquals(firstQueryHitCount, total.get("value")); - assertNotNull(total.get("relation")); - assertEquals(RELATION_EQUAL_TO, total.get("relation")); + List> hitsNestedList = getNestedHits(firstSearchResponseAsMap); + List ids = new ArrayList<>(); + List scores = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + ids.add((String) oneHit.get("_id")); + scores.add((Double) oneHit.get("_score")); + } - // second query is served from the cache - Map secondSearchResponseAsMap = search( - TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) - ); + // verify that scores are in desc order + assertTrue(IntStream.range(0, scores.size() - 1).noneMatch(idx -> scores.get(idx) < scores.get(idx + 1))); + // verify that all ids are unique + assertEquals(Set.copyOf(ids).size(), ids.size()); - assertEquals(firstQueryHitCount, getHitCount(secondSearchResponseAsMap)); + Map total = getTotalHits(firstSearchResponseAsMap); + assertNotNull(total.get("value")); + assertEquals(firstQueryHitCount, total.get("value")); + assertNotNull(total.get("relation")); + assertEquals(RELATION_EQUAL_TO, total.get("relation")); - List> hitsNestedListSecondQuery = getNestedHits(secondSearchResponseAsMap); - List idsSecondQuery = new ArrayList<>(); - List scoresSecondQuery = new ArrayList<>(); - for (Map oneHit : hitsNestedListSecondQuery) { - idsSecondQuery.add((String) oneHit.get("_id")); - scoresSecondQuery.add((Double) oneHit.get("_score")); - } + // second query is served from the cache + Map secondSearchResponseAsMap = search( + TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE, "request_cache", Boolean.TRUE.toString()) + ); - // verify that scores are in desc order - assertTrue( - IntStream.range(0, scoresSecondQuery.size() - 1) - .noneMatch(idx -> scoresSecondQuery.get(idx) < scoresSecondQuery.get(idx + 1)) - ); - // verify that all ids are unique - assertEquals(Set.copyOf(idsSecondQuery).size(), idsSecondQuery.size()); + assertEquals(firstQueryHitCount, getHitCount(secondSearchResponseAsMap)); - Map totalSecondQuery = getTotalHits(secondSearchResponseAsMap); - assertNotNull(totalSecondQuery.get("value")); - assertEquals(firstQueryHitCount, totalSecondQuery.get("value")); - assertNotNull(totalSecondQuery.get("relation")); - assertEquals(RELATION_EQUAL_TO, totalSecondQuery.get("relation")); - } finally { - wipeOfTestResources(TEST_INDEX_WITH_KEYWORDS_THREE_SHARDS, null, null, SEARCH_PIPELINE); + List> hitsNestedListSecondQuery = getNestedHits(secondSearchResponseAsMap); + List idsSecondQuery = new ArrayList<>(); + List scoresSecondQuery = new ArrayList<>(); + for (Map oneHit : hitsNestedListSecondQuery) { + idsSecondQuery.add((String) oneHit.get("_id")); + scoresSecondQuery.add((Double) oneHit.get("_score")); } + + // verify that scores are in desc order + assertTrue( + IntStream.range(0, scoresSecondQuery.size() - 1).noneMatch(idx -> scoresSecondQuery.get(idx) < scoresSecondQuery.get(idx + 1)) + ); + // verify that all ids are unique + assertEquals(Set.copyOf(idsSecondQuery).size(), idsSecondQuery.size()); + + Map totalSecondQuery = getTotalHits(secondSearchResponseAsMap); + assertNotNull(totalSecondQuery.get("value")); + assertEquals(firstQueryHitCount, totalSecondQuery.get("value")); + assertNotNull(totalSecondQuery.get("relation")); + assertEquals(RELATION_EQUAL_TO, totalSecondQuery.get("relation")); } @SneakyThrows @@ -658,7 +624,6 @@ 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, null, SEARCH_PIPELINE); } } @@ -695,7 +660,6 @@ 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, null, SEARCH_PIPELINE); } } @@ -741,7 +705,6 @@ public void testConcurrentSearchWithMultipleSlices_whenSingleShardIndex_thenSucc assertNotNull(total.get("relation")); assertEquals(RELATION_EQUAL_TO, total.get("relation")); } finally { - wipeOfTestResources(TEST_INDEX_DOC_QTY_ONE_SHARD, null, null, SEARCH_PIPELINE); updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); } } @@ -789,57 +752,40 @@ public void testConcurrentSearchWithMultipleSlices_whenMultipleShardsIndex_thenS assertNotNull(total.get("relation")); assertEquals(RELATION_EQUAL_TO, total.get("relation")); } finally { - wipeOfTestResources(TEST_INDEX_DOC_QTY_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); } } @SneakyThrows public void testPaginationOnSingleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); } @SneakyThrows public void testPaginationOnSingleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); } @SneakyThrows public void testPaginationOnMultipleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME); } @SneakyThrows public void testPaginationOnMultipleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSuccessful(TEST_MULTI_DOC_INDEX_NAME); } @SneakyThrows @@ -872,74 +818,66 @@ public void testHybridQuery_whenFromAndPaginationDepthIsGreaterThanZero_thenSucc @SneakyThrows public void testHybridQuery_whenFromIsGreaterThanTotalResultCount_thenFail() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - HybridQueryBuilder hybridQueryBuilderOnlyMatchAll = new HybridQueryBuilder(); - hybridQueryBuilderOnlyMatchAll.add(new MatchAllQueryBuilder()); - hybridQueryBuilderOnlyMatchAll.paginationDepth(10); - - ResponseException responseException = assertThrows( - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, - hybridQueryBuilderOnlyMatchAll, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 5 - ) - ); + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + HybridQueryBuilder hybridQueryBuilderOnlyMatchAll = new HybridQueryBuilder(); + hybridQueryBuilderOnlyMatchAll.add(new MatchAllQueryBuilder()); + hybridQueryBuilderOnlyMatchAll.paginationDepth(10); - org.hamcrest.MatcherAssert.assertThat( - responseException.getMessage(), - allOf(containsString("Reached end of search result, increase pagination_depth value to see more results")) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + ResponseException responseException = assertThrows( + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, + hybridQueryBuilderOnlyMatchAll, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 5 + ) + ); + + org.hamcrest.MatcherAssert.assertThat( + responseException.getMessage(), + allOf(containsString("Reached end of search result, increase pagination_depth value to see more results")) + ); } @SneakyThrows public void testHybridQuery_whenPaginationDepthIsOutOfRange_thenFail() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); - createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - HybridQueryBuilder hybridQueryBuilderOnlyMatchAll = new HybridQueryBuilder(); - hybridQueryBuilderOnlyMatchAll.add(new MatchAllQueryBuilder()); - hybridQueryBuilderOnlyMatchAll.paginationDepth(100001); - - ResponseException responseException = assertThrows( - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, - hybridQueryBuilderOnlyMatchAll, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - null, - false, - null, - 0 - ) - ); + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + HybridQueryBuilder hybridQueryBuilderOnlyMatchAll = new HybridQueryBuilder(); + hybridQueryBuilderOnlyMatchAll.add(new MatchAllQueryBuilder()); + hybridQueryBuilderOnlyMatchAll.paginationDepth(100001); - org.hamcrest.MatcherAssert.assertThat( - responseException.getMessage(), - allOf(containsString("pagination_depth should be less than or equal to index.max_result_window setting")) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, null, null, SEARCH_PIPELINE); - } + ResponseException responseException = assertThrows( + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_NAME_ONE_SHARD, + hybridQueryBuilderOnlyMatchAll, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ) + ); + + org.hamcrest.MatcherAssert.assertThat( + responseException.getMessage(), + allOf(containsString("pagination_depth should be less than or equal to index.max_result_window setting")) + ); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryPostFilterIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryPostFilterIT.java index 82b8ef6ec..0f9b08434 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryPostFilterIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryPostFilterIT.java @@ -67,54 +67,38 @@ public static void setUpCluster() { @SneakyThrows public void testPostFilterOnIndexWithSingleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testPostFilterOnIndexWithSingleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testPostFilterOnIndexWithMultipleShards_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } @SneakyThrows public void testPostFilterOnIndexWithMultipleShards_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testPostFilterRangeQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testPostFilterBoolQuery(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testPostFilterMatchAllAndMatchNoneQueries(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } /*{ diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java index 875c66310..f9cad4092 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java @@ -71,53 +71,37 @@ public static void setUpCluster() { @SneakyThrows public void testSortOnSingleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testSortOnSingleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testSortOnMultipleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } @SneakyThrows public void testSortOnMultipleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testSingleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testScoreSort_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } @SneakyThrows @@ -182,119 +166,95 @@ private void testMultipleFieldSort_whenMultipleSubQueries_thenSuccessful(String @SneakyThrows public void testSingleFieldSort_whenTrackScoresIsEnabled_thenFail() { - try { - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( - "mission", - "part", - LTE_OF_RANGE_IN_HYBRID_QUERY, - GTE_OF_RANGE_IN_HYBRID_QUERY - ); - Map fieldSortOrderMap = new HashMap<>(); - fieldSortOrderMap.put("stock", SortOrder.DESC); - assertThrows( - "Hybrid search results when sorted by any field, docId or _id, track_scores must be set to false.", - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - createSortBuilders(fieldSortOrderMap, false), - true, - null, - 0 - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( + "mission", + "part", + LTE_OF_RANGE_IN_HYBRID_QUERY, + GTE_OF_RANGE_IN_HYBRID_QUERY + ); + Map fieldSortOrderMap = new HashMap<>(); + fieldSortOrderMap.put("stock", SortOrder.DESC); + assertThrows( + "Hybrid search results when sorted by any field, docId or _id, track_scores must be set to false.", + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + createSortBuilders(fieldSortOrderMap, false), + true, + null, + 0 + ) + ); } @SneakyThrows public void testSingleFieldSort_whenSortCriteriaIsByScoreAndField_thenFail() { - try { - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( - "mission", - "part", - LTE_OF_RANGE_IN_HYBRID_QUERY, - GTE_OF_RANGE_IN_HYBRID_QUERY - ); - Map fieldSortOrderMap = new LinkedHashMap<>(); - fieldSortOrderMap.put("stock", SortOrder.DESC); - fieldSortOrderMap.put("_score", SortOrder.DESC); - assertThrows( - "_score sort criteria cannot be applied with any other criteria. Please select one sort criteria out of them.", - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - createSortBuilders(fieldSortOrderMap, false), - true, - null, - 0 - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( + "mission", + "part", + LTE_OF_RANGE_IN_HYBRID_QUERY, + GTE_OF_RANGE_IN_HYBRID_QUERY + ); + Map fieldSortOrderMap = new LinkedHashMap<>(); + fieldSortOrderMap.put("stock", SortOrder.DESC); + fieldSortOrderMap.put("_score", SortOrder.DESC); + assertThrows( + "_score sort criteria cannot be applied with any other criteria. Please select one sort criteria out of them.", + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + createSortBuilders(fieldSortOrderMap, false), + true, + null, + 0 + ) + ); } @SneakyThrows public void testSearchAfterWithSortOnSingleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testSearchAfterWithSortOnSingleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); - testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_SINGLE_NODE_CLUSTER); + testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); + testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD); } @SneakyThrows public void testSearchAfterWithSortOnMultipleShard_whenConcurrentSearchEnabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, true); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } @SneakyThrows public void testSearchAfterWithSortOnMultipleShard_whenConcurrentSearchDisabled_thenSuccessful() { - try { - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + testSearchAfter_whenSingleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); + testSearchAfter_whenMultipleFieldSort_thenSuccessful(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); } @SneakyThrows @@ -399,75 +359,67 @@ private void testScoreSort_whenSingleFieldSort_thenSuccessful(String indexName) @SneakyThrows public void testSort_whenSortFieldsSizeNotEqualToSearchAfterSize_thenFail() { - try { - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( - "mission", - "part", - LTE_OF_RANGE_IN_HYBRID_QUERY, - GTE_OF_RANGE_IN_HYBRID_QUERY - ); - Map fieldSortOrderMap = new LinkedHashMap<>(); - fieldSortOrderMap.put("stock", SortOrder.DESC); - List searchAfter = new ArrayList<>(); - searchAfter.add(25); - searchAfter.add(0); - assertThrows( - "after.fields has 2 values but sort has 1", - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - createSortBuilders(fieldSortOrderMap, false), - true, - searchAfter, - 0 - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( + "mission", + "part", + LTE_OF_RANGE_IN_HYBRID_QUERY, + GTE_OF_RANGE_IN_HYBRID_QUERY + ); + Map fieldSortOrderMap = new LinkedHashMap<>(); + fieldSortOrderMap.put("stock", SortOrder.DESC); + List searchAfter = new ArrayList<>(); + searchAfter.add(25); + searchAfter.add(0); + assertThrows( + "after.fields has 2 values but sort has 1", + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + createSortBuilders(fieldSortOrderMap, false), + true, + searchAfter, + 0 + ) + ); } @SneakyThrows public void testSearchAfter_whenAfterFieldIsNotPassed_thenFail() { - try { - prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( - "mission", - "part", - LTE_OF_RANGE_IN_HYBRID_QUERY, - GTE_OF_RANGE_IN_HYBRID_QUERY - ); - Map fieldSortOrderMap = new LinkedHashMap<>(); - fieldSortOrderMap.put("stock", SortOrder.DESC); - List searchAfter = new ArrayList<>(); - searchAfter.add(null); - assertThrows( - "after.fields wasn't set; you must pass fillFields=true for the previous search", - ResponseException.class, - () -> search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - null, - null, - createSortBuilders(fieldSortOrderMap, false), - true, - searchAfter, - 0 - ) - ); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResourcesBeforeTestExecution(SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( + "mission", + "part", + LTE_OF_RANGE_IN_HYBRID_QUERY, + GTE_OF_RANGE_IN_HYBRID_QUERY + ); + Map fieldSortOrderMap = new LinkedHashMap<>(); + fieldSortOrderMap.put("stock", SortOrder.DESC); + List searchAfter = new ArrayList<>(); + searchAfter.add(null); + assertThrows( + "after.fields wasn't set; you must pass fillFields=true for the previous search", + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + createSortBuilders(fieldSortOrderMap, false), + true, + searchAfter, + 0 + ) + ); } @SneakyThrows @@ -528,144 +480,135 @@ public void testSortingWithRescoreWhenConcurrentSegmentSearchEnabledAndDisabled_ ) ); } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); } } @SneakyThrows public void testExplainAndSort_whenIndexWithMultipleShards_thenSuccessful() { - try { - // Setup - updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + // Setup + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); - initializeIndexIfNotExists(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); - createSearchPipeline(SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); - // Assert - // scores for search hits - HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( - "mission", - "part", - LTE_OF_RANGE_IN_HYBRID_QUERY, - GTE_OF_RANGE_IN_HYBRID_QUERY - ); + initializeIndexIfNotExists(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SHARDS_COUNT_IN_MULTI_NODE_CLUSTER); + createSearchPipeline(SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, DEFAULT_COMBINATION_METHOD, Map.of(), true); + // Assert + // scores for search hits + HybridQueryBuilder hybridQueryBuilder = createHybridQueryBuilderWithMatchTermAndRangeQuery( + "mission", + "part", + LTE_OF_RANGE_IN_HYBRID_QUERY, + GTE_OF_RANGE_IN_HYBRID_QUERY + ); - Map fieldSortOrderMap = new HashMap<>(); - fieldSortOrderMap.put("stock", SortOrder.DESC); + Map fieldSortOrderMap = new HashMap<>(); + fieldSortOrderMap.put("stock", SortOrder.DESC); - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilder, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()), - null, - null, - createSortBuilders(fieldSortOrderMap, false), - false, - null, - 0 - ); - List> nestedHits = validateHitsCountAndFetchNestedHits(searchResponseAsMap, 6, 6); - assertStockValueWithSortOrderInHybridQueryResults(nestedHits, SortOrder.DESC, LARGEST_STOCK_VALUE_IN_QUERY_RESULT, true, true); - - // explain - Map searchHit1 = nestedHits.get(0); - Map explanationForHit1 = (Map) searchHit1.get("_explanation"); - assertNotNull(explanationForHit1); - assertNull(searchHit1.get("_score")); - String expectedGeneralCombineScoreDescription = "arithmetic_mean combination of:"; - assertEquals(expectedGeneralCombineScoreDescription, explanationForHit1.get("description")); - List> hit1Details = getListOfValues(explanationForHit1, "details"); - assertEquals(2, hit1Details.size()); - Map hit1DetailsForHit1 = hit1Details.get(0); - assertEquals(1.0, hit1DetailsForHit1.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit1.get("description")); - List> hit1DetailsForHit1Details = getListOfValues(hit1DetailsForHit1, "details"); - assertEquals(1, hit1DetailsForHit1Details.size()); - - Map hit1DetailsForHit1DetailsForHit1 = hit1DetailsForHit1Details.get(0); - assertEquals("weight(name:mission in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit1.get("description")); - assertTrue((double) hit1DetailsForHit1DetailsForHit1.get("value") > 0.0f); - assertEquals(1, getListOfValues(hit1DetailsForHit1DetailsForHit1, "details").size()); - - Map hit1DetailsForHit1DetailsForHit1DetailsForHit1 = getListOfValues( - hit1DetailsForHit1DetailsForHit1, - "details" - ).get(0); - assertEquals( - "score(freq=1.0), computed as boost * idf * tf from:", - hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("description") - ); - assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("value") > 0.0f); - assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").size()); - - assertEquals("boost", getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("description")); - assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("value") > 0.0f); - assertEquals( - "idf, computed as log(1 + (N - n + 0.5) / (n + 0.5)) from:", - getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("description") - ); - assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("value") > 0.0f); - assertEquals( - "tf, computed as freq / (freq + k1 * (1 - b + b * dl / avgdl)) from:", - getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("description") - ); - assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("value") > 0.0f); - - // hit 4 - Map searchHit4 = nestedHits.get(3); - Map explanationForHit4 = (Map) searchHit4.get("_explanation"); - assertNotNull(explanationForHit4); - assertNull(searchHit4.get("_score")); - assertEquals(expectedGeneralCombineScoreDescription, explanationForHit4.get("description")); - List> hit4Details = getListOfValues(explanationForHit4, "details"); - assertEquals(2, hit4Details.size()); - Map hit1DetailsForHit4 = hit4Details.get(0); - assertEquals(1.0, hit1DetailsForHit4.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit4.get("description")); - assertEquals(1, ((List) hit1DetailsForHit4.get("details")).size()); - List> hit1DetailsForHit4Details = getListOfValues(hit1DetailsForHit4, "details"); - assertEquals(1, hit1DetailsForHit4Details.size()); - - Map hit1DetailsForHit1DetailsForHit4 = hit1DetailsForHit4Details.get(0); - assertEquals("weight(name:part in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit4.get("description")); - assertTrue((double) hit1DetailsForHit1DetailsForHit4.get("value") > 0.0f); - assertEquals(1, getListOfValues(hit1DetailsForHit1DetailsForHit4, "details").size()); - - Map hit1DetailsForHit1DetailsForHit1DetailsForHit4 = getListOfValues( - hit1DetailsForHit1DetailsForHit4, - "details" - ).get(0); - assertEquals( - "score(freq=1.0), computed as boost * idf * tf from:", - hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("description") - ); - assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("value") > 0.0f); - assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit4, "details").size()); - - // hit 6 - Map searchHit6 = nestedHits.get(5); - Map explanationForHit6 = (Map) searchHit6.get("_explanation"); - assertNotNull(explanationForHit6); - assertNull(searchHit6.get("_score")); - assertEquals(expectedGeneralCombineScoreDescription, explanationForHit6.get("description")); - List> hit6Details = getListOfValues(explanationForHit6, "details"); - assertEquals(1, hit6Details.size()); - Map hit1DetailsForHit6 = hit6Details.get(0); - assertEquals(1.0, hit1DetailsForHit6.get("value")); - assertEquals("min_max normalization of:", hit1DetailsForHit6.get("description")); - assertEquals(1, ((List) hit1DetailsForHit6.get("details")).size()); - List> hit1DetailsForHit6Details = getListOfValues(hit1DetailsForHit6, "details"); - assertEquals(1, hit1DetailsForHit6Details.size()); - - Map hit1DetailsForHit1DetailsForHit6 = hit1DetailsForHit6Details.get(0); - assertEquals("weight(name:part in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit4.get("description")); - assertTrue((double) hit1DetailsForHit1DetailsForHit6.get("value") > 0.0f); - assertEquals(0, getListOfValues(hit1DetailsForHit1DetailsForHit6, "details").size()); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE, "explain", Boolean.TRUE.toString()), + null, + null, + createSortBuilders(fieldSortOrderMap, false), + false, + null, + 0 + ); + List> nestedHits = validateHitsCountAndFetchNestedHits(searchResponseAsMap, 6, 6); + assertStockValueWithSortOrderInHybridQueryResults(nestedHits, SortOrder.DESC, LARGEST_STOCK_VALUE_IN_QUERY_RESULT, true, true); + + // explain + Map searchHit1 = nestedHits.get(0); + Map explanationForHit1 = (Map) searchHit1.get("_explanation"); + assertNotNull(explanationForHit1); + assertNull(searchHit1.get("_score")); + String expectedGeneralCombineScoreDescription = "arithmetic_mean combination of:"; + assertEquals(expectedGeneralCombineScoreDescription, explanationForHit1.get("description")); + List> hit1Details = getListOfValues(explanationForHit1, "details"); + assertEquals(2, hit1Details.size()); + Map hit1DetailsForHit1 = hit1Details.get(0); + assertEquals(1.0, hit1DetailsForHit1.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit1.get("description")); + List> hit1DetailsForHit1Details = getListOfValues(hit1DetailsForHit1, "details"); + assertEquals(1, hit1DetailsForHit1Details.size()); + + Map hit1DetailsForHit1DetailsForHit1 = hit1DetailsForHit1Details.get(0); + assertEquals("weight(name:mission in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit1.get("description")); + assertTrue((double) hit1DetailsForHit1DetailsForHit1.get("value") > 0.0f); + assertEquals(1, getListOfValues(hit1DetailsForHit1DetailsForHit1, "details").size()); + + Map hit1DetailsForHit1DetailsForHit1DetailsForHit1 = getListOfValues(hit1DetailsForHit1DetailsForHit1, "details") + .get(0); + assertEquals( + "score(freq=1.0), computed as boost * idf * tf from:", + hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("description") + ); + assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("value") > 0.0f); + assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").size()); + + assertEquals("boost", getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("description")); + assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("value") > 0.0f); + assertEquals( + "idf, computed as log(1 + (N - n + 0.5) / (n + 0.5)) from:", + getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("description") + ); + assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("value") > 0.0f); + assertEquals( + "tf, computed as freq / (freq + k1 * (1 - b + b * dl / avgdl)) from:", + getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("description") + ); + assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("value") > 0.0f); + + // hit 4 + Map searchHit4 = nestedHits.get(3); + Map explanationForHit4 = (Map) searchHit4.get("_explanation"); + assertNotNull(explanationForHit4); + assertNull(searchHit4.get("_score")); + assertEquals(expectedGeneralCombineScoreDescription, explanationForHit4.get("description")); + List> hit4Details = getListOfValues(explanationForHit4, "details"); + assertEquals(2, hit4Details.size()); + Map hit1DetailsForHit4 = hit4Details.get(0); + assertEquals(1.0, hit1DetailsForHit4.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit4.get("description")); + assertEquals(1, ((List) hit1DetailsForHit4.get("details")).size()); + List> hit1DetailsForHit4Details = getListOfValues(hit1DetailsForHit4, "details"); + assertEquals(1, hit1DetailsForHit4Details.size()); + + Map hit1DetailsForHit1DetailsForHit4 = hit1DetailsForHit4Details.get(0); + assertEquals("weight(name:part in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit4.get("description")); + assertTrue((double) hit1DetailsForHit1DetailsForHit4.get("value") > 0.0f); + assertEquals(1, getListOfValues(hit1DetailsForHit1DetailsForHit4, "details").size()); + + Map hit1DetailsForHit1DetailsForHit1DetailsForHit4 = getListOfValues(hit1DetailsForHit1DetailsForHit4, "details") + .get(0); + assertEquals( + "score(freq=1.0), computed as boost * idf * tf from:", + hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("description") + ); + assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("value") > 0.0f); + assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit4, "details").size()); + + // hit 6 + Map searchHit6 = nestedHits.get(5); + Map explanationForHit6 = (Map) searchHit6.get("_explanation"); + assertNotNull(explanationForHit6); + assertNull(searchHit6.get("_score")); + assertEquals(expectedGeneralCombineScoreDescription, explanationForHit6.get("description")); + List> hit6Details = getListOfValues(explanationForHit6, "details"); + assertEquals(1, hit6Details.size()); + Map hit1DetailsForHit6 = hit6Details.get(0); + assertEquals(1.0, hit1DetailsForHit6.get("value")); + assertEquals("min_max normalization of:", hit1DetailsForHit6.get("description")); + assertEquals(1, ((List) hit1DetailsForHit6.get("details")).size()); + List> hit1DetailsForHit6Details = getListOfValues(hit1DetailsForHit6, "details"); + assertEquals(1, hit1DetailsForHit6Details.size()); + + Map hit1DetailsForHit1DetailsForHit6 = hit1DetailsForHit6Details.get(0); + assertEquals("weight(name:part in 0) [PerFieldSimilarity], result of:", hit1DetailsForHit1DetailsForHit4.get("description")); + assertTrue((double) hit1DetailsForHit1DetailsForHit6.get("value") > 0.0f); + assertEquals(0, getListOfValues(hit1DetailsForHit1DetailsForHit6, "details").size()); } private HybridQueryBuilder createHybridQueryBuilderWithMatchTermAndRangeQuery(String text, String value, int lte, int gte) { diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java index 37d1320d4..de3313654 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java @@ -99,93 +99,81 @@ public void setUp() throws Exception { @SneakyThrows public void testQueryWithBoostAndImageQueryAndRadialQuery() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareModel(); - 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); - Map searchResponseAsMapTextQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderTextQuery, 1); - Map firstInnerHitTextQuery = getFirstInnerHit(searchResponseAsMapTextQuery); - - assertEquals("1", firstInnerHitTextQuery.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHitTextQuery.get("_score")), DELTA_FOR_SCORE_ASSERTION); - - 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); - - 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 - ); - - // To save test resources, IT tests for radial search are added below. - // Context: https://github.com/opensearch-project/neural-search/pull/697#discussion_r1571549776 - - // Test radial search max distance query - 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, - neuralQueryWithMaxDistanceBuilder, - 1 - ); - Map firstInnerHitWithMaxDistanceQuery = getFirstInnerHit(searchResponseAsMapWithMaxDistanceQuery); - - assertEquals("1", firstInnerHitWithMaxDistanceQuery.get("_id")); - float expectedScoreWithMaxDistanceQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals( - expectedScoreWithMaxDistanceQuery, - objectToFloat(firstInnerHitWithMaxDistanceQuery.get("_score")), - DELTA_FOR_SCORE_ASSERTION - ); - - // Test radial search min score query - 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); - - assertEquals("1", firstInnerHitWithMinScoreQuery.get("_id")); - float expectedScoreWithMinScoreQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertEquals( - expectedScoreWithMinScoreQuery, - objectToFloat(firstInnerHitWithMinScoreQuery.get("_score")), - DELTA_FOR_SCORE_ASSERTION - ); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + modelId = prepareModel(); + 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); + Map searchResponseAsMapTextQuery = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilderTextQuery, 1); + Map firstInnerHitTextQuery = getFirstInnerHit(searchResponseAsMapTextQuery); + + assertEquals("1", firstInnerHitTextQuery.get("_id")); + float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHitTextQuery.get("_score")), DELTA_FOR_SCORE_ASSERTION); + + 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); + + 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); + + // To save test resources, IT tests for radial search are added below. + // Context: https://github.com/opensearch-project/neural-search/pull/697#discussion_r1571549776 + + // Test radial search max distance query + 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, neuralQueryWithMaxDistanceBuilder, 1); + Map firstInnerHitWithMaxDistanceQuery = getFirstInnerHit(searchResponseAsMapWithMaxDistanceQuery); + + assertEquals("1", firstInnerHitWithMaxDistanceQuery.get("_id")); + float expectedScoreWithMaxDistanceQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertEquals( + expectedScoreWithMaxDistanceQuery, + objectToFloat(firstInnerHitWithMaxDistanceQuery.get("_score")), + DELTA_FOR_SCORE_ASSERTION + ); + + // Test radial search min score query + 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); + + assertEquals("1", firstInnerHitWithMinScoreQuery.get("_id")); + float expectedScoreWithMinScoreQuery = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertEquals( + expectedScoreWithMinScoreQuery, + objectToFloat(firstInnerHitWithMinScoreQuery.get("_score")), + DELTA_FOR_SCORE_ASSERTION + ); } /** @@ -211,26 +199,22 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { @SneakyThrows public void testRescoreQuery() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareModel(); - MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); - 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); - - 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); - } + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + modelId = prepareModel(); + MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); + 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); + + 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); } /** @@ -282,64 +266,60 @@ public void testRescoreQuery() { @SneakyThrows public void testBooleanQuery_withMultipleNeuralQueries() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME); - modelId = prepareModel(); - // verify two neural queries wrapped into bool - BoolQueryBuilder boolQueryBuilderTwoNeuralQueries = new BoolQueryBuilder(); - 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); - - Map searchResponseAsMapTwoNeuralQueries = search( - TEST_MULTI_VECTOR_FIELD_INDEX_NAME, - boolQueryBuilderTwoNeuralQueries, - 1 - ); - Map firstInnerHitTwoNeuralQueries = getFirstInnerHit(searchResponseAsMapTwoNeuralQueries); - - assertEquals("1", firstInnerHitTwoNeuralQueries.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - 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 = 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); - - boolQueryBuilderMixOfNeuralAndBM25.should(neuralQueryBuilder).should(matchQueryBuilder); - - Map searchResponseAsMapMixOfNeuralAndBM25 = search( - TEST_MULTI_VECTOR_FIELD_INDEX_NAME, - boolQueryBuilderMixOfNeuralAndBM25, - 1 - ); - Map firstInnerHitMixOfNeuralAndBM25 = getFirstInnerHit(searchResponseAsMapMixOfNeuralAndBM25); - - assertEquals("1", firstInnerHitMixOfNeuralAndBM25.get("_id")); - float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); - assertTrue(minExpectedScore < objectToFloat(firstInnerHitMixOfNeuralAndBM25.get("_score"))); - } finally { - wipeOfTestResources(TEST_MULTI_VECTOR_FIELD_INDEX_NAME, null, modelId, null); - } + initializeIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME); + modelId = prepareModel(); + // verify two neural queries wrapped into bool + BoolQueryBuilder boolQueryBuilderTwoNeuralQueries = new BoolQueryBuilder(); + 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); + + Map searchResponseAsMapTwoNeuralQueries = search( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + boolQueryBuilderTwoNeuralQueries, + 1 + ); + Map firstInnerHitTwoNeuralQueries = getFirstInnerHit(searchResponseAsMapTwoNeuralQueries); + + assertEquals("1", firstInnerHitTwoNeuralQueries.get("_id")); + float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + 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 = 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); + + boolQueryBuilderMixOfNeuralAndBM25.should(neuralQueryBuilder).should(matchQueryBuilder); + + Map searchResponseAsMapMixOfNeuralAndBM25 = search( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + boolQueryBuilderMixOfNeuralAndBM25, + 1 + ); + Map firstInnerHitMixOfNeuralAndBM25 = getFirstInnerHit(searchResponseAsMapMixOfNeuralAndBM25); + + assertEquals("1", firstInnerHitMixOfNeuralAndBM25.get("_id")); + float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertTrue(minExpectedScore < objectToFloat(firstInnerHitMixOfNeuralAndBM25.get("_score"))); } /** @@ -363,25 +343,21 @@ public void testBooleanQuery_withMultipleNeuralQueries() { @SneakyThrows public void testNestedQuery() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_NESTED_INDEX_NAME); - modelId = prepareModel(); - 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); - - 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_NESTED_INDEX_NAME, null, modelId, null); - } + initializeIndexIfNotExist(TEST_NESTED_INDEX_NAME); + modelId = prepareModel(); + 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); + + 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); } /** @@ -408,26 +384,22 @@ public void testNestedQuery() { @SneakyThrows public void testFilterQuery() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); - modelId = prepareModel(); - 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); - assertEquals("3", 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_MULTI_DOC_INDEX_NAME, null, modelId, null); - } + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + modelId = prepareModel(); + 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); + assertEquals("3", firstInnerHit.get("_id")); + float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA_FOR_SCORE_ASSERTION); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java index 4790169e5..7135bd09d 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java @@ -58,22 +58,18 @@ public void setUp() throws Exception { @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) - .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId) + .boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); - } + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -97,21 +93,17 @@ public void testBasicQueryUsingQueryText() { */ @SneakyThrows public void testBasicQueryUsingQueryTokens() { - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - Map queryTokens = createRandomTokenWeightMap(TEST_TOKENS); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryTokensSupplier(() -> queryTokens) - .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + Map queryTokens = createRandomTokenWeightMap(TEST_TOKENS); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryTokensSupplier(() -> queryTokens) + .boost(2.0f); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, sparseEncodingQueryBuilder.queryTokensSupplier().get()); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); - } + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, sparseEncodingQueryBuilder.queryTokensSupplier().get()); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -137,22 +129,18 @@ public void testBasicQueryUsingQueryTokens() { @SneakyThrows public void testRescoreQuery() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, matchAllQueryBuilder, sparseEncodingQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, matchAllQueryBuilder, 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); - } + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -181,29 +169,25 @@ public void testRescoreQuery() { @SneakyThrows public void testBooleanQuery_withMultipleSparseEncodingQueries() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId); - boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); + boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); - Map searchResponseAsMap = search(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + Map searchResponseAsMap = search(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); - } finally { - wipeOfTestResources(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, null, modelId, null); - } + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } /** @@ -229,45 +213,34 @@ public void testBooleanQuery_withMultipleSparseEncodingQueries() { @SneakyThrows public void testBooleanQuery_withSparseEncodingAndBM25Queries() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); - boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId); + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); - Map searchResponseAsMap = search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); - Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + Map searchResponseAsMap = search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); - assertEquals("1", firstInnerHit.get("_id")); - float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); - assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); - } finally { - wipeOfTestResources(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, null, modelId, null); - } + assertEquals("1", firstInnerHit.get("_id")); + float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); } @SneakyThrows public void testBasicQueryUsingQueryText_whenQueryWrongFieldType_thenFail() { String modelId = null; - try { - initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); - modelId = prepareSparseEncodingModel(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_TEXT_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId); + initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); + modelId = prepareSparseEncodingModel(); + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_TEXT_FIELD_NAME_1) + .queryText(TEST_QUERY_TEXT) + .modelId(modelId); - expectThrows( - ResponseException.class, - () -> search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, sparseEncodingQueryBuilder, 1) - ); - } finally { - wipeOfTestResources(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, null, modelId, null); - } + expectThrows(ResponseException.class, () -> search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, sparseEncodingQueryBuilder, 1)); } @SneakyThrows diff --git a/src/test/java/org/opensearch/neuralsearch/query/aggregation/BucketAggregationsWithHybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/aggregation/BucketAggregationsWithHybridQueryIT.java index 7385f48e5..be47ffb42 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/aggregation/BucketAggregationsWithHybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/aggregation/BucketAggregationsWithHybridQueryIT.java @@ -231,282 +231,256 @@ public void testWithConcurrentSegmentSearch_whenSignificantTermsAggs_thenSuccess } private void testAvgNestedIntoFilter() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.filter( - GENERIC_AGGREGATION_NAME, - QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).lte(3000) - ).subAggregation(AggregationBuilders.avg(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.filter( + GENERIC_AGGREGATION_NAME, + QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).lte(3000) + ).subAggregation(AggregationBuilders.avg(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - double avgValue = getAggregationValue(getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME), AVG_AGGREGATION_NAME); - assertEquals(1789.5, avgValue, DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + double avgValue = getAggregationValue(getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME), AVG_AGGREGATION_NAME); + assertEquals(1789.5, avgValue, DELTA_FOR_SCORE_ASSERTION); } private void testSumNestedIntoFilters() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.filters( - GENERIC_AGGREGATION_NAME, - QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).lte(3000), - QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE) - ).otherBucket(true).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.filters( + GENERIC_AGGREGATION_NAME, + QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).lte(3000), + QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE) + ).otherBucket(true).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(buckets); - assertEquals(3, buckets.size()); + List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(buckets); + assertEquals(3, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(2, firstBucket.size()); - assertEquals(2, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(3579.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + Map firstBucket = buckets.get(0); + assertEquals(2, firstBucket.size()); + assertEquals(2, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(3579.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - Map secondBucket = buckets.get(1); - assertEquals(2, secondBucket.size()); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(1234.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + Map secondBucket = buckets.get(1); + assertEquals(2, secondBucket.size()); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(1234.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - Map thirdBucket = buckets.get(2); - assertEquals(2, thirdBucket.size()); - assertEquals(1, thirdBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(3456.0, getAggregationValue(thirdBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.size()); + assertEquals(1, thirdBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(3456.0, getAggregationValue(thirdBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); } private void testGlobalAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); - AggregationBuilder aggsBuilder = AggregationBuilders.global(GENERIC_AGGREGATION_NAME) - .subAggregation(AggregationBuilders.sum(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + AggregationBuilder aggsBuilder = AggregationBuilders.global(GENERIC_AGGREGATION_NAME) + .subAggregation(AggregationBuilders.sum(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggsBuilder), - hybridQueryBuilderNeuralThenTerm, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggsBuilder), + hybridQueryBuilderNeuralThenTerm, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - double avgValue = getAggregationValue(getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME), AVG_AGGREGATION_NAME); - assertEquals(15058.0, avgValue, DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + double avgValue = getAggregationValue(getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME), AVG_AGGREGATION_NAME); + assertEquals(15058.0, avgValue, DELTA_FOR_SCORE_ASSERTION); } private void testHistogramAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.histogram(GENERIC_AGGREGATION_NAME) - .field(INTEGER_FIELD_PRICE) - .interval(100); + AggregationBuilder aggsBuilder = AggregationBuilders.histogram(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_PRICE).interval(100); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(buckets); - assertEquals(2, buckets.size()); + List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(buckets); + assertEquals(2, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(2, firstBucket.size()); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0, (Double) firstBucket.get(KEY), DELTA_FOR_SCORE_ASSERTION); + Map firstBucket = buckets.get(0); + assertEquals(2, firstBucket.size()); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0, (Double) firstBucket.get(KEY), DELTA_FOR_SCORE_ASSERTION); - Map secondBucket = buckets.get(1); - assertEquals(2, secondBucket.size()); - assertEquals(2, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(100.0, (Double) secondBucket.get(KEY), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map secondBucket = buckets.get(1); + assertEquals(2, secondBucket.size()); + assertEquals(2, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(100.0, (Double) secondBucket.get(KEY), DELTA_FOR_SCORE_ASSERTION); } private void testNestedAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.nested(GENERIC_AGGREGATION_NAME, NESTED_TYPE_FIELD_USER) - .subAggregation( - AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1) - .field(String.join(".", NESTED_TYPE_FIELD_USER, NESTED_FIELD_FIRSTNAME)) - ); - - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.nested(GENERIC_AGGREGATION_NAME, NESTED_TYPE_FIELD_USER) + .subAggregation( + AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1) + .field(String.join(".", NESTED_TYPE_FIELD_USER, NESTED_FIELD_FIRSTNAME)) ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - Map nestedAgg = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(nestedAgg); + Map nestedAgg = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(nestedAgg); - assertEquals(3, nestedAgg.get(BUCKET_AGG_DOC_COUNT_FIELD)); - List> buckets = getAggregationBuckets(nestedAgg, BUCKETS_AGGREGATION_NAME_1); + assertEquals(3, nestedAgg.get(BUCKET_AGG_DOC_COUNT_FIELD)); + List> buckets = getAggregationBuckets(nestedAgg, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(buckets); - assertEquals(3, buckets.size()); + assertNotNull(buckets); + assertEquals(3, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(2, firstBucket.size()); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(NESTED_FIELD_FIRSTNAME_FRODO, firstBucket.get(KEY)); + Map firstBucket = buckets.get(0); + assertEquals(2, firstBucket.size()); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(NESTED_FIELD_FIRSTNAME_FRODO, firstBucket.get(KEY)); - Map secondBucket = buckets.get(1); - assertEquals(2, secondBucket.size()); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(NESTED_FIELD_FIRSTNAME_JOHN, secondBucket.get(KEY)); + Map secondBucket = buckets.get(1); + assertEquals(2, secondBucket.size()); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(NESTED_FIELD_FIRSTNAME_JOHN, secondBucket.get(KEY)); - Map thirdBucket = buckets.get(2); - assertEquals(2, thirdBucket.size()); - assertEquals(1, thirdBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(NESTED_FIELD_FIRSTNAME_SUN, thirdBucket.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map thirdBucket = buckets.get(2); + assertEquals(2, thirdBucket.size()); + assertEquals(1, thirdBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(NESTED_FIELD_FIRSTNAME_SUN, thirdBucket.get(KEY)); } private void testDiversifiedSampler() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.diversifiedSampler(GENERIC_AGGREGATION_NAME) - .field(KEYWORD_FIELD_DOCKEYWORD) - .shardSize(2) - .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_DOCKEYWORD)); + AggregationBuilder aggsBuilder = AggregationBuilders.diversifiedSampler(GENERIC_AGGREGATION_NAME) + .field(KEYWORD_FIELD_DOCKEYWORD) + .shardSize(2) + .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_DOCKEYWORD)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertEquals(2, aggValue.size()); - assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); - Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(nestedAggs); - assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); - List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertEquals(2, buckets.size()); + Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertEquals(2, aggValue.size()); + assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); + Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertNotNull(nestedAggs); + assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); + List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertEquals(2, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("likeable", firstBucket.get(KEY)); + Map firstBucket = buckets.get(0); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("likeable", firstBucket.get(KEY)); - Map secondBucket = buckets.get(1); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("workable", secondBucket.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map secondBucket = buckets.get(1); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("workable", secondBucket.get(KEY)); } private void testAdjacencyMatrixAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.adjacencyMatrix( - GENERIC_AGGREGATION_NAME, - Map.of( - "grpA", - QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE), - "grpB", - QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY), - "grpC", - QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_LIKABLE) - ) - ); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.adjacencyMatrix( + GENERIC_AGGREGATION_NAME, + Map.of( + "grpA", + QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE), + "grpB", + QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY), + "grpC", + QueryBuilders.matchQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_LIKABLE) + ) + ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(buckets); - assertEquals(2, buckets.size()); - Map grpA = buckets.get(0); - assertEquals(1, grpA.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("grpA", grpA.get(KEY)); - Map grpC = buckets.get(1); - assertEquals(1, grpC.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("grpC", grpC.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(buckets); + assertEquals(2, buckets.size()); + Map grpA = buckets.get(0); + assertEquals(1, grpA.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("grpA", grpA.get(KEY)); + Map grpC = buckets.get(1); + assertEquals(1, grpC.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("grpC", grpC.get(KEY)); } private void testDateBucketedSumsPipelinedToBucketMinMaxSumAvgAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - BucketMetricsPipelineAggregationBuilder aggAvgBucket = PipelineAggregatorBuilders - .avgBucket(BUCKETS_AGGREGATION_NAME_1, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggAvgBucket = PipelineAggregatorBuilders.avgBucket( + BUCKETS_AGGREGATION_NAME_1, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggSumBucket = PipelineAggregatorBuilders - .sumBucket(BUCKETS_AGGREGATION_NAME_2, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggSumBucket = PipelineAggregatorBuilders.sumBucket( + BUCKETS_AGGREGATION_NAME_2, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggMinBucket = PipelineAggregatorBuilders - .minBucket(BUCKETS_AGGREGATION_NAME_3, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggMinBucket = PipelineAggregatorBuilders.minBucket( + BUCKETS_AGGREGATION_NAME_3, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - BucketMetricsPipelineAggregationBuilder aggMaxBucket = PipelineAggregatorBuilders - .maxBucket(BUCKETS_AGGREGATION_NAME_4, GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME); + BucketMetricsPipelineAggregationBuilder aggMaxBucket = PipelineAggregatorBuilders.maxBucket( + BUCKETS_AGGREGATION_NAME_4, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggDateHisto, aggAvgBucket, aggSumBucket, aggMinBucket, aggMaxBucket), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggDateHisto, aggAvgBucket, aggSumBucket, aggMinBucket, aggMaxBucket), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - assertResultsOfPipelineSumtoDateHistogramAggs(searchResponseAsMap); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + assertResultsOfPipelineSumtoDateHistogramAggs(searchResponseAsMap); } private void assertResultsOfPipelineSumtoDateHistogramAggs(Map searchResponseAsMap) { @@ -553,221 +527,199 @@ private void assertResultsOfPipelineSumtoDateHistogramAggs(Map s } private void testDateBucketedSumsPipelinedToBucketStatsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - StatsBucketPipelineAggregationBuilder aggStatsBucket = PipelineAggregatorBuilders.statsBucket( - BUCKETS_AGGREGATION_NAME_1, - GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME - ); + StatsBucketPipelineAggregationBuilder aggStatsBucket = PipelineAggregatorBuilders.statsBucket( + BUCKETS_AGGREGATION_NAME_1, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggDateHisto, aggStatsBucket), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggDateHisto, aggStatsBucket), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - Map statsAggs = getAggregationValues(aggregations, BUCKETS_AGGREGATION_NAME_1); + Map statsAggs = getAggregationValues(aggregations, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(statsAggs); + assertNotNull(statsAggs); - assertEquals(3517.5, (Double) statsAggs.get("avg"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(7035.0, (Double) statsAggs.get("sum"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, (Double) statsAggs.get("min"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(5801.0, (Double) statsAggs.get("max"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(2, (int) statsAggs.get("count")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + assertEquals(3517.5, (Double) statsAggs.get("avg"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(7035.0, (Double) statsAggs.get("sum"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, (Double) statsAggs.get("min"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(5801.0, (Double) statsAggs.get("max"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(2, (int) statsAggs.get("count")); } private void testDateBucketedSumsPipelinedToBucketScriptedAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregations( - new AggregatorFactories.Builder().addAggregator( - AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX) + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregations( + new AggregatorFactories.Builder().addAggregator(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)) + .addAggregator( + AggregationBuilders.filter( + GENERIC_AGGREGATION_NAME, + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) + ) + .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))) + ).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME_2).field(INTEGER_FIELD_PRICE)) ) - .addAggregator( - AggregationBuilders.filter( - GENERIC_AGGREGATION_NAME, - QueryBuilders.boolQuery() - .should( - QueryBuilders.boolQuery() - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) - ) - .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))) - ).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME_2).field(INTEGER_FIELD_PRICE)) - ) - .addPipelineAggregator( - PipelineAggregatorBuilders.bucketScript( - BUCKETS_AGGREGATION_NAME_1, - Map.of("docNum", GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME_2, "totalNum", SUM_AGGREGATION_NAME), - new Script("params.docNum / params.totalNum") - ) + .addPipelineAggregator( + PipelineAggregatorBuilders.bucketScript( + BUCKETS_AGGREGATION_NAME_1, + Map.of("docNum", GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME_2, "totalNum", SUM_AGGREGATION_NAME), + new Script("params.docNum / params.totalNum") ) - ); - - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ) ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); - - assertNotNull(buckets); - assertEquals(21, buckets.size()); - - // check content of few buckets - // first bucket have all the aggs values - Map firstBucket = buckets.get(0); - assertEquals(6, firstBucket.size()); - assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.1053, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(firstBucket.containsKey(KEY)); - - Map inBucketAggValues = getAggregationValues(firstBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inBucketAggValues); - assertEquals(1, inBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(130.0, getAggregationValue(inBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - - // second bucket is empty - Map secondBucket = buckets.get(1); - assertEquals(5, secondBucket.size()); - assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertFalse(secondBucket.containsKey(BUCKETS_AGGREGATION_NAME_1)); - assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(secondBucket.containsKey(KEY)); - - Map inSecondBucketAggValues = getAggregationValues(secondBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inSecondBucketAggValues); - assertEquals(0, inSecondBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0, getAggregationValue(inSecondBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - - // last bucket has values - Map lastBucket = buckets.get(buckets.size() - 1); - assertEquals(6, lastBucket.size()); - assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(2, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0172, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertEquals(5801.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(lastBucket.containsKey(KEY)); - - Map inLastBucketAggValues = getAggregationValues(lastBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inLastBucketAggValues); - assertEquals(1, inLastBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(100.0, getAggregationValue(inLastBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); + + assertNotNull(buckets); + assertEquals(21, buckets.size()); + + // check content of few buckets + // first bucket have all the aggs values + Map firstBucket = buckets.get(0); + assertEquals(6, firstBucket.size()); + assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.1053, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(firstBucket.containsKey(KEY)); + + Map inBucketAggValues = getAggregationValues(firstBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inBucketAggValues); + assertEquals(1, inBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(130.0, getAggregationValue(inBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); + + // second bucket is empty + Map secondBucket = buckets.get(1); + assertEquals(5, secondBucket.size()); + assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertFalse(secondBucket.containsKey(BUCKETS_AGGREGATION_NAME_1)); + assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(secondBucket.containsKey(KEY)); + + Map inSecondBucketAggValues = getAggregationValues(secondBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inSecondBucketAggValues); + assertEquals(0, inSecondBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0, getAggregationValue(inSecondBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); + + // last bucket has values + Map lastBucket = buckets.get(buckets.size() - 1); + assertEquals(6, lastBucket.size()); + assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(2, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0172, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertEquals(5801.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(lastBucket.containsKey(KEY)); + + Map inLastBucketAggValues = getAggregationValues(lastBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inLastBucketAggValues); + assertEquals(1, inLastBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(100.0, getAggregationValue(inLastBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); } private void testSampler() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.sampler(GENERIC_AGGREGATION_NAME) - .shardSize(2) - .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_DOCKEYWORD)); + AggregationBuilder aggsBuilder = AggregationBuilders.sampler(GENERIC_AGGREGATION_NAME) + .shardSize(2) + .subAggregation(AggregationBuilders.terms(BUCKETS_AGGREGATION_NAME_1).field(KEYWORD_FIELD_DOCKEYWORD)); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertEquals(2, aggValue.size()); - assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); - Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(nestedAggs); - assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); - List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); - assertEquals(2, buckets.size()); + Map aggValue = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertEquals(2, aggValue.size()); + assertEquals(3, aggValue.get(BUCKET_AGG_DOC_COUNT_FIELD)); + Map nestedAggs = getAggregationValues(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertNotNull(nestedAggs); + assertEquals(0, nestedAggs.get("doc_count_error_upper_bound")); + List> buckets = getAggregationBuckets(aggValue, BUCKETS_AGGREGATION_NAME_1); + assertEquals(2, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("likeable", firstBucket.get(KEY)); + Map firstBucket = buckets.get(0); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("likeable", firstBucket.get(KEY)); - Map secondBucket = buckets.get(1); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals("workable", secondBucket.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map secondBucket = buckets.get(1); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals("workable", secondBucket.get(KEY)); } private void testTermsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.terms(GENERIC_AGGREGATION_NAME).field(KEYWORD_FIELD_DOCKEYWORD); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + AggregationBuilder aggsBuilder = AggregationBuilders.terms(GENERIC_AGGREGATION_NAME).field(KEYWORD_FIELD_DOCKEYWORD); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - List> buckets = ((Map) getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME)).get( - "buckets" - ); - assertNotNull(buckets); - assertEquals(2, buckets.size()); - Map firstBucket = buckets.get(0); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(KEYWORD_FIELD_DOCKEYWORD_LIKABLE, firstBucket.get(KEY)); - Map secondBucket = buckets.get(1); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(KEYWORD_FIELD_DOCKEYWORD_WORKABLE, secondBucket.get(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + List> buckets = ((Map) getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME)).get( + "buckets" + ); + assertNotNull(buckets); + assertEquals(2, buckets.size()); + Map firstBucket = buckets.get(0); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(KEYWORD_FIELD_DOCKEYWORD_LIKABLE, firstBucket.get(KEY)); + Map secondBucket = buckets.get(1); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(KEYWORD_FIELD_DOCKEYWORD_WORKABLE, secondBucket.get(KEY)); } private void testSignificantTermsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.significantTerms(GENERIC_AGGREGATION_NAME).field(KEYWORD_FIELD_DOCKEYWORD); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + AggregationBuilder aggsBuilder = AggregationBuilders.significantTerms(GENERIC_AGGREGATION_NAME).field(KEYWORD_FIELD_DOCKEYWORD); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(buckets); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + List> buckets = getAggregationBuckets(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(buckets); - Map significantTermsAggregations = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + Map significantTermsAggregations = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(significantTermsAggregations); - assertEquals(3, (int) getAggregationValues(significantTermsAggregations, BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(11, (int) getAggregationValues(significantTermsAggregations, "bg_count")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + assertNotNull(significantTermsAggregations); + assertEquals(3, (int) getAggregationValues(significantTermsAggregations, BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(11, (int) getAggregationValues(significantTermsAggregations, "bg_count")); } private Map executeQueryAndGetAggsResults(final Object aggsBuilder, String indexName) { diff --git a/src/test/java/org/opensearch/neuralsearch/query/aggregation/MetricAggregationsWithHybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/aggregation/MetricAggregationsWithHybridQueryIT.java index 2466bb3d9..8d600205d 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/aggregation/MetricAggregationsWithHybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/aggregation/MetricAggregationsWithHybridQueryIT.java @@ -182,274 +182,230 @@ public void testWithConcurrentSegmentSearch_whenValueCount_thenSuccessful() { } private void testAvgAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.avg(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(AVG_AGGREGATION_NAME)); - double maxAggsValue = getAggregationValue(aggregations, AVG_AGGREGATION_NAME); - assertEquals(maxAggsValue, 2345.0, DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.avg(AVG_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(AVG_AGGREGATION_NAME)); + double maxAggsValue = getAggregationValue(aggregations, AVG_AGGREGATION_NAME); + assertEquals(maxAggsValue, 2345.0, DELTA_FOR_SCORE_ASSERTION); } private void testCardinalityAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.cardinality(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - int aggsValue = getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME); - assertEquals(aggsValue, 3); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.cardinality(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + int aggsValue = getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME); + assertEquals(aggsValue, 3); } private void testExtendedStatsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.extendedStats(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - Map extendedStatsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(extendedStatsValues); - - assertEquals((double) extendedStatsValues.get("max"), 3456.0, DELTA_FOR_SCORE_ASSERTION); - assertEquals((int) extendedStatsValues.get("count"), 3); - assertEquals((double) extendedStatsValues.get("sum"), 7035.0, DELTA_FOR_SCORE_ASSERTION); - assertEquals((double) extendedStatsValues.get("avg"), 2345.0, DELTA_FOR_SCORE_ASSERTION); - assertEquals((double) extendedStatsValues.get("variance"), 822880.666, DELTA_FOR_SCORE_ASSERTION); - assertEquals((double) extendedStatsValues.get("std_deviation"), 907.127, DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.extendedStats(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + Map extendedStatsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(extendedStatsValues); + + assertEquals((double) extendedStatsValues.get("max"), 3456.0, DELTA_FOR_SCORE_ASSERTION); + assertEquals((int) extendedStatsValues.get("count"), 3); + assertEquals((double) extendedStatsValues.get("sum"), 7035.0, DELTA_FOR_SCORE_ASSERTION); + assertEquals((double) extendedStatsValues.get("avg"), 2345.0, DELTA_FOR_SCORE_ASSERTION); + assertEquals((double) extendedStatsValues.get("variance"), 822880.666, DELTA_FOR_SCORE_ASSERTION); + assertEquals((double) extendedStatsValues.get("std_deviation"), 907.127, DELTA_FOR_SCORE_ASSERTION); } private void testTopHitsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.topHits(GENERIC_AGGREGATION_NAME).size(4); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - Map aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(aggsValues); - assertHitResultsFromQuery(3, aggsValues); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.topHits(GENERIC_AGGREGATION_NAME).size(4); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + Map aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(aggsValues); + assertHitResultsFromQuery(3, aggsValues); } private void testScriptedMetricsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - // compute sum of all int fields that are not blank - AggregationBuilder aggsBuilder = AggregationBuilders.scriptedMetric(GENERIC_AGGREGATION_NAME) - .initScript(new Script("state.price = []")) - .mapScript( - new Script( - "state.price.add(doc[\"" - + INTEGER_FIELD_DOCINDEX - + "\"].size() == 0 ? 0 : doc." - + INTEGER_FIELD_DOCINDEX - + ".value)" - ) + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + // compute sum of all int fields that are not blank + AggregationBuilder aggsBuilder = AggregationBuilders.scriptedMetric(GENERIC_AGGREGATION_NAME) + .initScript(new Script("state.price = []")) + .mapScript( + new Script( + "state.price.add(doc[\"" + INTEGER_FIELD_DOCINDEX + "\"].size() == 0 ? 0 : doc." + INTEGER_FIELD_DOCINDEX + ".value)" ) - .combineScript(new Script("state.price.stream().mapToInt(Integer::intValue).sum()")) - .reduceScript(new Script("states.stream().mapToInt(Integer::intValue).sum()")); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - int aggsValue = getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME); - assertEquals(7035, aggsValue); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + ) + .combineScript(new Script("state.price.stream().mapToInt(Integer::intValue).sum()")) + .reduceScript(new Script("states.stream().mapToInt(Integer::intValue).sum()")); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + int aggsValue = getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME); + assertEquals(7035, aggsValue); } private void testPercentileAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.percentiles(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - assertHitResultsFromQuery(3, searchResponseAsMap); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - Map> aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(aggsValues); - - Map values = aggsValues.get("values"); - assertNotNull(values); - assertEquals(7, values.size()); - assertEquals(1234.0, values.get("1.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, values.get("5.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, values.get("25.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(2345.0, values.get("50.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(3456.0, values.get("75.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(3456.0, values.get("95.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(3456.0, values.get("99.0"), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.percentiles(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + assertHitResultsFromQuery(3, searchResponseAsMap); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + Map> aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(aggsValues); + + Map values = aggsValues.get("values"); + assertNotNull(values); + assertEquals(7, values.size()); + assertEquals(1234.0, values.get("1.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, values.get("5.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, values.get("25.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(2345.0, values.get("50.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(3456.0, values.get("75.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(3456.0, values.get("95.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(3456.0, values.get("99.0"), DELTA_FOR_SCORE_ASSERTION); } private void testPercentileRankAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.percentileRanks(GENERIC_AGGREGATION_NAME, new double[] { 2000, 3000 }) - .field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - assertHitResultsFromQuery(3, searchResponseAsMap); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - Map> aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); - assertNotNull(aggsValues); - Map values = aggsValues.get("values"); - assertNotNull(values); - assertEquals(33.333, values.get("2000.0"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(66.666, values.get("3000.0"), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.percentileRanks(GENERIC_AGGREGATION_NAME, new double[] { 2000, 3000 }) + .field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + assertHitResultsFromQuery(3, searchResponseAsMap); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + Map> aggsValues = getAggregationValues(aggregations, GENERIC_AGGREGATION_NAME); + assertNotNull(aggsValues); + Map values = aggsValues.get("values"); + assertNotNull(values); + assertEquals(33.333, values.get("2000.0"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(66.666, values.get("3000.0"), DELTA_FOR_SCORE_ASSERTION); } private void testSumAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(SUM_AGGREGATION_NAME)); - double maxAggsValue = getAggregationValue(aggregations, SUM_AGGREGATION_NAME); - assertEquals(7035.0, maxAggsValue, DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(SUM_AGGREGATION_NAME)); + double maxAggsValue = getAggregationValue(aggregations, SUM_AGGREGATION_NAME); + assertEquals(7035.0, maxAggsValue, DELTA_FOR_SCORE_ASSERTION); } private void testValueCountAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggsBuilder = AggregationBuilders.count(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggsBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); - - assertHitResultsFromQuery(3, searchResponseAsMap); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); - assertEquals(3, (int) getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + AggregationBuilder aggsBuilder = AggregationBuilders.count(GENERIC_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggsBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + assertHitResultsFromQuery(3, searchResponseAsMap); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + assertTrue(aggregations.containsKey(GENERIC_AGGREGATION_NAME)); + assertEquals(3, (int) getAggregationValue(aggregations, GENERIC_AGGREGATION_NAME)); } private void testSumAggsAndRangePostFilter() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggsBuilder = AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); - - TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); - TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); - TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); - - HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); - hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder3); - - QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).gte(3000).lte(5000); - - Map searchResponseAsMap = search( - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, - hybridQueryBuilderNeuralThenTerm, - null, - 10, - Map.of("search_pipeline", SEARCH_PIPELINE), - List.of(aggsBuilder), - rangeFilterQuery, - null, - false, - null, - 0 - ); - - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - assertTrue(aggregations.containsKey(SUM_AGGREGATION_NAME)); - double maxAggsValue = getAggregationValue(aggregations, SUM_AGGREGATION_NAME); - assertEquals(11602.0, maxAggsValue, DELTA_FOR_SCORE_ASSERTION); - - assertHitResultsFromQuery(2, searchResponseAsMap); - - // assert post-filter - List> hitsNestedList = getNestedHits(searchResponseAsMap); - - List docIndexes = new ArrayList<>(); - for (Map oneHit : hitsNestedList) { - assertNotNull(oneHit.get("_source")); - Map source = (Map) oneHit.get("_source"); - int docIndex = (int) source.get(INTEGER_FIELD_DOCINDEX); - docIndexes.add(docIndex); - } - assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 3000 || docIndex > 5000).count()); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggsBuilder = AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX); + + TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); + TermQueryBuilder termQueryBuilder3 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT5); + + HybridQueryBuilder hybridQueryBuilderNeuralThenTerm = new HybridQueryBuilder(); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder1); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder2); + hybridQueryBuilderNeuralThenTerm.add(termQueryBuilder3); + + QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery(INTEGER_FIELD_DOCINDEX).gte(3000).lte(5000); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, + hybridQueryBuilderNeuralThenTerm, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + List.of(aggsBuilder), + rangeFilterQuery, + null, + false, + null, + 0 + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + assertTrue(aggregations.containsKey(SUM_AGGREGATION_NAME)); + double maxAggsValue = getAggregationValue(aggregations, SUM_AGGREGATION_NAME); + assertEquals(11602.0, maxAggsValue, DELTA_FOR_SCORE_ASSERTION); + + assertHitResultsFromQuery(2, searchResponseAsMap); + + // assert post-filter + List> hitsNestedList = getNestedHits(searchResponseAsMap); + + List docIndexes = new ArrayList<>(); + for (Map oneHit : hitsNestedList) { + assertNotNull(oneHit.get("_source")); + Map source = (Map) oneHit.get("_source"); + int docIndex = (int) source.get(INTEGER_FIELD_DOCINDEX); + docIndexes.add(docIndex); } + assertEquals(0, docIndexes.stream().filter(docIndex -> docIndex < 3000 || docIndex > 5000).count()); } private Map executeQueryAndGetAggsResults(final Object aggsBuilder, String indexName) { diff --git a/src/test/java/org/opensearch/neuralsearch/query/aggregation/PipelineAggregationsWithHybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/aggregation/PipelineAggregationsWithHybridQueryIT.java index fd118629b..d0ba802cf 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/aggregation/PipelineAggregationsWithHybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/aggregation/PipelineAggregationsWithHybridQueryIT.java @@ -100,265 +100,244 @@ public void testPipelineParentAggs_whenDateBucketedSumsPipelinedToCumulativeSumA } private void testDateBucketedSumsPipelinedToBucketStatsAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); + AggregationBuilder aggDateHisto = AggregationBuilders.dateHistogram(GENERIC_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)); - StatsBucketPipelineAggregationBuilder aggStatsBucket = PipelineAggregatorBuilders.statsBucket( - BUCKETS_AGGREGATION_NAME_1, - GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME - ); + StatsBucketPipelineAggregationBuilder aggStatsBucket = PipelineAggregatorBuilders.statsBucket( + BUCKETS_AGGREGATION_NAME_1, + GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME + ); - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggDateHisto, aggStatsBucket), - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS - ); + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggDateHisto, aggStatsBucket), + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); - Map statsAggs = getAggregationValues(aggregations, BUCKETS_AGGREGATION_NAME_1); + Map statsAggs = getAggregationValues(aggregations, BUCKETS_AGGREGATION_NAME_1); - assertNotNull(statsAggs); + assertNotNull(statsAggs); - assertEquals(3517.5, (Double) statsAggs.get("avg"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(7035.0, (Double) statsAggs.get("sum"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, (Double) statsAggs.get("min"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(5801.0, (Double) statsAggs.get("max"), DELTA_FOR_SCORE_ASSERTION); - assertEquals(2, (int) statsAggs.get("count")); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + assertEquals(3517.5, (Double) statsAggs.get("avg"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(7035.0, (Double) statsAggs.get("sum"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, (Double) statsAggs.get("min"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(5801.0, (Double) statsAggs.get("max"), DELTA_FOR_SCORE_ASSERTION); + assertEquals(2, (int) statsAggs.get("count")); } private void testDateBucketedSumsPipelinedToBucketScriptedAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregations( - new AggregatorFactories.Builder().addAggregator( - AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX) + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregations( + new AggregatorFactories.Builder().addAggregator(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)) + .addAggregator( + AggregationBuilders.filter( + GENERIC_AGGREGATION_NAME, + QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) + ) + .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))) + ).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME_2).field(INTEGER_FIELD_PRICE)) ) - .addAggregator( - AggregationBuilders.filter( - GENERIC_AGGREGATION_NAME, - QueryBuilders.boolQuery() - .should( - QueryBuilders.boolQuery() - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) - ) - .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))) - ).subAggregation(AggregationBuilders.sum(SUM_AGGREGATION_NAME_2).field(INTEGER_FIELD_PRICE)) - ) - .addPipelineAggregator( - PipelineAggregatorBuilders.bucketScript( - BUCKETS_AGGREGATION_NAME_1, - Map.of("docNum", GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME_2, "totalNum", SUM_AGGREGATION_NAME), - new Script("params.docNum / params.totalNum") - ) + .addPipelineAggregator( + PipelineAggregatorBuilders.bucketScript( + BUCKETS_AGGREGATION_NAME_1, + Map.of("docNum", GENERIC_AGGREGATION_NAME + ">" + SUM_AGGREGATION_NAME_2, "totalNum", SUM_AGGREGATION_NAME), + new Script("params.docNum / params.totalNum") ) - ); - - Map searchResponseAsMap = executeQueryAndGetAggsResults( - aggBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ) ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); - - assertNotNull(buckets); - assertEquals(21, buckets.size()); - - // check content of few buckets - // first bucket have all the aggs values - Map firstBucket = buckets.get(0); - assertEquals(6, firstBucket.size()); - assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.1053, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(firstBucket.containsKey(KEY)); - - Map inBucketAggValues = getAggregationValues(firstBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inBucketAggValues); - assertEquals(1, inBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(130.0, getAggregationValue(inBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - - // second bucket is empty - Map secondBucket = buckets.get(1); - assertEquals(5, secondBucket.size()); - assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertFalse(secondBucket.containsKey(BUCKETS_AGGREGATION_NAME_1)); - assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(secondBucket.containsKey(KEY)); - - Map inSecondBucketAggValues = getAggregationValues(secondBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inSecondBucketAggValues); - assertEquals(0, inSecondBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0, getAggregationValue(inSecondBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - - // last bucket has values - Map lastBucket = buckets.get(buckets.size() - 1); - assertEquals(6, lastBucket.size()); - assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(2, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0172, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertEquals(5801.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(lastBucket.containsKey(KEY)); - - Map inLastBucketAggValues = getAggregationValues(lastBucket, GENERIC_AGGREGATION_NAME); - assertNotNull(inLastBucketAggValues); - assertEquals(1, inLastBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(100.0, getAggregationValue(inLastBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + Map searchResponseAsMap = executeQueryAndGetAggsResults( + aggBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); + + assertNotNull(buckets); + assertEquals(21, buckets.size()); + + // check content of few buckets + // first bucket have all the aggs values + Map firstBucket = buckets.get(0); + assertEquals(6, firstBucket.size()); + assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.1053, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(firstBucket.containsKey(KEY)); + + Map inBucketAggValues = getAggregationValues(firstBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inBucketAggValues); + assertEquals(1, inBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(130.0, getAggregationValue(inBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); + + // second bucket is empty + Map secondBucket = buckets.get(1); + assertEquals(5, secondBucket.size()); + assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertFalse(secondBucket.containsKey(BUCKETS_AGGREGATION_NAME_1)); + assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(secondBucket.containsKey(KEY)); + + Map inSecondBucketAggValues = getAggregationValues(secondBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inSecondBucketAggValues); + assertEquals(0, inSecondBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0, getAggregationValue(inSecondBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); + + // last bucket has values + Map lastBucket = buckets.get(buckets.size() - 1); + assertEquals(6, lastBucket.size()); + assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(2, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0172, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertEquals(5801.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(lastBucket.containsKey(KEY)); + + Map inLastBucketAggValues = getAggregationValues(lastBucket, GENERIC_AGGREGATION_NAME); + assertNotNull(inLastBucketAggValues); + assertEquals(1, inLastBucketAggValues.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(100.0, getAggregationValue(inLastBucketAggValues, SUM_AGGREGATION_NAME_2), DELTA_FOR_SCORE_ASSERTION); } private void testDateBucketedSumsPipelinedToBucketSortAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregations( - new AggregatorFactories.Builder().addAggregator( - AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX) + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregations( + new AggregatorFactories.Builder().addAggregator(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)) + .addPipelineAggregator( + PipelineAggregatorBuilders.bucketSort( + BUCKETS_AGGREGATION_NAME_1, + List.of(new FieldSortBuilder(SUM_AGGREGATION_NAME).order(SortOrder.DESC)) + ).size(5) ) - .addPipelineAggregator( - PipelineAggregatorBuilders.bucketSort( - BUCKETS_AGGREGATION_NAME_1, - List.of(new FieldSortBuilder(SUM_AGGREGATION_NAME).order(SortOrder.DESC)) - ).size(5) - ) - ); - - QueryBuilder queryBuilder = QueryBuilders.boolQuery() - .should( - QueryBuilders.boolQuery() - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) - ) - .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))); - - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggBuilder), - queryBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); - - assertNotNull(buckets); - assertEquals(3, buckets.size()); - - // check content of few buckets - Map firstBucket = buckets.get(0); - assertEquals(4, firstBucket.size()); - assertEquals("01/01/2015", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(2345.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(firstBucket.containsKey(KEY)); - - // second bucket is empty - Map secondBucket = buckets.get(1); - assertEquals(4, secondBucket.size()); - assertEquals("01/01/1995", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(1234.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(secondBucket.containsKey(KEY)); - - // last bucket has values - Map lastBucket = buckets.get(buckets.size() - 1); - assertEquals(4, lastBucket.size()); - assertEquals("01/01/2007", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertTrue(lastBucket.containsKey(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + QueryBuilder queryBuilder = QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) + ) + .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))); + + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggBuilder), + queryBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); + + assertNotNull(buckets); + assertEquals(3, buckets.size()); + + // check content of few buckets + Map firstBucket = buckets.get(0); + assertEquals(4, firstBucket.size()); + assertEquals("01/01/2015", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(2345.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(firstBucket.containsKey(KEY)); + + // second bucket is empty + Map secondBucket = buckets.get(1); + assertEquals(4, secondBucket.size()); + assertEquals("01/01/1995", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(1234.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(secondBucket.containsKey(KEY)); + + // last bucket has values + Map lastBucket = buckets.get(buckets.size() - 1); + assertEquals(4, lastBucket.size()); + assertEquals("01/01/2007", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertTrue(lastBucket.containsKey(KEY)); } private void testDateBucketedSumsPipelinedToCumulativeSumAggs() throws IOException { - try { - prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); - - AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) - .calendarInterval(DateHistogramInterval.YEAR) - .field(DATE_FIELD) - .subAggregations( - new AggregatorFactories.Builder().addAggregator( - AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX) - ).addPipelineAggregator(PipelineAggregatorBuilders.cumulativeSum(BUCKETS_AGGREGATION_NAME_1, SUM_AGGREGATION_NAME)) - ); - - QueryBuilder queryBuilder = QueryBuilders.boolQuery() - .should( - QueryBuilders.boolQuery() - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) - .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) - ) - .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))); - - Map searchResponseAsMap = executeQueryAndGetAggsResults( - List.of(aggBuilder), - queryBuilder, - TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); + + AggregationBuilder aggBuilder = AggregationBuilders.dateHistogram(DATE_AGGREGATION_NAME) + .calendarInterval(DateHistogramInterval.YEAR) + .field(DATE_FIELD) + .subAggregations( + new AggregatorFactories.Builder().addAggregator(AggregationBuilders.sum(SUM_AGGREGATION_NAME).field(INTEGER_FIELD_DOCINDEX)) + .addPipelineAggregator(PipelineAggregatorBuilders.cumulativeSum(BUCKETS_AGGREGATION_NAME_1, SUM_AGGREGATION_NAME)) ); - Map aggregations = getAggregations(searchResponseAsMap); - assertNotNull(aggregations); - - List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); - - assertNotNull(buckets); - assertEquals(21, buckets.size()); - - // check content of few buckets - Map firstBucket = buckets.get(0); - assertEquals(5, firstBucket.size()); - assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertTrue(firstBucket.containsKey(KEY)); - - Map secondBucket = buckets.get(1); - assertEquals(5, secondBucket.size()); - assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertEquals(1234.0, getAggregationValue(secondBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertTrue(secondBucket.containsKey(KEY)); - - // last bucket is empty - Map lastBucket = buckets.get(buckets.size() - 1); - assertEquals(5, lastBucket.size()); - assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); - assertEquals(1, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); - assertEquals(2345.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); - assertEquals(3579.0, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); - assertTrue(lastBucket.containsKey(KEY)); - } finally { - wipeOfTestResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); - } + QueryBuilder queryBuilder = QueryBuilders.boolQuery() + .should( + QueryBuilders.boolQuery() + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_WORKABLE)) + .should(QueryBuilders.termQuery(KEYWORD_FIELD_DOCKEYWORD, KEYWORD_FIELD_DOCKEYWORD_ANGRY)) + ) + .should(QueryBuilders.boolQuery().mustNot(QueryBuilders.existsQuery(KEYWORD_FIELD_DOCKEYWORD))); + + Map searchResponseAsMap = executeQueryAndGetAggsResults( + List.of(aggBuilder), + queryBuilder, + TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS + ); + + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + List> buckets = getAggregationBuckets(aggregations, DATE_AGGREGATION_NAME); + + assertNotNull(buckets); + assertEquals(21, buckets.size()); + + // check content of few buckets + Map firstBucket = buckets.get(0); + assertEquals(5, firstBucket.size()); + assertEquals("01/01/1995", firstBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, firstBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(1234.0, getAggregationValue(firstBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, getAggregationValue(firstBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertTrue(firstBucket.containsKey(KEY)); + + Map secondBucket = buckets.get(1); + assertEquals(5, secondBucket.size()); + assertEquals("01/01/1996", secondBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(0, secondBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(0.0, getAggregationValue(secondBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertEquals(1234.0, getAggregationValue(secondBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertTrue(secondBucket.containsKey(KEY)); + + // last bucket is empty + Map lastBucket = buckets.get(buckets.size() - 1); + assertEquals(5, lastBucket.size()); + assertEquals("01/01/2015", lastBucket.get(BUCKET_AGG_KEY_AS_STRING)); + assertEquals(1, lastBucket.get(BUCKET_AGG_DOC_COUNT_FIELD)); + assertEquals(2345.0, getAggregationValue(lastBucket, SUM_AGGREGATION_NAME), DELTA_FOR_SCORE_ASSERTION); + assertEquals(3579.0, getAggregationValue(lastBucket, BUCKETS_AGGREGATION_NAME_1), DELTA_FOR_SCORE_ASSERTION); + assertTrue(lastBucket.containsKey(KEY)); } private Map executeQueryAndGetAggsResults(final Object aggsBuilder, String indexName) { diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index 54edfd260..a8abb41fa 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -4,13 +4,23 @@ */ package org.opensearch.neuralsearch; +import org.junit.After; +import org.opensearch.action.search.SearchResponse; import org.opensearch.client.ResponseException; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.MediaType; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; import org.opensearch.ml.common.model.MLModelState; + +import static org.opensearch.knn.common.KNNConstants.MODEL_INDEX_NAME; import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -53,6 +63,7 @@ import org.opensearch.neuralsearch.processor.ExplanationResponseProcessor; import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import org.opensearch.neuralsearch.util.TokenWeightUtil; +import org.opensearch.search.SearchHit; import org.opensearch.search.sort.SortBuilder; import org.opensearch.test.ClusterServiceUtils; import org.opensearch.threadpool.TestThreadPool; @@ -62,14 +73,20 @@ import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; import com.google.common.collect.ImmutableList; +import static org.opensearch.neuralsearch.util.TestUtils.INGEST_PIPELINE_TYPE; import static org.opensearch.neuralsearch.util.TestUtils.MAX_TASK_RESULT_QUERY_TIME_IN_SECOND; import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_TASK_RESULT_QUERY_INTERVAL_IN_MILLISECOND; import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_USER_AGENT; import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_NORMALIZATION_METHOD; import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_COMBINATION_METHOD; +import static org.opensearch.neuralsearch.util.TestUtils.ML_PLUGIN_SYSTEM_INDEX_PREFIX; +import static org.opensearch.neuralsearch.util.TestUtils.OPENDISTRO_SECURITY; +import static org.opensearch.neuralsearch.util.TestUtils.OPENSEARCH_SYSTEM_INDEX_PREFIX; import static org.opensearch.neuralsearch.util.TestUtils.PARAM_NAME_WEIGHTS; import static org.opensearch.neuralsearch.util.TestUtils.MAX_RETRY; import static org.opensearch.neuralsearch.util.TestUtils.MAX_TIME_OUT_INTERVAL; +import static org.opensearch.neuralsearch.util.TestUtils.SEARCH_PIPELINE_TYPE; +import static org.opensearch.neuralsearch.util.TestUtils.SECURITY_AUDITLOG_PREFIX; import lombok.AllArgsConstructor; import lombok.Getter; @@ -96,6 +113,12 @@ public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase { protected static final String CONCURRENT_SEGMENT_SEARCH_ENABLED = "search.concurrent_segment_search.enabled"; protected static final String RRF_SEARCH_PIPELINE = "rrf-search-pipeline"; + private final Set IMMUTABLE_INDEX_PREFIXES = Set.of( + SECURITY_AUDITLOG_PREFIX, + OPENSEARCH_SYSTEM_INDEX_PREFIX, + ML_PLUGIN_SYSTEM_INDEX_PREFIX + ); + protected final ClassLoader classLoader = this.getClass().getClassLoader(); protected ThreadPool threadPool; @@ -111,6 +134,11 @@ public void setupSettings() { NeuralSearchClusterUtil.instance().initialize(clusterService); } + @After + public void cleanUp() throws IOException { + wipeOfTestResources(); + } + protected ThreadPool setUpThreadPool() { return new TestThreadPool(getClass().getName(), threadPoolSettings()); } @@ -1243,18 +1271,6 @@ protected void createSearchPipelineWithDefaultResultsPostProcessor(final String ); } - @SneakyThrows - protected void deleteSearchPipeline(final String pipelineId) { - makeRequest( - client(), - "DELETE", - String.format(LOCALE, "/_search/pipeline/%s", pipelineId), - null, - toHttpEntity(""), - ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT)) - ); - } - @SneakyThrows private String getModelGroupId() { String modelGroupRegisterRequestBody = Files.readString( @@ -1411,6 +1427,80 @@ protected void reindex(String fromIndexName, String toIndexName) throws Exceptio assertEquals(0, ((List) map.get("failures")).size()); } + /** + * Retrieve all available index + * @return index as a list of map object + */ + protected List> retrieveIndices() throws IOException { + Request request = new Request("GET", "/_cat/indices?format=json&expand_wildcards=all"); + Response response = client().performRequest(request); + MediaType mediaType = MediaType.fromMediaType(response.getEntity().getContentType()); + try ( + XContentParser parser = mediaType.xContent() + .createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + response.getEntity().getContent() + ) + ) { + XContentParser.Token token = parser.nextToken(); + List> parserList = null; + if (token == XContentParser.Token.START_ARRAY) { + parserList = parser.listOrderedMap().stream().map(obj -> (Map) obj).collect(Collectors.toList()); + } else { + parserList = Collections.singletonList(parser.mapOrdered()); + } + + return parserList; + } + } + + /** + * Fetch all existing indices and clean up, note that some system generated indices are skipped + */ + protected void deleteExistingIndices() throws IOException { + List> indices = retrieveIndices(); + for (Map index : indices) { + final String indexName = (String) index.get("index"); + if (!shouldDeleteIndex(indexName)) { + deleteIndex(indexName); + } + } + } + + private boolean shouldDeleteIndex(String indexName) { + return indexName == null + || OPENDISTRO_SECURITY.equals(indexName) + || IMMUTABLE_INDEX_PREFIXES.stream().anyMatch(indexName::startsWith) + || MODEL_INDEX_NAME.equals(indexName); + } + + /** + * Retrieve all existing pipelines or a specific pipeline + * @param pipelineType _ingest for retrieving ingest pipelines, _search for retrieving search pipelines + * @param pipelineName a specific pipeline to retrieve, if value is empty, it returns all pipelines + * @return get pipeline response as a map object + */ + @SneakyThrows(ParseException.class) + protected Map retrievePipelines(final String pipelineType, final String pipelineName) throws IOException { + Request request = new Request("GET", "/" + pipelineType + "/pipeline/" + (StringUtils.isEmpty(pipelineName) ? "" : pipelineName)); + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + String responseBody = EntityUtils.toString(response.getEntity()); + Map responseMap = createParser(XContentType.JSON.xContent(), responseBody).map(); + return StringUtils.isEmpty(pipelineName) ? responseMap : (Map) responseMap.get(pipelineName); + } + + private void deleteExistingIngestionPipelines() throws IOException { + Map pipelines = retrievePipelines(INGEST_PIPELINE_TYPE, null); + pipelines.keySet().forEach(this::deleteIngestPipeline); + } + + private void deleteExistingSearchPipelines() throws IOException { + Map pipelines = retrievePipelines(SEARCH_PIPELINE_TYPE, null); + pipelines.keySet().forEach(this::deleteSearchPipeline); + } + /** * Get ingest pipeline * @param pipelineName of the ingest pipeline @@ -1419,24 +1509,20 @@ protected void reindex(String fromIndexName, String toIndexName) throws Exceptio */ @SneakyThrows protected Map getIngestionPipeline(final String pipelineName) { - Request request = new Request("GET", "/_ingest/pipeline/" + pipelineName); - Response response = client().performRequest(request); - assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); - String responseBody = EntityUtils.toString(response.getEntity()); - Map responseMap = createParser(XContentType.JSON.xContent(), responseBody).map(); - return (Map) responseMap.get(pipelineName); + return retrievePipelines(INGEST_PIPELINE_TYPE, pipelineName); } /** * Delete pipeline * + * @param pipelineType _ingest for ingest pipelines, _search for search pipelines * @param pipelineName of the pipeline * * @return delete pipeline response as a map object */ @SneakyThrows - protected Map deletePipeline(final String pipelineName) { - Request request = new Request("DELETE", "/_ingest/pipeline/" + pipelineName); + protected Map deletePipeline(final String pipelineType, final String pipelineName) { + Request request = new Request("DELETE", "/" + pipelineType + "/pipeline/" + pipelineName); Response response = client().performRequest(request); assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); String responseBody = EntityUtils.toString(response.getEntity()); @@ -1444,6 +1530,72 @@ protected Map deletePipeline(final String pipelineName) { return responseMap; } + /** + * Delete an ingest pipeline + * + * @param pipelineName of the pipeline + * + * @return delete pipeline response as a map object + */ + @SneakyThrows + protected Map deleteIngestPipeline(final String pipelineName) { + return deletePipeline(INGEST_PIPELINE_TYPE, pipelineName); + } + + /** + * Delete a search pipeline + * + * @param pipelineName of the pipeline + * + * @return delete pipeline response as a map object + */ + @SneakyThrows + protected Map deleteSearchPipeline(final String pipelineName) { + return deletePipeline(SEARCH_PIPELINE_TYPE, pipelineName); + } + + /** + * Retrieves all existing models + * @return models as a map object + */ + @SneakyThrows(ParseException.class) + protected List retrieveModels() throws IOException { + Response response = makeRequest( + client(), + "POST", + "/_plugins/_ml/models/_search", + null, + toHttpEntity("{\"query\":{\"match_all\":{}}}"), + ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, DEFAULT_USER_AGENT)) + ); + + assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + + final String responseBody = EntityUtils.toString(response.getEntity()); + assertNotNull(responseBody); + + final XContentParser parser = createParser(MediaTypeRegistry.getDefaultMediaType().xContent(), responseBody); + final SearchResponse searchResponse = SearchResponse.fromXContent(parser); + + return Arrays.stream(searchResponse.getHits().getHits()) + .filter(h -> !h.getSourceAsMap().containsKey("chunk_number")) + .map(SearchHit::getId) + .toList(); + } + + private void deleteExistingModels() throws IOException { + List modelIds = retrieveModels(); + modelIds.forEach(m -> { + try { + deleteModel(m); + } catch (AssertionError e) { + // sometimes we have flaky test that the model state doesn't change after call undeploy api + // for this case we can call undeploy api one more time + deleteModel(m); + } + }); + } + protected float computeExpectedScore(final String modelId, final Map tokenWeightMap, final String queryText) { Map queryTokens = runSparseModelInference(modelId, queryText); return computeExpectedScore(tokenWeightMap, queryTokens); @@ -1499,36 +1651,38 @@ protected Float getFeatureFieldCompressedNumber(final Float originNumber) { } // Wipe of all the resources after execution of the tests. - protected void wipeOfTestResources( - final String indexName, - final String ingestPipeline, - final String modelId, - final String searchPipeline - ) throws IOException { + protected void wipeOfTestResources() throws IOException { try { - if (ingestPipeline != null) { - deletePipeline(ingestPipeline); - } - if (searchPipeline != null) { - deleteSearchPipeline(searchPipeline); - } - if (modelId != null) { - try { - deleteModel(modelId); - } catch (AssertionError e) { - // sometimes we have flaky test that the model state doesn't change after call undeploy api - // for this case we can call undeploy api one more time - deleteModel(modelId); - } - } - if (indexName != null) { - deleteIndex(indexName); - } + deleteExistingIngestionPipelines(); } catch (ResponseException e) { // It's possible that test fails during resources (index, model, pipeline, etc.) creation, when clean up // these resources after test run, the delete methods will throw ResponseException with 404 Not Found code // In this case, we just need to ignore this exception, for other exceptions, continue throwing - if (RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode()) != RestStatus.NOT_FOUND) { + if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { + throw e; + } + } + + try { + deleteExistingSearchPipelines(); + } catch (ResponseException e) { + if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { + throw e; + } + } + + try { + deleteExistingModels(); + } catch (ResponseException e) { + if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { + throw e; + } + } + + try { + deleteExistingIndices(); + } catch (ResponseException e) { + if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { throw e; } } diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java index bb072ab55..281ac8ab3 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java @@ -54,6 +54,9 @@ public class TestUtils { public static final String SKIP_DELETE_MODEL_INDEX = "tests.skip_delete_model_index"; public static final String SECURITY_AUDITLOG_PREFIX = "security-auditlog"; public static final String OPENSEARCH_SYSTEM_INDEX_PREFIX = ".opensearch"; + public static final String ML_PLUGIN_SYSTEM_INDEX_PREFIX = ".plugins-ml"; + public static final String INGEST_PIPELINE_TYPE = "_ingest"; + public static final String SEARCH_PIPELINE_TYPE = "_search"; public static final String TEXT_EMBEDDING_PROCESSOR = "text_embedding"; public static final String TEXT_IMAGE_EMBEDDING_PROCESSOR = "text_image_embedding"; public static final int MAX_TASK_RESULT_QUERY_TIME_IN_SECOND = 60 * 5; From 3780933c5b31d14c1ef2422fea3a898733781e8a Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Wed, 5 Feb 2025 16:58:45 -0800 Subject: [PATCH 3/5] Code clean up Signed-off-by: Weijia Zhao --- .../neuralsearch/BaseNeuralSearchIT.java | 104 ++++++++---------- 1 file changed, 43 insertions(+), 61 deletions(-) diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index a8abb41fa..72b9604c9 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -134,9 +134,13 @@ public void setupSettings() { NeuralSearchClusterUtil.instance().initialize(clusterService); } + // Wipe of all the resources after execution of the tests. @After - public void cleanUp() throws IOException { - wipeOfTestResources(); + public void cleanUp() { + deleteExistingIngestionPipelines(); + deleteExistingSearchPipelines(); + deleteExistingModels(); + deleteExistingIndices(); } protected ThreadPool setUpThreadPool() { @@ -1431,7 +1435,8 @@ protected void reindex(String fromIndexName, String toIndexName) throws Exceptio * Retrieve all available index * @return index as a list of map object */ - protected List> retrieveIndices() throws IOException { + @SneakyThrows + protected List> retrieveIndices() { Request request = new Request("GET", "/_cat/indices?format=json&expand_wildcards=all"); Response response = client().performRequest(request); MediaType mediaType = MediaType.fromMediaType(response.getEntity().getContentType()); @@ -1458,45 +1463,60 @@ protected List> retrieveIndices() throws IOException { /** * Fetch all existing indices and clean up, note that some system generated indices are skipped */ - protected void deleteExistingIndices() throws IOException { + @SneakyThrows + protected void deleteExistingIndices() { List> indices = retrieveIndices(); for (Map index : indices) { final String indexName = (String) index.get("index"); - if (!shouldDeleteIndex(indexName)) { + if (shouldDeleteIndex(indexName)) { deleteIndex(indexName); } } } private boolean shouldDeleteIndex(String indexName) { - return indexName == null - || OPENDISTRO_SECURITY.equals(indexName) - || IMMUTABLE_INDEX_PREFIXES.stream().anyMatch(indexName::startsWith) - || MODEL_INDEX_NAME.equals(indexName); + return indexName != null + && !OPENDISTRO_SECURITY.equals(indexName) + && IMMUTABLE_INDEX_PREFIXES.stream().noneMatch(indexName::startsWith) + && !MODEL_INDEX_NAME.equals(indexName); } /** * Retrieve all existing pipelines or a specific pipeline * @param pipelineType _ingest for retrieving ingest pipelines, _search for retrieving search pipelines - * @param pipelineName a specific pipeline to retrieve, if value is empty, it returns all pipelines + * @param pipelineName a specific pipeline to retrieve, if the value is null, it returns all pipelines * @return get pipeline response as a map object */ - @SneakyThrows(ParseException.class) - protected Map retrievePipelines(final String pipelineType, final String pipelineName) throws IOException { - Request request = new Request("GET", "/" + pipelineType + "/pipeline/" + (StringUtils.isEmpty(pipelineName) ? "" : pipelineName)); - Response response = client().performRequest(request); - assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); - String responseBody = EntityUtils.toString(response.getEntity()); - Map responseMap = createParser(XContentType.JSON.xContent(), responseBody).map(); - return StringUtils.isEmpty(pipelineName) ? responseMap : (Map) responseMap.get(pipelineName); + @SneakyThrows + protected Map retrievePipelines(final String pipelineType, final String pipelineName) { + try { + Request request = new Request( + "GET", + String.format(LOCALE, "/%s/pipeline/%s", pipelineType, Optional.ofNullable(pipelineName).orElse("")) + ); + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + String responseBody = EntityUtils.toString(response.getEntity()); + Map responseMap = createParser(XContentType.JSON.xContent(), responseBody).map(); + return StringUtils.isEmpty(pipelineName) ? responseMap : (Map) responseMap.get(pipelineName); + } catch (ResponseException e) { + // If no pipeline exits, we will probably receive a 404 NOT Found exception in the above GET call, + // see issue: https://github.com/opensearch-project/OpenSearch/issues/15917, + // we can just ignore the 404 exception, and rethrow the exception for other code + if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { + throw e; + } + } + + return Map.of(); } - private void deleteExistingIngestionPipelines() throws IOException { + private void deleteExistingIngestionPipelines() { Map pipelines = retrievePipelines(INGEST_PIPELINE_TYPE, null); pipelines.keySet().forEach(this::deleteIngestPipeline); } - private void deleteExistingSearchPipelines() throws IOException { + private void deleteExistingSearchPipelines() { Map pipelines = retrievePipelines(SEARCH_PIPELINE_TYPE, null); pipelines.keySet().forEach(this::deleteSearchPipeline); } @@ -1558,8 +1578,8 @@ protected Map deleteSearchPipeline(final String pipelineName) { * Retrieves all existing models * @return models as a map object */ - @SneakyThrows(ParseException.class) - protected List retrieveModels() throws IOException { + @SneakyThrows + protected List retrieveModels() { Response response = makeRequest( client(), "POST", @@ -1583,7 +1603,7 @@ protected List retrieveModels() throws IOException { .toList(); } - private void deleteExistingModels() throws IOException { + private void deleteExistingModels() { List modelIds = retrieveModels(); modelIds.forEach(m -> { try { @@ -1650,44 +1670,6 @@ protected Float getFeatureFieldCompressedNumber(final Float originNumber) { return Float.intBitsToFloat(freqBits); } - // Wipe of all the resources after execution of the tests. - protected void wipeOfTestResources() throws IOException { - try { - deleteExistingIngestionPipelines(); - } catch (ResponseException e) { - // It's possible that test fails during resources (index, model, pipeline, etc.) creation, when clean up - // these resources after test run, the delete methods will throw ResponseException with 404 Not Found code - // In this case, we just need to ignore this exception, for other exceptions, continue throwing - if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { - throw e; - } - } - - try { - deleteExistingSearchPipelines(); - } catch (ResponseException e) { - if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { - throw e; - } - } - - try { - deleteExistingModels(); - } catch (ResponseException e) { - if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { - throw e; - } - } - - try { - deleteExistingIndices(); - } catch (ResponseException e) { - if (RestStatus.NOT_FOUND != RestStatus.fromCode(e.getResponse().getStatusLine().getStatusCode())) { - throw e; - } - } - } - protected Object validateDocCountAndInfo( String indexName, int expectedDocCount, From 2f601fbf2ed4ee502f506060783139264e1e8ba9 Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Fri, 7 Feb 2025 17:09:38 -0800 Subject: [PATCH 4/5] Fixing failing bwc tests Signed-off-by: Weijia Zhao --- .../AbstractRestartUpgradeRestTestCase.java | 36 +++++++- .../bwc/restart/BatchIngestionIT.java | 24 ++--- .../bwc/restart/HybridSearchIT.java | 27 +++--- .../restart/HybridSearchWithRescoreIT.java | 25 ++--- .../bwc/restart/KnnRadialSearchIT.java | 24 ++--- .../bwc/restart/MultiModalSearchIT.java | 24 ++--- .../NeuralQueryEnricherProcessorIT.java | 67 +++++++------- .../bwc/restart/NeuralSparseSearchIT.java | 24 ++--- .../NeuralSparseTwoPhaseProcessorIT.java | 27 +++--- .../bwc/restart/SemanticSearchIT.java | 23 ++--- .../bwc/restart/TextChunkingProcessorIT.java | 13 +-- .../AbstractRollingUpgradeTestCase.java | 35 ++++++- .../bwc/rolling/BatchIngestionIT.java | 30 +++--- .../bwc/rolling/HybridSearchIT.java | 44 ++++----- .../rolling/HybridSearchWithRescoreIT.java | 44 ++++----- .../bwc/rolling/KnnRadialSearchIT.java | 35 +++---- .../bwc/rolling/MultiModalSearchIT.java | 33 +++---- .../NeuralQueryEnricherProcessorIT.java | 92 +++++++++---------- .../bwc/rolling/NeuralSparseSearchIT.java | 34 +++---- .../NeuralSparseTwoPhaseProcessorIT.java | 35 +++---- .../bwc/rolling/SemanticSearchIT.java | 33 +++---- .../bwc/rolling/TextChunkingProcessorIT.java | 3 +- .../neuralsearch/BaseNeuralSearchIT.java | 41 ++++++++- 23 files changed, 443 insertions(+), 330 deletions(-) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java index ca5a3e34e..151a8adf8 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java @@ -8,6 +8,9 @@ import java.nio.file.Path; import java.util.Locale; import java.util.Optional; + +import lombok.SneakyThrows; +import org.junit.After; import org.junit.Before; import org.opensearch.common.settings.Settings; import org.opensearch.neuralsearch.BaseNeuralSearchIT; @@ -20,11 +23,30 @@ public abstract class AbstractRestartUpgradeRestTestCase extends BaseNeuralSearchIT { + // Resources to be cleaned up after each test, need to assign the actual values in the test itself + protected String modelId; + protected String ingestPipelineName; + protected String searchPipelineName; + protected String indexName; + @Before - protected String getIndexNameForTest() { + public void initialize() { + // Initialize variables + this.modelId = null; + this.ingestPipelineName = null; + this.searchPipelineName = null; + // Creating index name by concatenating "neural-bwc-" prefix with test method name // for all the tests in this sub-project - return NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + this.indexName = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + } + + @SneakyThrows + @After + public void cleanUpResources() { + if (!isRunningAgainstOldCluster()) { + wipeOfTestResources(this.indexName, this.ingestPipelineName, this.modelId, this.searchPipelineName); + } } @Override @@ -53,6 +75,16 @@ protected final Settings restClientSettings() { .build(); } + @Override + protected boolean shouldCleanUpResources() { + // All NEW CLUSTER tests depend on resources created in OLD CLUSTER test cases + // Before NEW CLUSTER tests run, all OLD CLUSTER test cases will be run first + // We only want to clean up resources in NEW CLUSTER tests, also we don't want to clean up after each test case finishes + // this is because the cleanup method will pull every resource and delete, which will impact other tests + // Overriding the method in base class so that resources won't be accidentally clean up + return false; + } + protected static final boolean isRunningAgainstOldCluster() { return Boolean.parseBoolean(System.getProperty(RESTART_UPGRADE_OLD_CLUSTER)); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java index 4ad0bdc6d..4ee4f7bc6 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java @@ -23,26 +23,26 @@ public class BatchIngestionIT extends AbstractRestartUpgradeRestTestCase { public void testBatchIngestionWithNeuralSparseProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - String indexName = getIndexNameForTest(); + super.ingestPipelineName = PIPELINE_NAME; + if (isRunningAgainstOldCluster()) { - String modelId = uploadSparseEncodingModel(); - loadModel(modelId); - createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME, batchSize); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME, batchSize); createIndexWithConfiguration( - indexName, + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); List> docs = prepareDataForBulkIngestion(0, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); - validateDocCountAndInfo(indexName, 5, () -> getDocById(indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); + validateDocCountAndInfo(super.indexName, 5, () -> getDocById(super.indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); } else { - String modelId = null; - modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); List> docs = prepareDataForBulkIngestion(5, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); - validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); + validateDocCountAndInfo(super.indexName, 10, () -> getDocById(super.indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java index d42a29c59..022d87ce8 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java @@ -55,25 +55,28 @@ public void testNormalizationProcessor_whenIndexWithSingleShard_E2EFlow() throws private void validateNormalizationProcessor(final String fileName, final String pipelineName, final String searchPipelineName) throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = pipelineName; + super.searchPipelineName = searchPipelineName; + if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, pipelineName); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, pipelineName); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource(fileName).toURI())), pipelineName ); - addDocuments(getIndexNameForTest(), true); + addDocuments(super.indexName, true); createSearchPipeline(searchPipelineName); } else { - String modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocuments(getIndexNameForTest(), false); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); - hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); + super.modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + addDocuments(super.indexName, false); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); + validateTestIndex(super.indexName, searchPipelineName, hybridQueryBuilder); + hybridQueryBuilder = getQueryBuilder(super.modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(super.indexName, searchPipelineName, hybridQueryBuilder); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java index 2ca84745b..2357a1062 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java @@ -41,17 +41,18 @@ public class HybridSearchWithRescoreIT extends AbstractRestartUpgradeRestTestCas */ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -59,14 +60,14 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f })) ); } else { - String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); + validateTestIndex(super.indexName, hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(super.modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(super.indexName, hybridQueryBuilder, rescorer); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java index 02e52ce9d..16fd26a26 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java @@ -26,22 +26,24 @@ public class KnnRadialSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate radial query, pipeline and document count in restart-upgrade scenario public void testKnnRadialSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateIndexQuery(modelId); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + + loadModel(super.modelId); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateIndexQuery(super.modelId); } } @@ -54,7 +56,7 @@ private void validateIndexQuery(final String modelId) { .minScore(0.01f) .build(); - Map responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); + Map responseWithMinScoreQuery = search(super.indexName, neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScoreQuery); NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() @@ -64,7 +66,7 @@ private void validateIndexQuery(final String modelId) { .modelId(modelId) .maxDistance(100000f) .build(); - Map responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); + Map responseWithMaxDistanceQuery = search(super.indexName, neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxDistanceQuery); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java index 4b89eb901..98c215349 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java @@ -10,6 +10,7 @@ import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR; import static org.opensearch.neuralsearch.util.TestUtils.getModelId; + import org.opensearch.neuralsearch.query.NeuralQueryBuilder; public class MultiModalSearchIT extends AbstractRestartUpgradeRestTestCase { @@ -26,27 +27,28 @@ public class MultiModalSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateTestIndex(modelId); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateTestIndex(super.modelId); } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(2, docCount); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() .fieldName("passage_embedding") @@ -55,7 +57,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId) .k(1) .build(); - Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); + Map response = search(super.indexName, neuralQueryBuilder, 1); assertNotNull(response); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java index eb29e909e..6c25db994 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java @@ -35,36 +35,35 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws // will set the model_id after we obtain the id NeuralSparseQueryBuilder sparseEncodingQueryBuilderWithModelId = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1); + super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; + super.searchPipelineName = SPARSE_SEARCH_PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadSparseEncodingModel(); - loadModel(modelId); - sparseEncodingQueryBuilderWithModelId.modelId(modelId); - createPipelineForSparseEncodingProcessor(modelId, SPARSE_INGEST_PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); - createSearchRequestProcessor(modelId, SPARSE_SEARCH_PIPELINE_NAME); - updateIndexSettings( - getIndexNameForTest(), - Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME) - ); + createSearchRequestProcessor(super.modelId, SPARSE_SEARCH_PIPELINE_NAME); + updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME)); assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); } else { - String modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); - sparseEncodingQueryBuilderWithModelId.modelId(modelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); + sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); } } @@ -79,34 +78,36 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep .fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1) .build(); + this.ingestPipelineName = DENSE_INGEST_PIPELINE_NAME; + this.searchPipelineName = DENSE_SEARCH_PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - neuralQueryBuilderWithModelId.modelId(modelId); - createPipelineProcessor(modelId, DENSE_INGEST_PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + neuralQueryBuilderWithModelId.modelId(super.modelId); + createPipelineProcessor(super.modelId, DENSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), DENSE_INGEST_PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_TEXT_FIELD, TEXT_1, null, null); + addDocument(super.indexName, "0", TEST_TEXT_FIELD, TEXT_1, null, null); - createSearchRequestProcessor(modelId, DENSE_SEARCH_PIPELINE_NAME); - updateIndexSettings(getIndexNameForTest(), Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); + createSearchRequestProcessor(super.modelId, DENSE_SEARCH_PIPELINE_NAME); + updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") ); } else { - String modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - neuralQueryBuilderWithModelId.modelId(modelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + neuralQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") ); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java index 893da4c87..0d40b18fc 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java @@ -32,18 +32,20 @@ public class NeuralSparseSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testSparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; + if (isRunningAgainstOldCluster()) { - String modelId = uploadSparseEncodingModel(); - loadModel(modelId); - createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); addSparseEncodingDoc( - getIndexNameForTest(), + super.indexName, "0", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc1), @@ -51,23 +53,23 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEXT_1) ); } else { - String modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); addSparseEncodingDoc( - getIndexNameForTest(), + super.indexName, "1", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc2), List.of(TEST_TEXT_FIELD), List.of(TEXT_2) ); - validateTestIndex(modelId); + validateTestIndex(super.modelId); } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(2, docCount); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_SPARSE_ENCODING_FIELD) @@ -75,7 +77,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId); MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD, TEXT_1); boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); - Map response = search(getIndexNameForTest(), boolQueryBuilder, 1); + Map response = search(super.indexName, boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(response); assertEquals("0", firstInnerHit.get("_id")); diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java index 3ee154d8d..febca719f 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java @@ -26,29 +26,32 @@ public class NeuralSparseTwoPhaseProcessorIT extends AbstractRestartUpgradeRestT public void testNeuralSparseQueryTwoPhaseProcessor_NeuralSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); + super.ingestPipelineName = NEURAL_SPARSE_INGEST_PIPELINE_NAME; + super.searchPipelineName = NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME; + if (isRunningAgainstOldCluster()) { - String modelId = uploadSparseEncodingModel(); - loadModel(modelId); - neuralSparseQueryBuilder.modelId(modelId); - createPipelineForSparseEncodingProcessor(modelId, NEURAL_SPARSE_INGEST_PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + neuralSparseQueryBuilder.modelId(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, NEURAL_SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), NEURAL_SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); createNeuralSparseTwoPhaseSearchProcessor(NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME); updateIndexSettings( - getIndexNameForTest(), + super.indexName, Settings.builder().put("index.search.default_pipeline", NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME) ); - Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); + Object resultWith2PhasePipeline = search(super.indexName, neuralSparseQueryBuilder, 1).get("hits"); assertNotNull(resultWith2PhasePipeline); } else { - String modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(modelId); - neuralSparseQueryBuilder.modelId(modelId); - Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); + super.modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); + neuralSparseQueryBuilder.modelId(super.modelId); + Object resultWith2PhasePipeline = search(super.indexName, neuralSparseQueryBuilder, 1).get("hits"); assertNotNull(resultWith2PhasePipeline); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java index d378ab12f..984906b66 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java @@ -24,27 +24,28 @@ public class SemanticSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testTextEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - String modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); } else { - String modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(modelId); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null); - validateTestIndex(modelId); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, null, null); + validateTestIndex(super.modelId); } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(2, docCount); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() .fieldName("passage_embedding") @@ -52,7 +53,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId) .k(1) .build(); - Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); + Map response = search(super.indexName, neuralQueryBuilder, 1); assertNotNull(response); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java index 254d49041..3a32201a5 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java @@ -33,15 +33,16 @@ public class TextChunkingProcessorIT extends AbstractRestartUpgradeRestTestCase // Validate process, pipeline and document count in restart-upgrade scenario public void testTextChunkingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - String indexName = getIndexNameForTest(); + super.ingestPipelineName = PIPELINE_NAME; + if (isRunningAgainstOldCluster()) { createPipelineForTextChunkingProcessor(PIPELINE_NAME); - createChunkingIndex(indexName); - addDocument(indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(indexName, OUTPUT_FIELD, 1, expectedPassages); + createChunkingIndex(super.indexName); + addDocument(super.indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(super.indexName, OUTPUT_FIELD, 1, expectedPassages); } else { - addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages); + addDocument(super.indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(super.indexName, OUTPUT_FIELD, 2, expectedPassages); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java index a1ad7178c..ceb8b795f 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java @@ -8,6 +8,9 @@ import java.nio.file.Path; import java.util.Locale; import java.util.Optional; + +import lombok.SneakyThrows; +import org.junit.After; import org.junit.Before; import org.opensearch.common.settings.Settings; import org.opensearch.neuralsearch.BaseNeuralSearchIT; @@ -22,12 +25,30 @@ import org.opensearch.test.rest.OpenSearchRestTestCase; public abstract class AbstractRollingUpgradeTestCase extends BaseNeuralSearchIT { + // Resources to be cleaned up after each test, need to assign the actual values in the test itself + protected String modelId; + protected String ingestPipelineName; + protected String searchPipelineName; + protected String indexName; @Before - protected String getIndexNameForTest() { + public void initialize() { + // Initialize variables + this.modelId = null; + this.ingestPipelineName = null; + this.searchPipelineName = null; + // Creating index name by concatenating "neural-bwc-" prefix with test method name // for all the tests in this sub-project - return NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + this.indexName = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + } + + @SneakyThrows + @After + public void cleanUpResources() { + if (getClusterType() == ClusterType.UPGRADED) { + wipeOfTestResources(this.indexName, this.ingestPipelineName, this.modelId, this.searchPipelineName); + } } @Override @@ -56,6 +77,16 @@ protected final Settings restClientSettings() { .build(); } + @Override + protected boolean shouldCleanUpResources() { + // All UPGRADE tests depend on resources created in OLD and MIXED test cases + // Before UPGRADE tests run, all OLD and MIXED test cases will be run first + // We only want to clean up resources in upgrade tests, also we don't want to clean up after each test case finishes + // this is because the cleanup method will pull every resource and delete, which will impact other tests + // Overriding the method in base class so that resources won't be accidentally clean up + return false; + } + protected enum ClusterType { OLD, MIXED, diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java index 6fc64fa96..20d2ba3e4 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java @@ -22,35 +22,35 @@ public class BatchIngestionIT extends AbstractRollingUpgradeTestCase { public void testBatchIngestion_SparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); - String indexName = getIndexNameForTest(); - String sparseModelId = null; + super.ingestPipelineName = SPARSE_PIPELINE; + switch (getClusterType()) { case OLD: - sparseModelId = uploadSparseEncodingModel(); - loadModel(sparseModelId); - createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_PIPELINE, 2); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_PIPELINE, 2); createIndexWithConfiguration( - indexName, + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_PIPELINE ); List> docs = prepareDataForBulkIngestion(0, 5); bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docs); - validateDocCountAndInfo(indexName, 5, () -> getDocById(indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); + validateDocCountAndInfo(super.indexName, 5, () -> getDocById(super.indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); break; case MIXED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); List> docsForMixed = prepareDataForBulkIngestion(5, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForMixed); - validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForMixed); + validateDocCountAndInfo(super.indexName, 10, () -> getDocById(super.indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); break; case UPGRADED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); List> docsForUpgraded = prepareDataForBulkIngestion(10, 5); - bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); - validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); + validateDocCountAndInfo(super.indexName, 15, () -> getDocById(super.indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java index c4beb3845..1f114710d 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java @@ -36,24 +36,26 @@ public class HybridSearchIT extends AbstractRollingUpgradeTestCase { private static final int NUM_DOCS_PER_ROUND = 1; private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding"; protected static final String RESCORE_QUERY = "hi"; - private static String modelId = ""; // Test rolling-upgrade normalization processor when index with multiple shards // Create Text Embedding Processor, Ingestion Pipeline, add document and search pipeline with noramlization processor // Validate process , pipeline and document count in rolling-upgrade scenario public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; + super.searchPipelineName = SEARCH_PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -62,28 +64,28 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr ); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); - hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); + loadModel(super.modelId); + addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, null); + hybridQueryBuilder = getQueryBuilder(super.modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, null); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -96,11 +98,11 @@ private void validateTestIndexOnUpgrade( HybridQueryBuilder hybridQueryBuilder, QueryBuilder rescorer ) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); Map searchResponseAsMap = search( - getIndexNameForTest(), + super.indexName, hybridQueryBuilder, rescorer, 1, diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java index afed627e8..a2233161a 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java @@ -37,24 +37,26 @@ public class HybridSearchWithRescoreIT extends AbstractRollingUpgradeTestCase { private static final int NUM_DOCS_PER_ROUND = 1; private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding"; protected static final String RESCORE_QUERY = "hi"; - private static String modelId = ""; /** * Test normalization with hybrid query and rescore. This test is required as rescore will not be compatible with version lower than 2.15 */ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; + super.searchPipelineName = SEARCH_PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -63,30 +65,30 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr ); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, rescorer); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, rescorer); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + loadModel(super.modelId); + addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(super.modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, rescorer); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -99,11 +101,11 @@ private void validateTestIndexOnUpgrade( HybridQueryBuilder hybridQueryBuilder, QueryBuilder rescorer ) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); Map searchResponseAsMap = search( - getIndexNameForTest(), + super.indexName, hybridQueryBuilder, rescorer, 1, diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java index e59f6e911..4ab6691f7 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java @@ -24,43 +24,44 @@ public class KnnRadialSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk"; private static final int NUM_DOCS_PER_ROUND = 1; - private static String modelId = ""; // Test rolling-upgrade with kNN radial search // Create Text Image Embedding Processor, Ingestion Pipeline and add document // Validate radial query, pipeline and document count in rolling-upgrade scenario public void testKnnRadialSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); + super.ingestPipelineName = PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadTextImageEmbeddingModel(); - loadModel(modelId); - createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextImageEmbeddingModel(); + loadModel(super.modelId); + createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); + validateIndexQueryOnUpgrade(totalDocsCountMixed, super.modelId, TEXT, TEST_IMAGE_TEXT); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); + validateIndexQueryOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateIndexQueryOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + loadModel(super.modelId); + addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateIndexQueryOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -69,7 +70,7 @@ public void testKnnRadialSearch_E2EFlow() throws Exception { private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String modelId, final String text, final String imageText) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); @@ -81,7 +82,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo .minScore(0.01f) .build(); - Map responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); + Map responseWithMinScore = search(super.indexName, neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScore); NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() @@ -92,7 +93,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo .maxDistance(100000f) .build(); - Map responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); + Map responseWithMaxScore = search(super.indexName, neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxScore); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java index 11be782a4..d9c15a780 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java @@ -24,43 +24,44 @@ public class MultiModalSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk"; private static final int NUM_DOCS_PER_ROUND = 1; - private static String modelId = ""; // Test rolling-upgrade test image embedding processor // Create Text Image Embedding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); + super.ingestPipelineName = PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadTextImageEmbeddingModel(); - loadModel(modelId); - createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextImageEmbeddingModel(); + loadModel(super.modelId); + createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT, TEST_IMAGE_TEXT); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + loadModel(super.modelId); + addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -69,7 +70,7 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId, final String text, final String imageText) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); NeuralQueryBuilder neuralQueryBuilderWithKQuery = NeuralQueryBuilder.builder() @@ -80,7 +81,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .k(1) .build(); - Map responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1); + Map responseWithKQuery = search(super.indexName, neuralQueryBuilderWithKQuery, 1); assertNotNull(responseWithKQuery); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java index 6f66d0f72..0aaa74af0 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java @@ -26,8 +26,6 @@ public class NeuralQueryEnricherProcessorIT extends AbstractRollingUpgradeTestCa private static final String TEST_ENCODING_FIELD = "passage_embedding"; private static final String TEST_TEXT_FIELD = "passage_text"; private static final String TEXT_1 = "Hello world a b"; - private String sparseModelId = ""; - private String denseModelId = ""; // test of NeuralQueryEnricherProcessor supports neural_sparse query default model_id // the feature is introduced from 2.13 @@ -38,47 +36,46 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws // will set the model_id after we obtain the id NeuralSparseQueryBuilder sparseEncodingQueryBuilderWithModelId = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1); + super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; + super.searchPipelineName = SPARSE_SEARCH_PIPELINE_NAME; switch (getClusterType()) { case OLD: - sparseModelId = uploadSparseEncodingModel(); - loadModel(sparseModelId); - sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); - createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_INGEST_PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); - createSearchRequestProcessor(sparseModelId, SPARSE_SEARCH_PIPELINE_NAME); - updateIndexSettings( - getIndexNameForTest(), - Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME) - ); + addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + createSearchRequestProcessor(super.modelId, SPARSE_SEARCH_PIPELINE_NAME); + updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME)); assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); break; case MIXED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); + sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); break; case UPGRADED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); + sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); break; default: @@ -98,49 +95,48 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep .fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1) .build(); + super.ingestPipelineName = DENSE_INGEST_PIPELINE_NAME; + super.searchPipelineName = DENSE_SEARCH_PIPELINE_NAME; switch (getClusterType()) { case OLD: - denseModelId = uploadTextEmbeddingModel(); - loadModel(denseModelId); - neuralQueryBuilderWithModelId.modelId(denseModelId); - createPipelineProcessor(denseModelId, DENSE_INGEST_PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + neuralQueryBuilderWithModelId.modelId(super.modelId); + createPipelineProcessor(super.modelId, DENSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), DENSE_INGEST_PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_TEXT_FIELD, TEXT_1, null, null); + addDocument(super.indexName, "0", TEST_TEXT_FIELD, TEXT_1, null, null); - createSearchRequestProcessor(denseModelId, DENSE_SEARCH_PIPELINE_NAME); - updateIndexSettings( - getIndexNameForTest(), - Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME) - ); + createSearchRequestProcessor(super.modelId, DENSE_SEARCH_PIPELINE_NAME); + updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") ); break; case MIXED: - denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(denseModelId); - neuralQueryBuilderWithModelId.modelId(denseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + neuralQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") ); break; case UPGRADED: - denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(denseModelId); - neuralQueryBuilderWithModelId.modelId(denseModelId); + super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(super.modelId); + neuralQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") ); break; default: diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java index b91211859..8cea81b24 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java @@ -8,6 +8,7 @@ import java.nio.file.Path; import java.util.List; import java.util.Map; + import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.neuralsearch.util.TestUtils; @@ -32,25 +33,26 @@ public class NeuralSparseSearchIT extends AbstractRollingUpgradeTestCase { private final Map testRankFeaturesDoc2 = TestUtils.createRandomTokenWeightMap(TEST_TOKENS_2); private final Map testRankFeaturesDoc3 = TestUtils.createRandomTokenWeightMap(TEST_TOKENS_3); private static final int NUM_DOCS_PER_ROUND = 1; - private static String modelId = ""; // Test rolling-upgrade test sparse embedding processor // Create Sparse Encoding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testSparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); + super.ingestPipelineName = PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadSparseEncodingModel(); - loadModel(modelId); - createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); addSparseEncodingDoc( - getIndexNameForTest(), + super.indexName, "0", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc1), @@ -59,13 +61,13 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { ); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId); addSparseEncodingDoc( - getIndexNameForTest(), + super.indexName, "1", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc2), @@ -74,22 +76,22 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { ); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); + loadModel(super.modelId); addSparseEncodingDoc( - getIndexNameForTest(), + super.indexName, "2", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc3), List.of(TEST_TEXT_FIELD), List.of(TEXT_UPGRADED) ); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -97,7 +99,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { } private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); @@ -106,7 +108,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .modelId(modelId); MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD, TEXT); boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); - Map response = search(getIndexNameForTest(), boolQueryBuilder, 1); + Map response = search(super.indexName, boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(response); assertEquals("0", firstInnerHit.get("_id")); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java index 60dd9eac7..cecea39ee 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java @@ -22,7 +22,6 @@ public class NeuralSparseTwoPhaseProcessorIT extends AbstractRollingUpgradeTestC private static final String TEST_ENCODING_FIELD = "passage_embedding"; private static final String TEST_TEXT_FIELD = "passage_text"; private static final String TEXT_1 = "Hello world a b"; - private String sparseModelId = ""; // test of NeuralSparseTwoPhaseProcessor supports neural_sparse query's two phase speed up // the feature is introduced from 2.15 @@ -30,37 +29,33 @@ public void testNeuralSparseTwoPhaseProcessorIT_NeuralSparseSearch_E2EFlow() thr waitForClusterHealthGreen(NODES_BWC_CLUSTER); // will set the model_id after we obtain the id NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); + super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; + super.searchPipelineName = SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME; switch (getClusterType()) { case OLD: - sparseModelId = uploadSparseEncodingModel(); - loadModel(sparseModelId); - neuralSparseQueryBuilder.modelId(sparseModelId); - createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_INGEST_PIPELINE_NAME); + super.modelId = uploadSparseEncodingModel(); + loadModel(super.modelId); + neuralSparseQueryBuilder.modelId(super.modelId); + createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); createNeuralSparseTwoPhaseSearchProcessor(SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME); updateIndexSettings( - getIndexNameForTest(), + super.indexName, Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME) ); - assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); + assertNotNull(search(super.indexName, neuralSparseQueryBuilder, 1).get("hits")); break; - case MIXED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - neuralSparseQueryBuilder.modelId(sparseModelId); - assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); - break; - case UPGRADED: - sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(sparseModelId); - neuralSparseQueryBuilder.modelId(sparseModelId); - assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); + case MIXED, UPGRADED: + super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(super.modelId); + neuralSparseQueryBuilder.modelId(super.modelId); + assertNotNull(search(super.indexName, neuralSparseQueryBuilder, 1).get("hits")); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java index 99107e79c..e84adde55 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java @@ -19,43 +19,44 @@ public class SemanticSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEXT_MIXED = "Hello world mixed"; private static final String TEXT_UPGRADED = "Hello world upgraded"; private static final int NUM_DOCS_PER_ROUND = 1; - private static String modelId = ""; // Test rolling-upgrade Semantic Search // Create Text Embedding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testSemanticSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); + super.ingestPipelineName = PIPELINE_NAME; + switch (getClusterType()) { case OLD: - modelId = uploadTextEmbeddingModel(); - loadModel(modelId); - createPipelineProcessor(modelId, PIPELINE_NAME); + super.modelId = uploadTextEmbeddingModel(); + loadModel(super.modelId); + createPipelineProcessor(super.modelId, PIPELINE_NAME); createIndexWithConfiguration( - getIndexNameForTest(), + super.indexName, Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); + addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); break; case MIXED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT); - addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT); + addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED); } break; case UPGRADED: - modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(modelId); - addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED); + loadModel(super.modelId); + addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED); break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -64,7 +65,7 @@ public void testSemanticSearch_E2EFlow() throws Exception { } private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId, final String text) throws Exception { - int docCount = getDocCount(getIndexNameForTest()); + int docCount = getDocCount(super.indexName); assertEquals(numberOfDocs, docCount); loadModel(modelId); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() @@ -73,7 +74,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .queryText(text) .k(1) .build(); - Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); + Map response = search(super.indexName, neuralQueryBuilder, 1); assertNotNull(response); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java index 227b64f85..1733d511b 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java @@ -35,7 +35,8 @@ public class TextChunkingProcessorIT extends AbstractRollingUpgradeTestCase { // Validate process, pipeline and document count in rolling-upgrade scenario public void testTextChunkingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - String indexName = getIndexNameForTest(); + super.ingestPipelineName = PIPELINE_NAME; + switch (getClusterType()) { case OLD: createPipelineForTextChunkingProcessor(PIPELINE_NAME); diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index 72b9604c9..3b6143e09 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -137,10 +137,16 @@ public void setupSettings() { // Wipe of all the resources after execution of the tests. @After public void cleanUp() { - deleteExistingIngestionPipelines(); - deleteExistingSearchPipelines(); - deleteExistingModels(); - deleteExistingIndices(); + if (shouldCleanUpResources()) { + deleteExistingIngestionPipelines(); + deleteExistingSearchPipelines(); + deleteExistingModels(); + deleteExistingIndices(); + } + } + + protected boolean shouldCleanUpResources() { + return true; } protected ThreadPool setUpThreadPool() { @@ -1616,6 +1622,33 @@ private void deleteExistingModels() { }); } + @SneakyThrows + protected void wipeOfTestResources( + final String indexName, + final String ingestPipeline, + final String modelId, + final String searchPipeline + ) { + if (ingestPipeline != null) { + deleteIngestPipeline(ingestPipeline); + } + if (searchPipeline != null) { + deleteSearchPipeline(searchPipeline); + } + if (modelId != null) { + try { + deleteModel(modelId); + } catch (AssertionError e) { + // sometimes we have flaky test that the model state doesn't change after call undeploy api + // for this case we can call undeploy api one more time + deleteModel(modelId); + } + } + if (indexName != null) { + deleteIndex(indexName); + } + } + protected float computeExpectedScore(final String modelId, final Map tokenWeightMap, final String queryText) { Map queryTokens = runSparseModelInference(modelId, queryText); return computeExpectedScore(tokenWeightMap, queryTokens); From e296bc6046a619d57b9b65c438ee71df974223c3 Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Mon, 10 Feb 2025 13:15:12 -0800 Subject: [PATCH 5/5] Revert bwc test resource cleanup logic Signed-off-by: Weijia Zhao --- .../AbstractRestartUpgradeRestTestCase.java | 26 +---- .../bwc/restart/BatchIngestionIT.java | 30 ++--- .../bwc/restart/HybridSearchIT.java | 32 +++--- .../restart/HybridSearchWithRescoreIT.java | 32 +++--- .../bwc/restart/KnnRadialSearchIT.java | 29 ++--- .../bwc/restart/MultiModalSearchIT.java | 29 ++--- .../NeuralQueryEnricherProcessorIT.java | 85 ++++++++------ .../bwc/restart/NeuralSparseSearchIT.java | 43 +++---- .../NeuralSparseTwoPhaseProcessorIT.java | 39 ++++--- .../bwc/restart/SemanticSearchIT.java | 28 +++-- .../bwc/restart/TextChunkingProcessorIT.java | 17 +-- .../AbstractRollingUpgradeTestCase.java | 25 +--- .../bwc/rolling/BatchIngestionIT.java | 36 +++--- .../bwc/rolling/HybridSearchIT.java | 50 ++++---- .../rolling/HybridSearchWithRescoreIT.java | 52 +++++---- .../bwc/rolling/KnnRadialSearchIT.java | 41 ++++--- .../bwc/rolling/MultiModalSearchIT.java | 39 ++++--- .../NeuralQueryEnricherProcessorIT.java | 108 ++++++++++-------- .../bwc/rolling/NeuralSparseSearchIT.java | 54 ++++----- .../NeuralSparseTwoPhaseProcessorIT.java | 44 ++++--- .../bwc/rolling/SemanticSearchIT.java | 39 ++++--- .../bwc/rolling/TextChunkingProcessorIT.java | 13 ++- 22 files changed, 469 insertions(+), 422 deletions(-) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java index 151a8adf8..d7428d409 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/AbstractRestartUpgradeRestTestCase.java @@ -8,9 +8,6 @@ import java.nio.file.Path; import java.util.Locale; import java.util.Optional; - -import lombok.SneakyThrows; -import org.junit.After; import org.junit.Before; import org.opensearch.common.settings.Settings; import org.opensearch.neuralsearch.BaseNeuralSearchIT; @@ -23,30 +20,11 @@ public abstract class AbstractRestartUpgradeRestTestCase extends BaseNeuralSearchIT { - // Resources to be cleaned up after each test, need to assign the actual values in the test itself - protected String modelId; - protected String ingestPipelineName; - protected String searchPipelineName; - protected String indexName; - @Before - public void initialize() { - // Initialize variables - this.modelId = null; - this.ingestPipelineName = null; - this.searchPipelineName = null; - + protected String getIndexNameForTest() { // Creating index name by concatenating "neural-bwc-" prefix with test method name // for all the tests in this sub-project - this.indexName = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); - } - - @SneakyThrows - @After - public void cleanUpResources() { - if (!isRunningAgainstOldCluster()) { - wipeOfTestResources(this.indexName, this.ingestPipelineName, this.modelId, this.searchPipelineName); - } + return NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); } @Override diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java index 4ee4f7bc6..7fa0b1301 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/BatchIngestionIT.java @@ -23,26 +23,30 @@ public class BatchIngestionIT extends AbstractRestartUpgradeRestTestCase { public void testBatchIngestionWithNeuralSparseProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - + String indexName = getIndexNameForTest(); if (isRunningAgainstOldCluster()) { - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME, batchSize); + String modelId = uploadSparseEncodingModel(); + loadModel(modelId); + createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME, batchSize); createIndexWithConfiguration( - super.indexName, + indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); List> docs = prepareDataForBulkIngestion(0, 5); - bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); - validateDocCountAndInfo(super.indexName, 5, () -> getDocById(super.indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); + validateDocCountAndInfo(indexName, 5, () -> getDocById(indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); } else { - super.modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - List> docs = prepareDataForBulkIngestion(5, 5); - bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); - validateDocCountAndInfo(super.indexName, 10, () -> getDocById(super.indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); + String modelId = null; + modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + try { + List> docs = prepareDataForBulkIngestion(5, 5); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, PIPELINE_NAME, docs); + validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); + } finally { + wipeOfTestResources(indexName, PIPELINE_NAME, modelId, null); + } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java index 022d87ce8..d08d208da 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchIT.java @@ -55,28 +55,30 @@ public void testNormalizationProcessor_whenIndexWithSingleShard_E2EFlow() throws private void validateNormalizationProcessor(final String fileName, final String pipelineName, final String searchPipelineName) throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = pipelineName; - super.searchPipelineName = searchPipelineName; - if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, pipelineName); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, pipelineName); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource(fileName).toURI())), pipelineName ); - addDocuments(super.indexName, true); + addDocuments(getIndexNameForTest(), true); createSearchPipeline(searchPipelineName); } else { - super.modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - addDocuments(super.indexName, false); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); - validateTestIndex(super.indexName, searchPipelineName, hybridQueryBuilder); - hybridQueryBuilder = getQueryBuilder(super.modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(super.indexName, searchPipelineName, hybridQueryBuilder); + String modelId = null; + try { + modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocuments(getIndexNameForTest(), false); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); + } finally { + wipeOfTestResources(getIndexNameForTest(), pipelineName, modelId, searchPipelineName); + } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java index 2357a1062..9329a934b 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/HybridSearchWithRescoreIT.java @@ -41,18 +41,17 @@ public class HybridSearchWithRescoreIT extends AbstractRestartUpgradeRestTestCas */ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, PIPELINE_NAME); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -60,14 +59,19 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr Map.of(PARAM_NAME_WEIGHTS, Arrays.toString(new float[] { 0.3f, 0.7f })) ); } else { - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); - QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndex(super.indexName, hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(super.modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndex(super.indexName, hybridQueryBuilder, rescorer); + String modelId = null; + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndex(getIndexNameForTest(), hybridQueryBuilder, rescorer); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java index 16fd26a26..2e3ac34ab 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/KnnRadialSearchIT.java @@ -26,24 +26,27 @@ public class KnnRadialSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate radial query, pipeline and document count in restart-upgrade scenario public void testKnnRadialSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - - loadModel(super.modelId); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateIndexQuery(super.modelId); + String modelId = null; + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateIndexQuery(modelId); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } } } @@ -56,7 +59,7 @@ private void validateIndexQuery(final String modelId) { .minScore(0.01f) .build(); - Map responseWithMinScoreQuery = search(super.indexName, neuralQueryBuilderWithMinScoreQuery, 1); + Map responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScoreQuery); NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() @@ -66,7 +69,7 @@ private void validateIndexQuery(final String modelId) { .modelId(modelId) .maxDistance(100000f) .build(); - Map responseWithMaxDistanceQuery = search(super.indexName, neuralQueryBuilderWithMaxDistanceQuery, 1); + Map responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxDistanceQuery); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java index 98c215349..df3f8e94e 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/MultiModalSearchIT.java @@ -10,7 +10,6 @@ import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR; import static org.opensearch.neuralsearch.util.TestUtils.getModelId; - import org.opensearch.neuralsearch.query.NeuralQueryBuilder; public class MultiModalSearchIT extends AbstractRestartUpgradeRestTestCase { @@ -27,28 +26,32 @@ public class MultiModalSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); } else { - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); - validateTestIndex(super.modelId); + String modelId = null; + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1); + validateTestIndex(modelId); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(2, docCount); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() .fieldName("passage_embedding") @@ -57,7 +60,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId) .k(1) .build(); - Map response = search(super.indexName, neuralQueryBuilder, 1); + Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java index 6c25db994..97182bfa0 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralQueryEnricherProcessorIT.java @@ -35,36 +35,42 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws // will set the model_id after we obtain the id NeuralSparseQueryBuilder sparseEncodingQueryBuilderWithModelId = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1); - super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; - super.searchPipelineName = SPARSE_SEARCH_PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); + String modelId = uploadSparseEncodingModel(); + loadModel(modelId); + sparseEncodingQueryBuilderWithModelId.modelId(modelId); + createPipelineForSparseEncodingProcessor(modelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); - createSearchRequestProcessor(super.modelId, SPARSE_SEARCH_PIPELINE_NAME); - updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME)); - assertEquals( - search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") + createSearchRequestProcessor(modelId, SPARSE_SEARCH_PIPELINE_NAME); + updateIndexSettings( + getIndexNameForTest(), + Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME) ); - } else { - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); assertEquals( - search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); + } else { + String modelId = null; + try { + modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + sparseEncodingQueryBuilderWithModelId.modelId(modelId); + assertEquals( + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + ); + } finally { + wipeOfTestResources(getIndexNameForTest(), SPARSE_INGEST_PIPELINE_NAME, modelId, SPARSE_SEARCH_PIPELINE_NAME); + } } } @@ -78,37 +84,40 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep .fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1) .build(); - this.ingestPipelineName = DENSE_INGEST_PIPELINE_NAME; - this.searchPipelineName = DENSE_SEARCH_PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - neuralQueryBuilderWithModelId.modelId(super.modelId); - createPipelineProcessor(super.modelId, DENSE_INGEST_PIPELINE_NAME); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + neuralQueryBuilderWithModelId.modelId(modelId); + createPipelineProcessor(modelId, DENSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), DENSE_INGEST_PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_TEXT_FIELD, TEXT_1, null, null); + addDocument(getIndexNameForTest(), "0", TEST_TEXT_FIELD, TEXT_1, null, null); - createSearchRequestProcessor(super.modelId, DENSE_SEARCH_PIPELINE_NAME); - updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); + createSearchRequestProcessor(modelId, DENSE_SEARCH_PIPELINE_NAME); + updateIndexSettings(getIndexNameForTest(), Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); assertEquals( - search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") ); } else { - super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - neuralQueryBuilderWithModelId.modelId(super.modelId); + String modelId = null; + try { + modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + neuralQueryBuilderWithModelId.modelId(modelId); - assertEquals( - search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") - ); + assertEquals( + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + ); + } finally { + wipeOfTestResources(getIndexNameForTest(), DENSE_INGEST_PIPELINE_NAME, modelId, DENSE_SEARCH_PIPELINE_NAME); + } } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java index 0d40b18fc..5978f71ab 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseSearchIT.java @@ -32,20 +32,18 @@ public class NeuralSparseSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testSparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - if (isRunningAgainstOldCluster()) { - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME); + String modelId = uploadSparseEncodingModel(); + loadModel(modelId); + createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); addSparseEncodingDoc( - super.indexName, + getIndexNameForTest(), "0", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc1), @@ -53,23 +51,28 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEXT_1) ); } else { - super.modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - addSparseEncodingDoc( - super.indexName, - "1", - List.of(TEST_SPARSE_ENCODING_FIELD), - List.of(testRankFeaturesDoc2), - List.of(TEST_TEXT_FIELD), - List.of(TEXT_2) - ); - validateTestIndex(super.modelId); + String modelId = null; + try { + modelId = TestUtils.getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + addSparseEncodingDoc( + getIndexNameForTest(), + "1", + List.of(TEST_SPARSE_ENCODING_FIELD), + List.of(testRankFeaturesDoc2), + List.of(TEST_TEXT_FIELD), + List.of(TEXT_2) + ); + validateTestIndex(modelId); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(2, docCount); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_SPARSE_ENCODING_FIELD) @@ -77,7 +80,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId); MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD, TEXT_1); boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); - Map response = search(super.indexName, boolQueryBuilder, 1); + Map response = search(getIndexNameForTest(), boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(response); assertEquals("0", firstInnerHit.get("_id")); diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java index febca719f..6a6994809 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/NeuralSparseTwoPhaseProcessorIT.java @@ -26,33 +26,40 @@ public class NeuralSparseTwoPhaseProcessorIT extends AbstractRestartUpgradeRestT public void testNeuralSparseQueryTwoPhaseProcessor_NeuralSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); - super.ingestPipelineName = NEURAL_SPARSE_INGEST_PIPELINE_NAME; - super.searchPipelineName = NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME; - if (isRunningAgainstOldCluster()) { - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - neuralSparseQueryBuilder.modelId(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, NEURAL_SPARSE_INGEST_PIPELINE_NAME); + String modelId = uploadSparseEncodingModel(); + loadModel(modelId); + neuralSparseQueryBuilder.modelId(modelId); + createPipelineForSparseEncodingProcessor(modelId, NEURAL_SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), NEURAL_SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); createNeuralSparseTwoPhaseSearchProcessor(NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME); updateIndexSettings( - super.indexName, + getIndexNameForTest(), Settings.builder().put("index.search.default_pipeline", NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME) ); - Object resultWith2PhasePipeline = search(super.indexName, neuralSparseQueryBuilder, 1).get("hits"); + Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); assertNotNull(resultWith2PhasePipeline); } else { - super.modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - neuralSparseQueryBuilder.modelId(super.modelId); - Object resultWith2PhasePipeline = search(super.indexName, neuralSparseQueryBuilder, 1).get("hits"); - assertNotNull(resultWith2PhasePipeline); + String modelId = null; + try { + modelId = TestUtils.getModelId(getIngestionPipeline(NEURAL_SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(modelId); + neuralSparseQueryBuilder.modelId(modelId); + Object resultWith2PhasePipeline = search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits"); + assertNotNull(resultWith2PhasePipeline); + } finally { + wipeOfTestResources( + getIndexNameForTest(), + NEURAL_SPARSE_INGEST_PIPELINE_NAME, + modelId, + NEURAL_SPARSE_TWO_PHASE_SEARCH_PIPELINE_NAME + ); + } } } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java index 984906b66..9dc35c02d 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/SemanticSearchIT.java @@ -24,28 +24,32 @@ public class SemanticSearchIT extends AbstractRestartUpgradeRestTestCase { // Validate process , pipeline and document count in restart-upgrade scenario public void testTextEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; if (isRunningAgainstOldCluster()) { - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, PIPELINE_NAME); + String modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); } else { - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_1, null, null); - validateTestIndex(super.modelId); + String modelId = null; + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(modelId); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, null, null); + validateTestIndex(modelId); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } } } private void validateTestIndex(final String modelId) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(2, docCount); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() .fieldName("passage_embedding") @@ -53,7 +57,7 @@ private void validateTestIndex(final String modelId) throws Exception { .modelId(modelId) .k(1) .build(); - Map response = search(super.indexName, neuralQueryBuilder, 1); + Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } } diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java index 3a32201a5..d68903478 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/restart/TextChunkingProcessorIT.java @@ -33,16 +33,19 @@ public class TextChunkingProcessorIT extends AbstractRestartUpgradeRestTestCase // Validate process, pipeline and document count in restart-upgrade scenario public void testTextChunkingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - + String indexName = getIndexNameForTest(); if (isRunningAgainstOldCluster()) { createPipelineForTextChunkingProcessor(PIPELINE_NAME); - createChunkingIndex(super.indexName); - addDocument(super.indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(super.indexName, OUTPUT_FIELD, 1, expectedPassages); + createChunkingIndex(indexName); + addDocument(indexName, "0", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(indexName, OUTPUT_FIELD, 1, expectedPassages); } else { - addDocument(super.indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(super.indexName, OUTPUT_FIELD, 2, expectedPassages); + try { + addDocument(indexName, "1", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(indexName, OUTPUT_FIELD, 2, expectedPassages); + } finally { + wipeOfTestResources(indexName, PIPELINE_NAME, null, null); + } } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java index ceb8b795f..5aa243693 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/AbstractRollingUpgradeTestCase.java @@ -8,9 +8,6 @@ import java.nio.file.Path; import java.util.Locale; import java.util.Optional; - -import lombok.SneakyThrows; -import org.junit.After; import org.junit.Before; import org.opensearch.common.settings.Settings; import org.opensearch.neuralsearch.BaseNeuralSearchIT; @@ -25,30 +22,12 @@ import org.opensearch.test.rest.OpenSearchRestTestCase; public abstract class AbstractRollingUpgradeTestCase extends BaseNeuralSearchIT { - // Resources to be cleaned up after each test, need to assign the actual values in the test itself - protected String modelId; - protected String ingestPipelineName; - protected String searchPipelineName; - protected String indexName; @Before - public void initialize() { - // Initialize variables - this.modelId = null; - this.ingestPipelineName = null; - this.searchPipelineName = null; - + protected String getIndexNameForTest() { // Creating index name by concatenating "neural-bwc-" prefix with test method name // for all the tests in this sub-project - this.indexName = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); - } - - @SneakyThrows - @After - public void cleanUpResources() { - if (getClusterType() == ClusterType.UPGRADED) { - wipeOfTestResources(this.indexName, this.ingestPipelineName, this.modelId, this.searchPipelineName); - } + return NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); } @Override diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java index 20d2ba3e4..01fce83f0 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java @@ -22,35 +22,39 @@ public class BatchIngestionIT extends AbstractRollingUpgradeTestCase { public void testBatchIngestion_SparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); - super.ingestPipelineName = SPARSE_PIPELINE; - + String indexName = getIndexNameForTest(); + String sparseModelId = null; switch (getClusterType()) { case OLD: - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_PIPELINE, 2); + sparseModelId = uploadSparseEncodingModel(); + loadModel(sparseModelId); + createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_PIPELINE, 2); createIndexWithConfiguration( - super.indexName, + indexName, Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_PIPELINE ); List> docs = prepareDataForBulkIngestion(0, 5); bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docs); - validateDocCountAndInfo(super.indexName, 5, () -> getDocById(super.indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); + validateDocCountAndInfo(indexName, 5, () -> getDocById(indexName, "4"), EMBEDDING_FIELD_NAME, Map.class); break; case MIXED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); List> docsForMixed = prepareDataForBulkIngestion(5, 5); - bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForMixed); - validateDocCountAndInfo(super.indexName, 10, () -> getDocById(super.indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForMixed); + validateDocCountAndInfo(indexName, 10, () -> getDocById(indexName, "9"), EMBEDDING_FIELD_NAME, Map.class); break; case UPGRADED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - List> docsForUpgraded = prepareDataForBulkIngestion(10, 5); - bulkAddDocuments(super.indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); - validateDocCountAndInfo(super.indexName, 15, () -> getDocById(super.indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); + try { + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_PIPELINE), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + List> docsForUpgraded = prepareDataForBulkIngestion(10, 5); + bulkAddDocuments(indexName, TEXT_FIELD_NAME, SPARSE_PIPELINE, docsForUpgraded); + validateDocCountAndInfo(indexName, 15, () -> getDocById(indexName, "14"), EMBEDDING_FIELD_NAME, Map.class); + } finally { + wipeOfTestResources(indexName, SPARSE_PIPELINE, sparseModelId, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java index 1f114710d..de7ddef55 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchIT.java @@ -36,26 +36,24 @@ public class HybridSearchIT extends AbstractRollingUpgradeTestCase { private static final int NUM_DOCS_PER_ROUND = 1; private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding"; protected static final String RESCORE_QUERY = "hi"; + private static String modelId = ""; // Test rolling-upgrade normalization processor when index with multiple shards // Create Text Embedding Processor, Ingestion Pipeline, add document and search pipeline with noramlization processor // Validate process , pipeline and document count in rolling-upgrade scenario public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - super.searchPipelineName = SEARCH_PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -64,28 +62,32 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr ); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, null); - hybridQueryBuilder = getQueryBuilder(super.modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, null); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, null); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -98,11 +100,11 @@ private void validateTestIndexOnUpgrade( HybridQueryBuilder hybridQueryBuilder, QueryBuilder rescorer ) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); Map searchResponseAsMap = search( - super.indexName, + getIndexNameForTest(), hybridQueryBuilder, rescorer, 1, diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java index a2233161a..3cea2400d 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/HybridSearchWithRescoreIT.java @@ -37,26 +37,24 @@ public class HybridSearchWithRescoreIT extends AbstractRollingUpgradeTestCase { private static final int NUM_DOCS_PER_ROUND = 1; private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding"; protected static final String RESCORE_QUERY = "hi"; + private static String modelId = ""; /** * Test normalization with hybrid query and rescore. This test is required as rescore will not be compatible with version lower than 2.15 */ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - super.searchPipelineName = SEARCH_PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); createSearchPipeline( SEARCH_PIPELINE_NAME, DEFAULT_NORMALIZATION_METHOD, @@ -65,30 +63,34 @@ public void testHybridQueryWithRescore_whenIndexWithMultipleShards_E2EFlow() thr ); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, rescorer); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, rescorer); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, hybridQueryBuilder, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(super.modelId, null, null); - QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, rescorer); - hybridQueryBuilder = getQueryBuilder(super.modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, hybridQueryBuilder, rescorer); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -101,11 +103,11 @@ private void validateTestIndexOnUpgrade( HybridQueryBuilder hybridQueryBuilder, QueryBuilder rescorer ) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); Map searchResponseAsMap = search( - super.indexName, + getIndexNameForTest(), hybridQueryBuilder, rescorer, 1, diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java index 4ab6691f7..e9fb1a4a7 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/KnnRadialSearchIT.java @@ -24,44 +24,47 @@ public class KnnRadialSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk"; private static final int NUM_DOCS_PER_ROUND = 1; + private static String modelId = ""; // Test rolling-upgrade with kNN radial search // Create Text Image Embedding Processor, Ingestion Pipeline and add document // Validate radial query, pipeline and document count in rolling-upgrade scenario public void testKnnRadialSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadTextImageEmbeddingModel(); - loadModel(super.modelId); - createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadTextImageEmbeddingModel(); + loadModel(modelId); + createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateIndexQueryOnUpgrade(totalDocsCountMixed, super.modelId, TEXT, TEST_IMAGE_TEXT); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); + validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateIndexQueryOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); + validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateIndexQueryOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateIndexQueryOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -70,7 +73,7 @@ public void testKnnRadialSearch_E2EFlow() throws Exception { private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String modelId, final String text, final String imageText) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); @@ -82,7 +85,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo .minScore(0.01f) .build(); - Map responseWithMinScore = search(super.indexName, neuralQueryBuilderWithMinScoreQuery, 1); + Map responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); assertNotNull(responseWithMinScore); NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder() @@ -93,7 +96,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo .maxDistance(100000f) .build(); - Map responseWithMaxScore = search(super.indexName, neuralQueryBuilderWithMaxDistanceQuery, 1); + Map responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); assertNotNull(responseWithMaxScore); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java index d9c15a780..6ced20d95 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/MultiModalSearchIT.java @@ -24,44 +24,47 @@ public class MultiModalSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk"; private static final int NUM_DOCS_PER_ROUND = 1; + private static String modelId = ""; // Test rolling-upgrade test image embedding processor // Create Text Image Embedding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); - super.ingestPipelineName = PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadTextImageEmbeddingModel(); - loadModel(super.modelId); - createPipelineForTextImageProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadTextImageEmbeddingModel(); + loadModel(modelId); + createPipelineForTextImageProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT, TEST_IMAGE_TEXT); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -70,7 +73,7 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception { private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId, final String text, final String imageText) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); NeuralQueryBuilder neuralQueryBuilderWithKQuery = NeuralQueryBuilder.builder() @@ -81,7 +84,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .k(1) .build(); - Map responseWithKQuery = search(super.indexName, neuralQueryBuilderWithKQuery, 1); + Map responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1); assertNotNull(responseWithKQuery); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java index 0aaa74af0..e22e65e4a 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralQueryEnricherProcessorIT.java @@ -26,6 +26,8 @@ public class NeuralQueryEnricherProcessorIT extends AbstractRollingUpgradeTestCa private static final String TEST_ENCODING_FIELD = "passage_embedding"; private static final String TEST_TEXT_FIELD = "passage_text"; private static final String TEXT_1 = "Hello world a b"; + private String sparseModelId = ""; + private String denseModelId = ""; // test of NeuralQueryEnricherProcessor supports neural_sparse query default model_id // the feature is introduced from 2.13 @@ -36,47 +38,52 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws // will set the model_id after we obtain the id NeuralSparseQueryBuilder sparseEncodingQueryBuilderWithModelId = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1); - super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; - super.searchPipelineName = SPARSE_SEARCH_PIPELINE_NAME; switch (getClusterType()) { case OLD: - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); + sparseModelId = uploadSparseEncodingModel(); + loadModel(sparseModelId); + sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); + createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); - createSearchRequestProcessor(super.modelId, SPARSE_SEARCH_PIPELINE_NAME); - updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME)); + addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + createSearchRequestProcessor(sparseModelId, SPARSE_SEARCH_PIPELINE_NAME); + updateIndexSettings( + getIndexNameForTest(), + Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_PIPELINE_NAME) + ); assertEquals( - search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); break; case MIXED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); assertEquals( - search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") ); break; case UPGRADED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - sparseEncodingQueryBuilderWithModelId.modelId(super.modelId); - assertEquals( - search(super.indexName, sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, sparseEncodingQueryBuilderWithModelId, 1).get("hits") - ); + try { + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + sparseEncodingQueryBuilderWithModelId.modelId(sparseModelId); + assertEquals( + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), sparseEncodingQueryBuilderWithModelId, 1).get("hits") + ); + } finally { + wipeOfTestResources(getIndexNameForTest(), SPARSE_INGEST_PIPELINE_NAME, sparseModelId, SPARSE_SEARCH_PIPELINE_NAME); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -95,49 +102,54 @@ public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Excep .fieldName(TEST_ENCODING_FIELD) .queryText(TEXT_1) .build(); - super.ingestPipelineName = DENSE_INGEST_PIPELINE_NAME; - super.searchPipelineName = DENSE_SEARCH_PIPELINE_NAME; switch (getClusterType()) { case OLD: - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - neuralQueryBuilderWithModelId.modelId(super.modelId); - createPipelineProcessor(super.modelId, DENSE_INGEST_PIPELINE_NAME); + denseModelId = uploadTextEmbeddingModel(); + loadModel(denseModelId); + neuralQueryBuilderWithModelId.modelId(denseModelId); + createPipelineProcessor(denseModelId, DENSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), DENSE_INGEST_PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_TEXT_FIELD, TEXT_1, null, null); + addDocument(getIndexNameForTest(), "0", TEST_TEXT_FIELD, TEXT_1, null, null); - createSearchRequestProcessor(super.modelId, DENSE_SEARCH_PIPELINE_NAME); - updateIndexSettings(super.indexName, Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME)); + createSearchRequestProcessor(denseModelId, DENSE_SEARCH_PIPELINE_NAME); + updateIndexSettings( + getIndexNameForTest(), + Settings.builder().put("index.search.default_pipeline", DENSE_SEARCH_PIPELINE_NAME) + ); assertEquals( - search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") ); break; case MIXED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - neuralQueryBuilderWithModelId.modelId(super.modelId); + denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(denseModelId); + neuralQueryBuilderWithModelId.modelId(denseModelId); assertEquals( - search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") ); break; case UPGRADED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - loadModel(super.modelId); - neuralQueryBuilderWithModelId.modelId(super.modelId); + try { + denseModelId = TestUtils.getModelId(getIngestionPipeline(DENSE_INGEST_PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + loadModel(denseModelId); + neuralQueryBuilderWithModelId.modelId(denseModelId); - assertEquals( - search(super.indexName, neuralQueryBuilderWithoutModelId, 1).get("hits"), - search(super.indexName, neuralQueryBuilderWithModelId, 1).get("hits") - ); + assertEquals( + search(getIndexNameForTest(), neuralQueryBuilderWithoutModelId, 1).get("hits"), + search(getIndexNameForTest(), neuralQueryBuilderWithModelId, 1).get("hits") + ); + } finally { + wipeOfTestResources(getIndexNameForTest(), DENSE_INGEST_PIPELINE_NAME, denseModelId, DENSE_SEARCH_PIPELINE_NAME); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java index 8cea81b24..8cebae4a2 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseSearchIT.java @@ -8,7 +8,6 @@ import java.nio.file.Path; import java.util.List; import java.util.Map; - import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.neuralsearch.util.TestUtils; @@ -33,26 +32,25 @@ public class NeuralSparseSearchIT extends AbstractRollingUpgradeTestCase { private final Map testRankFeaturesDoc2 = TestUtils.createRandomTokenWeightMap(TEST_TOKENS_2); private final Map testRankFeaturesDoc3 = TestUtils.createRandomTokenWeightMap(TEST_TOKENS_3); private static final int NUM_DOCS_PER_ROUND = 1; + private static String modelId = ""; // Test rolling-upgrade test sparse embedding processor // Create Sparse Encoding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testSparseEncodingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); - super.ingestPipelineName = PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadSparseEncodingModel(); + loadModel(modelId); + createPipelineForSparseEncodingProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), PIPELINE_NAME ); addSparseEncodingDoc( - super.indexName, + getIndexNameForTest(), "0", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc1), @@ -61,13 +59,13 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { ); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId); addSparseEncodingDoc( - super.indexName, + getIndexNameForTest(), "1", List.of(TEST_SPARSE_ENCODING_FIELD), List.of(testRankFeaturesDoc2), @@ -76,22 +74,26 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { ); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addSparseEncodingDoc( - super.indexName, - "2", - List.of(TEST_SPARSE_ENCODING_FIELD), - List.of(testRankFeaturesDoc3), - List.of(TEST_TEXT_FIELD), - List.of(TEXT_UPGRADED) - ); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addSparseEncodingDoc( + getIndexNameForTest(), + "2", + List.of(TEST_SPARSE_ENCODING_FIELD), + List.of(testRankFeaturesDoc3), + List.of(TEST_TEXT_FIELD), + List.of(TEXT_UPGRADED) + ); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -99,7 +101,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { } private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); @@ -108,7 +110,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .modelId(modelId); MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD, TEXT); boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); - Map response = search(super.indexName, boolQueryBuilder, 1); + Map response = search(getIndexNameForTest(), boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(response); assertEquals("0", firstInnerHit.get("_id")); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java index cecea39ee..5338b777f 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/NeuralSparseTwoPhaseProcessorIT.java @@ -22,6 +22,7 @@ public class NeuralSparseTwoPhaseProcessorIT extends AbstractRollingUpgradeTestC private static final String TEST_ENCODING_FIELD = "passage_embedding"; private static final String TEST_TEXT_FIELD = "passage_text"; private static final String TEXT_1 = "Hello world a b"; + private String sparseModelId = ""; // test of NeuralSparseTwoPhaseProcessor supports neural_sparse query's two phase speed up // the feature is introduced from 2.15 @@ -29,33 +30,46 @@ public void testNeuralSparseTwoPhaseProcessorIT_NeuralSparseSearch_E2EFlow() thr waitForClusterHealthGreen(NODES_BWC_CLUSTER); // will set the model_id after we obtain the id NeuralSparseQueryBuilder neuralSparseQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1); - super.ingestPipelineName = SPARSE_INGEST_PIPELINE_NAME; - super.searchPipelineName = SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME; switch (getClusterType()) { case OLD: - super.modelId = uploadSparseEncodingModel(); - loadModel(super.modelId); - neuralSparseQueryBuilder.modelId(super.modelId); - createPipelineForSparseEncodingProcessor(super.modelId, SPARSE_INGEST_PIPELINE_NAME); + sparseModelId = uploadSparseEncodingModel(); + loadModel(sparseModelId); + neuralSparseQueryBuilder.modelId(sparseModelId); + createPipelineForSparseEncodingProcessor(sparseModelId, SPARSE_INGEST_PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/SparseIndexMappings.json").toURI())), SPARSE_INGEST_PIPELINE_NAME ); - addSparseEncodingDoc(super.indexName, "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); + addSparseEncodingDoc(getIndexNameForTest(), "0", List.of(), List.of(), List.of(TEST_TEXT_FIELD), List.of(TEXT_1)); createNeuralSparseTwoPhaseSearchProcessor(SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME); updateIndexSettings( - super.indexName, + getIndexNameForTest(), Settings.builder().put("index.search.default_pipeline", SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME) ); - assertNotNull(search(super.indexName, neuralSparseQueryBuilder, 1).get("hits")); + assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); break; - case MIXED, UPGRADED: - super.modelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); - loadModel(super.modelId); - neuralSparseQueryBuilder.modelId(super.modelId); - assertNotNull(search(super.indexName, neuralSparseQueryBuilder, 1).get("hits")); + case MIXED: + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + neuralSparseQueryBuilder.modelId(sparseModelId); + assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); + break; + case UPGRADED: + try { + sparseModelId = TestUtils.getModelId(getIngestionPipeline(SPARSE_INGEST_PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); + loadModel(sparseModelId); + neuralSparseQueryBuilder.modelId(sparseModelId); + assertNotNull(search(getIndexNameForTest(), neuralSparseQueryBuilder, 1).get("hits")); + } finally { + wipeOfTestResources( + getIndexNameForTest(), + SPARSE_INGEST_PIPELINE_NAME, + sparseModelId, + SPARSE_SEARCH_TWO_PHASE_PIPELINE_NAME + ); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java index e84adde55..989d53897 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/SemanticSearchIT.java @@ -19,44 +19,47 @@ public class SemanticSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEXT_MIXED = "Hello world mixed"; private static final String TEXT_UPGRADED = "Hello world upgraded"; private static final int NUM_DOCS_PER_ROUND = 1; + private static String modelId = ""; // Test rolling-upgrade Semantic Search // Create Text Embedding Processor, Ingestion Pipeline and add document // Validate process , pipeline and document count in rolling-upgrade scenario public void testSemanticSearch_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90); - super.ingestPipelineName = PIPELINE_NAME; - switch (getClusterType()) { case OLD: - super.modelId = uploadTextEmbeddingModel(); - loadModel(super.modelId); - createPipelineProcessor(super.modelId, PIPELINE_NAME); + modelId = uploadTextEmbeddingModel(); + loadModel(modelId); + createPipelineProcessor(modelId, PIPELINE_NAME); createIndexWithConfiguration( - super.indexName, + getIndexNameForTest(), Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), PIPELINE_NAME ); - addDocument(super.indexName, "0", TEST_FIELD, TEXT, null, null); + addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, null, null); break; case MIXED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT); - addDocument(super.indexName, "1", TEST_FIELD, TEXT_MIXED, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT); + addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - validateTestIndexOnUpgrade(totalDocsCountMixed, super.modelId, TEXT_MIXED); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED); } break; case UPGRADED: - super.modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - loadModel(super.modelId); - addDocument(super.indexName, "2", TEST_FIELD, TEXT_UPGRADED, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, super.modelId, TEXT_UPGRADED); + try { + modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_EMBEDDING_PROCESSOR); + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + loadModel(modelId); + addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED); + } finally { + wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType()); @@ -65,7 +68,7 @@ public void testSemanticSearch_E2EFlow() throws Exception { } private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId, final String text) throws Exception { - int docCount = getDocCount(super.indexName); + int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder() @@ -74,7 +77,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod .queryText(text) .k(1) .build(); - Map response = search(super.indexName, neuralQueryBuilder, 1); + Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); assertNotNull(response); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java index 1733d511b..29ec1fed5 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/TextChunkingProcessorIT.java @@ -35,8 +35,7 @@ public class TextChunkingProcessorIT extends AbstractRollingUpgradeTestCase { // Validate process, pipeline and document count in rolling-upgrade scenario public void testTextChunkingProcessor_E2EFlow() throws Exception { waitForClusterHealthGreen(NODES_BWC_CLUSTER); - super.ingestPipelineName = PIPELINE_NAME; - + String indexName = getIndexNameForTest(); switch (getClusterType()) { case OLD: createPipelineForTextChunkingProcessor(PIPELINE_NAME); @@ -55,9 +54,13 @@ public void testTextChunkingProcessor_E2EFlow() throws Exception { } break; case UPGRADED: - int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; - addDocument(indexName, "2", INPUT_FIELD, TEST_INGEST_TEXT, null, null); - validateTestIndex(indexName, OUTPUT_FIELD, totalDocsCountUpgraded, expectedPassages); + try { + int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; + addDocument(indexName, "2", INPUT_FIELD, TEST_INGEST_TEXT, null, null); + validateTestIndex(indexName, OUTPUT_FIELD, totalDocsCountUpgraded, expectedPassages); + } finally { + wipeOfTestResources(indexName, PIPELINE_NAME, null, null); + } break; default: throw new IllegalStateException("Unexpected value: " + getClusterType());