1// Copyright 2012-present Oliver Eilhard. All rights reserved.
2// Use of this source code is governed by a MIT-license.
3// See http://olivere.mit-license.org/license.txt for details.
4
5package elastic
6
7// SamplerAggregation is a filtering aggregation used to limit any
8// sub aggregations' processing to a sample of the top-scoring documents.
9// Optionally, diversity settings can be used to limit the number of matches
10// that share a common value such as an "author".
11//
12// See: https://www.elastic.co/guide/en/elasticsearch/reference/5.2/search-aggregations-bucket-sampler-aggregation.html
13type SamplerAggregation struct {
14	subAggregations map[string]Aggregation
15	meta            map[string]interface{}
16
17	shardSize       int
18	maxDocsPerValue int
19	executionHint   string
20}
21
22func NewSamplerAggregation() *SamplerAggregation {
23	return &SamplerAggregation{
24		shardSize:       -1,
25		maxDocsPerValue: -1,
26		subAggregations: make(map[string]Aggregation),
27	}
28}
29
30func (a *SamplerAggregation) SubAggregation(name string, subAggregation Aggregation) *SamplerAggregation {
31	a.subAggregations[name] = subAggregation
32	return a
33}
34
35// Meta sets the meta data to be included in the aggregation response.
36func (a *SamplerAggregation) Meta(metaData map[string]interface{}) *SamplerAggregation {
37	a.meta = metaData
38	return a
39}
40
41// ShardSize sets the maximum number of docs returned from each shard.
42func (a *SamplerAggregation) ShardSize(shardSize int) *SamplerAggregation {
43	a.shardSize = shardSize
44	return a
45}
46
47func (a *SamplerAggregation) MaxDocsPerValue(maxDocsPerValue int) *SamplerAggregation {
48	a.maxDocsPerValue = maxDocsPerValue
49	return a
50}
51
52func (a *SamplerAggregation) ExecutionHint(hint string) *SamplerAggregation {
53	a.executionHint = hint
54	return a
55}
56
57func (a *SamplerAggregation) Source() (interface{}, error) {
58	// Example:
59	// {
60	//     "aggs" : {
61	//         "sample" : {
62	//             "sampler" : {
63	//                 "shard_size" : 200
64	//             },
65	// 						 "aggs": {
66	//                 "keywords": {
67	//                     "significant_terms": {
68	//                         "field": "text"
69	//                      }
70	//                 }
71	//             }
72	//         }
73	//     }
74	// }
75	//
76	// This method returns only the { "sampler" : { ... } } part.
77
78	source := make(map[string]interface{})
79	opts := make(map[string]interface{})
80	source["sampler"] = opts
81
82	if a.shardSize >= 0 {
83		opts["shard_size"] = a.shardSize
84	}
85	if a.maxDocsPerValue >= 0 {
86		opts["max_docs_per_value"] = a.maxDocsPerValue
87	}
88	if a.executionHint != "" {
89		opts["execution_hint"] = a.executionHint
90	}
91
92	// AggregationBuilder (SubAggregations)
93	if len(a.subAggregations) > 0 {
94		aggsMap := make(map[string]interface{})
95		source["aggregations"] = aggsMap
96		for name, aggregate := range a.subAggregations {
97			src, err := aggregate.Source()
98			if err != nil {
99				return nil, err
100			}
101			aggsMap[name] = src
102		}
103	}
104
105	// Add Meta data if available
106	if len(a.meta) > 0 {
107		source["meta"] = a.meta
108	}
109
110	return source, nil
111}
112