github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/kv/zipfian.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // Copyright 2009 The Go Authors. All rights reserved.
    12  // Use of this source code is governed by a BSD-style
    13  // license that can be found in the LICENSE file.
    14  
    15  // W.Hormann, G.Derflinger:
    16  // "Rejection-Inversion to Generate Variates
    17  // from Monotone Discrete Distributions"
    18  // http://eeyore.wu-wien.ac.at/papers/96-04-04.wh-der.ps.gz
    19  
    20  package kv
    21  
    22  import (
    23  	"math"
    24  	"math/rand"
    25  )
    26  
    27  // A zipf generates Zipf distributed variates.
    28  // This was added here from math/rand so we could
    29  // have the Zipfian distribution give us a deterministic
    30  // result based on the read key so we don't read missing
    31  // entries.
    32  //
    33  // Our changes involve being supplied with
    34  // a seeded rand object during the time of retrieval instead
    35  // of a rand object during creation.
    36  type zipf struct {
    37  	imax         float64
    38  	v            float64
    39  	q            float64
    40  	s            float64
    41  	oneminusQ    float64
    42  	oneminusQinv float64
    43  	hxm          float64
    44  	hx0minusHxm  float64
    45  }
    46  
    47  func (z *zipf) h(x float64) float64 {
    48  	return math.Exp(z.oneminusQ*math.Log(z.v+x)) * z.oneminusQinv
    49  }
    50  
    51  func (z *zipf) hinv(x float64) float64 {
    52  	return math.Exp(z.oneminusQinv*math.Log(z.oneminusQ*x)) - z.v
    53  }
    54  
    55  // newZipf returns a Zipf variate generator.
    56  // The generator generates values k ∈ [0, imax]
    57  // such that P(k) is proportional to (v + k) ** (-s).
    58  // Requirements: s > 1 and v >= 1.
    59  func newZipf(s float64, v float64, imax uint64) *zipf {
    60  	z := new(zipf)
    61  	if s <= 1.0 || v < 1 {
    62  		return nil
    63  	}
    64  	z.imax = float64(imax)
    65  	z.v = v
    66  	z.q = s
    67  	z.oneminusQ = 1.0 - z.q
    68  	z.oneminusQinv = 1.0 / z.oneminusQ
    69  	z.hxm = z.h(z.imax + 0.5)
    70  	z.hx0minusHxm = z.h(0.5) - math.Exp(math.Log(z.v)*(-z.q)) - z.hxm
    71  	z.s = 1 - z.hinv(z.h(1.5)-math.Exp(-z.q*math.Log(z.v+1.0)))
    72  	return z
    73  }
    74  
    75  // Uint64 returns a value drawn from the Zipf distribution described
    76  // by the Zipf object.
    77  func (z *zipf) Uint64(random *rand.Rand) uint64 {
    78  	if z == nil {
    79  		panic("rand: nil Zipf")
    80  	}
    81  	k := 0.0
    82  
    83  	for {
    84  		r := random.Float64() // r in [0.0, 1.0]
    85  		ur := z.hxm + r*z.hx0minusHxm
    86  		x := z.hinv(ur)
    87  		k = math.Floor(x + 0.5)
    88  		if k-x <= z.s {
    89  			break
    90  		}
    91  		if ur >= z.h(k+0.5)-math.Exp(-math.Log(k+z.v)*z.q) {
    92  			break
    93  		}
    94  	}
    95  	return uint64(int64(k))
    96  }