Enable fenwick tree after some threshold
- Id
- 31ecb6ee68b594629e4405432ba129338cc1c4f5
- Author
- Vladimir Mihailenco
- Commit time
- 2018-11-08T12:19:16+02:00
Modified summary.go
)
type summary struct {
- means []float64
- counts []uint32
- bitree *fenwick.List
+ initialCapacity int
+ means []float64
+ counts []uint32
+ bitree *fenwick.List
}
-func newSummary(initialCapacity uint) *summary {
+func newSummary(initialCapacity int) *summary {
s := &summary{
- means: make([]float64, 0, initialCapacity),
- counts: make([]uint32, 0, initialCapacity),
+ initialCapacity: initialCapacity,
+ means: make([]float64, 0, initialCapacity),
+ counts: make([]uint32, 0, initialCapacity),
}
- s.rebuildFenwickTree()
+ s.rebuildFenwickTree(-1)
return s
}
s.means[idx] = key
s.counts[idx] = value
- // Reinitialize the prefixSum cache
- if s.bitree.Len() >= len(s.counts) {
- for i := idx; i < len(s.counts); i++ {
- s.bitree.Set(i, s.counts[i])
- }
- } else {
- s.rebuildFenwickTree()
- }
+ s.rebuildFenwickTree(idx)
return nil
}
-func (s *summary) rebuildFenwickTree() {
- s.bitree = fenwick.New(s.counts[:cap(s.counts)]...)
+func (s *summary) rebuildFenwickTree(idx int) {
+ if !s.useFenwickTree() {
+ return
+ }
+
+ if s.bitree == nil || s.bitree.Len() < len(s.counts) {
+ s.bitree = fenwick.New(s.counts[:cap(s.counts)]...)
+ return
+ }
+
+ for i := idx; i < len(s.counts); i++ {
+ s.bitree.Set(i, s.counts[i])
+ }
+}
+
+func (s *summary) useFenwickTree() bool {
+ return s.initialCapacity > 100 && len(s.counts) >= s.initialCapacity
}
func (s summary) Floor(x float64) int {
})
}
-func (s summary) HeadSum(index int) (sum float64) {
- return float64(s.bitree.Sum(index))
+// This method is the hotspot when calling Add(), which in turn is called by
+// Compress() and Merge().
+func (s summary) HeadSum(idx int) (sum float64) {
+ if s.bitree != nil {
+ return float64(s.bitree.Sum(idx))
+ }
+ return float64(sumUntilIndex(s.counts, idx))
}
func (s summary) FindIndex(x float64) int {
s.counts[index] = count
s.adjustRight(index)
s.adjustLeft(index)
- s.bitree.Set(index, count)
+ if s.bitree != nil {
+ s.bitree.Set(index, count)
+ }
}
func (s *summary) adjustRight(index int) {
means: append([]float64{}, s.means...),
counts: append([]uint32{}, s.counts...),
}
+}
+
+// A simple loop unroll saves a surprising amount of time.
+func sumUntilIndex(s []uint32, idx int) uint64 {
+ var cumSum uint64
+ var i int
+ for i = idx - 1; i >= 3; i -= 4 {
+ cumSum += uint64(s[i])
+ cumSum += uint64(s[i-1])
+ cumSum += uint64(s[i-2])
+ cumSum += uint64(s[i-3])
+ }
+ for ; i >= 0; i-- {
+ cumSum += uint64(s[i])
+ }
+ return cumSum
}
Modified tdigest.go
}
oldTree := t.summary
- t.summary = newSummary(uint(t.summary.Len()))
+ t.summary = newSummary(t.summary.Len())
t.count = 0
shuffle(oldTree.means, oldTree.counts, t.rng)
// Clone returns a deep copy of a TDigest.
func (t *TDigest) Clone() *TDigest {
summary := t.summary.Clone()
- summary.rebuildFenwickTree()
+ summary.rebuildFenwickTree(-1)
return &TDigest{
summary: summary,
compression: t.compression,
}
}
-func estimateCapacity(compression float64) uint {
- return uint(compression) * 10
+func estimateCapacity(compression float64) int {
+ return int(compression) * 10
}
Modified tdigest_test.go
"sort"
"testing"
+ "github.com/caio/go-tdigest/internal/fenwick"
"github.com/leesper/go_rng"
"gonum.org/v1/gonum/stat"
)
count: 1250,
rng: &globalRNG{},
}
- td.summary.rebuildFenwickTree()
+ td.summary.rebuildFenwickTree(-1)
if cdf := td.CDF(7.144560976650238e+06); cdf > 1 {
t.Fatalf("invalid: %v", cdf)
}
}
-func benchmarkAddCompression(compression uint32, b *testing.B) {
+var compressions = []uint32{1, 10, 20, 30, 50, 100}
+
+func BenchmarkTDigestAddOnce(b *testing.B) {
+ for _, compression := range compressions {
+ compression := compression
+ b.Run(fmt.Sprintf("compression=%d", compression), func(b *testing.B) {
+ benchmarkAddOnce(b, compression)
+ })
+ }
+}
+
+func benchmarkAddOnce(b *testing.B, compression uint32) {
t := uncheckedNew(Compression(compression))
data := make([]float64, b.N)
b.ReportAllocs()
b.ResetTimer()
for n := 0; n < b.N; n++ {
- err := t.AddWeighted(data[n], 1)
+ err := t.Add(data[n])
if err != nil {
b.Error(err)
}
b.StopTimer()
}
-func BenchmarkAddCompression1(b *testing.B) {
- benchmarkAddCompression(1, b)
+func BenchmarkTDigestAddMulti(b *testing.B) {
+ for _, compression := range compressions {
+ compression := compression
+ for _, n := range []int{10, 100, 1000, 10000} {
+ n := n
+ name := fmt.Sprintf("compression=%d n=%d", compression, n)
+ b.Run(name, func(b *testing.B) {
+ benchmarkAddMulti(b, compression, n)
+ })
+ }
+ }
}
-func BenchmarkAddCompression10(b *testing.B) {
- benchmarkAddCompression(10, b)
-}
-
-func BenchmarkAddCompression100(b *testing.B) {
- benchmarkAddCompression(100, b)
-}
-
-func benchmarkAddMultipleTimes(b *testing.B, times int) {
+func benchmarkAddMulti(b *testing.B, compression uint32, times int) {
data := make([]float64, times)
for i := 0; i < times; i++ {
data[i] = rand.Float64()
b.ReportAllocs()
b.ResetTimer()
for n := 0; n < b.N; n++ {
- t := uncheckedNew()
+ t := uncheckedNew(Compression(compression))
for i := 0; i < times; i++ {
err := t.AddWeighted(data[i], 1)
if err != nil {
b.StopTimer()
}
-func BenchmarkAddMultipleTimes(b *testing.B) {
- for _, n := range []int{10, 100, 1000, 10000} {
- n := n
- b.Run(fmt.Sprint(n), func(b *testing.B) {
- benchmarkAddMultipleTimes(b, n)
+func BenchmarkTDigestMerge(b *testing.B) {
+ for _, compression := range compressions {
+ compression := compression
+ for _, n := range []int{1, 10, 100} {
+ name := fmt.Sprintf("compression=%d n=%d", compression, n)
+ b.Run(name, func(b *testing.B) {
+ benchmarkMerge(b, compression, n)
+ })
+ }
+ }
+}
+
+func benchmarkMerge(b *testing.B, compression uint32, times int) {
+ ts := make([]*TDigest, times)
+ for i := 0; i < times; i++ {
+ ts[i] = randomTDigest(compression)
+ }
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for n := 0; n < b.N; n++ {
+ dst := uncheckedNew(Compression(compression))
+
+ for i := 0; i < times; i++ {
+ err := dst.Merge(ts[i])
+ if err != nil {
+ b.Fatal(err)
+ }
+ }
+
+ err := dst.Compress()
+ if err != nil {
+ b.Fatal(err)
+ }
+ }
+}
+
+func randomTDigest(compression uint32) *TDigest {
+ t := uncheckedNew(Compression(compression))
+ n := 20 * int(compression)
+ for i := 0; i < n; i++ {
+ err := t.Add(rand.Float64())
+ if err != nil {
+ panic(err)
+ }
+ }
+ return t
+}
+
+var sumSizes = []int{10, 100, 1000, 10000}
+
+func BenchmarkSumFenwickTree(b *testing.B) {
+ for _, size := range sumSizes {
+ size := size
+ b.Run(fmt.Sprint(size), func(b *testing.B) {
+ benchmarkSumFenwickTree(b, size)
})
}
+}
+
+func benchmarkSumFenwickTree(b *testing.B, size int) {
+ counts := generateCounts(size)
+ indexes := generateIndexes(size)
+ tree := fenwick.New(counts...)
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for n := 0; n < b.N; n++ {
+ for _, idx := range indexes {
+ _ = tree.Sum(idx)
+ }
+ }
+}
+
+func BenchmarkSumLoopSimple(b *testing.B) {
+ for _, size := range sumSizes {
+ size := size
+ b.Run(fmt.Sprint(size), func(b *testing.B) {
+ benchmarkSumLoopSimple(b, size)
+ })
+ }
+}
+
+func benchmarkSumLoopSimple(b *testing.B, size int) {
+ counts := generateCounts(size)
+ indexes := generateIndexes(size)
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for n := 0; n < b.N; n++ {
+ for _, idx := range indexes {
+ _ = sumUntilIndexSimple(counts, idx)
+ }
+ }
+}
+
+func BenchmarkSumLoopUnrolled(b *testing.B) {
+ for _, size := range sumSizes {
+ size := size
+ b.Run(fmt.Sprint(size), func(b *testing.B) {
+ benchmarkSumLoopUnrolled(b, size)
+ })
+ }
+}
+
+func benchmarkSumLoopUnrolled(b *testing.B, size int) {
+ counts := generateCounts(size)
+ indexes := generateIndexes(size)
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for n := 0; n < b.N; n++ {
+ for _, idx := range indexes {
+ _ = sumUntilIndex(counts, idx)
+ }
+ }
+}
+
+func generateCounts(size int) []uint32 {
+ counts := make([]uint32, size)
+ for i := 0; i < size; i++ {
+ counts[i] = rand.Uint32()
+ }
+ return counts
+}
+
+func generateIndexes(size int) []int {
+ const num = 100
+
+ indexes := make([]int, num)
+ for i := 0; i < num; i++ {
+ indexes[i] = rand.Intn(size)
+ }
+ return indexes
+}
+
+func sumUntilIndexSimple(counts []uint32, idx int) uint64 {
+ var sum uint64
+ for _, c := range counts {
+ sum += uint64(c)
+ }
+ return sum
}