diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7b3e5278..070d7ea4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -103,7 +103,7 @@ jobs: - name: netbsd version: '11.0' - name: openbsd - version: '7.9' + version: '8.0' tflags: '-test.short' - name: illumos action: omnios @@ -130,7 +130,7 @@ jobs: TESTFLAGS: ${{ matrix.os.tflags }} run: .github/workflows/build-test.sh - - uses: cross-platform-actions/action@v1.6.0 + - uses: cross-platform-actions/action@v1.7.0 with: operating_system: ${{ matrix.os.action || matrix.os.name }} architecture: ${{ matrix.os.arch }} diff --git a/ext/README.md b/ext/README.md index 5d042887..f0e472c7 100644 --- a/ext/README.md +++ b/ext/README.md @@ -5,6 +5,8 @@ you can load into your database connections. ### Extensions +- [`github.com/ncruces/go-sqlite3/ext/approx`](https://pkg.go.dev/github.com/ncruces/go-sqlite3/ext/approx) + provides approximate aggregate functions. - [`github.com/ncruces/go-sqlite3/ext/array`](https://pkg.go.dev/github.com/ncruces/go-sqlite3/ext/array) provides the [`array`](https://sqlite.org/carray.html) table-valued function. - [`github.com/ncruces/go-sqlite3/ext/blobio`](https://pkg.go.dev/github.com/ncruces/go-sqlite3/ext/blobio) diff --git a/ext/approx/approx.go b/ext/approx/approx.go new file mode 100644 index 00000000..ffe6a1d4 --- /dev/null +++ b/ext/approx/approx.go @@ -0,0 +1,35 @@ +// Package approx provides approximate aggregate functions. +// +// Provided functions: +// - approx_count_distinct(X, ...): approximate count of distinct elements using HyperLogLog. +// - approx_count_distinct_accumulate(X, ...): accumulates elements into a HyperLogLog sketch BLOB. +// - approx_count_distinct_combine(B): merges HyperLogLog sketch BLOBs into a single sketch BLOB. +// - approx_count_distinct_estimate(B): estimates the cardinality from a HyperLogLog sketch BLOB. +// +// Sketch BLOBs use the [Aggregate Knowledge HLL Storage Specification] (version 1), +// compatible with PostgreSQL's postgresql-hll extension. +// +// [Aggregate Knowledge HLL Storage Specification]: https://github.com/aggregateknowledge/hll-storage-spec +package approx + +import ( + "errors" + + "github.com/ncruces/go-sqlite3" + "github.com/ncruces/go-sqlite3/util/sql3util" +) + +// Register registers approximate aggregate functions. +func Register(db *sqlite3.Conn) error { + const flags = sqlite3.DETERMINISTIC | sqlite3.INNOCUOUS + return errors.Join( + db.CreateWindowFunction("approx_count_distinct", -1, flags, + sql3util.WindowFromCombiner[hllSketch](1e4)), + db.CreateWindowFunction("approx_count_distinct_accumulate", -1, flags, + sql3util.WindowFromCombiner[hllAccumulate](1e4)), + db.CreateWindowFunction("approx_count_distinct_combine", 1, flags, + sql3util.WindowFromCombiner[hllCombine](1e4)), + db.CreateFunction("approx_count_distinct_estimate", 1, flags, + hllEstimate), + ) +} diff --git a/ext/approx/hll.go b/ext/approx/hll.go new file mode 100644 index 00000000..a06da995 --- /dev/null +++ b/ext/approx/hll.go @@ -0,0 +1,213 @@ +package approx + +import ( + "encoding/binary" + "fmt" + "math" + + "github.com/segmentio/go-hll" + "github.com/twmb/murmur3" + + "github.com/ncruces/go-sqlite3" +) + +var emptyHLL, _ = hll.NewHll(hll.Settings{ + Log2m: 11, + Regwidth: 5, + ExplicitThreshold: hll.AutoExplicitThreshold, + SparseEnabled: true, +}) + +type hllSketch struct { + sketch *hll.Hll + hash uint64 +} + +func (b *hllSketch) Value(ctx sqlite3.Context) { + switch { + case b.sketch != nil: + ctx.ResultInt64(int64(b.sketch.Cardinality())) + case b.hash != 0: + ctx.ResultInt64(1) + default: + ctx.ResultInt64(0) + } +} + +func (b *hllSketch) Step(ctx sqlite3.Context, arg ...sqlite3.Value) { + var hash uint64 + for _, a := range arg { + if hash = hashValue(hash, a); hash == 0 { + return + } + } + b.insert(hash) +} + +func (b *hllSketch) Combine(other hllSketch) { + if other.hash != 0 { + b.insert(other.hash) + return + } + if other.sketch != nil { + if b.sketch != nil { + b.sketch.Union(*other.sketch) + return + } + h := emptyHLL + h.Union(*other.sketch) + if b.hash != 0 { + h.AddRaw(b.hash) + b.hash = 0 + } + b.sketch = &h + } +} + +func (b *hllSketch) insert(hash uint64) { + switch { + case hash == 0 || hash == b.hash: + return + case b.sketch != nil: + b.sketch.AddRaw(hash) + case b.hash == 0: + b.hash = hash + default: + h := emptyHLL + h.AddRaw(b.hash) + h.AddRaw(hash) + b.sketch = &h + b.hash = 0 + } +} + +type hllAccumulate struct{ hllSketch } + +func (b *hllAccumulate) Value(ctx sqlite3.Context) { + switch { + case b.hash != 0: + h := emptyHLL + h.AddRaw(b.hash) + ctx.ResultBlob(h.ToBytes()) + case b.sketch != nil: + ctx.ResultBlob(b.sketch.ToBytes()) + default: + ctx.ResultBlob(emptyHLL.ToBytes()) + } +} + +func (b *hllAccumulate) Combine(other hllAccumulate) { + b.hllSketch.Combine(other.hllSketch) +} + +type hllCombine struct { + sketch *hll.Hll + err error +} + +func (b *hllCombine) Value(ctx sqlite3.Context) { + switch { + case b.err != nil: + ctx.ResultError(b.err) + case b.sketch != nil: + ctx.ResultBlob(b.sketch.ToBytes()) + } +} + +func (b *hllCombine) Step(ctx sqlite3.Context, arg ...sqlite3.Value) { + switch t := arg[0].Type(); t { + case sqlite3.NULL: + return + case sqlite3.BLOB: + other, err := hll.FromBytes(arg[0].RawBlob()) + if err != nil { + ctx.ResultError(fmt.Errorf("approx: %w", err)) + return + } + if b.sketch == nil { + b.sketch = &other + return + } + if err := b.sketch.StrictUnion(other); err != nil { + ctx.ResultError(fmt.Errorf("approx: %w", err)) + } + default: + ctx.ResultError(fmt.Errorf("approx: invalid type: %v", t)) + } +} + +func (b *hllCombine) Combine(other hllCombine) { + if b.err != nil { + return + } + if other.err != nil { + b.err = other.err + return + } + if other.sketch == nil { + return + } + if b.sketch == nil { + h := *other.sketch + h.Clear() + h.Union(*other.sketch) + b.sketch = &h + return + } + if err := b.sketch.StrictUnion(*other.sketch); err != nil { + b.err = fmt.Errorf("approx: %w", err) + } +} + +func hllEstimate(ctx sqlite3.Context, arg ...sqlite3.Value) { + switch t := arg[0].Type(); t { + case sqlite3.NULL: + return + case sqlite3.BLOB: + sketch, err := hll.FromBytes(arg[0].RawBlob()) + if err != nil { + ctx.ResultError(fmt.Errorf("approx: %w", err)) + return + } + ctx.ResultInt64(int64(sketch.Cardinality())) + default: + ctx.ResultError(fmt.Errorf("approx: invalid type: %v", t)) + } +} + +func hashValue(seed uint64, val sqlite3.Value) uint64 { + var ( + buf [8]byte + data []byte + ) + switch val.Type() { + default: // sqlite3.NULL + return 0 + case sqlite3.INTEGER: + data = binary.LittleEndian.AppendUint64(buf[:0], + uint64(val.Int64())) + case sqlite3.FLOAT: + var u uint64 + f := val.Float() + if i, ok := floatIsInt64(f); ok { + u = uint64(i) + } else { + seed ^= 2 + u = math.Float64bits(f) + } + data = binary.LittleEndian.AppendUint64(buf[:0], u) + case sqlite3.TEXT: + data = val.RawText() + case sqlite3.BLOB: + seed ^= 1 + data = val.RawBlob() + } + return murmur3.SeedSum64(seed, data) +} + +func floatIsInt64(f float64) (int64, bool) { + // math.MaxInt64 is not exactly representable as a float64 (it rounds to 2^63), + // but int64(2^63) may saturate back to math.MaxInt64. + i := int64(f) + return i, float64(i) == f && i != math.MaxInt64 +} diff --git a/ext/approx/hll_test.go b/ext/approx/hll_test.go new file mode 100644 index 00000000..a2429372 --- /dev/null +++ b/ext/approx/hll_test.go @@ -0,0 +1,242 @@ +package approx_test + +import ( + "math" + "os" + "testing" + + "github.com/ncruces/go-sqlite3" + "github.com/ncruces/go-sqlite3/ext/approx" + "github.com/ncruces/go-sqlite3/internal/testcfg" +) + +func TestMain(m *testing.M) { + sqlite3.AutoExtension(approx.Register) + os.Exit(m.Run()) +} + +func TestRegister_hll(t *testing.T) { + t.Parallel() + + db, err := sqlite3.OpenContext(testcfg.Context(t), ":memory:") + if err != nil { + t.Fatal(err) + } + defer db.Close() + + err = db.Exec(`CREATE TABLE data (x, y)`) + if err != nil { + t.Fatal(err) + } + + // Empty table. + stmt, _, err := db.Prepare(` + SELECT + approx_count_distinct(x), + approx_count_distinct_estimate(approx_count_distinct_accumulate(x)), + approx_count_distinct_combine(x), + approx_count_distinct_estimate(NULL) + FROM data`) + if err != nil { + t.Fatal(err) + } + if !stmt.Step() { + t.Fatal(stmt.Err()) + } + if got := stmt.ColumnInt64(0); got != 0 { + t.Errorf("got %v, want 0", got) + } + if got := stmt.ColumnInt64(1); got != 0 { + t.Errorf("got %v, want 0", got) + } + if got := stmt.ColumnType(2); got != sqlite3.NULL { + t.Errorf("got %v, want NULL", got) + } + if got := stmt.ColumnType(3); got != sqlite3.NULL { + t.Errorf("got %v, want NULL", got) + } + stmt.Close() + + // Mixed types: NULLs ignored, 1 == 1.0, '1' != 1, X'31' != '1', large ints, tuples. + err = db.Exec(` + INSERT INTO data (x, y) VALUES + (NULL, 1), + (0, NULL), + (-0.0, 1), + (1, 2), + (1.0, 2.0), + ('1', 2), + (X'31', 2), + (2.5, 1), + (2.5, 1), + (9007199254740992, 1), + (9007199254740993, 1), + (9007199254740992.0, 1) + `) + if err != nil { + t.Fatal(err) + } + + stmt, _, err = db.Prepare(` + SELECT + count(DISTINCT x), + approx_count_distinct(x), + approx_count_distinct(x, y), + approx_count_distinct() + FROM data`) + if err != nil { + t.Fatal(err) + } + if !stmt.Step() { + t.Fatal(stmt.Err()) + } + if got, want := stmt.ColumnInt64(1), stmt.ColumnInt64(0); got != want { + t.Errorf("got %v, want %v", got, want) + } + if got := stmt.ColumnInt64(2); got != 7 { + t.Errorf("got %v, want 7", got) + } + if got := stmt.ColumnInt64(3); got != 0 { + t.Errorf("got %v, want 0", got) + } + stmt.Close() + + // Sliding window. + stmt, _, err = db.Prepare(` + SELECT + approx_count_distinct(x) OVER w, + approx_count_distinct_estimate(approx_count_distinct_accumulate(x) OVER w) + FROM data + WINDOW w AS (ROWS BETWEEN 2 PRECEDING AND 1 FOLLOWING)`) + if err != nil { + t.Fatal(err) + } + want := [...]int64{1, 1, 2, 2, 3, 3, 4, 3, 3, 3, 3, 2} + for i := 0; stmt.Step(); i++ { + if got := stmt.ColumnInt64(0); got != want[i] { + t.Errorf("row %d: got %v, want %v", i, got, want[i]) + } + if got := stmt.ColumnInt64(1); got != want[i] { + t.Errorf("row %d: got %v, want %v", i, got, want[i]) + } + } + stmt.Close() +} + +func TestRegister_large(t *testing.T) { + t.Parallel() + + db, err := sqlite3.OpenContext(testcfg.Context(t), ":memory:") + if err != nil { + t.Fatal(err) + } + defer db.Close() + + err = db.Exec(` + CREATE TABLE sketches (grp, sketch); + INSERT INTO sketches VALUES (0, NULL); + INSERT INTO sketches SELECT 1, approx_count_distinct_accumulate(NULL); + INSERT INTO sketches + SELECT value % 4 + 2, approx_count_distinct_accumulate(value % 25000) + FROM generate_series(1, 50000) + GROUP BY value % 4; + `) + if err != nil { + t.Fatal(err) + } + + stmt, _, err := db.Prepare(` + SELECT + approx_count_distinct_estimate(approx_count_distinct_combine(sketch) OVER ()), + approx_count_distinct_estimate(approx_count_distinct_combine(sketch) OVER (ORDER BY grp ROWS 2 PRECEDING)) + FROM sketches`) + if err != nil { + t.Fatal(err) + } + want := [...]float64{0, 0, 6250, 12500, 18750, 18750} + for i := 0; stmt.Step(); i++ { + if got := stmt.ColumnFloat(0); math.Abs(got-25000)/25000 > 0.05 { + t.Errorf("row %d: got %.0f, want ~25000", i, got) + } + if i == 0 && stmt.ColumnType(1) != sqlite3.NULL { + t.Errorf("row 0: got %v, want NULL", stmt.ColumnType(1)) + } + if got := stmt.ColumnFloat(1); math.Abs(got-want[i]) > 0.05*want[i] { + t.Errorf("row %d: got %.0f, want ~%.0f", i, got, want[i]) + } + } + stmt.Close() +} + +func TestRegister_errors(t *testing.T) { + t.Parallel() + + db, err := sqlite3.OpenContext(testcfg.Context(t), ":memory:") + if err != nil { + t.Fatal(err) + } + defer db.Close() + + tests := []string{ + `SELECT approx_count_distinct_estimate('not a blob')`, + `SELECT approx_count_distinct_estimate(X'00')`, + `SELECT approx_count_distinct_combine('not a blob')`, + `SELECT approx_count_distinct_combine(X'00')`, + `SELECT approx_count_distinct_combine(column1) FROM (VALUES (X'118A7F'), (X'118B7F'))`, + `SELECT approx_count_distinct_combine(column1) OVER (ROWS 1 PRECEDING) FROM (VALUES (X'118A7F'), (X'118B7F'))`, + } + for _, sql := range tests { + if err := db.Exec(sql); err == nil { + t.Errorf("expected error for %q", sql) + } + } +} + +func Benchmark_hll(b *testing.B) { + db, err := sqlite3.OpenContext(b.Context(), ":memory:") + if err != nil { + b.Fatal(err) + } + defer db.Close() + + err = db.Exec(` + CREATE TABLE data (grp10 INT, grp2 INT, x INT); + INSERT INTO data SELECT value / 10, value / 2, value FROM generate_series(0, 49999); + CREATE INDEX idx_grp10 ON data (grp10, x); + CREATE INDEX idx_grp2 ON data (grp2, x); + `) + if err != nil { + b.Fatal(err) + } + + queries := []struct { + name, sql string + }{ + {"agg_5k", `SELECT approx_count_distinct(x) FROM data WHERE rowid <= 5000`}, + {"agg_50k", `SELECT approx_count_distinct(x) FROM data`}, + {"group_2k_x2", `SELECT approx_count_distinct(x) FROM data WHERE grp2 < 2000 GROUP BY grp2`}, + {"group_1k_x10", `SELECT approx_count_distinct(x) FROM data WHERE grp10 < 1000 GROUP BY grp10`}, + {"win_20_explicit", `SELECT approx_count_distinct(x) OVER (ROWS 20 PRECEDING) FROM data WHERE rowid <= 2000`}, + {"win_250_sparse", `SELECT approx_count_distinct(x) OVER (ROWS 250 PRECEDING) FROM data WHERE rowid <= 2000`}, + {"win_1000_dense", `SELECT approx_count_distinct(x) OVER (ROWS 1000 PRECEDING) FROM data WHERE rowid <= 2000`}, + } + + for _, q := range queries { + b.Run(q.name, func(b *testing.B) { + b.ReportAllocs() + stmt, _, err := db.Prepare(q.sql) + if err != nil { + b.Fatal(err) + } + defer stmt.Close() + + for b.Loop() { + for stmt.Step() { + } + if err := stmt.Reset(); err != nil { + b.Fatal(err) + } + } + }) + } +} diff --git a/go.mod b/go.mod index 3c51b658..988cc330 100644 --- a/go.mod +++ b/go.mod @@ -14,12 +14,16 @@ require ( github.com/dchest/siphash v1.2.3 // ext/bloom github.com/google/uuid v1.6.0 // ext/uuid github.com/psanford/httpreadat v0.1.0 // example + github.com/segmentio/go-hll v1.0.1 // ext/approx + github.com/twmb/murmur3 v1.1.8 // ext/approx golang.org/x/crypto v0.58.0 // vfs/adiantum vfs/xts golang.org/x/sync v0.24.0 // test golang.org/x/text v0.43.0 // ext/unicode lukechampine.com/adiantum v1.1.1 // vfs/adiantum ) +require github.com/pkg/errors v0.8.0 // indirect + retract ( v0.23.2 // tagged from the wrong branch v0.4.0 // tagged from the wrong branch diff --git a/go.sum b/go.sum index 07676368..790206aa 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dchest/siphash v1.2.3 h1:QXwFc8cFOR2dSa/gE6o/HokBMWtLUaNDVd+22aKHeEA= github.com/dchest/siphash v1.2.3/go.mod h1:0NvQU092bT0ipiFN++/rXm69QG9tVxLAlQHIXMPAkHc= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -10,8 +12,18 @@ github.com/ncruces/sort v1.0.3 h1:2RpWKWvfftMmb3RbUIJF3yVeV/ksHc2HmiW0M7/PTnY= github.com/ncruces/sort v1.0.3/go.mod h1:yD/qToXSDwliVMWAdN02/6+58XrEKBcN0Uyr+KeahFM= github.com/ncruces/wbt v1.0.0 h1:8iBE7UPjTLUpzu3/FCRjAmuQjWzgxo10RGBgt3ooLSc= github.com/ncruces/wbt v1.0.0/go.mod h1:DtF92amvMxH69EmBFUSFWRDAlo6hOEfoNQnClxj9C/c= +github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/psanford/httpreadat v0.1.0 h1:VleW1HS2zO7/4c7c7zNl33fO6oYACSagjJIyMIwZLUE= github.com/psanford/httpreadat v0.1.0/go.mod h1:Zg7P+TlBm3bYbyHTKv/EdtSJZn3qwbPwpfZ/I9GKCRE= +github.com/segmentio/go-hll v1.0.1 h1:ph2Dy18eNdzzT6H2UsEacnWvDnHY6fbsmAnqDWNJeoQ= +github.com/segmentio/go-hll v1.0.1/go.mod h1:cJhmOoNwstKKd+IWBWRlctKkitBI8XHjw9NgyE27Xdk= +github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg= +github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ= golang.org/x/crypto v0.58.0 h1:COkYLr4k7nDI4r6QtMML2AfzeIswWHAJthLS6K/M/54= golang.org/x/crypto v0.58.0/go.mod h1:Xh+kl5A+M0gMiWDytH2Hhr5Z7WbNCpv/eB/EFgfSk+w= golang.org/x/sync v0.24.0 h1:iqiDrGoDgyuvAKrhLQthUEzU/R378N78gXbBjxed7oA= diff --git a/util/sql3util/combiner.go b/util/sql3util/combiner.go index 8b5b1a57..9f26e8e1 100644 --- a/util/sql3util/combiner.go +++ b/util/sql3util/combiner.go @@ -27,6 +27,9 @@ type WindowCombiner[T any] interface { // // Both receiver and other may be empty (the zero value): // x.Combine(empty) must be equivalent to x. + // + // Implementations must not modify or retain mutable state + // referenced by other. Combine(other T) }