eintopf.info@v0.13.16/service/search/aggregation.go

// Copyright (C) 2022 The Eintopf authors
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package search

import (
	"fmt"
	"sort"
	"time"
)

type AggregationType string

const (
	TermsAggregation     = "terms"
	ObjectsAggregation   = "objects"
	DateRangeAggregation = "daterange"
)

type Aggregation struct {
	Type    AggregationType `json:"type"`
	Field   string          `json:"field"`
	Filters []Filter        `json:"filters"`
}

// CacheKey returns the same string for every identical aggregations. All
// aggregations resulting in the same bucket (given the same index) should
// return an equal cache key.
func (a Aggregation) CacheKey() string {
	cacheKey := fmt.Sprint(a.Type) + a.Field
	for _, filter := range a.Filters {
		cacheKey += filter.CacheKey()
	}
	return cacheKey
}

// aggregator aggregates values into a bucket.
type aggregator interface {
	// aggregate aggregates a single value.
	// It may return an error if the type or value is invalid.
	aggregate(value interface{}) error
	// bucket returns the resulting bucket.
	bucket() Bucket
}

type termsAggregator struct {
	// terms maps unique terms to its count.
	terms map[string]int
}

// aggregate takes a value with one of the following types:
//   - string
//   - []string
//   - []interface{}
//
// Returns an error if the value has a different type.
func (t *termsAggregator) aggregate(value interface{}) error {
	// Determine the term value and increase the count in the terms map.
	switch v := value.(type) {
	case string:
		t.terms[v]++
	case []string:
		for _, vv := range v {
			t.terms[vv]++
		}
	case []interface{}:
		for _, vv := range v {
			t.aggregate(vv)
		}
	default:
		return fmt.Errorf("invalid type: %T", v)
	}
	return nil
}

// bucket returns a new TermsBucket.
func (t *termsAggregator) bucket() Bucket {
	terms := make(TermsBucket, 0, len(t.terms))
	for term, count := range t.terms {
		terms = append(terms, Term{Term: term, Count: count})
	}
	// Sort the terms slice, to provide a stable output.
	sort.Slice(terms, func(i, j int) bool { return terms[i].Term < terms[j].Term })
	return terms
}

// DateLayout is the date layout used internally for date aggregations.
const DateLayout = "2006-01-02T15:04:05Z07:00"

type dateRangeAggregator struct {
	min time.Time
	max time.Time
}

func (d *dateRangeAggregator) aggregate(value interface{}) error {
	switch v := value.(type) {
	case string:
		return d.aggregateString(v)
	case []interface{}:
		for _, vv := range v {
			switch vvv := vv.(type) {
			case string:
				err := d.aggregateString(vvv)
				if err != nil {
					return err
				}
			default:
				return fmt.Errorf("invalid type: %T", v)
			}
		}
		return nil
	default:
		return fmt.Errorf("invalid type: %T", v)
	}
}

func (d *dateRangeAggregator) aggregateString(value string) error {
	date, err := time.Parse(DateLayout, value)
	if err != nil {
		return fmt.Errorf("invalid date format: %s", err)
	}

	if date.Unix() < d.min.Unix() {
		d.min = date
	}
	if date.Unix() > d.max.Unix() {
		d.max = date
	}
	return nil
}

func (d *dateRangeAggregator) bucket() Bucket {
	return DateRangeBucket{Min: d.min, Max: d.max}
}

type Bucket interface {
	BucketType() AggregationType
}

type TermsBucket []Term

func (t TermsBucket) BucketType() AggregationType {
	return TermsAggregation
}

type Term struct {
	Term  string `json:"term"`
	Count int    `json:"count"`
}

type DateRangeBucket struct {
	Min time.Time `json:"min"`
	Max time.Time `json:"max"`
}

func (d DateRangeBucket) BucketType() AggregationType {
	return DateRangeAggregation
}