github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/internal/randvar/zipf.go

// 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
}