From 7dc84b50d3b1727cafb4b07342df9ddff2561ffd Mon Sep 17 00:00:00 2001 From: Weijia Zhao Date: Mon, 10 Feb 2025 15:02:56 -0800 Subject: [PATCH] Handle exceptions gracefully when delete non-existent resources during integ test resource clean up (#1154) * Handle exceptions gracefully when delete non-existent resources during integ test resource clean up Signed-off-by: Weijia Zhao * Clean up test resource dynamically after each test case Signed-off-by: Weijia Zhao * Code clean up Signed-off-by: Weijia Zhao * Fixing failing bwc tests Signed-off-by: Weijia Zhao * Revert bwc test resource cleanup logic Signed-off-by: Weijia Zhao --------- Signed-off-by: Weijia Zhao --- .../AbstractRestartUpgradeRestTestCase.java | 10 + .../AbstractRollingUpgradeTestCase.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 | 273 +++- .../neuralsearch/util/TestUtils.java | 3 + 26 files changed, 3761 insertions(+), 4146 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..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 @@ -53,6 +53,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/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..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 @@ -56,6 +56,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/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 2182d8d79..3b6143e09 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -4,12 +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; @@ -52,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; @@ -61,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; @@ -95,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; @@ -110,6 +134,21 @@ public void setupSettings() { NeuralSearchClusterUtil.instance().initialize(clusterService); } + // Wipe of all the resources after execution of the tests. + @After + public void cleanUp() { + if (shouldCleanUpResources()) { + deleteExistingIngestionPipelines(); + deleteExistingSearchPipelines(); + deleteExistingModels(); + deleteExistingIndices(); + } + } + + protected boolean shouldCleanUpResources() { + return true; + } + protected ThreadPool setUpThreadPool() { return new TestThreadPool(getClass().getName(), threadPoolSettings()); } @@ -1242,18 +1281,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( @@ -1410,6 +1437,96 @@ 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 + */ + @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()); + 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 + */ + @SneakyThrows + protected void deleteExistingIndices() { + 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().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 the value is null, it returns all pipelines + * @return get pipeline response as a map object + */ + @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() { + Map pipelines = retrievePipelines(INGEST_PIPELINE_TYPE, null); + pipelines.keySet().forEach(this::deleteIngestPipeline); + } + + private void deleteExistingSearchPipelines() { + Map pipelines = retrievePipelines(SEARCH_PIPELINE_TYPE, null); + pipelines.keySet().forEach(this::deleteSearchPipeline); + } + /** * Get ingest pipeline * @param pipelineName of the ingest pipeline @@ -1418,24 +1535,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()); @@ -1443,6 +1556,99 @@ 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 + protected List retrieveModels() { + 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() { + 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); + } + }); + } + + @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); @@ -1497,33 +1703,6 @@ protected Float getFeatureFieldCompressedNumber(final Float originNumber) { return Float.intBitsToFloat(freqBits); } - // 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 { - 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); - } - } - protected Object validateDocCountAndInfo( String indexName, int expectedDocCount, 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;