Files
sqlite3/ext/stats/percentile.go
2025-01-20 14:39:36 +00:00

110 lines
2.1 KiB
Go

package stats
import (
"encoding/json"
"fmt"
"math"
"slices"
"github.com/ncruces/go-sqlite3"
"github.com/ncruces/go-sqlite3/internal/util"
"github.com/ncruces/sort/quick"
)
// Compatible with:
// https://sqlite.org/src/file/ext/misc/percentile.c
const (
median = iota
percentile_100
percentile_cont
percentile_disc
)
func newPercentile(kind int) func() sqlite3.AggregateFunction {
return func() sqlite3.AggregateFunction { return &percentile{kind: kind} }
}
type percentile struct {
nums []float64
arg1 []byte
kind int
}
func (q *percentile) Step(ctx sqlite3.Context, arg ...sqlite3.Value) {
a := arg[0]
f := a.Float()
if f != 0.0 || a.NumericType() != sqlite3.NULL {
q.nums = append(q.nums, f)
}
if q.kind != median && q.arg1 == nil {
q.arg1 = append(q.arg1, arg[1].RawText()...)
}
}
func (q *percentile) Inverse(ctx sqlite3.Context, arg ...sqlite3.Value) {
a := arg[0]
f := a.Float()
if f != 0.0 || a.NumericType() != sqlite3.NULL {
i := slices.Index(q.nums, f)
l := len(q.nums) - 1
q.nums[i] = q.nums[l]
q.nums = q.nums[:l]
}
}
func (q *percentile) Value(ctx sqlite3.Context) {
if len(q.nums) == 0 {
return
}
var (
err error
float float64
floats []float64
)
if q.kind == median {
float, err = q.at(0.5)
ctx.ResultFloat(float)
} else if err = json.Unmarshal(q.arg1, &float); err == nil {
float, err = q.at(float)
ctx.ResultFloat(float)
} else if err = json.Unmarshal(q.arg1, &floats); err == nil {
err = q.atMore(floats)
ctx.ResultJSON(floats)
}
if err != nil {
ctx.ResultError(fmt.Errorf("percentile: %w", err)) // notest
}
}
func (q *percentile) at(pos float64) (float64, error) {
if q.kind == percentile_100 {
pos = pos / 100
}
if pos < 0 || pos > 1 {
return 0, util.ErrorString("invalid pos")
}
i, f := math.Modf(pos * float64(len(q.nums)-1))
m0 := quick.Select(q.nums, int(i))
if f == 0 || q.kind == percentile_disc {
return m0, nil
}
m1 := slices.Min(q.nums[int(i)+1:])
return util.Lerp(m0, m1, f), nil
}
func (q *percentile) atMore(pos []float64) error {
for i := range pos {
v, err := q.at(pos[i])
if err != nil {
return err
}
pos[i] = v
}
return nil
}