/
fmsketch.go
137 lines (125 loc) · 3.37 KB
/
fmsketch.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
//
// This source code is a modified form of original source from the TiDB project, which has the following copyright header(s):
//
// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.
package statistics
import (
"hash"
"github.com/squareup/pranadb/errors"
"github.com/squareup/pranadb/tidb/sessionctx/stmtctx"
"github.com/squareup/pranadb/tidb/types"
"github.com/squareup/pranadb/tidb/util/codec"
"github.com/twmb/murmur3"
)
// FMSketch is used to count the number of distinct elements in a set.
type FMSketch struct {
hashset map[uint64]bool
mask uint64
maxSize int
hashFunc hash.Hash64
}
// Copy makes a copy for current FMSketch.
func (s *FMSketch) Copy() *FMSketch {
if s == nil {
return nil
}
hashset := make(map[uint64]bool)
for key, value := range s.hashset {
hashset[key] = value
}
return &FMSketch{
hashset: hashset,
mask: s.mask,
maxSize: s.maxSize,
hashFunc: murmur3.New64(),
}
}
// NDV returns the ndv of the sketch.
func (s *FMSketch) NDV() int64 {
if s == nil {
return 0
}
return int64(s.mask+1) * int64(len(s.hashset))
}
func (s *FMSketch) insertHashValue(hashVal uint64) {
if (hashVal & s.mask) != 0 {
return
}
s.hashset[hashVal] = true
if len(s.hashset) > s.maxSize {
s.mask = s.mask*2 + 1
for key := range s.hashset {
if (key & s.mask) != 0 {
delete(s.hashset, key)
}
}
}
}
// InsertValue inserts a value into the FM sketch.
func (s *FMSketch) InsertValue(sc *stmtctx.StatementContext, value types.Datum) error {
bytes, err := codec.EncodeValue(sc, nil, value)
if err != nil {
return errors.Trace(err)
}
s.hashFunc.Reset()
_, err = s.hashFunc.Write(bytes)
if err != nil {
return errors.Trace(err)
}
s.insertHashValue(s.hashFunc.Sum64())
return nil
}
// InsertRowValue inserts multi-column values to the sketch.
func (s *FMSketch) InsertRowValue(sc *stmtctx.StatementContext, values []types.Datum) error {
b := make([]byte, 0, 8)
s.hashFunc.Reset()
for _, v := range values {
b = b[:0]
b, err := codec.EncodeValue(sc, b, v)
if err != nil {
return err
}
_, err = s.hashFunc.Write(b)
if err != nil {
return err
}
}
s.insertHashValue(s.hashFunc.Sum64())
return nil
}
// MergeFMSketch merges two FM Sketch.
func (s *FMSketch) MergeFMSketch(rs *FMSketch) {
if s == nil || rs == nil {
return
}
if s.mask < rs.mask {
s.mask = rs.mask
for key := range s.hashset {
if (key & s.mask) != 0 {
delete(s.hashset, key)
}
}
}
for key := range rs.hashset {
s.insertHashValue(key)
}
}
// MemoryUsage returns the total memory usage of a FMSketch.
func (s *FMSketch) MemoryUsage() (sum int64) {
// In FMSketch, we will ignore the memory usage of `hashFunc`.
// As for the variables mask(uint64) and maxSize(int) each will consume 8 bytes. This is the origin of the constant 16.
// And for the variables hashset(map[uint64]bool), each element in map will consume 9 bytes(8[uint64] + 1[bool]).
sum = int64(16 + 9*len(s.hashset))
return
}