Repository: opensearch-project/opensearch-py
Pull request #827 (2024-09-25 12:46:27)

Collapse support added to AsyncSearch helper (#769)
### Description
This PR adds collapse support to AsyncSearch helper

### Issues Resolved
#769 

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check [here](https://github.com/opensearch-project/OpenSearch/blob/main/CONTRIBUTING.md#developer-certificate-of-origin).


CorneeSean: To maintainers: the same tests seem to fail on other PRs as well :)
dblock: > To maintainers: the same tests seem to fail on other PRs as well :)

I debugged this but couldn't find a solution. You can rebase with the workaround in https://github.com/opensearch-project/opensearch-py/pull/828 for now (or maybe you know how to fix that failure?).
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/827?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 70.34%. Comparing base [(`ba715b9`)](https://app.codecov.io/gh/opensearch-project/opensearch-py/commit/ba715b9d3f0609d25bfb0210b4b52234601e4dd8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) to head [(`26ad804`)](https://app.codecov.io/gh/opensearch-project/opensearch-py/commit/26ad8048c77169bf45bf3a343f996bc34b3f7121?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).
> Report is 60 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #827      +/-   ##
==========================================
- Coverage   71.95%   70.34%   -1.61%     
==========================================
  Files          91      113      +22     
  Lines        8001     8890     +889     
==========================================
+ Hits         5757     6254     +497     
- Misses       2244     2636     +392     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/827?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 672fd065..44af4328 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased]
 ### Added
+- Added `AsyncSearch#collapse` ([827](https://github.com/opensearch-project/opensearch-py/pull/827))
 ### Changed
 ### Deprecated
 ### Removed
diff --git a/opensearchpy/_async/helpers/search.py b/opensearchpy/_async/helpers/search.py
index e6a52428..c82999ce 100644
--- a/opensearchpy/_async/helpers/search.py
+++ b/opensearchpy/_async/helpers/search.py
@@ -8,7 +8,7 @@
 # GitHub history for details.
 
 import copy
-from typing import Any, Sequence
+from typing import Any, Dict, Sequence, cast
 
 from opensearchpy._async.helpers.actions import aiter, async_scan
 from opensearchpy.connection.async_connections import get_connection
@@ -39,6 +39,7 @@ class AsyncSearch(Request):
 
         self.aggs = AggsProxy(self)
         self._sort: Sequence[Any] = []
+        self._collapse: Dict[str, Any] = {}
         self._source: Any = None
         self._highlight: Any = {}
         self._highlight_opts: Any = {}
@@ -111,13 +112,13 @@ class AsyncSearch(Request):
         s.update_from_dict(d)
         return s
 
-    def _clone(self) -> Any:
+    def _clone(self) -> "AsyncSearch":
         """
         Return a clone of the current search request. Performs a shallow copy
         of all the underlying objects. Used internally by most state modifying
         APIs.
         """
-        s = super()._clone()
+        s = cast(AsyncSearch, super()._clone())
 
         s._response_class = self._response_class
         s._sort = self._sort[:]
@@ -126,6 +127,7 @@ class AsyncSearch(Request):
         s._highlight_opts = self._highlight_opts.copy()
         s._suggest = self._suggest.copy()
         s._script_fields = self._script_fields.copy()
+        s._collapse = self._collapse.copy()
         for x in ("query", "post_filter"):
             getattr(s, x)._proxied = getattr(self, x)._proxied
 
@@ -281,6 +283,34 @@ class AsyncSearch(Request):
             s._sort.append(k)
         return s
 
+    def collapse(
+        self,
+        field: Any = None,
+        inner_hits: Any = None,
+        max_concurrent_group_searches: Any = None,
+    ) -> "AsyncSearch":
+        """
+        Add collapsing information to the search request.
+
+        If called without providing ``field``, it will remove all collapse
+        requirements, otherwise it will replace them with the provided
+        arguments.
+
+        The API returns a copy of the AsyncSearch object and can thus be chained.
+        """
+        s = self._clone()
+        s._collapse = {}
+
+        if field is None:
+            return s
+
+        s._collapse["field"] = field
+        if inner_hits:
+            s._collapse["inner_hits"] = inner_hits
+        if max_concurrent_group_searches:
+            s._collapse["max_concurrent_group_searches"] = max_concurrent_group_searches
+        return s
+
     def highlight_options(self, **kwargs: Any) -> Any:
         """
         Update the global highlighting options used for this request. For
@@ -376,6 +406,9 @@ class AsyncSearch(Request):
             if self._sort:
                 d["sort"] = self._sort
 
+            if self._collapse:
+                d["collapse"] = self._collapse
+
             d.update(recursive_to_dict(self._extra))
 
             if self._source not in (None, {}):


Test patch:
diff --git a/test_opensearchpy/test_async/test_helpers/test_search.py b/test_opensearchpy/test_async/test_helpers/test_search.py
index d01f0b80..81ea75a4 100644
--- a/test_opensearchpy/test_async/test_helpers/test_search.py
+++ b/test_opensearchpy/test_async/test_helpers/test_search.py
@@ -240,6 +240,40 @@ async def test_sort_by_score() -> None:
         s.sort("-_score")
 
 
+def test_collapse() -> None:
+    s = search.AsyncSearch()
+
+    inner_hits = {"name": "most_recent", "size": 5, "sort": [{"@timestamp": "desc"}]}
+    s = s.collapse(
+        field="user.id", inner_hits=inner_hits, max_concurrent_group_searches=4
+    )
+
+    assert {
+        "field": "user.id",
+        "inner_hits": {
+            "name": "most_recent",
+            "size": 5,
+            "sort": [{"@timestamp": "desc"}],
+        },
+        "max_concurrent_group_searches": 4,
+    } == s._collapse
+    assert {
+        "collapse": {
+            "field": "user.id",
+            "inner_hits": {
+                "name": "most_recent",
+                "size": 5,
+                "sort": [{"@timestamp": "desc"}],
+            },
+            "max_concurrent_group_searches": 4,
+        }
+    } == s.to_dict()
+
+    s = s.collapse()
+    assert {} == s._collapse
+    assert search.AsyncSearch().to_dict() == s.to_dict()
+
+
 async def test_slice() -> None:
     s = search.AsyncSearch()
     assert {"from": 3, "size": 7} == s[3:10].to_dict()
@@ -546,3 +580,19 @@ async def test_rescore_query_to_dict() -> None:
             },
         },
     }
+
+
+def test_collapse_chaining() -> None:
+    s = search.AsyncSearch(index="index_name")
+    s = s.filter("term", color="red")
+    s = s.collapse(field="category")
+    s = s.filter("term", brand="something")
+
+    assert {
+        "query": {
+            "bool": {
+                "filter": [{"term": {"color": "red"}}, {"term": {"brand": "something"}}]
+            }
+        },
+        "collapse": {"field": "category"},
+    } == s.to_dict()