caio.co/de/go-tdigest


Get rid of the centroid abstraction 💬 by Caio 8 years ago (log)
This was only being used to pack {float64,uint32}, all the
other functionality was skipped or became unused over
time for performance reasons. Away it goes.

Blob summary.go

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
package tdigest

import (
	"fmt"
	"math"
	"sort"
)

type summary struct {
	means  []float64
	counts []uint32
}

func newSummary(initialCapacity uint) *summary {
	return &summary{
		means:  make([]float64, 0, initialCapacity),
		counts: make([]uint32, 0, initialCapacity),
	}
}

func (s summary) Len() int {
	return len(s.means)
}

func (s *summary) Add(key float64, value uint32) error {

	if math.IsNaN(key) {
		return fmt.Errorf("Key must not be NaN")
	}

	if value == 0 {
		return fmt.Errorf("Count must be >0")
	}

	idx := s.FindInsertionIndex(key)

	s.means = append(s.means, math.NaN())
	s.counts = append(s.counts, 0)

	copy(s.means[idx+1:], s.means[idx:])
	copy(s.counts[idx+1:], s.counts[idx:])

	s.means[idx] = key
	s.counts[idx] = value

	return nil
}

func (s summary) Floor(x float64) int {
	return sort.Search(len(s.means), func(i int) bool {
		return s.means[i] >= x
	}) - 1
}

// Always insert to the right
func (s summary) FindInsertionIndex(x float64) int {
	return sort.Search(len(s.means), func(i int) bool {
		return s.means[i] > x
	})
}

func (s summary) HeadSum(index int) (sum float64) {
	for i := 0; i < index; i++ {
		sum += float64(s.counts[i])
	}
	return sum
}

func (s summary) FindIndex(x float64) int {
	idx := sort.Search(len(s.means), func(i int) bool {
		return s.means[i] >= x
	})
	if idx < s.Len() && s.means[idx] == x {
		return idx
	}
	return s.Len()
}

func (s summary) Mean(uncheckedIndex int) float64 {
	return s.means[uncheckedIndex]
}

func (s summary) Count(uncheckedIndex int) uint32 {
	return s.counts[uncheckedIndex]
}

// return the index of the last item which the sum of counts
// of items before it is less than or equal to `sum`. -1 in
// case no centroid satisfies the requirement.
// Since it's cheap, this also returns the `HeadSum` until
// the found index (i.e. cumSum = HeadSum(FloorSum(x)))
func (s summary) FloorSum(sum float64) (index int, cumSum float64) {
	index = -1
	for i := 0; i < s.Len(); i++ {
		if cumSum <= sum {
			index = i
		} else {
			break
		}
		cumSum += float64(s.counts[i])
	}
	if index != -1 {
		cumSum -= float64(s.counts[index])
	}
	return index, cumSum
}

func (s *summary) setAt(index int, mean float64, count uint32) {
	s.means[index] = mean
	s.counts[index] = count
	s.adjustRight(index)
	s.adjustLeft(index)
}

func (s *summary) adjustRight(index int) {
	for i := index + 1; i < len(s.means) && s.means[i-1] > s.means[i]; i++ {
		s.means[i-1], s.means[i] = s.means[i], s.means[i-1]
		s.counts[i-1], s.counts[i] = s.counts[i], s.counts[i-1]
	}
}

func (s *summary) adjustLeft(index int) {
	for i := index - 1; i >= 0 && s.means[i] > s.means[i+1]; i-- {
		s.means[i], s.means[i+1] = s.means[i+1], s.means[i]
		s.counts[i], s.counts[i+1] = s.counts[i+1], s.counts[i]
	}
}

func (s summary) ForEach(f func(float64, uint32) bool) {
	for i := 0; i < len(s.means); i++ {
		if !f(s.means[i], s.counts[i]) {
			break
		}
	}
}

func (s summary) Clone() *summary {
	return &summary{
		means:  append([]float64{}, s.means...),
		counts: append([]uint32{}, s.counts...),
	}
}