-
Notifications
You must be signed in to change notification settings - Fork 0
/
sketch.go
77 lines (69 loc) · 1.84 KB
/
sketch.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
/*
Count-Min Sketch, an approximate counting data structure for summarizing data streams
for more information see https://github.com/jehiah/countmin
*/
package countmin
import (
"encoding/binary"
"hash/crc32"
"hash/fnv"
)
type Sketch interface {
Add([]byte, uint32) uint32
Query([]byte) uint32
AddString(string, uint32) uint32
QueryString(string) uint32
}
type countMinSketch struct {
Hashes int
Columns int
Data []uint32
}
// Create a new Sketch. Settings for hashes and columns affect performance
// of Adding and Querying items, but also accuracy.
func NewCountMinSketch(hashes int, columns int) Sketch {
s := countMinSketch{
Hashes: hashes,
Columns: columns,
Data: make([]uint32, hashes*columns),
}
return &s
}
func (s *countMinSketch) AddString(key string, count uint32) uint32 {
return s.Add([]byte(key), count)
}
func (s *countMinSketch) QueryString(key string) uint32 {
return s.Query([]byte(key))
}
func (s *countMinSketch) Add(key []byte, count uint32) uint32 {
// TODO: this is a bad implementation because we hash all twice in worst case.
newValue := s.Query(key) + count
h := fnv.New64a()
h.Write(key)
columns := uint32(s.Columns)
var b []byte
for base := uint32(0); base < uint32(s.Hashes)*columns; base += columns {
binary.Write(h, binary.LittleEndian, uint32(base))
index := crc32.ChecksumIEEE(h.Sum(b)) % columns
if s.Data[base+index] <= newValue {
s.Data[base+index] = newValue
}
}
return newValue
}
func (s *countMinSketch) Query(key []byte) uint32 {
h := fnv.New64a()
h.Write(key)
var min uint32
var b []byte
columns := uint32(s.Columns)
for base := uint32(0); base < uint32(s.Hashes)*columns; base += columns {
binary.Write(h, binary.LittleEndian, uint32(base))
index := crc32.ChecksumIEEE(h.Sum(b)) % columns
v := s.Data[base+index]
if base == 0 || v < min {
min = v
}
}
return min
}