Add Adjacency Matrix aggregation

This commit adds the Adjacency Matrix aggregation to Elastic v6 (https://www.elastic.co/guide/en/elasticsearch/reference/6.2/search-aggregations-bucket-adjacency-matrix-aggregation.html).

See #725

Author: Zyqsempai

README.md

@@ -182,7 +182,7 @@ There are also [some recipes](https://github.com/olivere/elastic/tree/release-br
   - [x] Top Hits
   - [x] Value Count
 - Bucket Aggregations
-  - [ ] Adjacency Matrix
+  - [x] Adjacency Matrix
   - [x] Children
   - [x] Date Histogram
   - [x] Date Range

search_aggs.go

@@ -248,6 +248,21 @@ func (a Aggregations) Filters(name string) (*AggregationBucketFilters, bool) {
 	return nil, false
 }
 
+// AdjacencyMatrix returning a form of adjacency matrix.
+// See: https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-adjacency-matrix-aggregation.html
+func (a Aggregations) AdjacencyMatrix(name string) (*AggregationBucketAdjacencyMatrix, bool) {
+	if raw, found := a[name]; found {
+		agg := new(AggregationBucketAdjacencyMatrix)
+		if raw == nil {
+			return agg, true
+		}
+		if err := json.Unmarshal(*raw, agg); err == nil {
+			return agg, true
+		}
+	}
+	return nil, false
+}
+
 // Missing returns missing results.
 // See: https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-missing-aggregation.html
 func (a Aggregations) Missing(name string) (*AggregationSingleBucket, bool) {
@@ -1215,6 +1230,33 @@ func (a *AggregationBucketFilters) UnmarshalJSON(data []byte) error {
 	return nil
 }
 
+// -- Bucket AdjacencyMatrix --
+
+// AggregationBucketAdjacencyMatrix is a multi-bucket aggregation that is returned
+// with a AdjacencyMatrix aggregation.
+type AggregationBucketAdjacencyMatrix struct {
+	Aggregations
+
+	Buckets []*AggregationBucketKeyItem //`json:"buckets"`
+	Meta    map[string]interface{}      // `json:"meta,omitempty"`
+}
+
+// UnmarshalJSON decodes JSON data and initializes an AggregationBucketAdjacencyMatrix structure.
+func (a *AggregationBucketAdjacencyMatrix) UnmarshalJSON(data []byte) error {
+	var aggs map[string]*json.RawMessage
+	if err := json.Unmarshal(data, &aggs); err != nil {
+		return err
+	}
+	if v, ok := aggs["buckets"]; ok && v != nil {
+		json.Unmarshal(*v, &a.Buckets)
+	}
+	if v, ok := aggs["meta"]; ok && v != nil {
+		json.Unmarshal(*v, &a.Meta)
+	}
+	a.Aggregations = aggs
+	return nil
+}
+
 // -- Bucket histogram items --
 
 // AggregationBucketHistogramItems is a bucket aggregation that is returned

search_aggs_bucket_adjacency_matrix.go

@@ -0,0 +1,96 @@
+// Copyright 2012-present Oliver Eilhard. All rights reserved.
+// Use of this source code is governed by a MIT-license.
+// See http://olivere.mit-license.org/license.txt for details.
+
+package elastic
+
+// AdjacencyMatrixAggregation  returning a form of adjacency matrix.
+// The request provides a collection of named filter expressions,
+// similar to the filters aggregation request. Each bucket in the
+// response represents a non-empty cell in the matrix of intersecting filters.
+//
+// For details, see
+// https://www.elastic.co/guide/en/elasticsearch/reference/6.0/search-aggregations-bucket-adjacency-matrix-aggregation.html
+type AdjacencyMatrixAggregation struct {
+	filters         map[string]Query
+	subAggregations map[string]Aggregation
+	meta            map[string]interface{}
+}
+
+// NewAdjacencyMatrixAggregation initializes a new AdjacencyMatrixAggregation.
+func NewAdjacencyMatrixAggregation() *AdjacencyMatrixAggregation {
+	return &AdjacencyMatrixAggregation{
+		filters:         make(map[string]Query),
+		subAggregations: make(map[string]Aggregation),
+	}
+}
+
+// Filters adds the filter
+func (a *AdjacencyMatrixAggregation) Filters(name string, filter Query) *AdjacencyMatrixAggregation {
+	a.filters[name] = filter
+	return a
+}
+
+// SubAggregation adds a sub-aggregation to this aggregation.
+func (a *AdjacencyMatrixAggregation) SubAggregation(name string, subAggregation Aggregation) *AdjacencyMatrixAggregation {
+	a.subAggregations[name] = subAggregation
+	return a
+}
+
+// Meta sets the meta data to be included in the aggregation response.
+func (a *AdjacencyMatrixAggregation) Meta(metaData map[string]interface{}) *AdjacencyMatrixAggregation {
+	a.meta = metaData
+	return a
+}
+
+// Source returns the a JSON-serializable interface.
+func (a *AdjacencyMatrixAggregation) Source() (interface{}, error) {
+	// Example:
+	//	{
+	//  "aggs" : {
+	//		"interactions" : {
+	//			"adjacency_matrix" : {
+	//				"filters" : {
+	//					"grpA" : { "terms" : { "accounts" : ["hillary", "sidney"] }},
+	//					"grpB" : { "terms" : { "accounts" : ["donald", "mitt"] }},
+	//					"grpC" : { "terms" : { "accounts" : ["vladimir", "nigel"] }}
+	//				}
+	//			}
+	//		}
+	//	}
+	// This method returns only the (outer) { "adjacency_matrix" : {} } part.
+
+	source := make(map[string]interface{})
+	adjacencyMatrix := make(map[string]interface{})
+	source["adjacency_matrix"] = adjacencyMatrix
+
+	dict := make(map[string]interface{})
+	for key, filter := range a.filters {
+		src, err := filter.Source()
+		if err != nil {
+			return nil, err
+		}
+		dict[key] = src
+	}
+	adjacencyMatrix["filters"] = dict
+
+	// AggregationBuilder (SubAggregations)
+	if len(a.subAggregations) > 0 {
+		aggsMap := make(map[string]interface{})
+		source["aggregations"] = aggsMap
+		for name, aggregate := range a.subAggregations {
+			src, err := aggregate.Source()
+			if err != nil {
+				return nil, err
+			}
+			aggsMap[name] = src
+		}
+	}
+
+	// Add Meta data if available
+	if len(a.meta) > 0 {
+		source["meta"] = a.meta
+	}
+
+	return source, nil
+}

search_aggs_bucket_adjacency_matrix_test.go

@@ -0,0 +1,71 @@
+// Copyright 2012-present Oliver Eilhard. All rights reserved.
+// Use of this source code is governed by a MIT-license.
+// See http://olivere.mit-license.org/license.txt for details.
+
+package elastic
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestAdjacencyMatrixAggregationFilters(t *testing.T) {
+	f1 := NewTermQuery("accounts", "sydney")
+	f2 := NewTermQuery("accounts", "mitt")
+	f3 := NewTermQuery("accounts", "nigel")
+	agg := NewAdjacencyMatrixAggregation().Filters("grpA", f1).Filters("grpB", f2).Filters("grpC", f3)
+	src, err := agg.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"adjacency_matrix":{"filters":{"grpA":{"term":{"accounts":"sydney"}},"grpB":{"term":{"accounts":"mitt"}},"grpC":{"term":{"accounts":"nigel"}}}}}`
+	if got != expected {
+		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+	}
+}
+
+func TestAdjacencyMatrixAggregationWithSubAggregation(t *testing.T) {
+	avgPriceAgg := NewAvgAggregation().Field("price")
+	f1 := NewTermQuery("accounts", "sydney")
+	f2 := NewTermQuery("accounts", "mitt")
+	f3 := NewTermQuery("accounts", "nigel")
+	agg := NewAdjacencyMatrixAggregation().SubAggregation("avg_price", avgPriceAgg).Filters("grpA", f1).Filters("grpB", f2).Filters("grpC", f3)
+	src, err := agg.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"adjacency_matrix":{"filters":{"grpA":{"term":{"accounts":"sydney"}},"grpB":{"term":{"accounts":"mitt"}},"grpC":{"term":{"accounts":"nigel"}}}},"aggregations":{"avg_price":{"avg":{"field":"price"}}}}`
+	if got != expected {
+		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+	}
+}
+
+func TestAdjacencyMatrixAggregationWithMetaData(t *testing.T) {
+	f1 := NewTermQuery("accounts", "sydney")
+	f2 := NewTermQuery("accounts", "mitt")
+	f3 := NewTermQuery("accounts", "nigel")
+	agg := NewAdjacencyMatrixAggregation().Filters("grpA", f1).Filters("grpB", f2).Filters("grpC", f3).Meta(map[string]interface{}{"name": "Oliver"})
+	src, err := agg.Source()
+	if err != nil {
+		t.Fatal(err)
+	}
+	data, err := json.Marshal(src)
+	if err != nil {
+		t.Fatalf("marshaling to JSON failed: %v", err)
+	}
+	got := string(data)
+	expected := `{"adjacency_matrix":{"filters":{"grpA":{"term":{"accounts":"sydney"}},"grpB":{"term":{"accounts":"mitt"}},"grpC":{"term":{"accounts":"nigel"}}}},"meta":{"name":"Oliver"}}`
+	if got != expected {
+		t.Errorf("expected\n%s\n,got:\n%s", expected, got)
+	}
+}

search_aggs_test.go

@@ -148,6 +148,11 @@ func TestAggs(t *testing.T) {
 		FilterWithName("olivere", NewTermQuery("user", "olivere")).
 		FilterWithName("sandrae", NewTermQuery("user", "sandrae"))
 	builder = builder.Aggregation("countByUser2", countByUserAgg2)
+	// AdjacencyMatrix
+	adjacencyMatrixAgg := NewAdjacencyMatrixAggregation().
+		Filters("groupA", NewTermQuery("user", "olivere")).
+		Filters("groupB", NewTermQuery("user", "sandrae"))
+	builder = builder.Aggregation("interactions", adjacencyMatrixAgg)
 	// AvgBucket
 	dateHisto := NewDateHistogramAggregation().Field("created").Interval("year")
 	dateHisto = dateHisto.SubAggregation("sumOfRetweets", NewSumAggregation().Field("retweets"))
@@ -978,6 +983,24 @@ func TestAggs(t *testing.T) {
 		t.Errorf("expected %d; got: %d", 1, b.DocCount)
 	}
 
+	// AdjacencyMatrix agg "adjacencyMatrixAgg" (named)
+	adjacencyMatrixAggRes, found := agg.AdjacencyMatrix("interactions")
+	if !found {
+		t.Errorf("expected %v; got: %v", true, found)
+	}
+	if adjacencyMatrixAggRes == nil {
+		t.Fatalf("expected != nil; got: nil")
+	}
+	if len(adjacencyMatrixAggRes.Buckets) != 2 {
+		t.Fatalf("expected %d; got: %d", 2, len(adjacencyMatrixAggRes.Buckets))
+	}
+	if adjacencyMatrixAggRes.Buckets[0].DocCount != 2 {
+		t.Errorf("expected %d; got: %d", 2, adjacencyMatrixAggRes.Buckets[0].DocCount)
+	}
+	if adjacencyMatrixAggRes.Buckets[1].DocCount != 1 {
+		t.Errorf("expected %d; got: %d", 1, adjacencyMatrixAggRes.Buckets[1].DocCount)
+	}
+
 	compositeAggRes, found := agg.Composite("composite")
 	if !found {
 		t.Errorf("expected %v; got: %v", true, found)
@@ -1912,6 +1935,71 @@ func TestAggsBucketFiltersWithNamedBuckets(t *testing.T) {
 	}
 }
 
+func TestAggsBucketAdjacencyMatrix(t *testing.T) {
+	s := `{
+	"interactions": {
+		"buckets": [
+			{
+				"key": "grpA",
+				"doc_count": 2,
+				"monthly": {
+					"buckets": []
+				}
+			},
+			{
+				"key": "grpA&grpB",
+				"doc_count": 1,
+				"monthly": {
+					"buckets": []
+				}
+			}
+		]
+	}
+}`
+
+	aggs := new(Aggregations)
+	err := json.Unmarshal([]byte(s), &aggs)
+	if err != nil {
+		t.Fatalf("expected no error decoding; got: %v", err)
+	}
+
+	agg, found := aggs.AdjacencyMatrix("interactions")
+	if !found {
+		t.Fatalf("expected aggregation to be found; got: %v", found)
+	}
+	if agg == nil {
+		t.Fatalf("expected aggregation != nil; got: %v", agg)
+	}
+	if agg.Buckets == nil {
+		t.Fatalf("expected aggregation buckets != %v; got: %v", nil, agg.Buckets)
+	}
+	if len(agg.Buckets) != 2 {
+		t.Fatalf("expected %d buckets; got: %d", 2, len(agg.Buckets))
+	}
+
+	if agg.Buckets[0].DocCount != 2 {
+		t.Fatalf("expected DocCount = %d; got: %d", 2, agg.Buckets[0].DocCount)
+	}
+	subAgg, found := agg.Buckets[0].Histogram("monthly")
+	if !found {
+		t.Fatalf("expected sub aggregation to be found; got: %v", found)
+	}
+	if subAgg == nil {
+		t.Fatalf("expected sub aggregation != %v; got: %v", nil, subAgg)
+	}
+
+	if agg.Buckets[1].DocCount != 1 {
+		t.Fatalf("expected DocCount = %d; got: %d", 1, agg.Buckets[1].DocCount)
+	}
+	subAgg, found = agg.Buckets[1].Histogram("monthly")
+	if !found {
+		t.Fatalf("expected sub aggregation to be found; got: %v", found)
+	}
+	if subAgg == nil {
+		t.Fatalf("expected sub aggregation != %v; got: %v", nil, subAgg)
+	}
+}
+
 func TestAggsBucketMissing(t *testing.T) {
 	s := `{
 	"products_without_a_price" : {