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	"encoding/json"
19	"reflect"
20	"sort"
21
22	"github.com/blevesearch/bleve/v2/size"
23	index "github.com/blevesearch/bleve_index_api"
24)
25
26var reflectStaticSizeFacetsBuilder int
27var reflectStaticSizeFacetResult int
28var reflectStaticSizeTermFacet int
29var reflectStaticSizeNumericRangeFacet int
30var reflectStaticSizeDateRangeFacet int
31
32func init() {
33	var fb FacetsBuilder
34	reflectStaticSizeFacetsBuilder = int(reflect.TypeOf(fb).Size())
35	var fr FacetResult
36	reflectStaticSizeFacetResult = int(reflect.TypeOf(fr).Size())
37	var tf TermFacet
38	reflectStaticSizeTermFacet = int(reflect.TypeOf(tf).Size())
39	var nrf NumericRangeFacet
40	reflectStaticSizeNumericRangeFacet = int(reflect.TypeOf(nrf).Size())
41	var drf DateRangeFacet
42	reflectStaticSizeDateRangeFacet = int(reflect.TypeOf(drf).Size())
43}
44
45type FacetBuilder interface {
46	StartDoc()
47	UpdateVisitor(term []byte)
48	EndDoc()
49
50	Result() *FacetResult
51	Field() string
52
53	Size() int
54}
55
56type FacetsBuilder struct {
57	indexReader   index.IndexReader
58	facetNames    []string
59	facets        []FacetBuilder
60	facetsByField map[string][]FacetBuilder
61	fields        []string
62}
63
64func NewFacetsBuilder(indexReader index.IndexReader) *FacetsBuilder {
65	return &FacetsBuilder{
66		indexReader: indexReader,
67	}
68}
69
70func (fb *FacetsBuilder) Size() int {
71	sizeInBytes := reflectStaticSizeFacetsBuilder + size.SizeOfPtr
72
73	for k, v := range fb.facets {
74		sizeInBytes += size.SizeOfString + v.Size() + len(fb.facetNames[k])
75	}
76
77	for _, entry := range fb.fields {
78		sizeInBytes += size.SizeOfString + len(entry)
79	}
80
81	return sizeInBytes
82}
83
84func (fb *FacetsBuilder) Add(name string, facetBuilder FacetBuilder) {
85	if fb.facetsByField == nil {
86		fb.facetsByField = map[string][]FacetBuilder{}
87	}
88
89	fb.facetNames = append(fb.facetNames, name)
90	fb.facets = append(fb.facets, facetBuilder)
91	fb.facetsByField[facetBuilder.Field()] = append(fb.facetsByField[facetBuilder.Field()], facetBuilder)
92	fb.fields = append(fb.fields, facetBuilder.Field())
93}
94
95func (fb *FacetsBuilder) RequiredFields() []string {
96	return fb.fields
97}
98
99func (fb *FacetsBuilder) StartDoc() {
100	for _, facetBuilder := range fb.facets {
101		facetBuilder.StartDoc()
102	}
103}
104
105func (fb *FacetsBuilder) EndDoc() {
106	for _, facetBuilder := range fb.facets {
107		facetBuilder.EndDoc()
108	}
109}
110
111func (fb *FacetsBuilder) UpdateVisitor(field string, term []byte) {
112	if facetBuilders, ok := fb.facetsByField[field]; ok {
113		for _, facetBuilder := range facetBuilders {
114			facetBuilder.UpdateVisitor(term)
115		}
116	}
117}
118
119type TermFacet struct {
120	Term  string `json:"term"`
121	Count int    `json:"count"`
122}
123
124type TermFacets struct {
125	termFacets []*TermFacet
126	termLookup map[string]*TermFacet
127}
128
129func (tf *TermFacets) Terms() []*TermFacet {
130	if tf == nil {
131		return []*TermFacet{}
132	}
133	return tf.termFacets
134}
135
136func (tf *TermFacets) TrimToTopN(n int) {
137	tf.termFacets = tf.termFacets[:n]
138}
139
140func (tf *TermFacets) Add(termFacets ...*TermFacet) {
141	for _, termFacet := range termFacets {
142		if tf.termLookup == nil {
143			tf.termLookup = map[string]*TermFacet{}
144		}
145
146		if term, ok := tf.termLookup[termFacet.Term]; ok {
147			term.Count += termFacet.Count
148			return
149		}
150
151		// if we got here it wasn't already in the existing terms
152		tf.termFacets = append(tf.termFacets, termFacet)
153		tf.termLookup[termFacet.Term] = termFacet
154	}
155}
156
157func (tf *TermFacets) Len() int {
158	// Handle case where *TermFacets is not fully initialized in index_impl.go.init()
159	if tf == nil {
160		return 0
161	}
162
163	return len(tf.termFacets)
164}
165func (tf *TermFacets) Swap(i, j int) {
166	tf.termFacets[i], tf.termFacets[j] = tf.termFacets[j], tf.termFacets[i]
167}
168func (tf *TermFacets) Less(i, j int) bool {
169	if tf.termFacets[i].Count == tf.termFacets[j].Count {
170		return tf.termFacets[i].Term < tf.termFacets[j].Term
171	}
172	return tf.termFacets[i].Count > tf.termFacets[j].Count
173}
174
175// TermFacets used to be a type alias for []*TermFacet.
176// To maintain backwards compatibility, we have to implement custom
177// JSON marshalling.
178func (tf *TermFacets) MarshalJSON() ([]byte, error) {
179	return json.Marshal(tf.termFacets)
180}
181
182func (tf *TermFacets) UnmarshalJSON(b []byte) error {
183	termFacets := []*TermFacet{}
184	err := json.Unmarshal(b, &termFacets)
185	if err != nil {
186		return err
187	}
188
189	for _, termFacet := range termFacets {
190		tf.Add(termFacet)
191	}
192
193	return nil
194}
195
196type NumericRangeFacet struct {
197	Name  string   `json:"name"`
198	Min   *float64 `json:"min,omitempty"`
199	Max   *float64 `json:"max,omitempty"`
200	Count int      `json:"count"`
201}
202
203func (nrf *NumericRangeFacet) Same(other *NumericRangeFacet) bool {
204	if nrf.Min == nil && other.Min != nil {
205		return false
206	}
207	if nrf.Min != nil && other.Min == nil {
208		return false
209	}
210	if nrf.Min != nil && other.Min != nil && *nrf.Min != *other.Min {
211		return false
212	}
213	if nrf.Max == nil && other.Max != nil {
214		return false
215	}
216	if nrf.Max != nil && other.Max == nil {
217		return false
218	}
219	if nrf.Max != nil && other.Max != nil && *nrf.Max != *other.Max {
220		return false
221	}
222
223	return true
224}
225
226type NumericRangeFacets []*NumericRangeFacet
227
228func (nrf NumericRangeFacets) Add(numericRangeFacet *NumericRangeFacet) NumericRangeFacets {
229	for _, existingNr := range nrf {
230		if numericRangeFacet.Same(existingNr) {
231			existingNr.Count += numericRangeFacet.Count
232			return nrf
233		}
234	}
235	// if we got here it wasn't already in the existing terms
236	nrf = append(nrf, numericRangeFacet)
237	return nrf
238}
239
240func (nrf NumericRangeFacets) Len() int      { return len(nrf) }
241func (nrf NumericRangeFacets) Swap(i, j int) { nrf[i], nrf[j] = nrf[j], nrf[i] }
242func (nrf NumericRangeFacets) Less(i, j int) bool {
243	if nrf[i].Count == nrf[j].Count {
244		return nrf[i].Name < nrf[j].Name
245	}
246	return nrf[i].Count > nrf[j].Count
247}
248
249type DateRangeFacet struct {
250	Name  string  `json:"name"`
251	Start *string `json:"start,omitempty"`
252	End   *string `json:"end,omitempty"`
253	Count int     `json:"count"`
254}
255
256func (drf *DateRangeFacet) Same(other *DateRangeFacet) bool {
257	if drf.Start == nil && other.Start != nil {
258		return false
259	}
260	if drf.Start != nil && other.Start == nil {
261		return false
262	}
263	if drf.Start != nil && other.Start != nil && *drf.Start != *other.Start {
264		return false
265	}
266	if drf.End == nil && other.End != nil {
267		return false
268	}
269	if drf.End != nil && other.End == nil {
270		return false
271	}
272	if drf.End != nil && other.End != nil && *drf.End != *other.End {
273		return false
274	}
275
276	return true
277}
278
279type DateRangeFacets []*DateRangeFacet
280
281func (drf DateRangeFacets) Add(dateRangeFacet *DateRangeFacet) DateRangeFacets {
282	for _, existingDr := range drf {
283		if dateRangeFacet.Same(existingDr) {
284			existingDr.Count += dateRangeFacet.Count
285			return drf
286		}
287	}
288	// if we got here it wasn't already in the existing terms
289	drf = append(drf, dateRangeFacet)
290	return drf
291}
292
293func (drf DateRangeFacets) Len() int      { return len(drf) }
294func (drf DateRangeFacets) Swap(i, j int) { drf[i], drf[j] = drf[j], drf[i] }
295func (drf DateRangeFacets) Less(i, j int) bool {
296	if drf[i].Count == drf[j].Count {
297		return drf[i].Name < drf[j].Name
298	}
299	return drf[i].Count > drf[j].Count
300}
301
302type FacetResult struct {
303	Field         string             `json:"field"`
304	Total         int                `json:"total"`
305	Missing       int                `json:"missing"`
306	Other         int                `json:"other"`
307	Terms         *TermFacets        `json:"terms,omitempty"`
308	NumericRanges NumericRangeFacets `json:"numeric_ranges,omitempty"`
309	DateRanges    DateRangeFacets    `json:"date_ranges,omitempty"`
310}
311
312func (fr *FacetResult) Size() int {
313	return reflectStaticSizeFacetResult + size.SizeOfPtr +
314		len(fr.Field) +
315		fr.Terms.Len()*(reflectStaticSizeTermFacet+size.SizeOfPtr) +
316		len(fr.NumericRanges)*(reflectStaticSizeNumericRangeFacet+size.SizeOfPtr) +
317		len(fr.DateRanges)*(reflectStaticSizeDateRangeFacet+size.SizeOfPtr)
318}
319
320func (fr *FacetResult) Merge(other *FacetResult) {
321	fr.Total += other.Total
322	fr.Missing += other.Missing
323	fr.Other += other.Other
324	if fr.Terms != nil && other.Terms != nil {
325		for _, term := range other.Terms.termFacets {
326			fr.Terms.Add(term)
327		}
328	}
329	if fr.NumericRanges != nil && other.NumericRanges != nil {
330		for _, nr := range other.NumericRanges {
331			fr.NumericRanges = fr.NumericRanges.Add(nr)
332		}
333	}
334	if fr.DateRanges != nil && other.DateRanges != nil {
335		for _, dr := range other.DateRanges {
336			fr.DateRanges = fr.DateRanges.Add(dr)
337		}
338	}
339}
340
341func (fr *FacetResult) Fixup(size int) {
342	if fr.Terms != nil {
343		sort.Sort(fr.Terms)
344		if fr.Terms.Len() > size {
345			moveToOther := fr.Terms.termFacets[size:]
346			for _, mto := range moveToOther {
347				fr.Other += mto.Count
348			}
349			fr.Terms.termFacets = fr.Terms.termFacets[0:size]
350		}
351	} else if fr.NumericRanges != nil {
352		sort.Sort(fr.NumericRanges)
353		if len(fr.NumericRanges) > size {
354			moveToOther := fr.NumericRanges[size:]
355			for _, mto := range moveToOther {
356				fr.Other += mto.Count
357			}
358			fr.NumericRanges = fr.NumericRanges[0:size]
359		}
360	} else if fr.DateRanges != nil {
361		sort.Sort(fr.DateRanges)
362		if len(fr.DateRanges) > size {
363			moveToOther := fr.DateRanges[size:]
364			for _, mto := range moveToOther {
365				fr.Other += mto.Count
366			}
367			fr.DateRanges = fr.DateRanges[0:size]
368		}
369	}
370}
371
372type FacetResults map[string]*FacetResult
373
374func (fr FacetResults) Merge(other FacetResults) {
375	for name, oFacetResult := range other {
376		facetResult, ok := fr[name]
377		if ok {
378			facetResult.Merge(oFacetResult)
379		} else {
380			fr[name] = oFacetResult
381		}
382	}
383}
384
385func (fr FacetResults) Fixup(name string, size int) {
386	facetResult, ok := fr[name]
387	if ok {
388		facetResult.Fixup(size)
389	}
390}
391
392func (fb *FacetsBuilder) Results() FacetResults {
393	fr := make(FacetResults)
394	for i, facetBuilder := range fb.facets {
395		facetResult := facetBuilder.Result()
396		fr[fb.facetNames[i]] = facetResult
397	}
398	return fr
399}
400