ewma: add a per-bye EWMA estimator

This will be used to estimate compression ratio of blocks based on the
compression ratios of recently sampled blocks.

Author: RaduBerinde

internal/ewma/ewma_bytes.go

@@ -0,0 +1,99 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package ewma
+
+import (
+	"math"
+
+	"github.com/cockroachdb/pebble/internal/invariants"
+)
+
+// Bytes is an estimator for an arbitrary value that is sampled from byte
+// blocks.
+//
+// Consider a stream of data which is divided into blocks of varying size. We
+// want to estimate a value (like compression ratio) based on the values from
+// recent blocks.
+//
+// Bytes implements a per-byte exponential moving average (EWMA) estimator: let
+// pos_i and val_i be the position and value of each byte for which we have
+// data; the estimate at position p is the weighted sum:
+//
+//		Sum_i val_i*(1-alpha)^(p-pos_i)
+//	  -------------------------------
+//	     Sum_i (1-alpha)^(p-pos_i)
+type Bytes struct {
+	alpha       float64
+	sum         float64
+	totalWeight float64
+	gap         int64
+}
+
+// Init the estimator such that a block sampled <half-life> bytes ago has half
+// the weight compared to a block sampled now.
+//
+// Intuitively, half of the estimate comes from the values within the half-life
+// window; and 75% of the estimate comes from values within 2x half-life.
+func (b *Bytes) Init(halfLife int64) {
+	*b = Bytes{}
+	// Exact value is 1 - 2^(-1/H). The straightforward calculation suffers from
+	// precision loss as H grows (we are subtracting two nearly equal numbers). We
+	// use a numerically stable alternative:
+	//   1 - 2^(-1/H) = 1 - e^(-ln(2)/H) = -expm1(-ln(2)/H)
+	b.alpha = -math.Expm1(-math.Ln2 / float64(halfLife))
+}
+
+// Estimate returns the current estimate of the value, based on the recent
+// SampledBlock() calls. Returns NaN if no blocks have been sampled yet.
+func (b *Bytes) Estimate() float64 {
+	return b.sum / b.totalWeight
+}
+
+// NoSample informs the estimator that a block of the given length was not
+// sampled.
+func (b *Bytes) NoSample(numBytes int64) {
+	if numBytes < 0 {
+		if invariants.Enabled {
+			panic("invalid numBytes")
+		}
+		return
+	}
+	// It would be equivalent (but less efficient) to multiply both sum and
+	// totalWeight by (1-alpha)^numBytes instead of keeping track of the gap.
+	b.gap += numBytes
+}
+
+// SampledBlock informs the estimator that a block of the given length was
+// sampled.
+func (b *Bytes) SampledBlock(numBytes int64, value float64) {
+	if numBytes < 1 {
+		if invariants.Enabled {
+			panic("invalid numBytes")
+		}
+		return
+	}
+	decay := b.decay(b.gap + numBytes)
+	b.sum *= decay
+	b.totalWeight *= decay
+	b.gap = 0
+
+	// The sum of weights for the new bytes is:
+	//
+	//                                   1 - (1 - alpha)^numBytes
+	//       Sum       (1 - alpha)^i  =  ------------------------
+	//   0≤i<numBytes                             alpha
+	//
+	// We can drop the 1/alpha factor from all weights (it cancels out).
+	w := 1 - b.decay(numBytes)
+	b.sum += value * w
+	b.totalWeight += w
+}
+
+// decay returns (1 - alpha)^n for the given n.
+func (b *Bytes) decay(n int64) float64 {
+	// (1 - alpha)^n = e^(n * log(1 - alpha)). Using Exp and Log1p is stable for
+	// very small alpha (unlike math.Pow).
+	return math.Exp(float64(n) * math.Log1p(-b.alpha))
+}

internal/ewma/ewma_bytes_test.go

@@ -0,0 +1,64 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package ewma
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBytes(t *testing.T) {
+	const eps = 0.1
+	var b Bytes
+
+	b.Init(1000)
+
+	require.True(t, math.IsNaN(b.Estimate()))
+
+	b.SampledBlock(100, 1.0)
+	require.InEpsilon(t, 1.0, b.Estimate(), eps)
+
+	b.NoSample(10000)
+	require.InEpsilon(t, 1.0, b.Estimate(), eps)
+	b.SampledBlock(100, 10.0)
+	require.InEpsilon(t, 10.0, b.Estimate(), eps)
+
+	b.NoSample(10000)
+	b.SampledBlock(10, 0.0)
+	b.SampledBlock(10, 10.0)
+	b.SampledBlock(10, 0.0)
+	b.SampledBlock(10, 10.0)
+	require.InEpsilon(t, 5.0, b.Estimate(), eps)
+	b.NoSample(1000)
+	b.SampledBlock(10, 0.0)
+	require.InEpsilon(t, 3.3, b.Estimate(), eps)
+
+	b.Init(1000)
+	b.SampledBlock(1, 0.0)
+	b.NoSample(1000)
+	b.SampledBlock(1, 1.0)
+	// The byte 1 half-life ago matters 1/2 as much, so the estimate is 2/3.
+	require.InEpsilon(t, 0.66, b.Estimate(), eps)
+}
+
+// TestBytesHalfLife verifies that the alpha and decay calculations are accurate.
+func TestBytesHalfLife(t *testing.T) {
+	for _, n := range []int64{1, 2, 3, 10, 100, 1000, 10_000, 1 << 20, 128 << 20, 1 << 30} {
+		t.Run(fmt.Sprint(n), func(t *testing.T) {
+			var b Bytes
+			b.Init(n)
+			const eps = 1e-8
+			require.InEpsilon(t, 1.0/2, b.decay(n), eps)
+			require.InDelta(t, 1.0/4, b.decay(2*n), eps)
+			require.InDelta(t, 1.0/8, b.decay(3*n), eps)
+			require.InDelta(t, 1.0/16, b.decay(4*n), eps)
+			require.InDelta(t, 1.0/32, b.decay(5*n), eps)
+			require.InDelta(t, 1.0/64, b.decay(6*n), eps)
+		})
+	}
+}