/
generator.go
84 lines (71 loc) · 1.26 KB
/
generator.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
package main
import (
"math/rand"
/* for Zipfian distribution */
"github.com/pingcap/go-ycsb/pkg/generator"
)
const (
OP_RD = iota
OP_WR
OP_SCAN
)
const (
DIST_UNIFORM = iota
DIST_ZIPFIAN
)
type Generator struct {
rd *rand.Rand
nKeys int
rKeys uint64
rdRatio uint64
zipfian *generator.Zipfian
dist int
}
func NewGenerator(
seed int,
nKeys int, rKeys, rdRatio uint64,
theta float64,
) *Generator {
rd := rand.New(rand.NewSource(int64(seed)))
var zipfian *generator.Zipfian
var dist int
if theta == -1 {
dist = DIST_UNIFORM
} else {
dist = DIST_ZIPFIAN
zipfian = generator.NewZipfianWithItems(int64(rKeys), theta)
}
gen := &Generator {
rd : rd,
nKeys : nKeys,
rKeys : rKeys,
rdRatio : rdRatio,
zipfian : zipfian,
dist : dist,
}
return gen
}
func (g *Generator) PickOp() int {
x := g.rd.Uint64() % 100
if x < g.rdRatio {
return OP_RD
} else {
return OP_WR
}
}
func (g *Generator) pickKeyUniform() uint64 {
return g.rd.Uint64() % g.rKeys
}
func (g *Generator) pickKeyZipfian() uint64 {
return uint64(g.zipfian.Next(g.rd))
}
func (g *Generator) PickKey() uint64 {
if g.dist == DIST_ZIPFIAN {
return g.pickKeyZipfian()
} else {
return g.pickKeyUniform()
}
}
func (g *Generator) NKeys() int {
return g.nKeys
}