1//  Copyright (c) 2014 Couchbase, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// 		http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package search
16
17import (
18	"reflect"
19	"sort"
20
21	"github.com/blevesearch/bleve/index"
22	"github.com/blevesearch/bleve/size"
23)
24
25var reflectStaticSizeFacetsBuilder int
26var reflectStaticSizeFacetResult int
27var reflectStaticSizeTermFacet int
28var reflectStaticSizeNumericRangeFacet int
29var reflectStaticSizeDateRangeFacet int
30
31func init() {
32	var fb FacetsBuilder
33	reflectStaticSizeFacetsBuilder = int(reflect.TypeOf(fb).Size())
34	var fr FacetResult
35	reflectStaticSizeFacetResult = int(reflect.TypeOf(fr).Size())
36	var tf TermFacet
37	reflectStaticSizeTermFacet = int(reflect.TypeOf(tf).Size())
38	var nrf NumericRangeFacet
39	reflectStaticSizeNumericRangeFacet = int(reflect.TypeOf(nrf).Size())
40	var drf DateRangeFacet
41	reflectStaticSizeDateRangeFacet = int(reflect.TypeOf(drf).Size())
42}
43
44type FacetBuilder interface {
45	StartDoc()
46	UpdateVisitor(field string, term []byte)
47	EndDoc()
48
49	Result() *FacetResult
50	Field() string
51
52	Size() int
53}
54
55type FacetsBuilder struct {
56	indexReader index.IndexReader
57	facets      map[string]FacetBuilder
58	fields      []string
59}
60
61func NewFacetsBuilder(indexReader index.IndexReader) *FacetsBuilder {
62	return &FacetsBuilder{
63		indexReader: indexReader,
64		facets:      make(map[string]FacetBuilder, 0),
65	}
66}
67
68func (fb *FacetsBuilder) Size() int {
69	sizeInBytes := reflectStaticSizeFacetsBuilder + size.SizeOfPtr
70
71	for k, v := range fb.facets {
72		sizeInBytes += size.SizeOfString + len(k) +
73			v.Size()
74	}
75
76	for _, entry := range fb.fields {
77		sizeInBytes += size.SizeOfString + len(entry)
78	}
79
80	return sizeInBytes
81}
82
83func (fb *FacetsBuilder) Add(name string, facetBuilder FacetBuilder) {
84	fb.facets[name] = facetBuilder
85	fb.fields = append(fb.fields, facetBuilder.Field())
86}
87
88func (fb *FacetsBuilder) RequiredFields() []string {
89	return fb.fields
90}
91
92func (fb *FacetsBuilder) StartDoc() {
93	for _, facetBuilder := range fb.facets {
94		facetBuilder.StartDoc()
95	}
96}
97
98func (fb *FacetsBuilder) EndDoc() {
99	for _, facetBuilder := range fb.facets {
100		facetBuilder.EndDoc()
101	}
102}
103
104func (fb *FacetsBuilder) UpdateVisitor(field string, term []byte) {
105	for _, facetBuilder := range fb.facets {
106		facetBuilder.UpdateVisitor(field, term)
107	}
108}
109
110type TermFacet struct {
111	Term  string `json:"term"`
112	Count int    `json:"count"`
113}
114
115type TermFacets []*TermFacet
116
117func (tf TermFacets) Add(termFacet *TermFacet) TermFacets {
118	for _, existingTerm := range tf {
119		if termFacet.Term == existingTerm.Term {
120			existingTerm.Count += termFacet.Count
121			return tf
122		}
123	}
124	// if we got here it wasn't already in the existing terms
125	tf = append(tf, termFacet)
126	return tf
127}
128
129func (tf TermFacets) Len() int      { return len(tf) }
130func (tf TermFacets) Swap(i, j int) { tf[i], tf[j] = tf[j], tf[i] }
131func (tf TermFacets) Less(i, j int) bool {
132	if tf[i].Count == tf[j].Count {
133		return tf[i].Term < tf[j].Term
134	}
135	return tf[i].Count > tf[j].Count
136}
137
138type NumericRangeFacet struct {
139	Name  string   `json:"name"`
140	Min   *float64 `json:"min,omitempty"`
141	Max   *float64 `json:"max,omitempty"`
142	Count int      `json:"count"`
143}
144
145func (nrf *NumericRangeFacet) Same(other *NumericRangeFacet) bool {
146	if nrf.Min == nil && other.Min != nil {
147		return false
148	}
149	if nrf.Min != nil && other.Min == nil {
150		return false
151	}
152	if nrf.Min != nil && other.Min != nil && *nrf.Min != *other.Min {
153		return false
154	}
155	if nrf.Max == nil && other.Max != nil {
156		return false
157	}
158	if nrf.Max != nil && other.Max == nil {
159		return false
160	}
161	if nrf.Max != nil && other.Max != nil && *nrf.Max != *other.Max {
162		return false
163	}
164
165	return true
166}
167
168type NumericRangeFacets []*NumericRangeFacet
169
170func (nrf NumericRangeFacets) Add(numericRangeFacet *NumericRangeFacet) NumericRangeFacets {
171	for _, existingNr := range nrf {
172		if numericRangeFacet.Same(existingNr) {
173			existingNr.Count += numericRangeFacet.Count
174			return nrf
175		}
176	}
177	// if we got here it wasn't already in the existing terms
178	nrf = append(nrf, numericRangeFacet)
179	return nrf
180}
181
182func (nrf NumericRangeFacets) Len() int      { return len(nrf) }
183func (nrf NumericRangeFacets) Swap(i, j int) { nrf[i], nrf[j] = nrf[j], nrf[i] }
184func (nrf NumericRangeFacets) Less(i, j int) bool {
185	if nrf[i].Count == nrf[j].Count {
186		return nrf[i].Name < nrf[j].Name
187	}
188	return nrf[i].Count > nrf[j].Count
189}
190
191type DateRangeFacet struct {
192	Name  string  `json:"name"`
193	Start *string `json:"start,omitempty"`
194	End   *string `json:"end,omitempty"`
195	Count int     `json:"count"`
196}
197
198func (drf *DateRangeFacet) Same(other *DateRangeFacet) bool {
199	if drf.Start == nil && other.Start != nil {
200		return false
201	}
202	if drf.Start != nil && other.Start == nil {
203		return false
204	}
205	if drf.Start != nil && other.Start != nil && *drf.Start != *other.Start {
206		return false
207	}
208	if drf.End == nil && other.End != nil {
209		return false
210	}
211	if drf.End != nil && other.End == nil {
212		return false
213	}
214	if drf.End != nil && other.End != nil && *drf.End != *other.End {
215		return false
216	}
217
218	return true
219}
220
221type DateRangeFacets []*DateRangeFacet
222
223func (drf DateRangeFacets) Add(dateRangeFacet *DateRangeFacet) DateRangeFacets {
224	for _, existingDr := range drf {
225		if dateRangeFacet.Same(existingDr) {
226			existingDr.Count += dateRangeFacet.Count
227			return drf
228		}
229	}
230	// if we got here it wasn't already in the existing terms
231	drf = append(drf, dateRangeFacet)
232	return drf
233}
234
235func (drf DateRangeFacets) Len() int      { return len(drf) }
236func (drf DateRangeFacets) Swap(i, j int) { drf[i], drf[j] = drf[j], drf[i] }
237func (drf DateRangeFacets) Less(i, j int) bool {
238	if drf[i].Count == drf[j].Count {
239		return drf[i].Name < drf[j].Name
240	}
241	return drf[i].Count > drf[j].Count
242}
243
244type FacetResult struct {
245	Field         string             `json:"field"`
246	Total         int                `json:"total"`
247	Missing       int                `json:"missing"`
248	Other         int                `json:"other"`
249	Terms         TermFacets         `json:"terms,omitempty"`
250	NumericRanges NumericRangeFacets `json:"numeric_ranges,omitempty"`
251	DateRanges    DateRangeFacets    `json:"date_ranges,omitempty"`
252}
253
254func (fr *FacetResult) Size() int {
255	return reflectStaticSizeFacetResult + size.SizeOfPtr +
256		len(fr.Field) +
257		len(fr.Terms)*(reflectStaticSizeTermFacet+size.SizeOfPtr) +
258		len(fr.NumericRanges)*(reflectStaticSizeNumericRangeFacet+size.SizeOfPtr) +
259		len(fr.DateRanges)*(reflectStaticSizeDateRangeFacet+size.SizeOfPtr)
260}
261
262func (fr *FacetResult) Merge(other *FacetResult) {
263	fr.Total += other.Total
264	fr.Missing += other.Missing
265	fr.Other += other.Other
266	if fr.Terms != nil && other.Terms != nil {
267		for _, term := range other.Terms {
268			fr.Terms = fr.Terms.Add(term)
269		}
270	}
271	if fr.NumericRanges != nil && other.NumericRanges != nil {
272		for _, nr := range other.NumericRanges {
273			fr.NumericRanges = fr.NumericRanges.Add(nr)
274		}
275	}
276	if fr.DateRanges != nil && other.DateRanges != nil {
277		for _, dr := range other.DateRanges {
278			fr.DateRanges = fr.DateRanges.Add(dr)
279		}
280	}
281}
282
283func (fr *FacetResult) Fixup(size int) {
284	if fr.Terms != nil {
285		sort.Sort(fr.Terms)
286		if len(fr.Terms) > size {
287			moveToOther := fr.Terms[size:]
288			for _, mto := range moveToOther {
289				fr.Other += mto.Count
290			}
291			fr.Terms = fr.Terms[0:size]
292		}
293	} else if fr.NumericRanges != nil {
294		sort.Sort(fr.NumericRanges)
295		if len(fr.NumericRanges) > size {
296			moveToOther := fr.NumericRanges[size:]
297			for _, mto := range moveToOther {
298				fr.Other += mto.Count
299			}
300			fr.NumericRanges = fr.NumericRanges[0:size]
301		}
302	} else if fr.DateRanges != nil {
303		sort.Sort(fr.DateRanges)
304		if len(fr.DateRanges) > size {
305			moveToOther := fr.DateRanges[size:]
306			for _, mto := range moveToOther {
307				fr.Other += mto.Count
308			}
309			fr.DateRanges = fr.DateRanges[0:size]
310		}
311	}
312}
313
314type FacetResults map[string]*FacetResult
315
316func (fr FacetResults) Merge(other FacetResults) {
317	for name, oFacetResult := range other {
318		facetResult, ok := fr[name]
319		if ok {
320			facetResult.Merge(oFacetResult)
321		} else {
322			fr[name] = oFacetResult
323		}
324	}
325}
326
327func (fr FacetResults) Fixup(name string, size int) {
328	facetResult, ok := fr[name]
329	if ok {
330		facetResult.Fixup(size)
331	}
332}
333
334func (fb *FacetsBuilder) Results() FacetResults {
335	fr := make(FacetResults)
336	for facetName, facetBuilder := range fb.facets {
337		facetResult := facetBuilder.Result()
338		fr[facetName] = facetResult
339	}
340	return fr
341}
342