From 88cf8456513ddd4e37969e0fa618d66fe6b83277 Mon Sep 17 00:00:00 2001
From: Nuno Cruces <ncruces@users.noreply.github.com>
Date: Thu, 25 Jan 2024 11:03:33 +0000
Subject: [PATCH] JSON stats.

---
 ext/stats/stats.go      |  6 ++++++
 ext/stats/stats_test.go |  8 +++++++-
 ext/stats/welford.go    | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/ext/stats/stats.go b/ext/stats/stats.go
index 6a566bb..c1d0e5a 100644
--- a/ext/stats/stats.go
+++ b/ext/stats/stats.go
@@ -17,6 +17,7 @@
 //   - regr_count: count non-null pairs of variables
 //   - regr_slope: slope of the least-squares-fit linear equation
 //   - regr_intercept: y-intercept of the least-squares-fit linear equation
+//   - regr_json: all regr stats in a JSON object
 //
 // These join the [Built-in Aggregate Functions]:
 //   - count: count rows/values
@@ -52,6 +53,7 @@ func Register(db *sqlite3.Conn) {
 	db.CreateWindowFunction("regr_slope", 2, flags, newCovariance(regr_slope))
 	db.CreateWindowFunction("regr_intercept", 2, flags, newCovariance(regr_intercept))
 	db.CreateWindowFunction("regr_count", 2, flags, newCovariance(regr_count))
+	db.CreateWindowFunction("regr_json", 2, flags, newCovariance(regr_json))
 }
 
 const (
@@ -69,6 +71,7 @@ const (
 	regr_slope
 	regr_intercept
 	regr_count
+	regr_json
 )
 
 func newVariance(kind int) func() sqlite3.AggregateFunction {
@@ -144,6 +147,9 @@ func (fn *covariance) Value(ctx sqlite3.Context) {
 	case regr_count:
 		ctx.ResultInt64(fn.regr_count())
 		return
+	case regr_json:
+		ctx.ResultText(fn.regr_json())
+		return
 	}
 	ctx.ResultFloat(r)
 }
diff --git a/ext/stats/stats_test.go b/ext/stats/stats_test.go
index 53df0af..411f1da 100644
--- a/ext/stats/stats_test.go
+++ b/ext/stats/stats_test.go
@@ -107,7 +107,7 @@ func TestRegister_covariance(t *testing.T) {
 		regr_avgy(y, x), regr_avgx(y, x),
 		regr_syy(y, x), regr_sxx(y, x), regr_sxy(y, x),
 		regr_slope(y, x), regr_intercept(y, x), regr_r2(y, x),
-		regr_count(y, x)
+		regr_count(y, x), regr_json(y, x)
 		FROM data`)
 	if err != nil {
 		t.Fatal(err)
@@ -151,6 +151,12 @@ func TestRegister_covariance(t *testing.T) {
 		if got := stmt.ColumnInt(11); got != 5 {
 			t.Errorf("got %v, want 5", got)
 		}
+		var a map[string]float64
+		if err := stmt.ColumnJSON(12, &a); err != nil {
+			t.Error(err)
+		} else if got := a["count"]; got != 5 {
+			t.Errorf("got %v, want 5", got)
+		}
 	}
 
 	{
diff --git a/ext/stats/welford.go b/ext/stats/welford.go
index 562549e..d2d74ad 100644
--- a/ext/stats/welford.go
+++ b/ext/stats/welford.go
@@ -1,6 +1,10 @@
 package stats
 
-import "math"
+import (
+	"math"
+	"strconv"
+	"strings"
+)
 
 // Welford's algorithm with Kahan summation:
 // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
@@ -108,6 +112,32 @@ func (w welford2) regr_r2() float64 {
 	return w.cov.hi * w.cov.hi / (w.m2y.hi * w.m2x.hi)
 }
 
+func (w welford2) regr_json() string {
+	var json strings.Builder
+	var num [32]byte
+	json.Grow(128)
+	json.WriteString(`{"count":`)
+	json.Write(strconv.AppendInt(num[:0], w.regr_count(), 10))
+	json.WriteString(`,"avgy":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_avgy(), 'g', -1, 64))
+	json.WriteString(`,"avgx":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_avgx(), 'g', -1, 64))
+	json.WriteString(`,"syy":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_syy(), 'g', -1, 64))
+	json.WriteString(`,"sxx":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_sxx(), 'g', -1, 64))
+	json.WriteString(`,"sxy":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_sxy(), 'g', -1, 64))
+	json.WriteString(`,"slope":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_slope(), 'g', -1, 64))
+	json.WriteString(`,"intercept":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_intercept(), 'g', -1, 64))
+	json.WriteString(`,"r2":`)
+	json.Write(strconv.AppendFloat(num[:0], w.regr_r2(), 'g', -1, 64))
+	json.WriteByte('}')
+	return json.String()
+}
+
 func (w *welford2) enqueue(y, x float64) {
 	w.n++
 	d1y := y - w.m1y.hi - w.m1y.lo