ceremonyclient/pebble/internal/randvar/zipf.go

148 lines
4.7 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
// Copyright 2017 The Cockroach Authors.
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License. See the AUTHORS file
// for names of contributors.
//
// ZipfGenerator implements the Incrementing Zipfian Random Number Generator from
// [1]: "Quickly Generating Billion-Record Synthetic Databases"
// by Gray, Sundaresan, Englert, Baclawski, and Weinberger, SIGMOD 1994.
package randvar
import (
"math"
"sync"
"github.com/cockroachdb/errors"
"golang.org/x/exp/rand"
)
const (
// See https://github.com/brianfrankcooper/YCSB/blob/f886c1e7988f8f4965cb88a1fe2f6bad2c61b56d/core/src/main/java/com/yahoo/ycsb/generator/ScrambledZipfianGenerator.java#L33-L35
defaultMax = 10000000000
defaultTheta = 0.99
defaultZetaN = 26.46902820178302
)
// Zipf is a random number generator that generates random numbers from a Zipf
// distribution. Unlike rand.Zipf, this generator supports incrementing the max
// parameter without performing an expensive recomputation of the underlying
// hidden parameters, which is a pattern used in [1] for efficiently generating
// large volumes of Zipf-distributed records for synthetic data. Second,
// rand.Zipf only supports theta <= 1, we suppose all values of theta.
type Zipf struct {
// Supplied constants.
theta float64
min uint64
// Internally computed constants.
alpha, zeta2 float64
halfPowTheta float64
// Mutable state.
mu struct {
sync.RWMutex
max uint64
eta float64
zetaN float64
}
}
// NewDefaultZipf constructs a new Zipf generator with the default parameters.
func NewDefaultZipf() (*Zipf, error) {
return NewZipf(1, defaultMax, defaultTheta)
}
// NewZipf constructs a new Zipf generator with the given parameters. Returns
// an error if the parameters are outside the accepted range.
func NewZipf(min, max uint64, theta float64) (*Zipf, error) {
if min > max {
return nil, errors.Errorf("min %d > max %d", errors.Safe(min), errors.Safe(max))
}
if theta < 0.0 || theta == 1.0 {
return nil, errors.New("0 < theta, and theta != 1")
}
z := &Zipf{
min: min,
theta: theta,
}
z.mu.max = max
// Compute hidden parameters.
z.zeta2 = computeZetaFromScratch(2, theta)
z.halfPowTheta = 1.0 + math.Pow(0.5, z.theta)
z.mu.zetaN = computeZetaFromScratch(max+1-min, theta)
z.alpha = 1.0 / (1.0 - theta)
z.mu.eta = (1 - math.Pow(2.0/float64(z.mu.max+1-z.min), 1.0-theta)) / (1.0 - z.zeta2/z.mu.zetaN)
return z, nil
}
// computeZetaIncrementally recomputes zeta(max, theta), assuming that sum =
// zeta(oldMax, theta). Returns zeta(max, theta), computed incrementally.
func computeZetaIncrementally(oldMax, max uint64, theta float64, sum float64) float64 {
if max < oldMax {
panic("unable to decrement max!")
}
for i := oldMax + 1; i <= max; i++ {
sum += 1.0 / math.Pow(float64(i), theta)
}
return sum
}
// The function zeta computes the value
// zeta(n, theta) = (1/1)^theta + (1/2)^theta + (1/3)^theta + ... + (1/n)^theta
func computeZetaFromScratch(n uint64, theta float64) float64 {
if n == defaultMax && theta == defaultTheta {
// Precomputed value, borrowed from ScrambledZipfianGenerator.java. This is
// quite slow to calculate from scratch due to the large n value.
return defaultZetaN
}
return computeZetaIncrementally(0, n, theta, 0.0)
}
// IncMax increments max and recomputes the internal values that depend on
// it. Returns an error if the recomputation failed.
func (z *Zipf) IncMax(delta int) {
z.mu.Lock()
oldMax := z.mu.max
z.mu.max += uint64(delta)
z.mu.zetaN = computeZetaIncrementally(oldMax+1-z.min, z.mu.max+1-z.min, z.theta, z.mu.zetaN)
z.mu.eta = (1 - math.Pow(2.0/float64(z.mu.max+1-z.min), 1.0-z.theta)) / (1.0 - z.zeta2/z.mu.zetaN)
z.mu.Unlock()
}
// Max returns the max.
func (z *Zipf) Max() uint64 {
z.mu.Lock()
defer z.mu.Unlock()
return z.mu.max
}
// Uint64 draws a new value between min and max, with probabilities according
// to the Zipf distribution.
func (z *Zipf) Uint64(rng *rand.Rand) uint64 {
u := rng.Float64()
z.mu.RLock()
uz := u * z.mu.zetaN
var result uint64
if uz < 1.0 {
result = z.min
} else if uz < z.halfPowTheta {
result = z.min + 1
} else {
spread := float64(z.mu.max + 1 - z.min)
result = z.min + uint64(spread*math.Pow(z.mu.eta*u-z.mu.eta+1.0, z.alpha))
}
z.mu.RUnlock()
return result
}