From db972e615b9156b4e364091d6a893d64fb3ef4f3 Mon Sep 17 00:00:00 2001 From: "C. Quentin" Date: Tue, 27 Jun 2023 01:19:15 +0200 Subject: [PATCH] feat: Add a `collapse` method to `opensearchpy.helpers.search.Search` (#409) Signed-off-by: Quentin Coumes Signed-off-by: Daniel (dB.) Doubrovkine Co-authored-by: Daniel (dB.) Doubrovkine --- CHANGELOG.md | 1 + opensearchpy/helpers/search.py | 27 +++++++++++++++ test_opensearchpy/test_helpers/test_search.py | 34 +++++++++++++++++++ 3 files changed, 62 insertions(+) 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, {}): 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()