add support for scored named queries

Signed-off-by: Dharin Shah <8616130+Dharin-shah@users.noreply.github.com>
opensearch-project · Dec 17, 2023 · 16b7c19 · 16b7c19
1 parent 332f4c5
commit 16b7c19
Show file tree

Hide file tree

Showing 11 changed files with 522 additions and 115 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -121,6 +121,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Adding slf4j license header to LoggerMessageFormat.java ([#11069](https://github.com/opensearch-project/OpenSearch/pull/11069))
 - [BWC and API enforcement] Introduce checks for enforcing the API restrictions ([#11175](https://github.com/opensearch-project/OpenSearch/pull/11175))
 - Create separate transport action for render search template action ([#11170](https://github.com/opensearch-project/OpenSearch/pull/11170))
+- Backwards compatible support for returning scores in matched queries ((https://github.com/opensearch-project/OpenSearch/pull/11549))
 
 ### Dependencies
 - Bump Lucene from 9.7.0 to 9.8.0 ([10276](https://github.com/opensearch-project/OpenSearch/pull/10276))

diff --git a/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java b/libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java
@@ -632,6 +632,77 @@ public String[] readOptionalStringArray() throws IOException {
         return null;
     }
 
+
+    /**
+     * Reads an ordered {@link Map} from a data source using provided key and value readers.
+     *
+     * This method creates a `LinkedHashMap`, ensuring that the iteration order of the map
+     * matches the order in which the entries were read from the data source. It uses the
+     * `keyReader` and `valueReader` to read each key-value pair and constructs the map
+     * with an initial capacity optimized based on the expected size.
+     *
+     * Usage example:
+     * <pre><code>
+     * Map&lt;Integer, String&gt; orderedMap = readOrderedMap(StreamInput::readInteger, StreamInput::readString);
+     * </code></pre>
+     *
+     * @param keyReader The {@link Writeable.Reader} used to read the keys of the map.
+     * @param valueReader The {@link Writeable.Reader} used to read the values of the map.
+     * @return A {@link LinkedHashMap} containing the keys and values read from the data source.
+     *         The map maintains the order in which keys and values were read.
+     * @throws IOException If an I/O error occurs during reading from the data source.
+     */
+    public <K, V> Map<K, V> readOrderedMap(Writeable.Reader<K> keyReader, Writeable.Reader<V> valueReader) throws IOException {
+        return readMap(keyReader, valueReader, (expectedSize) -> new LinkedHashMap<>(capacity(expectedSize)));
+    }
+
+    static int capacity(int expectedSize) {
+        assert expectedSize >= 0;
+        return expectedSize < 2 ? expectedSize + 1 : (int) (expectedSize / 0.75 + 1.0);
+    }
+
+    /**
+     * Reads a {@link Map} from a data source using provided key and value readers and a map constructor.
+     *
+     * This method is a flexible utility for reading a map from a data source, such as a file or network stream.
+     * It uses the provided `keyReader` and `valueReader` to read each key-value pair. The size of the map is
+     * determined first by reading the array size. A map constructor is also provided to create a map of the
+     * appropriate size, allowing for optimized performance based on the expected number of entries.
+     *
+     * If the read size is zero, an empty map is returned. Otherwise, the method iterates over the size,
+     * reading keys and values and adding them to the map.
+     *
+     * Usage example:
+     * <pre><code>
+     * Map&lt;Integer, String&gt; map = readMap(StreamInput::readInteger, StreamInput::readString, HashMap::new);
+     * </code></pre>
+     *
+     * Note: If the map is empty, an immutable empty map is returned. Otherwise, a mutable map is created.
+     *
+     * @param keyReader The {@link Writeable.Reader} used to read the keys of the map.
+     * @param valueReader The {@link Writeable.Reader} used to read the values of the map.
+     * @param constructor A function that takes an integer (the initial size) and returns a new map.
+     *                    This allows for the creation of a map with an initial capacity matching the
+     *                    expected size, optimizing performance.
+     * @return A {@link Map} containing the keys and values read from the data source. This map is either
+     *         empty (and immutable) or mutable and filled with the read data.
+     * @throws IOException If an I/O error occurs during reading from the data source.
+     */
+    private <K, V> Map<K, V> readMap(Writeable.Reader<K> keyReader, Writeable.Reader<V> valueReader, IntFunction<Map<K, V>> constructor)
+    throws IOException {
+        int size = readArraySize();
+        if (size == 0) {
+            return Collections.emptyMap();
+        }
+        Map<K, V> map = constructor.apply(size);
+        for (int i = 0; i < size; i++) {
+            K key = keyReader.read(this);
+            V value = valueReader.read(this);
+            map.put(key, value);
+        }
+        return map;
+    }
+
     /**
      * If the returned map contains any entries it will be mutable. If it is empty it might be immutable.
      */

diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/search.json b/rest-api-spec/src/main/resources/rest-api-spec/api/search.json
@@ -229,6 +229,11 @@
       "search_pipeline": {
         "type": "string",
         "description": "The search pipeline to use to execute this request"
+      },
+      "include_named_queries_score":{
+        "type": "boolean",
+        "description":"Indicates whether hit.matched_queries should be rendered as a map that includes the name of the matched query associated with its score (true) or as an array containing the name of the matched queries (false)",
+        "default":false
       }
     },
     "body":{

diff --git a/rest-api-spec/src/main/resources/rest-api-spec/test/search/350_matched_queries.yml b/rest-api-spec/src/main/resources/rest-api-spec/test/search/350_matched_queries.yml
@@ -0,0 +1,103 @@
+setup:
+  - skip:
+      version: " - 2.12.0"
+      reason: "implemented for versions post 2.12.0"
+
+---
+"matched queries":
+  - do:
+      indices.create:
+        index: test
+
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
+          - '{"field" : 1 }'
+          - '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
+          - '{"field" : [1, 2] }'
+
+  - do:
+      search:
+        index: test_1
+        body:
+          query:
+            bool: {
+              should: [
+                {
+                  match: {
+                    field: {
+                      query: 1,
+                      _name: match_field_1
+                    }
+                  }
+                },
+                {
+                  match: {
+                    field: {
+                      query: 2,
+                      _name: match_field_2,
+                      boost: 10
+                    }
+                  }
+                }
+              ]
+            }
+
+  - match: {hits.total.value: 2}
+  - length: {hits.hits.0.matched_queries: 2}
+  - match: {hits.hits.0.matched_queries: [ "match_field_1", "match_field_2" ]}
+  - length: {hits.hits.1.matched_queries: 1}
+  - match: {hits.hits.1.matched_queries: [ "match_field_1" ]}
+
+---
+
+"matched queries with scores":
+  - do:
+      indices.create:
+        index: test
+
+  - do:
+      bulk:
+        refresh: true
+        body:
+          - '{ "index" : { "_index" : "test_1", "_id" : "1" } }'
+          - '{"field" : 1 }'
+          - '{ "index" : { "_index" : "test_1", "_id" : "2" } }'
+          - '{"field" : [1, 2] }'
+
+  - do:
+      search:
+        include_named_queries_score: true
+        index: test_1
+        body:
+          query:
+            bool: {
+              should: [
+                {
+                  match: {
+                    field: {
+                      query: 1,
+                      _name: match_field_1
+                    }
+                  }
+                },
+                {
+                  match: {
+                    field: {
+                      query: 2,
+                      _name: match_field_2,
+                      boost: 10
+                    }
+                  }
+                }
+              ]
+            }
+
+  - match: { hits.total.value: 2 }
+  - length: { hits.hits.0.matched_queries: 2 }
+  - match: { hits.hits.0.matched_queries.match_field_1: 1 }
+  - match: { hits.hits.0.matched_queries.match_field_2: 10 }
+  - length: { hits.hits.1.matched_queries: 1 }
+  - match: { hits.hits.1.matched_queries.match_field_1: 1 }
diff --git a/...r/src/internalClusterTest/java/org/opensearch/search/fetch/subphase/MatchedQueriesIT.java b/...r/src/internalClusterTest/java/org/opensearch/search/fetch/subphase/MatchedQueriesIT.java
@@ -62,7 +62,9 @@
 import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.hasItemInArray;
+import static org.hamcrest.Matchers.hasKey;
 
 public class MatchedQueriesIT extends ParameterizedOpenSearchIntegTestCase {
 
@@ -105,11 +107,13 @@ public void testSimpleMatchedQueryFromFilteredQuery() throws Exception {
         assertHitCount(searchResponse, 3L);
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("3") || hit.getId().equals("2")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("test2"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("test2"));
+                assertThat(hit.getMatchedQueryScore("test2"), equalTo(1f));
             } else if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("test1"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("test1"));
+                assertThat(hit.getMatchedQueryScore("test1"), equalTo(1f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -123,11 +127,13 @@ public void testSimpleMatchedQueryFromFilteredQuery() throws Exception {
         assertHitCount(searchResponse, 3L);
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1") || hit.getId().equals("2")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("test1"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("test1"));
+                assertThat(hit.getMatchedQueryScore("test1"), equalTo(1f));
             } else if (hit.getId().equals("3")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("test2"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("test2"));
+                assertThat(hit.getMatchedQueryScore("test2"), equalTo(1f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -153,12 +159,15 @@ public void testSimpleMatchedQueryFromTopLevelFilter() throws Exception {
         assertHitCount(searchResponse, 3L);
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(2));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("name"));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("title"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(2));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("name"));
+                assertThat(hit.getMatchedQueryScore("name"), greaterThan(0f));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("title"));
+                assertThat(hit.getMatchedQueryScore("title"), greaterThan(0f));
             } else if (hit.getId().equals("2") || hit.getId().equals("3")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("name"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("name"));
+                assertThat(hit.getMatchedQueryScore("name"), greaterThan(0f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -174,12 +183,15 @@ public void testSimpleMatchedQueryFromTopLevelFilter() throws Exception {
         assertHitCount(searchResponse, 3L);
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(2));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("name"));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("title"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(2));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("name"));
+                assertThat(hit.getMatchedQueryScore("name"), greaterThan(0f));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("title"));
+                assertThat(hit.getMatchedQueryScore("title"), greaterThan(0f));
             } else if (hit.getId().equals("2") || hit.getId().equals("3")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("name"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("name"));
+                assertThat(hit.getMatchedQueryScore("name"), greaterThan(0f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -203,9 +215,11 @@ public void testSimpleMatchedQueryFromTopLevelFilterAndFilteredQuery() throws Ex
         assertHitCount(searchResponse, 3L);
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1") || hit.getId().equals("2") || hit.getId().equals("3")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(2));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("name"));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("title"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(2));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("name"));
+                assertThat(hit.getMatchedQueryScore("name"), greaterThan(0f));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("title"));
+                assertThat(hit.getMatchedQueryScore("title"), greaterThan(0f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -242,8 +256,9 @@ public void testRegExpQuerySupportsName() throws InterruptedException {
 
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("regex"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("regex"));
+                assertThat(hit.getMatchedQueryScore("regex"), equalTo(1f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -265,8 +280,9 @@ public void testPrefixQuerySupportsName() throws InterruptedException {
 
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("prefix"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("prefix"));
+                assertThat(hit.getMatchedQueryScore("prefix"), equalTo(1f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -288,8 +304,9 @@ public void testFuzzyQuerySupportsName() throws InterruptedException {
 
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("fuzzy"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("fuzzy"));
+                assertThat(hit.getMatchedQueryScore("fuzzy"), greaterThan(0f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -311,8 +328,9 @@ public void testWildcardQuerySupportsName() throws InterruptedException {
 
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("wildcard"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("wildcard"));
+                assertThat(hit.getMatchedQueryScore("wildcard"), equalTo(1f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -334,8 +352,9 @@ public void testSpanFirstQuerySupportsName() throws InterruptedException {
 
         for (SearchHit hit : searchResponse.getHits()) {
             if (hit.getId().equals("1")) {
-                assertThat(hit.getMatchedQueries().length, equalTo(1));
-                assertThat(hit.getMatchedQueries(), hasItemInArray("span"));
+                assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                assertThat(hit.getMatchedQueriesAndScores(), hasKey("span"));
+                assertThat(hit.getMatchedQueryScore("span"), greaterThan(0f));
             } else {
                 fail("Unexpected document returned with id " + hit.getId());
             }
@@ -369,11 +388,13 @@ public void testMatchedWithShould() throws Exception {
             assertHitCount(searchResponse, 2L);
             for (SearchHit hit : searchResponse.getHits()) {
                 if (hit.getId().equals("1")) {
-                    assertThat(hit.getMatchedQueries().length, equalTo(1));
-                    assertThat(hit.getMatchedQueries(), hasItemInArray("dolor"));
+                    assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                    assertThat(hit.getMatchedQueriesAndScores(), hasKey("dolor"));
+                    assertThat(hit.getMatchedQueryScore("dolor"), greaterThan(0f));
                 } else if (hit.getId().equals("2")) {
-                    assertThat(hit.getMatchedQueries().length, equalTo(1));
-                    assertThat(hit.getMatchedQueries(), hasItemInArray("elit"));
+                    assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+                    assertThat(hit.getMatchedQueriesAndScores(), hasKey("elit"));
+                    assertThat(hit.getMatchedQueryScore("elit"), greaterThan(0f));
                 } else {
                     fail("Unexpected document returned with id " + hit.getId());
                 }
@@ -397,7 +418,10 @@ public void testMatchedWithWrapperQuery() throws Exception {
         for (QueryBuilder query : queries) {
             SearchResponse searchResponse = client().prepareSearch().setQuery(query).get();
             assertHitCount(searchResponse, 1L);
-            assertThat(searchResponse.getHits().getAt(0).getMatchedQueries()[0], equalTo("abc"));
+            SearchHit hit = searchResponse.getHits().getAt(0);
+            assertThat(hit.getMatchedQueriesAndScores().size(), equalTo(1));
+            assertThat(hit.getMatchedQueriesAndScores(), hasKey("abc"));
+            assertThat(hit.getMatchedQueryScore("abc"), greaterThan(0f));
         }
     }
 }