From 09ed23b5e5abaa24ea1e193bc538540c1ffc725b Mon Sep 17 00:00:00 2001
From: Hardy <caoxiaozhen@secnium.com>
Date: Sat, 17 Jul 2021 16:22:41 +0800
Subject: [PATCH] feat: add collapse search result function

---
 search.go      | 15 ++++++++++++---
 search_test.go | 10 ++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/search.go b/search.go
index 5c661a2..0a08b4b 100644
--- a/search.go
+++ b/search.go
@@ -26,7 +26,7 @@ type SearchRequest struct {
 	sort        Sort
 	source      Source
 	timeout     *time.Duration
-
+	collapse    map[string]interface{}
 }
 
 // Search creates a new SearchRequest object, to be filled via method chaining.
@@ -76,6 +76,14 @@ func (req *SearchRequest) Sort(name string, order Order) *SearchRequest {
 	return req
 }
 
+// Collapse the results
+func (req *SearchRequest) Collapse(field string) *SearchRequest {
+	req.collapse = map[string]interface{}{
+		"field": field,
+	}
+	return req
+}
+
 // SearchAfter retrieve the sorted result
 func (req *SearchRequest) SearchAfter(s ...interface{}) *SearchRequest {
 	req.searchAfter = append(req.searchAfter, s...)
@@ -113,8 +121,6 @@ func (req *SearchRequest) Highlight(highlight Mappable) *SearchRequest {
 	return req
 }
 
-
-
 // Map implements the Mappable interface. It converts the request to into a
 // nested map[string]interface{}, as expected by the go-elasticsearch library.
 func (req *SearchRequest) Map() map[string]interface{} {
@@ -155,6 +161,9 @@ func (req *SearchRequest) Map() map[string]interface{} {
 		m["search_after"] = req.searchAfter
 	}
 
+	if req.collapse != nil {
+		m["collapse"] = req.collapse
+	}
 
 	source := req.source.Map()
 	if len(source) > 0 {
diff --git a/search_test.go b/search_test.go
index 1bcf948..a6378eb 100644
--- a/search_test.go
+++ b/search_test.go
@@ -14,6 +14,16 @@ func TestSearchMaps(t *testing.T) {
 				"search_after": []string{"_id", "name"},
 			},
 		},
+		{
+			"a simple query with collapse",
+			Search().Query(Term("user.id", "Tony")).Collapse("user.id"),
+			map[string]interface{}{
+				"query": Term("user.id", "Tony").Map(),
+				"collapse": map[string]interface{}{
+					"field": "user.id",
+				},
+			},
+		},
 		{
 			"a simple match_all query with a size and no aggs",
 			Search().Query(MatchAll()).Size(20),