Repository: opensearch-project/opensearch-py
Pull request #409 (2023-06-10 05:06:55)

feat: Add a `collapse` method to `opensearchpy.helpers.search.Search`
### Description
This PR adds a `collapse` method to `opensearchpy.helpers.search.Search` which allows adding a `collapse` requirement to a query without the need of `.extra()`.

This is achieved similarly to how `sort()` is done.

### Issues Resolved
#408 

qcoumes: I resolved the conflict.

I don't know where I should add some documentation in the user guide. The *Using DSL features from opensearch-dsl-py* chapter is mainly a redirection to `opensearch-dsl-py`'s documentation.

It seems weird to have a subchapter about `collapse` without some documentation about basic features.
dblock: > It seems weird to have a subchapter about `collapse` without some documentation about basic features.

Yes, you're of course right. Would love someone (hint hint) to contribute some basic features doc. I'll let it slide for this PR :)

Thank you!
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/409?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) Report
> Merging [#409](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/409?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) (f62d6dd) into [main](https://app.codecov.io/gh/opensearch-project/opensearch-py/commit/be57a4d604faf337a39f0d0734636bf9bd001dbf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) (be57a4d) will **decrease** coverage by `0.02%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##             main     #409      +/-   ##
==========================================
- Coverage   71.66%   71.65%   -0.02%     
==========================================
  Files          77       77              
  Lines        7190     7204      +14     
==========================================
+ Hits         5153     5162       +9     
- Misses       2037     2042       +5     
```


| [Impacted Files](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/409?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project) | Coverage Δ | |
|---|---|---|
| [opensearchpy/helpers/search.py](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/409?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=opensearch-project#diff-b3BlbnNlYXJjaHB5L2hlbHBlcnMvc2VhcmNoLnB5) | `82.07% <100.00%> (+0.67%)` | :arrow_up: |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/opensearch-project/opensearch-py/pull/409/indirect-changes?src=pr&el=tree-more&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 9a2afe06..89b9f894 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Compatibility with OpenSearch 2.1.0 - 2.6.0 ([#381](https://github.com/opensearch-project/opensearch-py/pull/381))
 - Added 'allow_redirects' parameter in perform_request function for RequestsHttpConnection ([#401](https://github.com/opensearch-project/opensearch-py/pull/401))
 - Enhanced YAML test runner to use OpenSearch rest-api-spec YAML tests ([#414](https://github.com/opensearch-project/opensearch-py/pull/414)
+- Added `Search#collapse` ([#409](https://github.com/opensearch-project/opensearch-py/issues/409))
 ### Changed
 - Upgrading pytest-asyncio to latest version - 0.21.0 ([#339](https://github.com/opensearch-project/opensearch-py/pull/339))
 - Fixed flaky CI tests by replacing httpbin with a simple http_server ([#395](https://github.com/opensearch-project/opensearch-py/pull/395))
diff --git a/opensearchpy/helpers/search.py b/opensearchpy/helpers/search.py
index 4fab50aa..0652b60a 100644
--- a/opensearchpy/helpers/search.py
+++ b/opensearchpy/helpers/search.py
@@ -331,6 +331,7 @@ class Search(Request):
 
         self.aggs = AggsProxy(self)
         self._sort = []
+        self._collapse = {}
         self._source = None
         self._highlight = {}
         self._highlight_opts = {}
@@ -579,6 +580,29 @@ class Search(Request):
             s._sort.append(k)
         return s
 
+    def collapse(self, field=None, inner_hits=None, max_concurrent_group_searches=None):
+        """
+        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 Search 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):
         """
         Update the global highlighting options used for this request. For
@@ -674,6 +698,9 @@ class Search(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_helpers/test_search.py b/test_opensearchpy/test_helpers/test_search.py
index afb032fe..91c7a709 100644
--- a/test_opensearchpy/test_helpers/test_search.py
+++ b/test_opensearchpy/test_helpers/test_search.py
@@ -266,6 +266,40 @@ def test_sort_by_score():
         s.sort("-_score")
 
 
+def test_collapse():
+    s = search.Search()
+
+    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.Search().to_dict() == s.to_dict()
+
+
 def test_slice():
     s = search.Search()
     assert {"from": 3, "size": 7} == s[3:10].to_dict()