feat: Add a collapse method to opensearchpy.helpers.search.Search (#409)

Signed-off-by: Quentin Coumes <[email protected]> Signed-off-by: Daniel (dB.) Doubrovkine <[email protected]> Co-authored-by: Daniel (dB.) Doubrovkine <[email protected]>
2023-06-26 19:19:15 -04:00
parent e5789c7b52
commit db972e615b
3 changed files with 62 additions and 0 deletions
@@ -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, {}):