diff --git a/release-notes/opensearch-neural-search.release-notes-2.18.0.0.md b/release-notes/opensearch-neural-search.release-notes-2.18.0.0.md index 298bd704b..189a62b97 100644 --- a/release-notes/opensearch-neural-search.release-notes-2.18.0.0.md +++ b/release-notes/opensearch-neural-search.release-notes-2.18.0.0.md @@ -3,6 +3,8 @@ Compatible with OpenSearch 2.18.0 +### Bug Fixes +- Fixed incorrect document order for nested aggregations in hybrid query ([#956](https://github.com/opensearch-project/neural-search/pull/956)) ### Enhancements - Implement `ignore_missing` field in text chunking processors ([#907](https://github.com/opensearch-project/neural-search/pull/907)) - Added rescorer in hybrid query ([#917](https://github.com/opensearch-project/neural-search/pull/917)) diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryScorer.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryScorer.java index 23dbd0e1d..eb410aa23 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryScorer.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryScorer.java @@ -97,8 +97,13 @@ public int advanceShallow(int target) throws IOException { */ @Override public float score() throws IOException { + return score(getSubMatches()); + } + + private float score(DisiWrapper topList) throws IOException { float totalScore = 0.0f; - for (DisiWrapper disiWrapper : subScorersPQ) { + for (DisiWrapper disiWrapper = topList; disiWrapper != null; disiWrapper = disiWrapper.next) { + // check if this doc has match in the subQuery. If not, add score as 0.0 and continue if (disiWrapper.scorer.docID() == DocIdSetIterator.NO_MORE_DOCS) { continue; } diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java index 3bb566aef..2df2c21a7 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryAggregationsIT.java @@ -12,8 +12,10 @@ import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.neuralsearch.BaseNeuralSearchIT; +import org.opensearch.script.Script; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregationBuilders; +import org.opensearch.search.aggregations.BucketOrder; import org.opensearch.search.aggregations.PipelineAggregatorBuilders; import org.opensearch.search.aggregations.bucket.histogram.DateHistogramInterval; import org.opensearch.search.aggregations.pipeline.AvgBucketPipelineAggregationBuilder; @@ -21,20 +23,25 @@ import org.opensearch.search.aggregations.pipeline.MaxBucketPipelineAggregationBuilder; import org.opensearch.search.aggregations.pipeline.MinBucketPipelineAggregationBuilder; import org.opensearch.search.aggregations.pipeline.SumBucketPipelineAggregationBuilder; +import org.opensearch.search.sort.SortBuilders; +import org.opensearch.search.sort.SortOrder; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; -import static org.opensearch.neuralsearch.util.TestUtils.DELTA_FOR_SCORE_ASSERTION; import static org.opensearch.neuralsearch.util.AggregationsTestUtils.getAggregationBuckets; import static org.opensearch.neuralsearch.util.AggregationsTestUtils.getAggregationValue; import static org.opensearch.neuralsearch.util.AggregationsTestUtils.getAggregationValues; import static org.opensearch.neuralsearch.util.AggregationsTestUtils.getAggregations; import static org.opensearch.neuralsearch.util.AggregationsTestUtils.getNestedHits; import static org.opensearch.neuralsearch.util.TestUtils.assertHitResultsFromQuery; +import static org.opensearch.neuralsearch.util.TestUtils.TEST_SPACE_TYPE; +import static org.opensearch.neuralsearch.util.TestUtils.DELTA_FOR_SCORE_ASSERTION; /** * Integration tests for base scenarios when aggregations are combined with hybrid query @@ -42,6 +49,7 @@ public class HybridQueryAggregationsIT extends BaseNeuralSearchIT { private static final String TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS = "test-hybrid-aggs-multi-doc-index-multiple-shards"; private static final String TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD = "test-hybrid-aggs-multi-doc-index-single-shard"; + private static final String TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS = "test-hybrid-nested-aggs-multi-doc-index"; private static final String TEST_QUERY_TEXT3 = "hello"; private static final String TEST_QUERY_TEXT4 = "everyone"; private static final String TEST_QUERY_TEXT5 = "welcome"; @@ -464,6 +472,166 @@ public void testPostFilterOnIndexWithSingleShards_WhenConcurrentSearchEnabled_th testPostFilterWithComplexHybridQuery(true, true); } + @SneakyThrows + public void testNestedAggs_whenMultipleShardsAndConcurrentSearchDisabled_thenSuccessful() { + updateClusterSettings(CONCURRENT_SEGMENT_SEARCH_ENABLED, false); + try { + if (!indexExists(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS)) { + createIndexWithConfiguration( + TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS, + buildIndexConfiguration( + List.of(new KNNFieldConfig("location", 2, TEST_SPACE_TYPE)), + List.of(), + List.of(), + List.of("imdb"), + List.of("actor"), + List.of(), + 3 + ), + "" + ); + + String ingestBulkPayload = Files.readString(Path.of(classLoader.getResource("processor/ingest_bulk.json").toURI())) + .replace("\"{indexname}\"", "\"" + TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS + "\""); + + bulkIngest(ingestBulkPayload, null); + } + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + + /* constructing following search query + { + "from": 0, + "aggs": { + "cardinality_of_unique_names": { + "cardinality": { + "field": "actor" + } + }, + "unique_names": { + "terms": { + "field": "actor", + "size": 10, + "order": { + "max_score": "desc" + } + }, + "aggs": { + "top_doc": { + "top_hits": { + "size": 1, + "sort": [ + { + "_score": { + "order": "desc" + } + } + ] + } + }, + "max_score": { + "max": { + "script": { + "source": "_score" + } + } + } + } + } + }, + "query": { + "hybrid": { + "queries": [ + { + "match": { + "actor": "anil" + } + }, + { + "range": { + "imdb": { + "gte": 1.0, + "lte": 10.0 + } + } + } + ]}}} + */ + + QueryBuilder rangeFilterQuery = QueryBuilders.rangeQuery("imdb").gte(1.0).lte(10.0); + QueryBuilder matchQuery = QueryBuilders.matchQuery("actor", "anil"); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(matchQuery).add(rangeFilterQuery); + + AggregationBuilder aggsBuilderCardinality = AggregationBuilders.cardinality("cardinality_of_unique_names").field("actor"); + AggregationBuilder aggsBuilderUniqueNames = AggregationBuilders.terms("unique_names") + .field("actor") + .size(10) + .order(BucketOrder.aggregation("max_score", false)) + .subAggregation(AggregationBuilders.topHits("top_doc").size(1).sort(SortBuilders.scoreSort().order(SortOrder.DESC))) + .subAggregation(AggregationBuilders.max("max_score").script(new Script("_score"))); + + Map searchResponseAsMap = search( + TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + List.of(aggsBuilderCardinality, aggsBuilderUniqueNames), + rangeFilterQuery, + null, + false, + null, + 0 + ); + assertNotNull(searchResponseAsMap); + + // assert actual results + // aggregations + Map aggregations = getAggregations(searchResponseAsMap); + assertNotNull(aggregations); + + int cardinalityValue = getAggregationValue(aggregations, "cardinality_of_unique_names"); + assertEquals(7, cardinalityValue); + + Map uniqueAggValue = getAggregationValues(aggregations, "unique_names"); + assertEquals(3, uniqueAggValue.size()); + assertEquals(0, uniqueAggValue.get("doc_count_error_upper_bound")); + assertEquals(0, uniqueAggValue.get("sum_other_doc_count")); + + List> buckets = getAggregationBuckets(aggregations, "unique_names"); + assertNotNull(buckets); + assertEquals(7, buckets.size()); + + // check content of few buckets + Map firstBucket = buckets.get(0); + assertEquals(4, firstBucket.size()); + assertEquals("anil", firstBucket.get(KEY)); + assertEquals(42, firstBucket.get("doc_count")); + assertNotNull(getAggregationValue(firstBucket, "max_score")); + assertTrue((double) getAggregationValue(firstBucket, "max_score") > 1.0f); + + Map secondBucket = buckets.get(1); + assertEquals(4, secondBucket.size()); + assertEquals("abhishek", secondBucket.get(KEY)); + assertEquals(8, secondBucket.get("doc_count")); + assertNotNull(getAggregationValue(secondBucket, "max_score")); + assertEquals(1.0, getAggregationValue(secondBucket, "max_score"), DELTA_FOR_SCORE_ASSERTION); + + Map lastBucket = buckets.get(buckets.size() - 1); + assertEquals(4, lastBucket.size()); + assertEquals("sanjay", lastBucket.get(KEY)); + assertEquals(7, lastBucket.get("doc_count")); + assertNotNull(getAggregationValue(lastBucket, "max_score")); + assertEquals(1.0, getAggregationValue(lastBucket, "max_score"), DELTA_FOR_SCORE_ASSERTION); + + // assert the hybrid query scores + assertHitResultsFromQuery(10, 92, searchResponseAsMap); + + } finally { + wipeOfTestResources(TEST_MULTI_DOC_INDEX_FOR_NESTED_AGGS_MULTIPLE_SHARDS, null, null, SEARCH_PIPELINE); + } + } + private void testMaxAggsOnSingleShardCluster() throws Exception { try { prepareResourcesForSingleShardIndex(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_SINGLE_SHARD, SEARCH_PIPELINE); @@ -501,8 +669,6 @@ private void testDateRange() throws IOException { try { initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); - // try { - // prepareResources(TEST_MULTI_DOC_INDEX_WITH_TEXT_AND_INT_MULTIPLE_SHARDS, SEARCH_PIPELINE); AggregationBuilder aggsBuilder = AggregationBuilders.dateRange(DATE_AGGREGATION_NAME) .field(DATE_FIELD_1) diff --git a/src/test/resources/processor/ingest_bulk.json b/src/test/resources/processor/ingest_bulk.json new file mode 100644 index 000000000..d7657b934 --- /dev/null +++ b/src/test/resources/processor/ingest_bulk.json @@ -0,0 +1,184 @@ +{"index": {"_index": "{indexname}"}} +{"id":"s9","passage_text":"hi there","imdb":7.7,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s10","passage_text":"hi world","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s11","passage_text":"you are hero","imdb":2.5,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s12","passage_text":"jalwa","imdb":8.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s13","passage_text":"suna hai","imdb":2.5,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s14","passage_text":"what is your name","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s15","passage_text":"people","imdb":8.0,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s16","passage_text":"peeking","imdb":6.6,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s17","passage_text":"hum","imdb":9.9,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s18","passage_text":"dillagi","imdb":9.9,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s19","passage_text":"beekhayali","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s20","passage_text":"bina","imdb":9.9,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s21","passage_text":"hum tum","imdb":9.8,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s22","passage_text":"like","imdb":9.8,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s23","passage_text":"like","imdb":6.6,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s24","passage_text":"liking","imdb":9.9,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s25","passage_text":"hike","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s26","passage_text":"hiking","imdb":2.5,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s27","passage_text":"hiker","imdb":8.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s28","passage_text":"trekker","imdb":3.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s29","passage_text":"trekking","imdb":7.7,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s30","passage_text":"mard","imdb":4.5,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s31","passage_text":"mard 1","imdb":10.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s32","passage_text":"do you like me","imdb":9.9,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s33","passage_text":"do you like me 1","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s34","passage_text":"do you like me 2","imdb":6.6,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s35","passage_text":"Avengers","imdb":9.9,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s36","passage_text":"Avengers 1","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s37","passage_text":"Avengers 3","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s38","passage_text":"Avengers 2","imdb":2.5,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s39","passage_text":"levitating","imdb":2.5,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s40","passage_text":"flowers","imdb":6.6,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s41","passage_text":"too sweet","imdb":10.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s42","passage_text":"kd pathak","imdb":2.5,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s43","passage_text":"tmkoc","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s44","passage_text":"wake up","imdb":7.7,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s45","passage_text":"sunrise","imdb":9.9,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s46","passage_text":"pretend","imdb":2.5,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s47","passage_text":"sunset","imdb":8.0,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s48","passage_text":"for me","imdb":8.0,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s49","passage_text":"I take my ","imdb":6.6,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s50","passage_text":"coffee","imdb":10.0,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s51","passage_text":"tuje dekha to jaana sanam","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s52","passage_text":"tumse hi","imdb":9.8,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s53","passage_text":"kareena","imdb":7.7,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s54","passage_text":"shahid","imdb":3.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s55","passage_text":"haider","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s56","passage_text":"tara rum pum","imdb":7.7,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s57","passage_text":"forever","imdb":9.9,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s58","passage_text":"tsa","imdb":4.5,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s59","passage_text":"I wish","imdb":2.5,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s60","passage_text":"could go","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s61","passage_text":"soft as a rain","imdb":8.0,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s62","passage_text":"pretty","imdb":7.7,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s63","passage_text":"take my whiskey","imdb":6.6,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s64","passage_text":"bed","imdb":9.9,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s65","passage_text":"cutu","imdb":2.5,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s66","passage_text":"pyar","imdb":8.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s67","passage_text":"amit","imdb":2.5,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s68","passage_text":"josh","imdb":10.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s69","passage_text":"owais","imdb":2.5,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s70","passage_text":"martin","imdb":9.8,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s71","passage_text":"navneet","imdb":3.0,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s72","passage_text":"vamshi","imdb":4.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s73","passage_text":"varun","imdb":3.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s74","passage_text":"amulya","imdb":8.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s75","passage_text":"minal","imdb":6.6,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s76","passage_text":"xin","imdb":9.9,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s77","passage_text":"xun","imdb":7.7,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s78","passage_text":"anandhi","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s79","passage_text":"mukul","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s80","passage_text":"north","imdb":7.7,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s81","passage_text":"restaurent","imdb":10.0,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s82","passage_text":"south","imdb":9.8,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s83","passage_text":"east","imdb":6.6,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s84","passage_text":"west","imdb":9.9,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s85","passage_text":"tree","imdb":2.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s86","passage_text":"house","imdb":8.0,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s87","passage_text":"jibhi","imdb":9.9,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s88","passage_text":"Popp","imdb":7.7,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s89","passage_text":"martins","imdb":9.8,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s90","passage_text":"marrige","imdb":10.0,"actor":"jackie"} +{"index": {"_index": "{indexname}"}} +{"id":"s91","passage_text":"airpod","imdb":4.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s92","passage_text":"mac","imdb":3.0,"actor":"ranveer"} +{"index": {"_index": "{indexname}"}} +{"id":"s93","passage_text":"apple","imdb":3.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s94","passage_text":"yellow","imdb":9.8,"actor":"salman"} +{"index": {"_index": "{indexname}"}} +{"id":"s95","passage_text":"green","imdb":4.5,"actor":"anil"} +{"index": {"_index": "{indexname}"}} +{"id":"s96","passage_text":"airport","imdb":8.0,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s97","passage_text":"mt rainier","imdb":6.6,"actor":"abhishek"} +{"index": {"_index": "{indexname}"}} +{"id":"s98","passage_text":"glacier","imdb":7.7,"actor":"ranbir"} +{"index": {"_index": "{indexname}"}} +{"id":"s99","passage_text":"alaska","imdb":8.0,"actor":"sanjay"} +{"index": {"_index": "{indexname}"}} +{"id":"s100","passage_text":"hawaii","imdb":9.8,"actor":"jackie"} diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index e6fb45d2a..afc545447 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -11,6 +11,7 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Locale; import java.util.Map; @@ -785,6 +786,39 @@ protected void bulkAddDocuments(final String index, final String textField, fina assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); } + @SneakyThrows + protected void bulkIngest(final String ingestBulkPayload, final String pipeline) { + Map params = new HashMap<>(); + params.put("refresh", "true"); + if (Objects.nonNull(pipeline)) { + params.put("pipeline", pipeline); + } + Response response = makeRequest( + client(), + "POST", + "_bulk", + params, + toHttpEntity(ingestBulkPayload), + ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana")) + ); + Map map = XContentHelper.convertToMap( + XContentType.JSON.xContent(), + EntityUtils.toString(response.getEntity()), + false + ); + + int failedDocCount = 0; + for (Object item : ((List) map.get("items"))) { + Map> itemMap = (Map>) item; + if (itemMap.get("index").get("error") != null) { + failedDocCount++; + } + } + assertEquals(0, failedDocCount); + + assertEquals("_bulk failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + /** * Parse the first returned hit from a search response as a map * @@ -924,6 +958,27 @@ protected String buildIndexConfiguration( final List keywordFields, final List dateFields, final int numberOfShards + ) { + return buildIndexConfiguration( + knnFieldConfigs, + nestedFields, + intFields, + Collections.emptyList(), + keywordFields, + dateFields, + numberOfShards + ); + } + + @SneakyThrows + protected String buildIndexConfiguration( + final List knnFieldConfigs, + final List nestedFields, + final List intFields, + final List floatFields, + final List keywordFields, + final List dateFields, + final int numberOfShards ) { XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() .startObject() @@ -964,6 +1019,10 @@ protected String buildIndexConfiguration( xContentBuilder.startObject(intField).field("type", "integer").endObject(); } + for (String floatField : floatFields) { + xContentBuilder.startObject(floatField).field("type", "float").endObject(); + } + for (String keywordField : keywordFields) { xContentBuilder.startObject(keywordField).field("type", "keyword").endObject(); } diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java index ab041c440..c10380e87 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/util/TestUtils.java @@ -304,6 +304,10 @@ public static void assertFetchResultScores(FetchSearchResult fetchSearchResult, } public static void assertHitResultsFromQuery(int expected, Map searchResponseAsMap) { + assertHitResultsFromQuery(expected, expected, searchResponseAsMap); + } + + public static void assertHitResultsFromQuery(int expected, int expectedTotal, Map searchResponseAsMap) { assertEquals(expected, getHitCount(searchResponseAsMap)); List> hitsNestedList = getNestedHits(searchResponseAsMap); @@ -321,7 +325,7 @@ public static void assertHitResultsFromQuery(int expected, Map s Map total = getTotalHits(searchResponseAsMap); assertNotNull(total.get("value")); - assertEquals(expected, total.get("value")); + assertEquals(expectedTotal, total.get("value")); assertNotNull(total.get("relation")); assertEquals(RELATION_EQUAL_TO, total.get("relation")); }