estimate.go

// Copyright 2024 Aerospike, Inc.
//
// 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.

package backup

import (
	"bytes"
	"fmt"
	"io"
	"log/slog"
	"math"
)

// zScore represents the number of standard deviations a given value is from the mean of a distribution.
// For a 99% confidence interval, the zScore is approximately 2.576.
// This means that there is a 99% probability that the true population parameter
// lies within the calculated confidence interval.
const zScore = 2.576

// estimateStats contains the mean and variance of the data.
type estimateStats struct {
	Mean     float64
	Variance float64
}

// getEstimate calculates the estimated backup file size.
func getEstimate(data []float64, total float64, logger *slog.Logger) float64 {
	stats := calculateStats(data)

	low, high := confidenceInterval(stats, len(data))
	logger.Debug("predicted record size:",
		slog.Float64("from", low),
		slog.Float64("to", high))

	avg := (low + high) / 2
	logger.Debug("average record size:", slog.Float64("average", avg))

	return total * avg
}

// calculateStats calculates the mean and variance of the data.
func calculateStats(data []float64) estimateStats {
	n := len(data)
	if n == 0 {
		return estimateStats{}
	}

	var sum float64
	for _, value := range data {
		sum += value
	}

	mean := sum / float64(n)

	// For n = 1, variance = 0
	if n == 1 {
		return estimateStats{
			mean,
			0.0,
		}
	}

	var varianceSum float64
	for _, value := range data {
		varianceSum += (value - mean) * (value - mean)
	}

	variance := varianceSum / float64(n-1)

	return estimateStats{
		Mean:     mean,
		Variance: variance,
	}
}

// confidenceInterval calculates the confidence interval of the data.
func confidenceInterval(stats estimateStats, sampleSize int) (low, high float64) {
	if sampleSize == 0 {
		return 0, 0
	}

	standardError := math.Sqrt(stats.Variance / float64(sampleSize))
	marginOfError := zScore * standardError

	low = stats.Mean - marginOfError
	high = stats.Mean + marginOfError

	return low, high
}

type nopWriteCloser struct{ io.Writer }

func (nopWriteCloser) Close() error { return nil }

// getCompressRatio calculates the compress ratio of the data.
func getCompressRatio(policy *CompressionPolicy, samplesData []byte) (float64, error) {
	// Create an io.WriteCloser from samplesData to calculate the compress ratio.
	var buf bytes.Buffer

	wc := nopWriteCloser{Writer: &buf} // satisfies io.WriteCloser
	// Create a compression writer similarly to the backup.
	encodedWriter, err := newCompressionWriter(policy, wc)
	if err != nil {
		return 0, err
	}

	if _, err := encodedWriter.Write(samplesData); err != nil {
		_ = encodedWriter.Close()
		return 0, err
	}

	if err := encodedWriter.Close(); err != nil {
		return 0, err
	}

	// Check buf.len for safety reasons.
	if buf.Len() == 0 {
		if len(samplesData) > 0 {
			return 0, fmt.Errorf("output length is zero, but samples data is not empty")
		}

		return 1, nil
	}

	return float64(len(samplesData)) / float64(buf.Len()), nil
}