From 0674222c66b34e4f55b72075280102333cb44889 Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Thu, 9 Jul 2026 16:45:35 -0700 Subject: [PATCH 1/4] First pass on sum/avg window function implementations --- server/analyzer/init.go | 16 + server/analyzer/type_sanitizer.go | 2 +- server/functions/aggregate/avg_aggregates.go | 805 ++++++++++++++++++ server/functions/aggregate/init.go | 2 + .../functions/aggregate/numeric_aggregates.go | 805 ++++++++++++++++++ server/functions/framework/catalog.go | 55 +- .../framework/compiled_aggregate_function.go | 110 ++- .../framework/compiled_window_function.go | 179 ++++ server/functions/framework/functions.go | 39 +- server/functions/window/init.go | 21 + server/functions/window/rank.go | 147 ++++ server/functions/window/row_number.go | 77 ++ server/initialization/initialization.go | 2 + testing/go/enginetest/doltgres_engine_test.go | 2 + testing/go/smoke_test.go | 4 +- testing/go/values_statement_test.go | 2 +- testing/go/window_test.go | 207 +++++ 17 files changed, 2451 insertions(+), 24 deletions(-) create mode 100644 server/functions/aggregate/avg_aggregates.go create mode 100644 server/functions/aggregate/numeric_aggregates.go create mode 100644 server/functions/framework/compiled_window_function.go create mode 100644 server/functions/window/init.go create mode 100644 server/functions/window/rank.go create mode 100644 server/functions/window/row_number.go create mode 100644 testing/go/window_test.go diff --git a/server/analyzer/init.go b/server/analyzer/init.go index 12601ec6bc..50794cbdf9 100644 --- a/server/analyzer/init.go +++ b/server/analyzer/init.go @@ -127,6 +127,7 @@ func initEngine() { plan.ValidateForeignKeyDefinition = validateForeignKeyDefinition planbuilder.IsAggregateFunc = IsAggregateFunc + planbuilder.IsWindowFunc = IsWindowFunc expression.DefaultExpressionFactory = pgexpression.PostgresExpressionFactory{} @@ -148,6 +149,21 @@ func IsAggregateFunc(name string) bool { return false } +// IsWindowFunc checks if the given function name is a window function. This is the entire set supported by +// MySQL plus some postgres specific ones. +func IsWindowFunc(name string) bool { + if planbuilder.IsMySQLWindowFuncName(name) { + return true + } + + switch name { + case "array_agg", "bool_and", "bool_or": + return true + } + + return false +} + // insertAnalyzerRules inserts the given rule(s) before or after the given analyzer.RuleId, returning an updated slice. func insertAnalyzerRules(rules []analyzer.Rule, id analyzer.RuleId, before bool, additionalRules ...analyzer.Rule) []analyzer.Rule { inserted := false diff --git a/server/analyzer/type_sanitizer.go b/server/analyzer/type_sanitizer.go index c9dd97369e..9f1b154fee 100644 --- a/server/analyzer/type_sanitizer.go +++ b/server/analyzer/type_sanitizer.go @@ -77,7 +77,7 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope if _, ok := expr.(framework.Function); !ok { // Some aggregation functions cannot be wrapped due to expectations in the analyzer, so we exclude them here. switch expr.FunctionName() { - case "Count", "CountDistinct", "group_concat", "JSONObjectAgg", "Sum": + case "Count", "CountDistinct", "group_concat", "JSONObjectAgg", "Sum", "Avg": case "coalesce": // Replace GMS Coalesce with a Doltgres-native implementation that uses // Postgres type-resolution rules (FindCommonType) to infer the result type. diff --git a/server/functions/aggregate/avg_aggregates.go b/server/functions/aggregate/avg_aggregates.go new file mode 100644 index 0000000000..c751d3050a --- /dev/null +++ b/server/functions/aggregate/avg_aggregates.go @@ -0,0 +1,805 @@ +// Copyright 2026 Dolthub, 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 aggregate + +import ( + "github.com/cockroachdb/apd/v3" + "github.com/cockroachdb/errors" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" + + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// avgDecimalCtx is used for the division in AVG's sum/count. sql.DecimalCtx (apd.BaseContext) has Precision: +// 0, which disables rounding - fine for the exact Add calls used to accumulate the sum, but Quo (true +// division) can be inexact and apd requires a nonzero precision for that. 20 digits is generous headroom +// over Postgres's own guard-digit scale for numeric division (real Postgres computes a much smaller +// dscale for this), and it's meaningfully cheaper than a very high precision: apd's Quo cost scales with +// precision (bigger intermediate big.Ints), and this runs once per GROUP BY group. +var avgDecimalCtx = sql.DecimalCtx.WithPrecision(16) + +// initAvgAggs registers the functions to the catalog. +func initAvgAggs() { + framework.RegisterAggregateFunction(avgInt16) + framework.RegisterAggregateFunction(avgInt32) + framework.RegisterAggregateFunction(avgInt64) + framework.RegisterAggregateFunction(avgNumeric) + framework.RegisterAggregateFunction(avgFloat32) + framework.RegisterAggregateFunction(avgFloat64) +} + +// avgInt16 represents the PostgreSQL avg(int2) function, which promotes to numeric. +var avgInt16 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int16, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgInt16Buffer, + NewAggWindowFunc: newAvgInt16WindowFunction, +} + +// avgInt16Buffer is an aggregation buffer for avg(int2). +type avgInt16Buffer struct { + expr sql.Expression + sum int64 + count int64 +} + +var _ sql.AggregationBuffer = (*avgInt16Buffer)(nil) + +// newAvgInt16Buffer creates a new aggregation buffer for the avg(int2) aggregate function. +func newAvgInt16Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgInt16Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgInt16Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgInt16Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return avgFromInt64Sum(b.sum, b.count) +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgInt16Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int16) + if !ok { + return errors.Errorf("avg: expected int16, got %T", v) + } + b.sum += int64(i) + b.count++ + return nil +} + +// avgInt16WindowFunction is the sql.WindowFunction used for avg(int2) within an OVER(...) clause. +type avgInt16WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgInt16WindowFunction)(nil) + +// newAvgInt16WindowFunction creates the sql.WindowFunction for avg(int2). +func newAvgInt16WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgInt16WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgInt16WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgInt16WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgInt16WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum, count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int16) + if !ok { + return nil, errors.Errorf("avg: expected int16, got %T", v) + } + sum += int64(iv) + count++ + } + if count == 0 { + return nil, nil + } + return avgFromInt64Sum(sum, count) +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgInt16WindowFunction) Dispose(ctx *sql.Context) {} + +// avgInt32 represents the PostgreSQL avg(int4) function, which promotes to numeric. +var avgInt32 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int32, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgInt32Buffer, + NewAggWindowFunc: newAvgInt32WindowFunction, +} + +// avgInt32Buffer is an aggregation buffer for avg(int4). +type avgInt32Buffer struct { + expr sql.Expression + sum int64 + count int64 +} + +var _ sql.AggregationBuffer = (*avgInt32Buffer)(nil) + +// newAvgInt32Buffer creates a new aggregation buffer for the avg(int4) aggregate function. +func newAvgInt32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgInt32Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgInt32Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgInt32Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return avgFromInt64Sum(b.sum, b.count) +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int32) + if !ok { + return errors.Errorf("avg: expected int32, got %T", v) + } + b.sum += int64(i) + b.count++ + return nil +} + +// avgFromInt64Sum returns the numeric average of sum/count. +func avgFromInt64Sum(sum, count int64) (*apd.Decimal, error) { + return quoAvg(apd.New(sum, 0), apd.New(count, 0)) +} + +// quoAvg divides dividend/divisor for use as an AVG result, reduced to strip the trailing zeros that Quo +// pads on to fill avgDecimalCtx's precision (e.g. 30/2 comes back as 15.00000...0 with 38 zeros, not 15). +func quoAvg(dividend, divisor *apd.Decimal) (*apd.Decimal, error) { + result := new(apd.Decimal) + if _, err := avgDecimalCtx.Quo(result, dividend, divisor); err != nil { + return nil, err + } + result.Reduce(result) + return result, nil +} + +// avgInt32WindowFunction is the sql.WindowFunction used for avg(int4) within an OVER(...) clause. +type avgInt32WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgInt32WindowFunction)(nil) + +// newAvgInt32WindowFunction creates the sql.WindowFunction for avg(int4). +func newAvgInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgInt32WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgInt32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgInt32WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum, count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int32) + if !ok { + return nil, errors.Errorf("avg: expected int32, got %T", v) + } + sum += int64(iv) + count++ + } + if count == 0 { + return nil, nil + } + return avgFromInt64Sum(sum, count) +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgInt32WindowFunction) Dispose(ctx *sql.Context) {} + +// avgInt64 represents the PostgreSQL avg(int8) function, which promotes to numeric. +var avgInt64 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int64, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgInt64Buffer, + NewAggWindowFunc: newAvgInt64WindowFunction, +} + +// avgInt64Buffer is an aggregation buffer for avg(int8). +type avgInt64Buffer struct { + expr sql.Expression + sum apd.Decimal + count int64 +} + +var _ sql.AggregationBuffer = (*avgInt64Buffer)(nil) + +// newAvgInt64Buffer creates a new aggregation buffer for the avg(int8) aggregate function. +func newAvgInt64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgInt64Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgInt64Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgInt64Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return quoAvg(&b.sum, apd.New(b.count, 0)) +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgInt64Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int64) + if !ok { + return errors.Errorf("avg: expected int64, got %T", v) + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, apd.New(i, 0)) + b.count++ + return err +} + +// avgInt64WindowFunction is the sql.WindowFunction used for avg(int8) within an OVER(...) clause. +type avgInt64WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgInt64WindowFunction)(nil) + +// newAvgInt64WindowFunction creates the sql.WindowFunction for avg(int8). +func newAvgInt64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgInt64WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgInt64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgInt64WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgInt64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum apd.Decimal + var count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int64) + if !ok { + return nil, errors.Errorf("avg: expected int64, got %T", v) + } + if _, err = sql.DecimalCtx.Add(&sum, &sum, apd.New(iv, 0)); err != nil { + return nil, err + } + count++ + } + if count == 0 { + return nil, nil + } + return quoAvg(&sum, apd.New(count, 0)) +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgInt64WindowFunction) Dispose(ctx *sql.Context) {} + +// avgNumeric represents the PostgreSQL avg(numeric) function, which stays numeric. +var avgNumeric = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Numeric, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgNumericBuffer, + NewAggWindowFunc: newAvgNumericWindowFunction, +} + +// avgNumericBuffer is an aggregation buffer for avg(numeric). +type avgNumericBuffer struct { + expr sql.Expression + sum apd.Decimal + count int64 +} + +var _ sql.AggregationBuffer = (*avgNumericBuffer)(nil) + +// newAvgNumericBuffer creates a new aggregation buffer for the avg(numeric) aggregate function. +func newAvgNumericBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgNumericBuffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgNumericBuffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgNumericBuffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return quoAvg(&b.sum, apd.New(b.count, 0)) +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + d, ok := v.(*apd.Decimal) + if !ok { + return errors.Errorf("avg: expected *apd.Decimal, got %T", v) + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, d) + b.count++ + return err +} + +// avgNumericWindowFunction is the sql.WindowFunction used for avg(numeric) within an OVER(...) clause. +type avgNumericWindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgNumericWindowFunction)(nil) + +// newAvgNumericWindowFunction creates the sql.WindowFunction for avg(numeric). +func newAvgNumericWindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgNumericWindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgNumericWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgNumericWindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgNumericWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum apd.Decimal + var count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + d, ok := v.(*apd.Decimal) + if !ok { + return nil, errors.Errorf("avg: expected *apd.Decimal, got %T", v) + } + if _, err = sql.DecimalCtx.Add(&sum, &sum, d); err != nil { + return nil, err + } + count++ + } + if count == 0 { + return nil, nil + } + return quoAvg(&sum, apd.New(count, 0)) +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgNumericWindowFunction) Dispose(ctx *sql.Context) {} + +// avgFloat32 represents the PostgreSQL avg(float4) function, which promotes to double precision. +var avgFloat32 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Float64, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Float32, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgFloat32Buffer, + NewAggWindowFunc: newAvgFloat32WindowFunction, +} + +// avgFloat32Buffer is an aggregation buffer for avg(float4). +type avgFloat32Buffer struct { + expr sql.Expression + sum float64 + count int64 +} + +var _ sql.AggregationBuffer = (*avgFloat32Buffer)(nil) + +// newAvgFloat32Buffer creates a new aggregation buffer for the avg(float4) aggregate function. +func newAvgFloat32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgFloat32Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgFloat32Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgFloat32Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return b.sum / float64(b.count), nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgFloat32Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + f, ok := v.(float32) + if !ok { + return errors.Errorf("avg: expected float32, got %T", v) + } + b.sum += float64(f) + b.count++ + return nil +} + +// avgFloat32WindowFunction is the sql.WindowFunction used for avg(float4) within an OVER(...) clause. +type avgFloat32WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgFloat32WindowFunction)(nil) + +// newAvgFloat32WindowFunction creates the sql.WindowFunction for avg(float4). +func newAvgFloat32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgFloat32WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgFloat32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgFloat32WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum float64 + var count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + fv, ok := v.(float32) + if !ok { + return nil, errors.Errorf("avg: expected float32, got %T", v) + } + sum += float64(fv) + count++ + } + if count == 0 { + return nil, nil + } + return sum / float64(count), nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgFloat32WindowFunction) Dispose(ctx *sql.Context) {} + +// avgFloat64 represents the PostgreSQL avg(float8) function, which stays double precision. +var avgFloat64 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "avg", + Return: pgtypes.Float64, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Float64, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newAvgFloat64Buffer, + NewAggWindowFunc: newAvgFloat64WindowFunction, +} + +// avgFloat64Buffer is an aggregation buffer for avg(float8). +type avgFloat64Buffer struct { + expr sql.Expression + sum float64 + count int64 +} + +var _ sql.AggregationBuffer = (*avgFloat64Buffer)(nil) + +// newAvgFloat64Buffer creates a new aggregation buffer for the avg(float8) aggregate function. +func newAvgFloat64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &avgFloat64Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *avgFloat64Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *avgFloat64Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return b.sum / float64(b.count), nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *avgFloat64Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + f, ok := v.(float64) + if !ok { + return errors.Errorf("avg: expected float64, got %T", v) + } + b.sum += f + b.count++ + return nil +} + +// avgFloat64WindowFunction is the sql.WindowFunction used for avg(float8) within an OVER(...) clause. +type avgFloat64WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*avgFloat64WindowFunction)(nil) + +// newAvgFloat64WindowFunction creates the sql.WindowFunction for avg(float8). +func newAvgFloat64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &avgFloat64WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *avgFloat64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *avgFloat64WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *avgFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum float64 + var count int64 + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + fv, ok := v.(float64) + if !ok { + return nil, errors.Errorf("avg: expected float64, got %T", v) + } + sum += fv + count++ + } + if count == 0 { + return nil, nil + } + return sum / float64(count), nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *avgFloat64WindowFunction) Dispose(ctx *sql.Context) {} diff --git a/server/functions/aggregate/init.go b/server/functions/aggregate/init.go index eb43c311c8..106797dc59 100755 --- a/server/functions/aggregate/init.go +++ b/server/functions/aggregate/init.go @@ -16,4 +16,6 @@ package aggregate func Init() { initBoolAggs() + initNumericAggs() + initAvgAggs() } diff --git a/server/functions/aggregate/numeric_aggregates.go b/server/functions/aggregate/numeric_aggregates.go new file mode 100644 index 0000000000..3c4ea639ff --- /dev/null +++ b/server/functions/aggregate/numeric_aggregates.go @@ -0,0 +1,805 @@ +// Copyright 2026 Dolthub, 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 aggregate + +import ( + "github.com/cockroachdb/apd/v3" + "github.com/cockroachdb/errors" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" + + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// initNumericAggs registers the functions to the catalog. +func initNumericAggs() { + framework.RegisterAggregateFunction(sumInt16) + framework.RegisterAggregateFunction(sumInt32) + framework.RegisterAggregateFunction(sumInt64) + framework.RegisterAggregateFunction(sumNumeric) + framework.RegisterAggregateFunction(sumFloat32) + framework.RegisterAggregateFunction(sumFloat64) +} + +// sumInt16 represents the PostgreSQL sum(int2) function, which promotes to bigint. +// +// Note that overload resolution for aggregates does not insert implicit widening casts the way it does for +// scalar functions (CompiledAggregateFunction hands NewBuffer/NewWindowFunc the raw, uncast Arguments) - so +// every distinct argument type sum/avg can be called on needs its own overload registered below; there's no +// way to cover e.g. both int2 and int4 with a single numeric-ish overload. +var sumInt16 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Int64, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int16, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumInt16Buffer, + NewAggWindowFunc: newSumInt16WindowFunction, +} + +// sumInt16Buffer is an aggregation buffer for sum(int2). +type sumInt16Buffer struct { + expr sql.Expression + sum int64 + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumInt16Buffer)(nil) + +// newSumInt16Buffer creates a new aggregation buffer for the sum(int2) aggregate function. +func newSumInt16Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumInt16Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumInt16Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumInt16Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumInt16Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int16) + if !ok { + return errors.Errorf("sum: expected int16, got %T", v) + } + b.sum += int64(i) + b.sawOne = true + return nil +} + +// sumInt16WindowFunction is the sql.WindowFunction used for sum(int2) within an OVER(...) clause. +type sumInt16WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumInt16WindowFunction)(nil) + +// newSumInt16WindowFunction creates the sql.WindowFunction for sum(int2). +func newSumInt16WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumInt16WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumInt16WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumInt16WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumInt16WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum int64 + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int16) + if !ok { + return nil, errors.Errorf("sum: expected int16, got %T", v) + } + sum += int64(iv) + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumInt16WindowFunction) Dispose(ctx *sql.Context) {} + +// sumInt32 represents the PostgreSQL sum(int4) function, which promotes to bigint. It supports both the +// GROUP BY buffer path and use as a window function within an OVER(...) clause, so that a query like +// `SUM(x) OVER (...)` returns a correctly-typed Postgres bigint directly, without any post-hoc type +// correction elsewhere in the analyzer. +var sumInt32 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Int64, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int32, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumInt32Buffer, + NewAggWindowFunc: newSumInt32WindowFunction, +} + +// sumInt32Buffer is an aggregation buffer for sum(int4). +type sumInt32Buffer struct { + expr sql.Expression + sum int64 + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumInt32Buffer)(nil) + +// newSumInt32Buffer creates a new aggregation buffer for the sum(int4) aggregate function. +func newSumInt32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumInt32Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumInt32Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumInt32Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int32) + if !ok { + return errors.Errorf("sum: expected int32, got %T", v) + } + b.sum += int64(i) + b.sawOne = true + return nil +} + +// sumInt32WindowFunction is the sql.WindowFunction used for sum(int4) within an OVER(...) clause. +type sumInt32WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumInt32WindowFunction)(nil) + +// newSumInt32WindowFunction creates the sql.WindowFunction for sum(int4), binding the frame declared on +// window (if any); with no explicit frame, DefaultFramer supplies Postgres's default (unbounded preceding to +// current row). +func newSumInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumInt32WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumInt32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumInt32WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum int64 + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int32) + if !ok { + return nil, errors.Errorf("sum: expected int32, got %T", v) + } + sum += int64(iv) + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumInt32WindowFunction) Dispose(ctx *sql.Context) {} + +// sumInt64 represents the PostgreSQL sum(int8) function. A running sum of bigints can itself overflow a +// bigint, so Postgres promotes this to numeric rather than bigint. +var sumInt64 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Int64, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumInt64Buffer, + NewAggWindowFunc: newSumInt64WindowFunction, +} + +// sumInt64Buffer is an aggregation buffer for sum(int8). +type sumInt64Buffer struct { + expr sql.Expression + sum apd.Decimal + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumInt64Buffer)(nil) + +// newSumInt64Buffer creates a new aggregation buffer for the sum(int8) aggregate function. +func newSumInt64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumInt64Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumInt64Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumInt64Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + result := b.sum + return &result, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumInt64Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + i, ok := v.(int64) + if !ok { + return errors.Errorf("sum: expected int64, got %T", v) + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, apd.New(i, 0)) + b.sawOne = true + return err +} + +// sumInt64WindowFunction is the sql.WindowFunction used for sum(int8) within an OVER(...) clause. +type sumInt64WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumInt64WindowFunction)(nil) + +// newSumInt64WindowFunction creates the sql.WindowFunction for sum(int8). +func newSumInt64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumInt64WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumInt64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumInt64WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumInt64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum apd.Decimal + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + iv, ok := v.(int64) + if !ok { + return nil, errors.Errorf("sum: expected int64, got %T", v) + } + if _, err = sql.DecimalCtx.Add(&sum, &sum, apd.New(iv, 0)); err != nil { + return nil, err + } + sawOne = true + } + if !sawOne { + return nil, nil + } + return &sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumInt64WindowFunction) Dispose(ctx *sql.Context) {} + +// sumNumeric represents the PostgreSQL sum(numeric) function, which stays numeric. It supports both the +// GROUP BY buffer path and use as a window function within an OVER(...) clause. +var sumNumeric = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Numeric, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Numeric, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumNumericBuffer, + NewAggWindowFunc: newSumNumericWindowFunction, +} + +// sumNumericBuffer is an aggregation buffer for sum(numeric). +type sumNumericBuffer struct { + expr sql.Expression + sum apd.Decimal + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumNumericBuffer)(nil) + +// newSumNumericBuffer creates a new aggregation buffer for the sum(numeric) aggregate function. +func newSumNumericBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumNumericBuffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumNumericBuffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumNumericBuffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + result := b.sum + return &result, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + d, ok := v.(*apd.Decimal) + if !ok { + return errors.Errorf("sum: expected *apd.Decimal, got %T", v) + } + if !b.sawOne { + b.sum.Set(d) + b.sawOne = true + return nil + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, d) + return err +} + +// sumNumericWindowFunction is the sql.WindowFunction used for sum(numeric) within an OVER(...) clause. +type sumNumericWindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumNumericWindowFunction)(nil) + +// newSumNumericWindowFunction creates the sql.WindowFunction for sum(numeric). +func newSumNumericWindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumNumericWindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumNumericWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumNumericWindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumNumericWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum apd.Decimal + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + d, ok := v.(*apd.Decimal) + if !ok { + return nil, errors.Errorf("sum: expected *apd.Decimal, got %T", v) + } + if !sawOne { + sum.Set(d) + sawOne = true + continue + } + if _, err = sql.DecimalCtx.Add(&sum, &sum, d); err != nil { + return nil, err + } + } + if !sawOne { + return nil, nil + } + return &sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumNumericWindowFunction) Dispose(ctx *sql.Context) {} + +// sumFloat32 represents the PostgreSQL sum(float4) function, which stays real (float32). +var sumFloat32 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Float32, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Float32, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumFloat32Buffer, + NewAggWindowFunc: newSumFloat32WindowFunction, +} + +// sumFloat32Buffer is an aggregation buffer for sum(float4). +type sumFloat32Buffer struct { + expr sql.Expression + sum float32 + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumFloat32Buffer)(nil) + +// newSumFloat32Buffer creates a new aggregation buffer for the sum(float4) aggregate function. +func newSumFloat32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumFloat32Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumFloat32Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumFloat32Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumFloat32Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + f, ok := v.(float32) + if !ok { + return errors.Errorf("sum: expected float32, got %T", v) + } + b.sum += f + b.sawOne = true + return nil +} + +// sumFloat32WindowFunction is the sql.WindowFunction used for sum(float4) within an OVER(...) clause. +type sumFloat32WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumFloat32WindowFunction)(nil) + +// newSumFloat32WindowFunction creates the sql.WindowFunction for sum(float4). +func newSumFloat32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumFloat32WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumFloat32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumFloat32WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum float32 + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + fv, ok := v.(float32) + if !ok { + return nil, errors.Errorf("sum: expected float32, got %T", v) + } + sum += fv + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumFloat32WindowFunction) Dispose(ctx *sql.Context) {} + +// sumFloat64 represents the PostgreSQL sum(float8) function, which stays double precision (float64). +var sumFloat64 = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "sum", + Return: pgtypes.Float64, + Parameters: [1]*pgtypes.DoltgresType{ + pgtypes.Float64, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newSumFloat64Buffer, + NewAggWindowFunc: newSumFloat64WindowFunction, +} + +// sumFloat64Buffer is an aggregation buffer for sum(float8). +type sumFloat64Buffer struct { + expr sql.Expression + sum float64 + sawOne bool +} + +var _ sql.AggregationBuffer = (*sumFloat64Buffer)(nil) + +// newSumFloat64Buffer creates a new aggregation buffer for the sum(float8) aggregate function. +func newSumFloat64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &sumFloat64Buffer{expr: exprs[0]}, nil +} + +// Dispose implements the sql.AggregationBuffer interface. +func (b *sumFloat64Buffer) Dispose(ctx *sql.Context) {} + +// Eval implements the sql.AggregationBuffer interface. +func (b *sumFloat64Buffer) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +// Update implements the sql.AggregationBuffer interface. +func (b *sumFloat64Buffer) Update(ctx *sql.Context, row sql.Row) error { + v, err := b.expr.Eval(ctx, row) + if err != nil { + return err + } + if v == nil { + return nil + } + f, ok := v.(float64) + if !ok { + return errors.Errorf("sum: expected float64, got %T", v) + } + b.sum += f + b.sawOne = true + return nil +} + +// sumFloat64WindowFunction is the sql.WindowFunction used for sum(float8) within an OVER(...) clause. +type sumFloat64WindowFunction struct { + expr sql.Expression + framer sql.WindowFramer +} + +var _ sql.WindowFunction = (*sumFloat64WindowFunction)(nil) + +// newSumFloat64WindowFunction creates the sql.WindowFunction for sum(float8). +func newSumFloat64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &sumFloat64WindowFunction{expr: exprs[0]} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer + } + return wf, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *sumFloat64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *sumFloat64WindowFunction) DefaultFramer() sql.WindowFramer { + if w.framer != nil { + return w.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *sumFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + var sum float64 + var sawOne bool + for i := interval.Start; i < interval.End; i++ { + v, err := w.expr.Eval(ctx, buf[i]) + if err != nil { + return nil, err + } + if v == nil { + continue + } + fv, ok := v.(float64) + if !ok { + return nil, errors.Errorf("sum: expected float64, got %T", v) + } + sum += fv + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *sumFloat64WindowFunction) Dispose(ctx *sql.Context) {} diff --git a/server/functions/framework/catalog.go b/server/functions/framework/catalog.go index 88df9ca946..eff3756832 100644 --- a/server/functions/framework/catalog.go +++ b/server/functions/framework/catalog.go @@ -34,6 +34,9 @@ var Catalog = map[string][]FunctionInterface{} // AggregateCatalog contains all of the PostgreSQL aggregate functions. var AggregateCatalog = map[string][]AggregateFunctionInterface{} +// WindowCatalog contains all of the PostgreSQL functions that may only be used as window functions. +var WindowCatalog = map[string][]WindowFunctionInterface{} + // initializedFunctions simply states whether Initialize has been called yet. var initializedFunctions = false @@ -97,6 +100,21 @@ func RegisterAggregateFunction(f AggregateFunctionInterface) { } } +// RegisterWindowFunction registers the given window-only function, so that it will be usable from a running +// server. This should be called from within an init(). +func RegisterWindowFunction(f WindowFunctionInterface) { + if initializedFunctions { + panic("attempted to register a function after the init() phase") + } + switch f := f.(type) { + case Func0Window: + name := strings.ToLower(f.Name) + WindowCatalog[name] = append(WindowCatalog[name], f) + default: + panic(fmt.Sprintf("unhandled function type %T", f)) + } +} + // Initialize handles the initialization of the catalog by overwriting the built-in GMS functions, since they do not // apply to PostgreSQL (and functions of the same name often have different behavior). func Initialize(astConvert func(parser.Statement) (sqlparser.Statement, error)) { @@ -114,6 +132,7 @@ func Initialize(astConvert func(parser.Statement) (sqlparser.Statement, error)) validateFunctions() compileFunctions() compileAggs() + compileWindowFuncs() } // replaceGmsBuiltIns replaces all GMS built-ins that have conflicting names with PostgreSQL functions. @@ -122,6 +141,12 @@ func replaceGmsBuiltIns() { for name := range Catalog { functionNames[strings.ToLower(name)] = struct{}{} } + for name := range AggregateCatalog { + functionNames[strings.ToLower(name)] = struct{}{} + } + for name := range WindowCatalog { + functionNames[strings.ToLower(name)] = struct{}{} + } var newBuiltIns []sql.Function for _, f := range function.BuiltIns { if _, ok := functionNames[strings.ToLower(f.FunctionName())]; !ok { @@ -193,10 +218,8 @@ func compileNonOperatorFunction(funcName string, overloads []FunctionInterface) // compileNonOperatorFunction creates a CompiledFunction for each overload of the given function. func compileAggFunction(funcName string, overloads []AggregateFunctionInterface) { - var newBuffer NewBufferFn overloadTree := NewOverloads() for _, functionOverload := range overloads { - newBuffer = functionOverload.NewBuffer if err := overloadTree.Add(functionOverload); err != nil { panic(err) } @@ -205,7 +228,27 @@ func compileAggFunction(funcName string, overloads []AggregateFunctionInterface) // Store the compiled function into the engine's built-in functions // TODO: don't do this, use an actual contract for communicating these functions to the engine catalog createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) { - return NewCompiledAggregateFunction(ctx, funcName, params, overloadTree, newBuffer), nil + return NewCompiledAggregateFunction(ctx, funcName, params, overloadTree), nil + } + function.BuiltIns = append(function.BuiltIns, sql.FunctionN{ + Name: funcName, + Fn: createFunc, + }) + compiledCatalog[funcName] = createFunc +} + +// compileWindowFunction creates a CompiledWindowFunction for each overload of the given window-only function. +func compileWindowFunction(funcName string, overloads []WindowFunctionInterface) { + overloadTree := NewOverloads() + for _, functionOverload := range overloads { + if err := overloadTree.Add(functionOverload); err != nil { + panic(err) + } + } + + // Store the compiled function into the engine's built-in functions + createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) { + return NewCompiledWindowFunction(ctx, funcName, params, overloadTree), nil } function.BuiltIns = append(function.BuiltIns, sql.FunctionN{ Name: funcName, @@ -260,3 +303,9 @@ func compileAggs() { compileAggFunction(funcName, overloads) } } + +func compileWindowFuncs() { + for funcName, overloads := range WindowCatalog { + compileWindowFunction(funcName, overloads) + } +} diff --git a/server/functions/framework/compiled_aggregate_function.go b/server/functions/framework/compiled_aggregate_function.go index e4eda47243..6a1c57b3c7 100644 --- a/server/functions/framework/compiled_aggregate_function.go +++ b/server/functions/framework/compiled_aggregate_function.go @@ -20,6 +20,7 @@ import ( cerrors "github.com/cockroachdb/errors" "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/expression" + "github.com/dolthub/go-mysql-server/sql/transform" ) // AggregateFunction is an expression that represents CompiledAggregateFunction @@ -34,27 +35,38 @@ type NewBufferFn func([]sql.Expression) (sql.AggregationBuffer, error) // CompiledAggregateFunction is an expression that represents a fully-analyzed PostgreSQL aggregate function. type CompiledAggregateFunction struct { *CompiledFunction - aggId sql.ColumnId - newBuffer NewBufferFn + aggId sql.ColumnId + window *sql.WindowDefinition } var _ AggregateFunction = (*CompiledAggregateFunction)(nil) // NewCompiledAggregateFunction returns a newly compiled function. -// TODO: newBuffer probably needs to be parameterized in the overloads -func NewCompiledAggregateFunction(ctx *sql.Context, name string, args []sql.Expression, functions *Overloads, newBuffer NewBufferFn) *CompiledAggregateFunction { - return newCompiledAggregateFunctionInternal(ctx, name, args, functions, functions.overloadsForParams(len(args)), newBuffer) +func NewCompiledAggregateFunction(ctx *sql.Context, name string, args []sql.Expression, functions *Overloads) *CompiledAggregateFunction { + return newCompiledAggregateFunctionInternal(ctx, name, args, functions, functions.overloadsForParams(len(args))) } // newCompiledAggregateFunctionInternal is called internally, which skips steps that may have already been processed. -func newCompiledAggregateFunctionInternal(ctx *sql.Context, name string, args []sql.Expression, overloads *Overloads, fnOverloads []Overload, newBuffer NewBufferFn) *CompiledAggregateFunction { +func newCompiledAggregateFunctionInternal(ctx *sql.Context, name string, args []sql.Expression, overloads *Overloads, fnOverloads []Overload) *CompiledAggregateFunction { cf := newCompiledFunctionInternal(ctx, name, args, overloads, fnOverloads, false, nil) - c := &CompiledAggregateFunction{ + return &CompiledAggregateFunction{ CompiledFunction: cf, - newBuffer: newBuffer, } +} - return c +// aggregateOverload returns the AggregateFunctionInterface that this function's overload resolution matched +// for its actual argument types. Each overload of a given function name carries its own NewBuffer/NewWindowFunc, +// so a name like "sum" that has separate int4/int8/numeric/etc. overloads gets the correct implementation for +// the arguments actually bound, rather than a single implementation shared across every overload of the name. +func (c *CompiledAggregateFunction) aggregateOverload() (AggregateFunctionInterface, error) { + if !c.overload.Valid() { + return nil, cerrors.Errorf("%s: no matching overload was resolved", c.Name) + } + agg, ok := c.overload.Function().(AggregateFunctionInterface) + if !ok { + return nil, cerrors.Errorf("%s: resolved overload is not an aggregate function", c.Name) + } + return agg, nil } // Eval implements the interface sql.Expression. @@ -62,14 +74,41 @@ func (c *CompiledAggregateFunction) Eval(ctx *sql.Context, row sql.Row) (interfa return nil, cerrors.New("Eval should not be called on CompiledAggregateFunction") } +// Children implements the interface sql.Expression. When this aggregate is bound to a window (via WithWindow), +// the window's PartitionBy/OrderBy expressions are included after the aggregate's own arguments, mirroring +// GMS's own unaryAggBase.Children(): analyzer passes such as column pruning only discover a node's column +// dependencies by walking Children(), and window.PartitionBy/OrderBy are otherwise invisible to them since +// they aren't part of Arguments. +func (c *CompiledAggregateFunction) Children() []sql.Expression { + children := append([]sql.Expression{}, c.Arguments...) + if c.window != nil { + children = append(children, c.window.ToExpressions()...) + } + return children +} + // WithChildren implements the interface sql.Expression. func (c *CompiledAggregateFunction) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) { - if len(children) != len(c.Arguments) { - return nil, sql.ErrInvalidChildrenNumber.New(len(children), len(c.Arguments)) + numArgs := len(c.Arguments) + if len(children) < numArgs { + return nil, sql.ErrInvalidChildrenNumber.New(len(children), numArgs) } // We have to re-resolve here, since the change in children may require it (e.g. we have more type info than we did) - return newCompiledAggregateFunctionInternal(ctx, c.Name, children, c.overloads, c.fnOverloads, c.newBuffer), nil + nc := newCompiledAggregateFunctionInternal(ctx, c.Name, children[:numArgs], c.overloads, c.fnOverloads) + // Preserve the aggregate/window identity assigned by the planbuilder: a later analyzer pass that + // rebuilds this expression via WithChildren (e.g. a generic bottom-up transform) must not silently + // lose the WithId/WithWindow state that was already bound onto c. + nc.aggId = c.aggId + nc.window = c.window + if len(children) > numArgs && c.window != nil { + w, err := c.window.FromExpressions(ctx, children[numArgs:]) + if err != nil { + return nil, err + } + nc.window = w + } + return nc, nil } // SetStatementRunner implements the interface analyzer.Interpreter. @@ -102,7 +141,15 @@ func (c *CompiledAggregateFunction) DebugString(ctx *sql.Context) string { // NewBuffer implements the interface sql.Aggregation. func (c *CompiledAggregateFunction) NewBuffer(ctx *sql.Context) (sql.AggregationBuffer, error) { - return c.newBuffer(c.Arguments) + agg, err := c.aggregateOverload() + if err != nil { + return nil, err + } + args, err := cloneArguments(ctx, c.Arguments) + if err != nil { + return nil, err + } + return agg.NewBuffer(args) } // Id implements the interface sql.Aggregation. @@ -119,15 +166,46 @@ func (c *CompiledAggregateFunction) WithId(id sql.ColumnId) sql.IdExpression { // NewWindowFunction implements the interface sql.WindowAdaptableExpression. func (c *CompiledAggregateFunction) NewWindowFunction(ctx *sql.Context) (sql.WindowFunction, error) { - panic("windows are not implemented yet") + agg, err := c.aggregateOverload() + if err != nil { + return nil, err + } + newWindowFunc := agg.NewWindowFunc() + if newWindowFunc == nil { + return nil, cerrors.Errorf("aggregate function %s cannot be used as a window function", c.Name) + } + args, err := cloneArguments(ctx, c.Arguments) + if err != nil { + return nil, err + } + return newWindowFunc(args, c.window) +} + +// cloneArguments returns a deep copy of args. Each partition/group gets its own AggregationBuffer or +// WindowFunction instance, but without cloning, they'd all share the same argument expressions - so a +// stateful expression like DISTINCT's dedup cache (which must reset per group/partition) would incorrectly +// carry state across them. This mirrors GMS's own aggregations (e.g. *aggregation.Sum.NewBuffer), which clone +// their child expression for the same reason. +func cloneArguments(ctx *sql.Context, args []sql.Expression) ([]sql.Expression, error) { + cloned := make([]sql.Expression, len(args)) + for i, arg := range args { + c, err := transform.Clone(ctx, arg) + if err != nil { + return nil, err + } + cloned[i] = c + } + return cloned, nil } // WithWindow implements the interface sql.WindowAdaptableExpression. func (c *CompiledAggregateFunction) WithWindow(ctx *sql.Context, window *sql.WindowDefinition) sql.WindowAdaptableExpression { - panic("windows are not implemented yet") + nc := *c + nc.window = window + return &nc } // Window implements the interface sql.WindowAdaptableExpression. func (c *CompiledAggregateFunction) Window() *sql.WindowDefinition { - panic("windows are not implemented yet") + return c.window } diff --git a/server/functions/framework/compiled_window_function.go b/server/functions/framework/compiled_window_function.go new file mode 100644 index 0000000000..1ef6fd045a --- /dev/null +++ b/server/functions/framework/compiled_window_function.go @@ -0,0 +1,179 @@ +// Copyright 2026 Dolthub, 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 framework + +import ( + "strings" + + cerrors "github.com/cockroachdb/errors" + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression" +) + +// WindowOnlyFunction is an expression that represents CompiledWindowFunction: a PostgreSQL function that may +// only be used as a window function (within an OVER(...) clause), such as row_number() or rank(). +type WindowOnlyFunction interface { + sql.FunctionExpression + sql.WindowAggregation + specificFuncImpl() +} + +// CompiledWindowFunction is an expression that represents a fully-analyzed PostgreSQL window-only function. +type CompiledWindowFunction struct { + *CompiledFunction + windowId sql.ColumnId + window *sql.WindowDefinition +} + +var _ WindowOnlyFunction = (*CompiledWindowFunction)(nil) + +// NewCompiledWindowFunction returns a newly compiled function. +func NewCompiledWindowFunction(ctx *sql.Context, name string, args []sql.Expression, functions *Overloads) *CompiledWindowFunction { + return newCompiledWindowFunctionInternal(ctx, name, args, functions, functions.overloadsForParams(len(args))) +} + +// newCompiledWindowFunctionInternal is called internally, which skips steps that may have already been processed. +func newCompiledWindowFunctionInternal(ctx *sql.Context, name string, args []sql.Expression, overloads *Overloads, fnOverloads []Overload) *CompiledWindowFunction { + cf := newCompiledFunctionInternal(ctx, name, args, overloads, fnOverloads, false, nil) + return &CompiledWindowFunction{ + CompiledFunction: cf, + } +} + +// windowOverload returns the WindowFunctionInterface that this function's overload resolution matched for its +// actual argument types, mirroring CompiledAggregateFunction.aggregateOverload. +func (c *CompiledWindowFunction) windowOverload() (WindowFunctionInterface, error) { + if !c.overload.Valid() { + return nil, cerrors.Errorf("%s: no matching overload was resolved", c.Name) + } + fn, ok := c.overload.Function().(WindowFunctionInterface) + if !ok { + return nil, cerrors.Errorf("%s: resolved overload is not a window function", c.Name) + } + return fn, nil +} + +// Eval implements the interface sql.Expression. +func (c *CompiledWindowFunction) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) { + return nil, cerrors.New("Eval should not be called on CompiledWindowFunction") +} + +// Children implements the interface sql.Expression. When this function is bound to a window (via +// WithWindow), the window's PartitionBy/OrderBy expressions are included after its own arguments, mirroring +// GMS's own unaryAggBase.Children(): analyzer passes such as column pruning only discover a node's column +// dependencies by walking Children(), and window.PartitionBy/OrderBy are otherwise invisible to them since +// they aren't part of Arguments. +func (c *CompiledWindowFunction) Children() []sql.Expression { + children := append([]sql.Expression{}, c.Arguments...) + if c.window != nil { + children = append(children, c.window.ToExpressions()...) + } + return children +} + +// WithChildren implements the interface sql.Expression. +func (c *CompiledWindowFunction) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) { + numArgs := len(c.Arguments) + if len(children) < numArgs { + return nil, sql.ErrInvalidChildrenNumber.New(len(children), numArgs) + } + + // We have to re-resolve here, since the change in children may require it (e.g. we have more type info than we did) + nc := newCompiledWindowFunctionInternal(ctx, c.Name, children[:numArgs], c.overloads, c.fnOverloads) + // Preserve the identity assigned by the planbuilder: a later analyzer pass that rebuilds this + // expression via WithChildren (e.g. a generic bottom-up transform) must not silently lose the + // WithId/WithWindow state that was already bound onto c. + nc.windowId = c.windowId + nc.window = c.window + if len(children) > numArgs && c.window != nil { + w, err := c.window.FromExpressions(ctx, children[numArgs:]) + if err != nil { + return nil, err + } + nc.window = w + } + return nc, nil +} + +// SetStatementRunner implements the interface analyzer.Interpreter. +func (c *CompiledWindowFunction) SetStatementRunner(ctx *sql.Context, runner sql.StatementRunner) sql.Expression { + nc := *c + nc.runner = runner + return &nc +} + +// specificFuncImpl implements the interface WindowOnlyFunction. +func (*CompiledWindowFunction) specificFuncImpl() {} + +func (c *CompiledWindowFunction) DebugString(ctx *sql.Context) string { + sb := strings.Builder{} + sb.WriteString("CompiledWindowFunction:") + sb.WriteString(c.Name + "(") + for i, param := range c.Arguments { + // Aliases will output the string "x as x", which is an artifact of how we build the AST, so we'll bypass it + if alias, ok := param.(*expression.Alias); ok { + param = alias.Child + } + if i > 0 { + sb.WriteString(", ") + } + sb.WriteString(sql.DebugString(ctx, param)) + } + sb.WriteString(")") + return sb.String() +} + +// Id implements the interface sql.IdExpression. +func (c *CompiledWindowFunction) Id() sql.ColumnId { + return c.windowId +} + +// WithId implements the interface sql.IdExpression. +func (c *CompiledWindowFunction) WithId(id sql.ColumnId) sql.IdExpression { + nc := *c + nc.windowId = id + return &nc +} + +// NewWindowFunction implements the interface sql.WindowAdaptableExpression. +func (c *CompiledWindowFunction) NewWindowFunction(ctx *sql.Context) (sql.WindowFunction, error) { + fn, err := c.windowOverload() + if err != nil { + return nil, err + } + newWindowFunc := fn.NewWindowFunc() + if newWindowFunc == nil { + return nil, cerrors.Errorf("function %s cannot be used as a window function", c.Name) + } + // See cloneArguments: each partition needs its own argument expression instances so stateful + // expressions (e.g. DISTINCT's dedup cache) don't leak state across partitions. + args, err := cloneArguments(ctx, c.Arguments) + if err != nil { + return nil, err + } + return newWindowFunc(args, c.window) +} + +// WithWindow implements the interface sql.WindowAdaptableExpression. +func (c *CompiledWindowFunction) WithWindow(ctx *sql.Context, window *sql.WindowDefinition) sql.WindowAdaptableExpression { + nc := *c + nc.window = window + return &nc +} + +// Window implements the interface sql.WindowAdaptableExpression. +func (c *CompiledWindowFunction) Window() *sql.WindowDefinition { + return c.window +} diff --git a/server/functions/framework/functions.go b/server/functions/framework/functions.go index 1765ca0523..54715701c5 100644 --- a/server/functions/framework/functions.go +++ b/server/functions/framework/functions.go @@ -56,6 +56,21 @@ type AggregateFunctionInterface interface { FunctionInterface // TODO: this maybe needs to take the place of the Callable function NewBuffer([]sql.Expression) (sql.AggregationBuffer, error) + // NewWindowFunc returns the constructor used to build this aggregate's runtime sql.WindowFunction for use + // within an OVER(...) clause, or nil if this aggregate does not yet support window use. + NewWindowFunc() NewWindowFunctionFn +} + +// NewWindowFunctionFn creates a sql.WindowFunction for the given arguments, for use within an OVER(...) clause +// bound to the given window definition. +type NewWindowFunctionFn func(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) + +// WindowFunctionInterface is an interface for PostgreSQL functions that may only be used as window functions +// (i.e. within an OVER(...) clause) and have no GROUP BY equivalent, such as row_number() or rank(). +type WindowFunctionInterface interface { + FunctionInterface + // NewWindowFunc returns the constructor used to build this function's runtime sql.WindowFunction. + NewWindowFunc() NewWindowFunctionFn } // Function0 is a function that does not take any parameters. @@ -605,10 +620,32 @@ func (f Function7) enforceInterfaceInheritance(error) {} type Func1Aggregate struct { Function1 NewAggBuffer func([]sql.Expression) (sql.AggregationBuffer, error) + // NewAggWindowFunc optionally builds this aggregate's runtime sql.WindowFunction for use within an + // OVER(...) clause. Leave nil if this aggregate does not yet support window use. + NewAggWindowFunc NewWindowFunctionFn } +var _ AggregateFunctionInterface = Func1Aggregate{} + func (f Func1Aggregate) NewBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { return f.NewAggBuffer(exprs) } -var _ AggregateFunctionInterface = Func1Aggregate{} +func (f Func1Aggregate) NewWindowFunc() NewWindowFunctionFn { + return f.NewAggWindowFunc +} + +// Func0Window is a PostgreSQL function that takes no parameters and may only be used as a window function +// (within an OVER(...) clause), such as row_number(). +type Func0Window struct { + Function0 + NewWinFunc func(window *sql.WindowDefinition) (sql.WindowFunction, error) +} + +var _ WindowFunctionInterface = Func0Window{} + +func (f Func0Window) NewWindowFunc() NewWindowFunctionFn { + return func(_ []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + return f.NewWinFunc(window) + } +} diff --git a/server/functions/window/init.go b/server/functions/window/init.go new file mode 100644 index 0000000000..2831b064b9 --- /dev/null +++ b/server/functions/window/init.go @@ -0,0 +1,21 @@ +// Copyright 2026 Dolthub, 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 window registers PostgreSQL functions that may only be used as window functions (i.e. within an +// OVER(...) clause) and have no GROUP BY equivalent, such as row_number() or rank(). +package window + +func Init() { + initRanking() +} diff --git a/server/functions/window/rank.go b/server/functions/window/rank.go new file mode 100644 index 0000000000..ed518d325b --- /dev/null +++ b/server/functions/window/rank.go @@ -0,0 +1,147 @@ +// Copyright 2026 Dolthub, 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 window + +import ( + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" + + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// rank represents the PostgreSQL rank() window function, returning a bigint directly rather than GMS's +// native uint64 representation. +var rank = framework.Func0Window{ + Function0: framework.Function0{ + Name: "rank", + Return: pgtypes.Int64, + }, + NewWinFunc: newRankWindowFunction, +} + +// denseRank represents the PostgreSQL dense_rank() window function, returning a bigint directly rather than +// GMS's native uint64 representation. +var denseRank = framework.Func0Window{ + Function0: framework.Function0{ + Name: "dense_rank", + Return: pgtypes.Int64, + }, + NewWinFunc: newDenseRankWindowFunction, +} + +// rankWindowFunction is the sql.WindowFunction used for rank() within an OVER(...) clause. Its algorithm +// mirrors GMS's own aggregation.rankBase.Compute exactly (see sql/expression/function/aggregation/ +// window_functions.go), just returning int64 instead of uint64: the rank of a row is the count of rows +// strictly before its peer group (rows sharing the same ORDER BY key), plus 1. Per the SQL standard, RANK +// ignores any explicit frame clause and groups rows into "peers" by the window's ORDER BY expressions, so +// DefaultFramer always returns a peer-group framer regardless of window.Frame. +type rankWindowFunction struct { + orderBy []sql.Expression + partitionStart, partitionEnd int + pos int +} + +var _ sql.WindowFunction = (*rankWindowFunction)(nil) + +// newRankWindowFunction creates the sql.WindowFunction for rank(). +func newRankWindowFunction(window *sql.WindowDefinition) (sql.WindowFunction, error) { + var orderBy []sql.Expression + if window != nil { + orderBy = window.OrderBy.ToExpressions() + } + return &rankWindowFunction{partitionStart: -1, partitionEnd: -1, pos: -1, orderBy: orderBy}, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *rankWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + w.partitionStart, w.partitionEnd = interval.Start, interval.End + w.pos = w.partitionStart + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *rankWindowFunction) DefaultFramer() sql.WindowFramer { + return aggregation.NewPeerGroupFramer(w.orderBy) +} + +// Compute implements the sql.WindowFunction interface. +func (w *rankWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End-interval.Start < 1 { + return nil, nil + } + defer func() { w.pos++ }() + switch { + case w.pos == 0: + return int64(1), nil + case w.partitionEnd-w.partitionStart == 1: + return int64(1), nil + default: + return int64(interval.Start-w.partitionStart) + 1, nil + } +} + +// Dispose implements the sql.WindowFunction interface. +func (w *rankWindowFunction) Dispose(ctx *sql.Context) {} + +// denseRankWindowFunction is the sql.WindowFunction used for dense_rank() within an OVER(...) clause. It +// wraps rankWindowFunction's peer-group counting the same way GMS's own DenseRank wraps rankBase: instead of +// counting all preceding rows, it counts preceding distinct peer groups. +type denseRankWindowFunction struct { + rank *rankWindowFunction + prevRank int64 + denseRank int64 +} + +var _ sql.WindowFunction = (*denseRankWindowFunction)(nil) + +// newDenseRankWindowFunction creates the sql.WindowFunction for dense_rank(). +func newDenseRankWindowFunction(window *sql.WindowDefinition) (sql.WindowFunction, error) { + r, err := newRankWindowFunction(window) + if err != nil { + return nil, err + } + return &denseRankWindowFunction{rank: r.(*rankWindowFunction)}, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *denseRankWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return w.rank.StartPartition(ctx, interval, buf) +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *denseRankWindowFunction) DefaultFramer() sql.WindowFramer { + return w.rank.DefaultFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *denseRankWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + r, err := w.rank.Compute(ctx, interval, buf) + if err != nil || r == nil { + return r, err + } + rankVal := r.(int64) + if rankVal == 1 { + w.prevRank = 1 + w.denseRank = 1 + } else if rankVal != w.prevRank { + w.prevRank = rankVal + w.denseRank++ + } + return w.denseRank, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *denseRankWindowFunction) Dispose(ctx *sql.Context) {} diff --git a/server/functions/window/row_number.go b/server/functions/window/row_number.go new file mode 100644 index 0000000000..f3a379ca22 --- /dev/null +++ b/server/functions/window/row_number.go @@ -0,0 +1,77 @@ +// Copyright 2026 Dolthub, 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 window + +import ( + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" + + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// initRanking registers the ranking window functions to the catalog. +func initRanking() { + framework.RegisterWindowFunction(rowNumber) + framework.RegisterWindowFunction(rank) + framework.RegisterWindowFunction(denseRank) +} + +// rowNumber represents the PostgreSQL row_number() window function, returning a bigint directly rather than +// GMS's native uint64/int representation. +var rowNumber = framework.Func0Window{ + Function0: framework.Function0{ + Name: "row_number", + Return: pgtypes.Int64, + }, + NewWinFunc: newRowNumberWindowFunction, +} + +// rowNumberWindowFunction is the sql.WindowFunction used for row_number() within an OVER(...) clause. Per the +// SQL standard, row_number() ignores any explicit frame clause and always numbers rows across the whole +// partition, so DefaultFramer always returns a partition-wide framer regardless of window.Frame. +type rowNumberWindowFunction struct { + pos int64 +} + +var _ sql.WindowFunction = (*rowNumberWindowFunction)(nil) + +// newRowNumberWindowFunction creates the sql.WindowFunction for row_number(). +func newRowNumberWindowFunction(_ *sql.WindowDefinition) (sql.WindowFunction, error) { + return &rowNumberWindowFunction{}, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *rowNumberWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + w.pos = 1 + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *rowNumberWindowFunction) DefaultFramer() sql.WindowFramer { + return aggregation.NewPartitionFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *rowNumberWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + if interval.End-interval.Start < 1 { + return nil, nil + } + defer func() { w.pos++ }() + return w.pos, nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *rowNumberWindowFunction) Dispose(ctx *sql.Context) {} diff --git a/server/initialization/initialization.go b/server/initialization/initialization.go index 03a50f0985..d1b936aaf2 100644 --- a/server/initialization/initialization.go +++ b/server/initialization/initialization.go @@ -34,6 +34,7 @@ import ( "github.com/dolthub/doltgresql/server/functions/binary" "github.com/dolthub/doltgresql/server/functions/framework" "github.com/dolthub/doltgresql/server/functions/unary" + "github.com/dolthub/doltgresql/server/functions/window" "github.com/dolthub/doltgresql/server/tables" "github.com/dolthub/doltgresql/server/tables/dtables" "github.com/dolthub/doltgresql/server/tables/information_schema" @@ -58,6 +59,7 @@ func Initialize(dEnv *env.DoltEnv, cfg *doltgresservercfg.DoltgresConfig) { unary.Init() functions.Init() aggregate.Init() + window.Init() builtInCasts := casts.Init(core.GetRunnerFromContext, func(f sql.Expression) bool { if cf, ok := f.(*framework.CompiledFunction); ok { return cf.IsStrict() diff --git a/testing/go/enginetest/doltgres_engine_test.go b/testing/go/enginetest/doltgres_engine_test.go index 0c5646a125..a152ec394d 100755 --- a/testing/go/enginetest/doltgres_engine_test.go +++ b/testing/go/enginetest/doltgres_engine_test.go @@ -540,6 +540,8 @@ func TestScripts(t *testing.T) { "INSERT IGNORE throws an error when json is badly formatted", // error messages don't match "identical expressions over different windows should produce different results", // ERROR: integer: unhandled type: float64 "windows without ORDER BY should be treated as RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", // ERROR: integer: unhandled type: float64 + "sum() and avg() on non-DECIMAL type column returns the DOUBLE type result", // MySQL-specific: MySQL widens FLOAT to DOUBLE for SUM; Postgres sum(real) stays real + "sum() and avg() on DECIMAL type column returns the DECIMAL type result", // MySQL-specific: MySQL truncates AVG's decimal scale; Postgres numeric division keeps full precision "division and int division operation on negative, small and big value for decimal type column of table", // numeric keys broken "update columns with default", // broken, see repro in update_test.go "select count(*) from t where (f in (null, cast(0.8 as float)));", // incorrect result, needs a fix diff --git a/testing/go/smoke_test.go b/testing/go/smoke_test.go index d86fca7ed4..54acced277 100644 --- a/testing/go/smoke_test.go +++ b/testing/go/smoke_test.go @@ -648,7 +648,7 @@ func TestSmokeTests(t *testing.T) { { Query: "SELECT SUM(v1) FROM test WHERE v1 BETWEEN 3 AND 5;", Expected: []sql.Row{ - {12.0}, + {int64(12)}, }, }, { @@ -658,7 +658,7 @@ func TestSmokeTests(t *testing.T) { { Query: "SELECT SUM(v1) FROM test WHERE v1 BETWEEN 3 AND 5;", Expected: []sql.Row{ - {12.0}, + {int64(12)}, }, }, }, diff --git a/testing/go/values_statement_test.go b/testing/go/values_statement_test.go index 9f8590df61..046dd7045c 100644 --- a/testing/go/values_statement_test.go +++ b/testing/go/values_statement_test.go @@ -242,7 +242,7 @@ var ValuesStatementTests = []ScriptTest{ { // AVG on mixed types Query: `SELECT AVG(n) FROM (VALUES(1),(2),(3),(4)) v(n);`, - Expected: []sql.Row{{2.5}}, + Expected: []sql.Row{{Numeric("2.5")}}, }, { // MIN/MAX on mixed types diff --git a/testing/go/window_test.go b/testing/go/window_test.go new file mode 100644 index 0000000000..b940595103 --- /dev/null +++ b/testing/go/window_test.go @@ -0,0 +1,207 @@ +// Copyright 2026 Dolthub, 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 _go + +import ( + "testing" + + "github.com/dolthub/go-mysql-server/sql" +) + +func TestWindowFunctions(t *testing.T) { + RunScripts(t, []ScriptTest{ + { + Name: "native sum and row_number as window functions", + SetUpScript: []string{ + "CREATE TABLE t (id INT PRIMARY KEY, grp INT, amt INT);", + "INSERT INTO t VALUES (1, 1, 10), (2, 1, 20), (3, 1, 30), (4, 2, 5), (5, 2, 15);", + }, + Assertions: []ScriptTestAssertion{ + { + // SUM(int4) OVER (...) must return a bigint (int64), not GMS's float64. + Query: "SELECT id, sum(amt) OVER (PARTITION BY grp ORDER BY id) FROM t ORDER BY id", + Expected: []sql.Row{ + {1, int64(10)}, + {2, int64(30)}, + {3, int64(60)}, + {4, int64(5)}, + {5, int64(20)}, + }, + }, + { + // row_number() must return a bigint directly and reset per partition. + Query: "SELECT id, row_number() OVER (PARTITION BY grp ORDER BY id) FROM t ORDER BY id", + Expected: []sql.Row{ + {1, int64(1)}, + {2, int64(2)}, + {3, int64(3)}, + {4, int64(1)}, + {5, int64(2)}, + }, + }, + { + // sum(amt) as a regular GROUP BY aggregate still works and is also bigint. + Query: "SELECT grp, sum(amt) FROM t GROUP BY grp ORDER BY grp", + Expected: []sql.Row{ + {1, int64(60)}, + {2, int64(20)}, + }, + }, + }, + }, + { + // https://github.com/dolthub/doltgresql/issues/1796 + Name: "basic window functions", + SetUpScript: []string{ + "CREATE TABLE c (c_id INT PRIMARY KEY, bill TEXT);", + "CREATE TABLE o (o_id INT PRIMARY KEY, c_id INT, ship TEXT);", + "INSERT INTO c VALUES (1, 'CA'), (2, 'TX'), (3, 'MA'), (4, 'TX'), (5, NULL), (6, 'FL');", + "INSERT INTO o VALUES (10, 1, 'CA'), (20, 1, 'CA'), (30, 1, 'CA'), (40, 2, 'CA'), (50, 2, 'TX'), (60, 2, NULL), (70, 4, 'WY'), (80, 4, NULL), (90, 6, 'WA');", + }, + Assertions: []ScriptTestAssertion{ + { + Query: "SELECT row_number() OVER () AS rn FROM o WHERE c_id=-999", + Expected: []sql.Row{}, + }, + { + Query: "SELECT row_number() OVER () AS rn FROM o WHERE c_id=1", + Expected: []sql.Row{ + {int64(1)}, + {int64(2)}, + {int64(3)}, + }, + }, + { + Query: "SELECT rank() OVER () AS rnk FROM o WHERE c_id=-999", + Expected: []sql.Row{}, + }, + { + Query: "SELECT o_id, c_id, rank() OVER (ORDER BY o_id) AS rnk FROM o WHERE c_id=1", + Expected: []sql.Row{ + {10, 1, int64(1)}, + {20, 1, int64(2)}, + {30, 1, int64(3)}, + }, + }, + { + Query: "SELECT dense_rank() OVER () AS drnk FROM o WHERE c_id=-999", + Expected: []sql.Row{}, + }, + { + // Doltgres inherits GMS/MySQL null-first sort order, so NULLs sort before + // non-NULLs in window ORDER BY and outer ORDER BY. Postgres default is NULLS LAST. + // TODO: update expected values once Doltgres adopts Postgres NULLS LAST ordering. + Query: "SELECT ship, dense_rank() OVER (ORDER BY ship) AS drnk FROM o WHERE c_id IN (1, 2) ORDER BY ship", + Expected: []sql.Row{ + {nil, int64(1)}, + {"CA", int64(2)}, + {"CA", int64(2)}, + {"CA", int64(2)}, + {"CA", int64(2)}, + {"TX", int64(3)}, + }, + }, + { + Query: "SELECT * FROM (SELECT c_id AS c_c_id, bill FROM c) sq1, LATERAL (SELECT row_number() OVER () AS rownum FROM o WHERE c_id = c_c_id) sq2 ORDER BY c_c_id, bill, rownum", + Expected: []sql.Row{ + {1, "CA", int64(1)}, + {1, "CA", int64(2)}, + {1, "CA", int64(3)}, + {2, "TX", int64(1)}, + {2, "TX", int64(2)}, + {2, "TX", int64(3)}, + {4, "TX", int64(1)}, + {4, "TX", int64(2)}, + {6, "FL", int64(1)}, + }, + }, + // ORDER BY on rank alias with LIMIT (multi-hop Limit→Sort→Window chain) + { + Query: "SELECT c_id, rank() OVER (ORDER BY c_id) AS rnk FROM c ORDER BY rnk DESC LIMIT 3", + Expected: []sql.Row{ + {6, int64(6)}, + {5, int64(5)}, + {4, int64(4)}, + }, + }, + // ORDER BY on rank alias (Sort→Window chain) + { + Query: "SELECT c_id, rank() OVER (ORDER BY c_id) AS r FROM c ORDER BY r", + Expected: []sql.Row{ + {1, int64(1)}, + {2, int64(2)}, + {3, int64(3)}, + {4, int64(4)}, + {5, int64(5)}, + {6, int64(6)}, + }, + }, + // DISTINCT + ORDER BY on rank alias + { + Query: "SELECT DISTINCT c_id, rank() OVER (ORDER BY c_id) AS r FROM c ORDER BY r", + Expected: []sql.Row{ + {1, int64(1)}, + {2, int64(2)}, + {3, int64(3)}, + {4, int64(4)}, + {5, int64(5)}, + {6, int64(6)}, + }, + }, + // CASE expression over rank() subquery (int64 rank values flow into Numeric cast) + { + Query: "SELECT sum(CASE WHEN r > 0 THEN 1 ELSE 0 END) FROM (SELECT rank() OVER (ORDER BY c_id) AS r FROM c) t", + Expected: []sql.Row{ + {int64(6)}, + }, + }, + // window SUM with non-empty result (GMS SumAgg.Compute returns float64, not int32) + { + Query: "SELECT c_id, SUM(o_id) OVER (PARTITION BY c_id) AS s FROM o WHERE c_id = 1 ORDER BY o_id", + Expected: []sql.Row{ + {1, int64(60)}, + {1, int64(60)}, + {1, int64(60)}, + }, + }, + }, + }, + { + Name: "window SUM/AVG wrapped in a subquery projection", + SetUpScript: []string{ + "CREATE TABLE wrapper_probe (grp INT, val INT);", + "INSERT INTO wrapper_probe VALUES (1, 10), (1, 20), (2, 5);", + }, + Assertions: []ScriptTestAssertion{ + { + Query: "SELECT grp, val, grp_total FROM (SELECT grp, val, SUM(val) OVER (PARTITION BY grp) AS grp_total FROM wrapper_probe) sub ORDER BY grp, val;", + Expected: []sql.Row{ + {1, 10, int64(30)}, + {1, 20, int64(30)}, + {2, 5, int64(5)}, + }, + }, + { + Query: "SELECT grp, val, grp_avg FROM (SELECT grp, val, AVG(val) OVER (PARTITION BY grp) AS grp_avg FROM wrapper_probe) sub ORDER BY grp, val;", + Expected: []sql.Row{ + {1, 10, Numeric("15")}, + {1, 20, Numeric("15")}, + {2, 5, Numeric("5")}, + }, + }, + }, + }, + }) +} From 6dc96cea33b06800ab2ccd1efdd201e98471ed05 Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Fri, 10 Jul 2026 13:14:34 -0700 Subject: [PATCH 2/4] Updating GMS to e06e93647d009ed3cbe362ee2af0203153dd0c41 --- go.mod | 2 +- go.sum | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/go.mod b/go.mod index b9ed822799..fc715b7ac0 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/dolthub/dolt/go v0.40.5-0.20260710160803-b101e404aa74 github.com/dolthub/eventsapi_schema v0.0.0-20260310172945-37a9265ade69 github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2 - github.com/dolthub/go-mysql-server v0.20.1-0.20260709224106-fecf6bb54eb0 + github.com/dolthub/go-mysql-server v0.20.1-0.20260709234404-e06e93647d00 github.com/dolthub/pg_query_go/v6 v6.0.0-20251215122834-fb20be4254d1 github.com/dolthub/sqllogictest/go v0.0.0-20260624223518-788480b24166 github.com/dolthub/vitess v0.0.0-20260624214226-81d034e0fde8 diff --git a/go.sum b/go.sum index 9be106897b..66adba99ea 100644 --- a/go.sum +++ b/go.sum @@ -258,6 +258,8 @@ github.com/dolthub/go-icu-regex v0.0.0-20260610153742-72563bc7ca83 h1:FEMjCGEroD github.com/dolthub/go-icu-regex v0.0.0-20260610153742-72563bc7ca83/go.mod h1:F3cnm+vMRK1HaU6+rNqQrOCyR03HHhR1GWG2gnPOqaE= github.com/dolthub/go-mysql-server v0.20.1-0.20260709224106-fecf6bb54eb0 h1:IUbkgDsRcl3pNZB3AhdLQo5DGPijmwa6f9f40Qi3WxY= github.com/dolthub/go-mysql-server v0.20.1-0.20260709224106-fecf6bb54eb0/go.mod h1:mj5/QX3V8i92REbA1w6CzyknJAFdKtdE7l931405C/E= +github.com/dolthub/go-mysql-server v0.20.1-0.20260709234404-e06e93647d00 h1:tQY5FOiUFKzcSqJQleh0NCVUXuPLvxl1iPQ5KtFvrbo= +github.com/dolthub/go-mysql-server v0.20.1-0.20260709234404-e06e93647d00/go.mod h1:mj5/QX3V8i92REbA1w6CzyknJAFdKtdE7l931405C/E= github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI= github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q= github.com/dolthub/ishell v0.0.0-20260414231531-5f031e3e9037 h1:oIW9HwuWrhxv+4HZxA+QQSKHLqWFyXZ2FmNjUYwkdiM= From a03e5966920fe150ef9379d5f9cbd0f4dbbba084 Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Fri, 10 Jul 2026 14:23:25 -0700 Subject: [PATCH 3/4] fixup: cleaning up typed agg function implementations --- server/functions/aggregate/avg_aggregates.go | 690 ++++-------------- .../functions/aggregate/numeric_aggregates.go | 667 ++++------------- 2 files changed, 267 insertions(+), 1090 deletions(-) diff --git a/server/functions/aggregate/avg_aggregates.go b/server/functions/aggregate/avg_aggregates.go index c751d3050a..6b0f3dc62e 100644 --- a/server/functions/aggregate/avg_aggregates.go +++ b/server/functions/aggregate/avg_aggregates.go @@ -24,190 +24,87 @@ import ( pgtypes "github.com/dolthub/doltgresql/server/types" ) -// avgDecimalCtx is used for the division in AVG's sum/count. sql.DecimalCtx (apd.BaseContext) has Precision: -// 0, which disables rounding - fine for the exact Add calls used to accumulate the sum, but Quo (true -// division) can be inexact and apd requires a nonzero precision for that. 20 digits is generous headroom -// over Postgres's own guard-digit scale for numeric division (real Postgres computes a much smaller -// dscale for this), and it's meaningfully cheaper than a very high precision: apd's Quo cost scales with -// precision (bigger intermediate big.Ints), and this runs once per GROUP BY group. -var avgDecimalCtx = sql.DecimalCtx.WithPrecision(16) - -// initAvgAggs registers the functions to the catalog. -func initAvgAggs() { - framework.RegisterAggregateFunction(avgInt16) - framework.RegisterAggregateFunction(avgInt32) - framework.RegisterAggregateFunction(avgInt64) - framework.RegisterAggregateFunction(avgNumeric) - framework.RegisterAggregateFunction(avgFloat32) - framework.RegisterAggregateFunction(avgFloat64) -} - -// avgInt16 represents the PostgreSQL avg(int2) function, which promotes to numeric. -var avgInt16 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int16, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newAvgInt16Buffer, - NewAggWindowFunc: newAvgInt16WindowFunction, -} - -// avgInt16Buffer is an aggregation buffer for avg(int2). -type avgInt16Buffer struct { - expr sql.Expression - sum int64 - count int64 -} - -var _ sql.AggregationBuffer = (*avgInt16Buffer)(nil) - -// newAvgInt16Buffer creates a new aggregation buffer for the avg(int2) aggregate function. -func newAvgInt16Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgInt16Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgInt16Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *avgInt16Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if b.count == 0 { - return nil, nil - } - return avgFromInt64Sum(b.sum, b.count) -} - -// Update implements the sql.AggregationBuffer interface. -func (b *avgInt16Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - i, ok := v.(int16) - if !ok { - return errors.Errorf("avg: expected int16, got %T", v) - } - b.sum += int64(i) - b.count++ - return nil -} - -// avgInt16WindowFunction is the sql.WindowFunction used for avg(int2) within an OVER(...) clause. -type avgInt16WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer -} - -var _ sql.WindowFunction = (*avgInt16WindowFunction)(nil) - -// newAvgInt16WindowFunction creates the sql.WindowFunction for avg(int2). -func newAvgInt16WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgInt16WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} - -// StartPartition implements the sql.WindowFunction interface. -func (w *avgInt16WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgInt16WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - -// Compute implements the sql.WindowFunction interface. -func (w *avgInt16WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil - } - var sum, count int64 - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - iv, ok := v.(int16) - if !ok { - return nil, errors.Errorf("avg: expected int16, got %T", v) - } - sum += int64(iv) - count++ +// avgGuardDigits is added on top of the dividend's own digit count when dividing for an AVG result, to get a +// meaningful number of fractional digits out of Quo (an integer dividend alone would otherwise round to an +// integer quotient). 16 matches Postgres's own displayed scale for avg() of integer types. +const avgGuardDigits = 16 + +// quoAvg divides dividend/divisor for use as an AVG result, then reduces the result to strip the trailing +// zeros Quo pads on to fill out the requested precision (e.g. 30/2 would otherwise come back as +// 15.00000...0). The precision is computed dynamically from the dividend's digit count, the same convention +// used elsewhere in this codebase for apd division/rounding (see div.go, round.go, sqrt.go, ln.go), rather +// than a single fixed precision for every call: sql.DecimalCtx (apd.BaseContext) has Precision: 0, which +// disables rounding entirely and Quo requires a nonzero value, but a fixed precision that's too low doesn't +// just lose fractional digits - if the dividend needs more significant digits than that to represent its +// *integer* part (e.g. avg() of a single huge bigint/numeric sum), apd rounds the integer part too, silently +// producing a wrong answer, not just an imprecise one. +func quoAvg(dividend, divisor *apd.Decimal) (*apd.Decimal, error) { + p := dividend.NumDigits() + if dividend.Exponent > 0 { + p += int64(dividend.Exponent) } - if count == 0 { - return nil, nil + p += avgGuardDigits + result := new(apd.Decimal) + if _, err := sql.DecimalCtx.WithPrecision(uint32(p)).Quo(result, dividend, divisor); err != nil { + return nil, err } - return avgFromInt64Sum(sum, count) + result.Reduce(result) + return result, nil } -// Dispose implements the sql.WindowFunction interface. -func (w *avgInt16WindowFunction) Dispose(ctx *sql.Context) {} - -// avgInt32 represents the PostgreSQL avg(int4) function, which promotes to numeric. -var avgInt32 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int32, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil +// initAvgAggs registers the functions to the catalog. See the comment on initNumericAggs for why avg needs a +// separate overload per input type rather than one generic/numeric-ish overload. +func initAvgAggs() { + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Int16, pgtypes.Numeric, newIntAvgBuffer[int16], newIntAvgWindowFunction[int16])) + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Int32, pgtypes.Numeric, newIntAvgBuffer[int32], newIntAvgWindowFunction[int32])) + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Int64, pgtypes.Numeric, newDecimalAvgBuffer(int64ToDecimal), newDecimalAvgWindowFunction(int64ToDecimal))) + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Numeric, pgtypes.Numeric, newDecimalAvgBuffer(decimalIdentity), newDecimalAvgWindowFunction(decimalIdentity))) + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Float32, pgtypes.Float64, newFloatAvgBuffer[float32], newFloatAvgWindowFunction[float32])) + framework.RegisterAggregateFunction(avgOverload("avg", pgtypes.Float64, pgtypes.Float64, newFloatAvgBuffer[float64], newFloatAvgWindowFunction[float64])) +} + +// avgOverload builds a single avg(...) overload; see sumOverload, which this mirrors. +func avgOverload(name string, paramType, returnType *pgtypes.DoltgresType, newBuffer framework.NewBufferFn, newWindowFunc framework.NewWindowFunctionFn) framework.Func1Aggregate { + return framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: name, + Return: returnType, + Parameters: [1]*pgtypes.DoltgresType{ + paramType, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, }, - }, - NewAggBuffer: newAvgInt32Buffer, - NewAggWindowFunc: newAvgInt32WindowFunction, + NewAggBuffer: newBuffer, + NewAggWindowFunc: newWindowFunc, + } } -// avgInt32Buffer is an aggregation buffer for avg(int4). -type avgInt32Buffer struct { +// intAvgBuffer is the GROUP BY buffer for avg(int2)/avg(int4), both of which promote to numeric. The running +// sum fits safely in an int64 accumulator; only the final division (in Eval) needs decimal arithmetic. +type intAvgBuffer[T int16 | int32] struct { expr sql.Expression sum int64 count int64 } -var _ sql.AggregationBuffer = (*avgInt32Buffer)(nil) +var _ sql.AggregationBuffer = (*intAvgBuffer[int32])(nil) -// newAvgInt32Buffer creates a new aggregation buffer for the avg(int4) aggregate function. -func newAvgInt32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgInt32Buffer{expr: exprs[0]}, nil +func newIntAvgBuffer[T int16 | int32](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &intAvgBuffer[T]{expr: exprs[0]}, nil } -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgInt32Buffer) Dispose(ctx *sql.Context) {} +func (b *intAvgBuffer[T]) Dispose(ctx *sql.Context) {} -// Eval implements the sql.AggregationBuffer interface. -func (b *avgInt32Buffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *intAvgBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if b.count == 0 { return nil, nil } - return avgFromInt64Sum(b.sum, b.count) + return quoAvg(apd.New(b.sum, 0), apd.New(b.count, 0)) } -// Update implements the sql.AggregationBuffer interface. -func (b *avgInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *intAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -215,42 +112,25 @@ func (b *avgInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - i, ok := v.(int32) + i, ok := v.(T) if !ok { - return errors.Errorf("avg: expected int32, got %T", v) + return errors.Errorf("avg: expected %T, got %T", i, v) } b.sum += int64(i) b.count++ return nil } -// avgFromInt64Sum returns the numeric average of sum/count. -func avgFromInt64Sum(sum, count int64) (*apd.Decimal, error) { - return quoAvg(apd.New(sum, 0), apd.New(count, 0)) -} - -// quoAvg divides dividend/divisor for use as an AVG result, reduced to strip the trailing zeros that Quo -// pads on to fill avgDecimalCtx's precision (e.g. 30/2 comes back as 15.00000...0 with 38 zeros, not 15). -func quoAvg(dividend, divisor *apd.Decimal) (*apd.Decimal, error) { - result := new(apd.Decimal) - if _, err := avgDecimalCtx.Quo(result, dividend, divisor); err != nil { - return nil, err - } - result.Reduce(result) - return result, nil -} - -// avgInt32WindowFunction is the sql.WindowFunction used for avg(int4) within an OVER(...) clause. -type avgInt32WindowFunction struct { +// intAvgWindowFunction is the sql.WindowFunction used for avg(int2)/avg(int4) within an OVER(...) clause. +type intAvgWindowFunction[T int16 | int32] struct { expr sql.Expression framer sql.WindowFramer } -var _ sql.WindowFunction = (*avgInt32WindowFunction)(nil) +var _ sql.WindowFunction = (*intAvgWindowFunction[int32])(nil) -// newAvgInt32WindowFunction creates the sql.WindowFunction for avg(int4). -func newAvgInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgInt32WindowFunction{expr: exprs[0]} +func newIntAvgWindowFunction[T int16 | int32](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &intAvgWindowFunction[T]{expr: exprs[0]} if window != nil && window.Frame != nil { framer, err := window.Frame.NewFramer(window) if err != nil { @@ -261,21 +141,18 @@ func newAvgInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinit return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *avgInt32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *intAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgInt32WindowFunction) DefaultFramer() sql.WindowFramer { +func (w *intAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *avgInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *intAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } @@ -288,9 +165,9 @@ func (w *avgInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowIn if v == nil { continue } - iv, ok := v.(int32) + iv, ok := v.(T) if !ok { - return nil, errors.Errorf("avg: expected int32, got %T", v) + return nil, errors.Errorf("avg: expected %T, got %T", iv, v) } sum += int64(iv) count++ @@ -298,181 +175,38 @@ func (w *avgInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowIn if count == 0 { return nil, nil } - return avgFromInt64Sum(sum, count) -} - -// Dispose implements the sql.WindowFunction interface. -func (w *avgInt32WindowFunction) Dispose(ctx *sql.Context) {} - -// avgInt64 represents the PostgreSQL avg(int8) function, which promotes to numeric. -var avgInt64 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int64, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newAvgInt64Buffer, - NewAggWindowFunc: newAvgInt64WindowFunction, -} - -// avgInt64Buffer is an aggregation buffer for avg(int8). -type avgInt64Buffer struct { - expr sql.Expression - sum apd.Decimal - count int64 -} - -var _ sql.AggregationBuffer = (*avgInt64Buffer)(nil) - -// newAvgInt64Buffer creates a new aggregation buffer for the avg(int8) aggregate function. -func newAvgInt64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgInt64Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgInt64Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *avgInt64Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if b.count == 0 { - return nil, nil - } - return quoAvg(&b.sum, apd.New(b.count, 0)) -} - -// Update implements the sql.AggregationBuffer interface. -func (b *avgInt64Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - i, ok := v.(int64) - if !ok { - return errors.Errorf("avg: expected int64, got %T", v) - } - _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, apd.New(i, 0)) - b.count++ - return err -} - -// avgInt64WindowFunction is the sql.WindowFunction used for avg(int8) within an OVER(...) clause. -type avgInt64WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer + return quoAvg(apd.New(sum, 0), apd.New(count, 0)) } -var _ sql.WindowFunction = (*avgInt64WindowFunction)(nil) - -// newAvgInt64WindowFunction creates the sql.WindowFunction for avg(int8). -func newAvgInt64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgInt64WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} +func (w *intAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} -// StartPartition implements the sql.WindowFunction interface. -func (w *avgInt64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil +// decimalAvgBuffer is the GROUP BY buffer for avg(int8)/avg(numeric), both of which stay numeric. convert +// adapts the buffer to either input type, same as decimalSumBuffer. +type decimalAvgBuffer[T int64 | *apd.Decimal] struct { + expr sql.Expression + sum apd.Decimal + count int64 + convert func(T) *apd.Decimal } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgInt64WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} +var _ sql.AggregationBuffer = (*decimalAvgBuffer[int64])(nil) -// Compute implements the sql.WindowFunction interface. -func (w *avgInt64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil +func newDecimalAvgBuffer[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewBufferFn { + return func(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &decimalAvgBuffer[T]{expr: exprs[0], convert: convert}, nil } - var sum apd.Decimal - var count int64 - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - iv, ok := v.(int64) - if !ok { - return nil, errors.Errorf("avg: expected int64, got %T", v) - } - if _, err = sql.DecimalCtx.Add(&sum, &sum, apd.New(iv, 0)); err != nil { - return nil, err - } - count++ - } - if count == 0 { - return nil, nil - } - return quoAvg(&sum, apd.New(count, 0)) } -// Dispose implements the sql.WindowFunction interface. -func (w *avgInt64WindowFunction) Dispose(ctx *sql.Context) {} +func (b *decimalAvgBuffer[T]) Dispose(ctx *sql.Context) {} -// avgNumeric represents the PostgreSQL avg(numeric) function, which stays numeric. -var avgNumeric = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Numeric, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newAvgNumericBuffer, - NewAggWindowFunc: newAvgNumericWindowFunction, -} - -// avgNumericBuffer is an aggregation buffer for avg(numeric). -type avgNumericBuffer struct { - expr sql.Expression - sum apd.Decimal - count int64 -} - -var _ sql.AggregationBuffer = (*avgNumericBuffer)(nil) - -// newAvgNumericBuffer creates a new aggregation buffer for the avg(numeric) aggregate function. -func newAvgNumericBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgNumericBuffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgNumericBuffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *avgNumericBuffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *decimalAvgBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if b.count == 0 { return nil, nil } return quoAvg(&b.sum, apd.New(b.count, 0)) } -// Update implements the sql.AggregationBuffer interface. -func (b *avgNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *decimalAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -480,51 +214,51 @@ func (b *avgNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - d, ok := v.(*apd.Decimal) + typedV, ok := v.(T) if !ok { - return errors.Errorf("avg: expected *apd.Decimal, got %T", v) + return errors.Errorf("avg: expected %T, got %T", typedV, v) } - _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, d) + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, b.convert(typedV)) b.count++ return err } -// avgNumericWindowFunction is the sql.WindowFunction used for avg(numeric) within an OVER(...) clause. -type avgNumericWindowFunction struct { - expr sql.Expression - framer sql.WindowFramer +// decimalAvgWindowFunction is the sql.WindowFunction used for avg(int8)/avg(numeric) within an OVER(...) +// clause. +type decimalAvgWindowFunction[T int64 | *apd.Decimal] struct { + expr sql.Expression + framer sql.WindowFramer + convert func(T) *apd.Decimal } -var _ sql.WindowFunction = (*avgNumericWindowFunction)(nil) +var _ sql.WindowFunction = (*decimalAvgWindowFunction[int64])(nil) -// newAvgNumericWindowFunction creates the sql.WindowFunction for avg(numeric). -func newAvgNumericWindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgNumericWindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err +func newDecimalAvgWindowFunction[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewWindowFunctionFn { + return func(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &decimalAvgWindowFunction[T]{expr: exprs[0], convert: convert} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer } - wf.framer = framer + return wf, nil } - return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *avgNumericWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *decimalAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgNumericWindowFunction) DefaultFramer() sql.WindowFramer { +func (w *decimalAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *avgNumericWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *decimalAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } @@ -538,11 +272,11 @@ func (w *avgNumericWindowFunction) Compute(ctx *sql.Context, interval sql.Window if v == nil { continue } - d, ok := v.(*apd.Decimal) + typedV, ok := v.(T) if !ok { - return nil, errors.Errorf("avg: expected *apd.Decimal, got %T", v) + return nil, errors.Errorf("avg: expected %T, got %T", typedV, v) } - if _, err = sql.DecimalCtx.Add(&sum, &sum, d); err != nil { + if _, err = sql.DecimalCtx.Add(&sum, &sum, w.convert(typedV)); err != nil { return nil, err } count++ @@ -553,52 +287,32 @@ func (w *avgNumericWindowFunction) Compute(ctx *sql.Context, interval sql.Window return quoAvg(&sum, apd.New(count, 0)) } -// Dispose implements the sql.WindowFunction interface. -func (w *avgNumericWindowFunction) Dispose(ctx *sql.Context) {} +func (w *decimalAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} -// avgFloat32 represents the PostgreSQL avg(float4) function, which promotes to double precision. -var avgFloat32 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Float64, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Float32, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newAvgFloat32Buffer, - NewAggWindowFunc: newAvgFloat32WindowFunction, -} - -// avgFloat32Buffer is an aggregation buffer for avg(float4). -type avgFloat32Buffer struct { +// floatAvgBuffer is the GROUP BY buffer for avg(float4)/avg(float8), both of which promote to double +// precision (unlike sum's float overloads, which preserve their input type). +type floatAvgBuffer[T float32 | float64] struct { expr sql.Expression sum float64 count int64 } -var _ sql.AggregationBuffer = (*avgFloat32Buffer)(nil) +var _ sql.AggregationBuffer = (*floatAvgBuffer[float64])(nil) -// newAvgFloat32Buffer creates a new aggregation buffer for the avg(float4) aggregate function. -func newAvgFloat32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgFloat32Buffer{expr: exprs[0]}, nil +func newFloatAvgBuffer[T float32 | float64](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &floatAvgBuffer[T]{expr: exprs[0]}, nil } -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgFloat32Buffer) Dispose(ctx *sql.Context) {} +func (b *floatAvgBuffer[T]) Dispose(ctx *sql.Context) {} -// Eval implements the sql.AggregationBuffer interface. -func (b *avgFloat32Buffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *floatAvgBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if b.count == 0 { return nil, nil } return b.sum / float64(b.count), nil } -// Update implements the sql.AggregationBuffer interface. -func (b *avgFloat32Buffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *floatAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -606,26 +320,26 @@ func (b *avgFloat32Buffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - f, ok := v.(float32) + f, ok := v.(T) if !ok { - return errors.Errorf("avg: expected float32, got %T", v) + return errors.Errorf("avg: expected %T, got %T", f, v) } b.sum += float64(f) b.count++ return nil } -// avgFloat32WindowFunction is the sql.WindowFunction used for avg(float4) within an OVER(...) clause. -type avgFloat32WindowFunction struct { +// floatAvgWindowFunction is the sql.WindowFunction used for avg(float4)/avg(float8) within an OVER(...) +// clause. +type floatAvgWindowFunction[T float32 | float64] struct { expr sql.Expression framer sql.WindowFramer } -var _ sql.WindowFunction = (*avgFloat32WindowFunction)(nil) +var _ sql.WindowFunction = (*floatAvgWindowFunction[float64])(nil) -// newAvgFloat32WindowFunction creates the sql.WindowFunction for avg(float4). -func newAvgFloat32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgFloat32WindowFunction{expr: exprs[0]} +func newFloatAvgWindowFunction[T float32 | float64](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &floatAvgWindowFunction[T]{expr: exprs[0]} if window != nil && window.Frame != nil { framer, err := window.Frame.NewFramer(window) if err != nil { @@ -636,21 +350,18 @@ func newAvgFloat32WindowFunction(exprs []sql.Expression, window *sql.WindowDefin return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *avgFloat32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *floatAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgFloat32WindowFunction) DefaultFramer() sql.WindowFramer { +func (w *floatAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *avgFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *floatAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } @@ -664,9 +375,9 @@ func (w *avgFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.Window if v == nil { continue } - fv, ok := v.(float32) + fv, ok := v.(T) if !ok { - return nil, errors.Errorf("avg: expected float32, got %T", v) + return nil, errors.Errorf("avg: expected %T, got %T", fv, v) } sum += float64(fv) count++ @@ -677,129 +388,4 @@ func (w *avgFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.Window return sum / float64(count), nil } -// Dispose implements the sql.WindowFunction interface. -func (w *avgFloat32WindowFunction) Dispose(ctx *sql.Context) {} - -// avgFloat64 represents the PostgreSQL avg(float8) function, which stays double precision. -var avgFloat64 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "avg", - Return: pgtypes.Float64, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Float64, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newAvgFloat64Buffer, - NewAggWindowFunc: newAvgFloat64WindowFunction, -} - -// avgFloat64Buffer is an aggregation buffer for avg(float8). -type avgFloat64Buffer struct { - expr sql.Expression - sum float64 - count int64 -} - -var _ sql.AggregationBuffer = (*avgFloat64Buffer)(nil) - -// newAvgFloat64Buffer creates a new aggregation buffer for the avg(float8) aggregate function. -func newAvgFloat64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &avgFloat64Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *avgFloat64Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *avgFloat64Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if b.count == 0 { - return nil, nil - } - return b.sum / float64(b.count), nil -} - -// Update implements the sql.AggregationBuffer interface. -func (b *avgFloat64Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - f, ok := v.(float64) - if !ok { - return errors.Errorf("avg: expected float64, got %T", v) - } - b.sum += f - b.count++ - return nil -} - -// avgFloat64WindowFunction is the sql.WindowFunction used for avg(float8) within an OVER(...) clause. -type avgFloat64WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer -} - -var _ sql.WindowFunction = (*avgFloat64WindowFunction)(nil) - -// newAvgFloat64WindowFunction creates the sql.WindowFunction for avg(float8). -func newAvgFloat64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &avgFloat64WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} - -// StartPartition implements the sql.WindowFunction interface. -func (w *avgFloat64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -// DefaultFramer implements the sql.WindowFunction interface. -func (w *avgFloat64WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - -// Compute implements the sql.WindowFunction interface. -func (w *avgFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil - } - var sum float64 - var count int64 - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - fv, ok := v.(float64) - if !ok { - return nil, errors.Errorf("avg: expected float64, got %T", v) - } - sum += fv - count++ - } - if count == 0 { - return nil, nil - } - return sum / float64(count), nil -} - -// Dispose implements the sql.WindowFunction interface. -func (w *avgFloat64WindowFunction) Dispose(ctx *sql.Context) {} +func (w *floatAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} diff --git a/server/functions/aggregate/numeric_aggregates.go b/server/functions/aggregate/numeric_aggregates.go index 3c4ea639ff..f551799092 100644 --- a/server/functions/aggregate/numeric_aggregates.go +++ b/server/functions/aggregate/numeric_aggregates.go @@ -24,191 +24,76 @@ import ( pgtypes "github.com/dolthub/doltgresql/server/types" ) -// initNumericAggs registers the functions to the catalog. -func initNumericAggs() { - framework.RegisterAggregateFunction(sumInt16) - framework.RegisterAggregateFunction(sumInt32) - framework.RegisterAggregateFunction(sumInt64) - framework.RegisterAggregateFunction(sumNumeric) - framework.RegisterAggregateFunction(sumFloat32) - framework.RegisterAggregateFunction(sumFloat64) -} +// int64ToDecimal converts an int64 sum to *apd.Decimal, for use as the decimalConvert of a decimalSumBuffer/ +// decimalAvgBuffer instantiated over int64 (i.e. sum(int8)/avg(int8), whose accumulator needs to be decimal +// since a running sum of bigints can itself overflow a bigint). +func int64ToDecimal(v int64) *apd.Decimal { return apd.New(v, 0) } -// sumInt16 represents the PostgreSQL sum(int2) function, which promotes to bigint. +// decimalIdentity is the decimalConvert for a decimalSumBuffer/decimalAvgBuffer instantiated over +// *apd.Decimal (i.e. sum(numeric)/avg(numeric)), whose input values are already decimal. +func decimalIdentity(v *apd.Decimal) *apd.Decimal { return v } + +// initNumericAggs registers the functions to the catalog. // // Note that overload resolution for aggregates does not insert implicit widening casts the way it does for // scalar functions (CompiledAggregateFunction hands NewBuffer/NewWindowFunc the raw, uncast Arguments) - so // every distinct argument type sum/avg can be called on needs its own overload registered below; there's no -// way to cover e.g. both int2 and int4 with a single numeric-ish overload. -var sumInt16 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Int64, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int16, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil +// way to cover e.g. both int2 and int4 with a single numeric-ish overload. This mirrors how other multi-type +// functions in this package (e.g. abs.go) already register one overload per Postgres type. +func initNumericAggs() { + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Int16, pgtypes.Int64, newIntSumBuffer[int16], newIntSumWindowFunction[int16])) + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Int32, pgtypes.Int64, newIntSumBuffer[int32], newIntSumWindowFunction[int32])) + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Int64, pgtypes.Numeric, newDecimalSumBuffer(int64ToDecimal), newDecimalSumWindowFunction(int64ToDecimal))) + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Numeric, pgtypes.Numeric, newDecimalSumBuffer(decimalIdentity), newDecimalSumWindowFunction(decimalIdentity))) + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Float32, pgtypes.Float32, newFloatSumBuffer[float32], newFloatSumWindowFunction[float32])) + framework.RegisterAggregateFunction(sumOverload("sum", pgtypes.Float64, pgtypes.Float64, newFloatSumBuffer[float64], newFloatSumWindowFunction[float64])) +} + +// sumOverload builds a single sum(...) overload: paramType is the Postgres type of the aggregated column, +// returnType is sum's result type for that input (e.g. sum(int4) promotes to bigint), and newBuffer/ +// newWindowFunc construct the GROUP BY and OVER(...) implementations respectively. +func sumOverload(name string, paramType, returnType *pgtypes.DoltgresType, newBuffer framework.NewBufferFn, newWindowFunc framework.NewWindowFunctionFn) framework.Func1Aggregate { + return framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: name, + Return: returnType, + Parameters: [1]*pgtypes.DoltgresType{ + paramType, + }, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { + return nil, nil + }, }, - }, - NewAggBuffer: newSumInt16Buffer, - NewAggWindowFunc: newSumInt16WindowFunction, -} - -// sumInt16Buffer is an aggregation buffer for sum(int2). -type sumInt16Buffer struct { - expr sql.Expression - sum int64 - sawOne bool -} - -var _ sql.AggregationBuffer = (*sumInt16Buffer)(nil) - -// newSumInt16Buffer creates a new aggregation buffer for the sum(int2) aggregate function. -func newSumInt16Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumInt16Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumInt16Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *sumInt16Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if !b.sawOne { - return nil, nil - } - return b.sum, nil -} - -// Update implements the sql.AggregationBuffer interface. -func (b *sumInt16Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - i, ok := v.(int16) - if !ok { - return errors.Errorf("sum: expected int16, got %T", v) - } - b.sum += int64(i) - b.sawOne = true - return nil -} - -// sumInt16WindowFunction is the sql.WindowFunction used for sum(int2) within an OVER(...) clause. -type sumInt16WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer -} - -var _ sql.WindowFunction = (*sumInt16WindowFunction)(nil) - -// newSumInt16WindowFunction creates the sql.WindowFunction for sum(int2). -func newSumInt16WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumInt16WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} - -// StartPartition implements the sql.WindowFunction interface. -func (w *sumInt16WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumInt16WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - -// Compute implements the sql.WindowFunction interface. -func (w *sumInt16WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil - } - var sum int64 - var sawOne bool - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - iv, ok := v.(int16) - if !ok { - return nil, errors.Errorf("sum: expected int16, got %T", v) - } - sum += int64(iv) - sawOne = true + NewAggBuffer: newBuffer, + NewAggWindowFunc: newWindowFunc, } - if !sawOne { - return nil, nil - } - return sum, nil } -// Dispose implements the sql.WindowFunction interface. -func (w *sumInt16WindowFunction) Dispose(ctx *sql.Context) {} - -// sumInt32 represents the PostgreSQL sum(int4) function, which promotes to bigint. It supports both the -// GROUP BY buffer path and use as a window function within an OVER(...) clause, so that a query like -// `SUM(x) OVER (...)` returns a correctly-typed Postgres bigint directly, without any post-hoc type -// correction elsewhere in the analyzer. -var sumInt32 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Int64, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int32, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newSumInt32Buffer, - NewAggWindowFunc: newSumInt32WindowFunction, -} - -// sumInt32Buffer is an aggregation buffer for sum(int4). -type sumInt32Buffer struct { +// intSumBuffer is the GROUP BY buffer for sum(int2)/sum(int4), which both promote to bigint. Their sum fits +// safely in an int64 accumulator (the widest int2/int4 value is nowhere near int64's range), so this is a +// simple integer sum, unlike sum(int8)/sum(numeric) (see decimalSumBuffer) which need decimal accumulation. +type intSumBuffer[T int16 | int32] struct { expr sql.Expression sum int64 sawOne bool } -var _ sql.AggregationBuffer = (*sumInt32Buffer)(nil) +var _ sql.AggregationBuffer = (*intSumBuffer[int32])(nil) -// newSumInt32Buffer creates a new aggregation buffer for the sum(int4) aggregate function. -func newSumInt32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumInt32Buffer{expr: exprs[0]}, nil +func newIntSumBuffer[T int16 | int32](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &intSumBuffer[T]{expr: exprs[0]}, nil } -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumInt32Buffer) Dispose(ctx *sql.Context) {} +func (b *intSumBuffer[T]) Dispose(ctx *sql.Context) {} -// Eval implements the sql.AggregationBuffer interface. -func (b *sumInt32Buffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *intSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if !b.sawOne { return nil, nil } return b.sum, nil } -// Update implements the sql.AggregationBuffer interface. -func (b *sumInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *intSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -216,28 +101,25 @@ func (b *sumInt32Buffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - i, ok := v.(int32) + i, ok := v.(T) if !ok { - return errors.Errorf("sum: expected int32, got %T", v) + return errors.Errorf("sum: expected %T, got %T", i, v) } b.sum += int64(i) b.sawOne = true return nil } -// sumInt32WindowFunction is the sql.WindowFunction used for sum(int4) within an OVER(...) clause. -type sumInt32WindowFunction struct { +// intSumWindowFunction is the sql.WindowFunction used for sum(int2)/sum(int4) within an OVER(...) clause. +type intSumWindowFunction[T int16 | int32] struct { expr sql.Expression framer sql.WindowFramer } -var _ sql.WindowFunction = (*sumInt32WindowFunction)(nil) +var _ sql.WindowFunction = (*intSumWindowFunction[int32])(nil) -// newSumInt32WindowFunction creates the sql.WindowFunction for sum(int4), binding the frame declared on -// window (if any); with no explicit frame, DefaultFramer supplies Postgres's default (unbounded preceding to -// current row). -func newSumInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumInt32WindowFunction{expr: exprs[0]} +func newIntSumWindowFunction[T int16 | int32](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &intSumWindowFunction[T]{expr: exprs[0]} if window != nil && window.Frame != nil { framer, err := window.Frame.NewFramer(window) if err != nil { @@ -248,21 +130,20 @@ func newSumInt32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinit return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *sumInt32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *intSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumInt32WindowFunction) DefaultFramer() sql.WindowFramer { +// DefaultFramer implements the sql.WindowFunction interface; with no explicit frame, this supplies +// Postgres's default (unbounded preceding to current row). +func (w *intSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *sumInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *intSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } @@ -276,9 +157,9 @@ func (w *sumInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowIn if v == nil { continue } - iv, ok := v.(int32) + iv, ok := v.(T) if !ok { - return nil, errors.Errorf("sum: expected int32, got %T", v) + return nil, errors.Errorf("sum: expected %T, got %T", iv, v) } sum += int64(iv) sawOne = true @@ -289,173 +170,30 @@ func (w *sumInt32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowIn return sum, nil } -// Dispose implements the sql.WindowFunction interface. -func (w *sumInt32WindowFunction) Dispose(ctx *sql.Context) {} - -// sumInt64 represents the PostgreSQL sum(int8) function. A running sum of bigints can itself overflow a -// bigint, so Postgres promotes this to numeric rather than bigint. -var sumInt64 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Int64, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newSumInt64Buffer, - NewAggWindowFunc: newSumInt64WindowFunction, -} - -// sumInt64Buffer is an aggregation buffer for sum(int8). -type sumInt64Buffer struct { - expr sql.Expression - sum apd.Decimal - sawOne bool -} - -var _ sql.AggregationBuffer = (*sumInt64Buffer)(nil) - -// newSumInt64Buffer creates a new aggregation buffer for the sum(int8) aggregate function. -func newSumInt64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumInt64Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumInt64Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *sumInt64Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if !b.sawOne { - return nil, nil - } - result := b.sum - return &result, nil -} - -// Update implements the sql.AggregationBuffer interface. -func (b *sumInt64Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - i, ok := v.(int64) - if !ok { - return errors.Errorf("sum: expected int64, got %T", v) - } - _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, apd.New(i, 0)) - b.sawOne = true - return err -} - -// sumInt64WindowFunction is the sql.WindowFunction used for sum(int8) within an OVER(...) clause. -type sumInt64WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer -} - -var _ sql.WindowFunction = (*sumInt64WindowFunction)(nil) - -// newSumInt64WindowFunction creates the sql.WindowFunction for sum(int8). -func newSumInt64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumInt64WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} +func (w *intSumWindowFunction[T]) Dispose(ctx *sql.Context) {} -// StartPartition implements the sql.WindowFunction interface. -func (w *sumInt64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil +// decimalSumBuffer is the GROUP BY buffer for sum(int8)/sum(numeric). sum(int8) promotes to numeric because +// a running sum of bigints can itself overflow a bigint; sum(numeric) stays numeric. convert adapts the +// buffer to either input type: int64 values are boxed via int64ToDecimal, while numeric values (already +// *apd.Decimal) pass through decimalIdentity unchanged. +type decimalSumBuffer[T int64 | *apd.Decimal] struct { + expr sql.Expression + sum apd.Decimal + sawOne bool + convert func(T) *apd.Decimal } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumInt64WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} +var _ sql.AggregationBuffer = (*decimalSumBuffer[int64])(nil) -// Compute implements the sql.WindowFunction interface. -func (w *sumInt64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil +func newDecimalSumBuffer[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewBufferFn { + return func(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &decimalSumBuffer[T]{expr: exprs[0], convert: convert}, nil } - var sum apd.Decimal - var sawOne bool - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - iv, ok := v.(int64) - if !ok { - return nil, errors.Errorf("sum: expected int64, got %T", v) - } - if _, err = sql.DecimalCtx.Add(&sum, &sum, apd.New(iv, 0)); err != nil { - return nil, err - } - sawOne = true - } - if !sawOne { - return nil, nil - } - return &sum, nil -} - -// Dispose implements the sql.WindowFunction interface. -func (w *sumInt64WindowFunction) Dispose(ctx *sql.Context) {} - -// sumNumeric represents the PostgreSQL sum(numeric) function, which stays numeric. It supports both the -// GROUP BY buffer path and use as a window function within an OVER(...) clause. -var sumNumeric = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Numeric, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Numeric, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newSumNumericBuffer, - NewAggWindowFunc: newSumNumericWindowFunction, } -// sumNumericBuffer is an aggregation buffer for sum(numeric). -type sumNumericBuffer struct { - expr sql.Expression - sum apd.Decimal - sawOne bool -} +func (b *decimalSumBuffer[T]) Dispose(ctx *sql.Context) {} -var _ sql.AggregationBuffer = (*sumNumericBuffer)(nil) - -// newSumNumericBuffer creates a new aggregation buffer for the sum(numeric) aggregate function. -func newSumNumericBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumNumericBuffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumNumericBuffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *sumNumericBuffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *decimalSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if !b.sawOne { return nil, nil } @@ -463,8 +201,7 @@ func (b *sumNumericBuffer) Eval(ctx *sql.Context) (interface{}, error) { return &result, nil } -// Update implements the sql.AggregationBuffer interface. -func (b *sumNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *decimalSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -472,10 +209,11 @@ func (b *sumNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - d, ok := v.(*apd.Decimal) + typedV, ok := v.(T) if !ok { - return errors.Errorf("sum: expected *apd.Decimal, got %T", v) + return errors.Errorf("sum: expected %T, got %T", typedV, v) } + d := b.convert(typedV) if !b.sawOne { b.sum.Set(d) b.sawOne = true @@ -485,42 +223,42 @@ func (b *sumNumericBuffer) Update(ctx *sql.Context, row sql.Row) error { return err } -// sumNumericWindowFunction is the sql.WindowFunction used for sum(numeric) within an OVER(...) clause. -type sumNumericWindowFunction struct { - expr sql.Expression - framer sql.WindowFramer +// decimalSumWindowFunction is the sql.WindowFunction used for sum(int8)/sum(numeric) within an OVER(...) +// clause. +type decimalSumWindowFunction[T int64 | *apd.Decimal] struct { + expr sql.Expression + framer sql.WindowFramer + convert func(T) *apd.Decimal } -var _ sql.WindowFunction = (*sumNumericWindowFunction)(nil) +var _ sql.WindowFunction = (*decimalSumWindowFunction[int64])(nil) -// newSumNumericWindowFunction creates the sql.WindowFunction for sum(numeric). -func newSumNumericWindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumNumericWindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err +func newDecimalSumWindowFunction[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewWindowFunctionFn { + return func(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &decimalSumWindowFunction[T]{expr: exprs[0], convert: convert} + if window != nil && window.Frame != nil { + framer, err := window.Frame.NewFramer(window) + if err != nil { + return nil, err + } + wf.framer = framer } - wf.framer = framer + return wf, nil } - return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *sumNumericWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *decimalSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumNumericWindowFunction) DefaultFramer() sql.WindowFramer { +func (w *decimalSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *sumNumericWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *decimalSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } @@ -534,10 +272,11 @@ func (w *sumNumericWindowFunction) Compute(ctx *sql.Context, interval sql.Window if v == nil { continue } - d, ok := v.(*apd.Decimal) + typedV, ok := v.(T) if !ok { - return nil, errors.Errorf("sum: expected *apd.Decimal, got %T", v) + return nil, errors.Errorf("sum: expected %T, got %T", typedV, v) } + d := w.convert(typedV) if !sawOne { sum.Set(d) sawOne = true @@ -553,176 +292,32 @@ func (w *sumNumericWindowFunction) Compute(ctx *sql.Context, interval sql.Window return &sum, nil } -// Dispose implements the sql.WindowFunction interface. -func (w *sumNumericWindowFunction) Dispose(ctx *sql.Context) {} - -// sumFloat32 represents the PostgreSQL sum(float4) function, which stays real (float32). -var sumFloat32 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Float32, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Float32, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newSumFloat32Buffer, - NewAggWindowFunc: newSumFloat32WindowFunction, -} - -// sumFloat32Buffer is an aggregation buffer for sum(float4). -type sumFloat32Buffer struct { - expr sql.Expression - sum float32 - sawOne bool -} - -var _ sql.AggregationBuffer = (*sumFloat32Buffer)(nil) - -// newSumFloat32Buffer creates a new aggregation buffer for the sum(float4) aggregate function. -func newSumFloat32Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumFloat32Buffer{expr: exprs[0]}, nil -} - -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumFloat32Buffer) Dispose(ctx *sql.Context) {} - -// Eval implements the sql.AggregationBuffer interface. -func (b *sumFloat32Buffer) Eval(ctx *sql.Context) (interface{}, error) { - if !b.sawOne { - return nil, nil - } - return b.sum, nil -} - -// Update implements the sql.AggregationBuffer interface. -func (b *sumFloat32Buffer) Update(ctx *sql.Context, row sql.Row) error { - v, err := b.expr.Eval(ctx, row) - if err != nil { - return err - } - if v == nil { - return nil - } - f, ok := v.(float32) - if !ok { - return errors.Errorf("sum: expected float32, got %T", v) - } - b.sum += f - b.sawOne = true - return nil -} - -// sumFloat32WindowFunction is the sql.WindowFunction used for sum(float4) within an OVER(...) clause. -type sumFloat32WindowFunction struct { - expr sql.Expression - framer sql.WindowFramer -} - -var _ sql.WindowFunction = (*sumFloat32WindowFunction)(nil) - -// newSumFloat32WindowFunction creates the sql.WindowFunction for sum(float4). -func newSumFloat32WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumFloat32WindowFunction{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer - } - return wf, nil -} - -// StartPartition implements the sql.WindowFunction interface. -func (w *sumFloat32WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumFloat32WindowFunction) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - -// Compute implements the sql.WindowFunction interface. -func (w *sumFloat32WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { - if interval.End <= interval.Start { - return nil, nil - } - var sum float32 - var sawOne bool - for i := interval.Start; i < interval.End; i++ { - v, err := w.expr.Eval(ctx, buf[i]) - if err != nil { - return nil, err - } - if v == nil { - continue - } - fv, ok := v.(float32) - if !ok { - return nil, errors.Errorf("sum: expected float32, got %T", v) - } - sum += fv - sawOne = true - } - if !sawOne { - return nil, nil - } - return sum, nil -} - -// Dispose implements the sql.WindowFunction interface. -func (w *sumFloat32WindowFunction) Dispose(ctx *sql.Context) {} - -// sumFloat64 represents the PostgreSQL sum(float8) function, which stays double precision (float64). -var sumFloat64 = framework.Func1Aggregate{ - Function1: framework.Function1{ - Name: "sum", - Return: pgtypes.Float64, - Parameters: [1]*pgtypes.DoltgresType{ - pgtypes.Float64, - }, - Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) { - return nil, nil - }, - }, - NewAggBuffer: newSumFloat64Buffer, - NewAggWindowFunc: newSumFloat64WindowFunction, -} +func (w *decimalSumWindowFunction[T]) Dispose(ctx *sql.Context) {} -// sumFloat64Buffer is an aggregation buffer for sum(float8). -type sumFloat64Buffer struct { +// floatSumBuffer is the GROUP BY buffer for sum(float4)/sum(float8), which (unlike the integer overloads) +// preserve their input type rather than promoting. +type floatSumBuffer[T float32 | float64] struct { expr sql.Expression - sum float64 + sum T sawOne bool } -var _ sql.AggregationBuffer = (*sumFloat64Buffer)(nil) +var _ sql.AggregationBuffer = (*floatSumBuffer[float64])(nil) -// newSumFloat64Buffer creates a new aggregation buffer for the sum(float8) aggregate function. -func newSumFloat64Buffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { - return &sumFloat64Buffer{expr: exprs[0]}, nil +func newFloatSumBuffer[T float32 | float64](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &floatSumBuffer[T]{expr: exprs[0]}, nil } -// Dispose implements the sql.AggregationBuffer interface. -func (b *sumFloat64Buffer) Dispose(ctx *sql.Context) {} +func (b *floatSumBuffer[T]) Dispose(ctx *sql.Context) {} -// Eval implements the sql.AggregationBuffer interface. -func (b *sumFloat64Buffer) Eval(ctx *sql.Context) (interface{}, error) { +func (b *floatSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { if !b.sawOne { return nil, nil } return b.sum, nil } -// Update implements the sql.AggregationBuffer interface. -func (b *sumFloat64Buffer) Update(ctx *sql.Context, row sql.Row) error { +func (b *floatSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { v, err := b.expr.Eval(ctx, row) if err != nil { return err @@ -730,26 +325,26 @@ func (b *sumFloat64Buffer) Update(ctx *sql.Context, row sql.Row) error { if v == nil { return nil } - f, ok := v.(float64) + f, ok := v.(T) if !ok { - return errors.Errorf("sum: expected float64, got %T", v) + return errors.Errorf("sum: expected %T, got %T", f, v) } b.sum += f b.sawOne = true return nil } -// sumFloat64WindowFunction is the sql.WindowFunction used for sum(float8) within an OVER(...) clause. -type sumFloat64WindowFunction struct { +// floatSumWindowFunction is the sql.WindowFunction used for sum(float4)/sum(float8) within an OVER(...) +// clause. +type floatSumWindowFunction[T float32 | float64] struct { expr sql.Expression framer sql.WindowFramer } -var _ sql.WindowFunction = (*sumFloat64WindowFunction)(nil) +var _ sql.WindowFunction = (*floatSumWindowFunction[float64])(nil) -// newSumFloat64WindowFunction creates the sql.WindowFunction for sum(float8). -func newSumFloat64WindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { - wf := &sumFloat64WindowFunction{expr: exprs[0]} +func newFloatSumWindowFunction[T float32 | float64](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &floatSumWindowFunction[T]{expr: exprs[0]} if window != nil && window.Frame != nil { framer, err := window.Frame.NewFramer(window) if err != nil { @@ -760,25 +355,22 @@ func newSumFloat64WindowFunction(exprs []sql.Expression, window *sql.WindowDefin return wf, nil } -// StartPartition implements the sql.WindowFunction interface. -func (w *sumFloat64WindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { +func (w *floatSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { return nil } -// DefaultFramer implements the sql.WindowFunction interface. -func (w *sumFloat64WindowFunction) DefaultFramer() sql.WindowFramer { +func (w *floatSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { if w.framer != nil { return w.framer } return aggregation.NewUnboundedPrecedingToCurrentRowFramer() } -// Compute implements the sql.WindowFunction interface. -func (w *sumFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { +func (w *floatSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil } - var sum float64 + var sum T var sawOne bool for i := interval.Start; i < interval.End; i++ { v, err := w.expr.Eval(ctx, buf[i]) @@ -788,9 +380,9 @@ func (w *sumFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.Window if v == nil { continue } - fv, ok := v.(float64) + fv, ok := v.(T) if !ok { - return nil, errors.Errorf("sum: expected float64, got %T", v) + return nil, errors.Errorf("sum: expected %T, got %T", fv, v) } sum += fv sawOne = true @@ -801,5 +393,4 @@ func (w *sumFloat64WindowFunction) Compute(ctx *sql.Context, interval sql.Window return sum, nil } -// Dispose implements the sql.WindowFunction interface. -func (w *sumFloat64WindowFunction) Dispose(ctx *sql.Context) {} +func (w *floatSumWindowFunction[T]) Dispose(ctx *sql.Context) {} From bd381c2601032094bb270de232e32ed891d37eff Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Fri, 10 Jul 2026 16:40:13 -0700 Subject: [PATCH 4/4] fixup: more simplifying --- postgres/parser/parser/sql.y | 2 +- server/analyzer/type_sanitizer.go | 12 +- server/functions/aggregate/avg_aggregates.go | 74 ++--------- .../functions/aggregate/numeric_aggregates.go | 115 ++++++++---------- server/functions/window/rank.go | 53 ++++++++ server/functions/window/row_number.go | 1 + testing/go/enginetest/doltgres_engine_test.go | 7 +- 7 files changed, 132 insertions(+), 132 deletions(-) diff --git a/postgres/parser/parser/sql.y b/postgres/parser/parser/sql.y index 680430a0f1..897ed5a4f0 100644 --- a/postgres/parser/parser/sql.y +++ b/postgres/parser/parser/sql.y @@ -13630,7 +13630,7 @@ over_clause: } | OVER window_name { - $$.val = &tree.WindowDef{Name: tree.Name($2)} + $$.val = &tree.WindowDef{RefName: tree.Name($2)} } | /* EMPTY */ { diff --git a/server/analyzer/type_sanitizer.go b/server/analyzer/type_sanitizer.go index 9f1b154fee..74af0d14b4 100644 --- a/server/analyzer/type_sanitizer.go +++ b/server/analyzer/type_sanitizer.go @@ -75,9 +75,17 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope case sql.FunctionExpression: // Compiled functions are Doltgres functions. We're only concerned with GMS functions. if _, ok := expr.(framework.Function); !ok { - // Some aggregation functions cannot be wrapped due to expectations in the analyzer, so we exclude them here. + // Aggregation/window-only expressions (Sum, Avg, Count, BitAnd, Rank, ...) can't be + // Eval()'d directly - only via NewBuffer/NewWindowFunction - so wrapping one in + // GMSCast (which evaluates its child directly) breaks it. + // sql.WindowAdaptableExpression is the common parent interface for both sql.Aggregation + // and sql.WindowAggregation, so checking it covers every current and future case in one + // shot. Only the *outer* reference to an aggregate's result (a GetField, handled + // elsewhere in this function) still needs its declared type corrected. + if _, ok := expr.(sql.WindowAdaptableExpression); ok { + return expr, transform.SameTree, nil + } switch expr.FunctionName() { - case "Count", "CountDistinct", "group_concat", "JSONObjectAgg", "Sum", "Avg": case "coalesce": // Replace GMS Coalesce with a Doltgres-native implementation that uses // Postgres type-resolution rules (FindCommonType) to infer the result type. diff --git a/server/functions/aggregate/avg_aggregates.go b/server/functions/aggregate/avg_aggregates.go index 6b0f3dc62e..458aa6e10d 100644 --- a/server/functions/aggregate/avg_aggregates.go +++ b/server/functions/aggregate/avg_aggregates.go @@ -18,7 +18,6 @@ import ( "github.com/cockroachdb/apd/v3" "github.com/cockroachdb/errors" "github.com/dolthub/go-mysql-server/sql" - "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" "github.com/dolthub/doltgresql/server/functions/framework" pgtypes "github.com/dolthub/doltgresql/server/types" @@ -123,35 +122,20 @@ func (b *intAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // intAvgWindowFunction is the sql.WindowFunction used for avg(int2)/avg(int4) within an OVER(...) clause. type intAvgWindowFunction[T int16 | int32] struct { - expr sql.Expression - framer sql.WindowFramer + windowFramerState + expr sql.Expression } var _ sql.WindowFunction = (*intAvgWindowFunction[int32])(nil) func newIntAvgWindowFunction[T int16 | int32](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &intAvgWindowFunction[T]{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } -func (w *intAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -func (w *intAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *intAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -178,8 +162,6 @@ func (w *intAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowI return quoAvg(apd.New(sum, 0), apd.New(count, 0)) } -func (w *intAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} - // decimalAvgBuffer is the GROUP BY buffer for avg(int8)/avg(numeric), both of which stay numeric. convert // adapts the buffer to either input type, same as decimalSumBuffer. type decimalAvgBuffer[T int64 | *apd.Decimal] struct { @@ -226,8 +208,8 @@ func (b *decimalAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // decimalAvgWindowFunction is the sql.WindowFunction used for avg(int8)/avg(numeric) within an OVER(...) // clause. type decimalAvgWindowFunction[T int64 | *apd.Decimal] struct { + windowFramerState expr sql.Expression - framer sql.WindowFramer convert func(T) *apd.Decimal } @@ -236,28 +218,13 @@ var _ sql.WindowFunction = (*decimalAvgWindowFunction[int64])(nil) func newDecimalAvgWindowFunction[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewWindowFunctionFn { return func(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &decimalAvgWindowFunction[T]{expr: exprs[0], convert: convert} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } } -func (w *decimalAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -func (w *decimalAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *decimalAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -287,8 +254,6 @@ func (w *decimalAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.Win return quoAvg(&sum, apd.New(count, 0)) } -func (w *decimalAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} - // floatAvgBuffer is the GROUP BY buffer for avg(float4)/avg(float8), both of which promote to double // precision (unlike sum's float overloads, which preserve their input type). type floatAvgBuffer[T float32 | float64] struct { @@ -332,35 +297,20 @@ func (b *floatAvgBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // floatAvgWindowFunction is the sql.WindowFunction used for avg(float4)/avg(float8) within an OVER(...) // clause. type floatAvgWindowFunction[T float32 | float64] struct { - expr sql.Expression - framer sql.WindowFramer + windowFramerState + expr sql.Expression } var _ sql.WindowFunction = (*floatAvgWindowFunction[float64])(nil) func newFloatAvgWindowFunction[T float32 | float64](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &floatAvgWindowFunction[T]{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } -func (w *floatAvgWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -func (w *floatAvgWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *floatAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -387,5 +337,3 @@ func (w *floatAvgWindowFunction[T]) Compute(ctx *sql.Context, interval sql.Windo } return sum / float64(count), nil } - -func (w *floatAvgWindowFunction[T]) Dispose(ctx *sql.Context) {} diff --git a/server/functions/aggregate/numeric_aggregates.go b/server/functions/aggregate/numeric_aggregates.go index f551799092..dd254e6b8d 100644 --- a/server/functions/aggregate/numeric_aggregates.go +++ b/server/functions/aggregate/numeric_aggregates.go @@ -24,6 +24,46 @@ import ( pgtypes "github.com/dolthub/doltgresql/server/types" ) +// windowFramerState holds the sql.WindowFramer setup shared by every native window-function implementation +// in this package, regardless of accumulator type: bind a framer from the window's explicit frame clause if +// one was given, otherwise fall back to Postgres's default (unbounded preceding to current row). None of +// this touches the per-row hot path (StartPartition/DefaultFramer/Dispose all run once per partition, not +// once per row), so unlike the accumulator itself, it's free to share via embedding across every T. +type windowFramerState struct { + framer sql.WindowFramer +} + +// bindFramer builds and stores this window's framer, if it declared an explicit frame clause; with no +// explicit frame, DefaultFramer's fallback applies instead. +func (s *windowFramerState) bindFramer(window *sql.WindowDefinition) error { + if window == nil || window.Frame == nil { + return nil + } + framer, err := window.Frame.NewFramer(window) + if err != nil { + return err + } + s.framer = framer + return nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (s *windowFramerState) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return nil +} + +// DefaultFramer implements the sql.WindowFunction interface; with no explicit frame, this supplies +// Postgres's default (unbounded preceding to current row). +func (s *windowFramerState) DefaultFramer() sql.WindowFramer { + if s.framer != nil { + return s.framer + } + return aggregation.NewUnboundedPrecedingToCurrentRowFramer() +} + +// Dispose implements the sql.WindowFunction interface. +func (s *windowFramerState) Dispose(ctx *sql.Context) {} + // int64ToDecimal converts an int64 sum to *apd.Decimal, for use as the decimalConvert of a decimalSumBuffer/ // decimalAvgBuffer instantiated over int64 (i.e. sum(int8)/avg(int8), whose accumulator needs to be decimal // since a running sum of bigints can itself overflow a bigint). @@ -112,37 +152,20 @@ func (b *intSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // intSumWindowFunction is the sql.WindowFunction used for sum(int2)/sum(int4) within an OVER(...) clause. type intSumWindowFunction[T int16 | int32] struct { - expr sql.Expression - framer sql.WindowFramer + windowFramerState + expr sql.Expression } var _ sql.WindowFunction = (*intSumWindowFunction[int32])(nil) func newIntSumWindowFunction[T int16 | int32](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &intSumWindowFunction[T]{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } -func (w *intSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -// DefaultFramer implements the sql.WindowFunction interface; with no explicit frame, this supplies -// Postgres's default (unbounded preceding to current row). -func (w *intSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *intSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -170,8 +193,6 @@ func (w *intSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowI return sum, nil } -func (w *intSumWindowFunction[T]) Dispose(ctx *sql.Context) {} - // decimalSumBuffer is the GROUP BY buffer for sum(int8)/sum(numeric). sum(int8) promotes to numeric because // a running sum of bigints can itself overflow a bigint; sum(numeric) stays numeric. convert adapts the // buffer to either input type: int64 values are boxed via int64ToDecimal, while numeric values (already @@ -226,8 +247,8 @@ func (b *decimalSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // decimalSumWindowFunction is the sql.WindowFunction used for sum(int8)/sum(numeric) within an OVER(...) // clause. type decimalSumWindowFunction[T int64 | *apd.Decimal] struct { + windowFramerState expr sql.Expression - framer sql.WindowFramer convert func(T) *apd.Decimal } @@ -236,28 +257,13 @@ var _ sql.WindowFunction = (*decimalSumWindowFunction[int64])(nil) func newDecimalSumWindowFunction[T int64 | *apd.Decimal](convert func(T) *apd.Decimal) framework.NewWindowFunctionFn { return func(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &decimalSumWindowFunction[T]{expr: exprs[0], convert: convert} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } } -func (w *decimalSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -func (w *decimalSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *decimalSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -292,8 +298,6 @@ func (w *decimalSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.Win return &sum, nil } -func (w *decimalSumWindowFunction[T]) Dispose(ctx *sql.Context) {} - // floatSumBuffer is the GROUP BY buffer for sum(float4)/sum(float8), which (unlike the integer overloads) // preserve their input type rather than promoting. type floatSumBuffer[T float32 | float64] struct { @@ -337,35 +341,20 @@ func (b *floatSumBuffer[T]) Update(ctx *sql.Context, row sql.Row) error { // floatSumWindowFunction is the sql.WindowFunction used for sum(float4)/sum(float8) within an OVER(...) // clause. type floatSumWindowFunction[T float32 | float64] struct { - expr sql.Expression - framer sql.WindowFramer + windowFramerState + expr sql.Expression } var _ sql.WindowFunction = (*floatSumWindowFunction[float64])(nil) func newFloatSumWindowFunction[T float32 | float64](exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { wf := &floatSumWindowFunction[T]{expr: exprs[0]} - if window != nil && window.Frame != nil { - framer, err := window.Frame.NewFramer(window) - if err != nil { - return nil, err - } - wf.framer = framer + if err := wf.bindFramer(window); err != nil { + return nil, err } return wf, nil } -func (w *floatSumWindowFunction[T]) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { - return nil -} - -func (w *floatSumWindowFunction[T]) DefaultFramer() sql.WindowFramer { - if w.framer != nil { - return w.framer - } - return aggregation.NewUnboundedPrecedingToCurrentRowFramer() -} - func (w *floatSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { if interval.End <= interval.Start { return nil, nil @@ -392,5 +381,3 @@ func (w *floatSumWindowFunction[T]) Compute(ctx *sql.Context, interval sql.Windo } return sum, nil } - -func (w *floatSumWindowFunction[T]) Dispose(ctx *sql.Context) {} diff --git a/server/functions/window/rank.go b/server/functions/window/rank.go index ed518d325b..144c6dfd7d 100644 --- a/server/functions/window/rank.go +++ b/server/functions/window/rank.go @@ -42,6 +42,15 @@ var denseRank = framework.Func0Window{ NewWinFunc: newDenseRankWindowFunction, } +// percentRank represents the PostgreSQL percent_rank() window function. +var percentRank = framework.Func0Window{ + Function0: framework.Function0{ + Name: "percent_rank", + Return: pgtypes.Float64, + }, + NewWinFunc: newPercentRankWindowFunction, +} + // rankWindowFunction is the sql.WindowFunction used for rank() within an OVER(...) clause. Its algorithm // mirrors GMS's own aggregation.rankBase.Compute exactly (see sql/expression/function/aggregation/ // window_functions.go), just returning int64 instead of uint64: the rank of a row is the count of rows @@ -145,3 +154,47 @@ func (w *denseRankWindowFunction) Compute(ctx *sql.Context, interval sql.WindowI // Dispose implements the sql.WindowFunction interface. func (w *denseRankWindowFunction) Dispose(ctx *sql.Context) {} + +// percentRankWindowFunction is the sql.WindowFunction used for percent_rank() within an OVER(...) clause. It +// reuses rankWindowFunction's peer-group counting the same way GMS's own PercentRank wraps rankBase: the +// result is (rank-1)/(partition size-1), or 0 for a single-row partition (rather than a divide-by-zero). +type percentRankWindowFunction struct { + rank *rankWindowFunction +} + +var _ sql.WindowFunction = (*percentRankWindowFunction)(nil) + +// newPercentRankWindowFunction creates the sql.WindowFunction for percent_rank(). +func newPercentRankWindowFunction(window *sql.WindowDefinition) (sql.WindowFunction, error) { + r, err := newRankWindowFunction(window) + if err != nil { + return nil, err + } + return &percentRankWindowFunction{rank: r.(*rankWindowFunction)}, nil +} + +// StartPartition implements the sql.WindowFunction interface. +func (w *percentRankWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + return w.rank.StartPartition(ctx, interval, buf) +} + +// DefaultFramer implements the sql.WindowFunction interface. +func (w *percentRankWindowFunction) DefaultFramer() sql.WindowFramer { + return w.rank.DefaultFramer() +} + +// Compute implements the sql.WindowFunction interface. +func (w *percentRankWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) (interface{}, error) { + partitionSize := w.rank.partitionEnd - w.rank.partitionStart + r, err := w.rank.Compute(ctx, interval, buf) + if err != nil || r == nil { + return r, err + } + if partitionSize <= 1 { + return float64(0), nil + } + return float64(r.(int64)-1) / float64(partitionSize-1), nil +} + +// Dispose implements the sql.WindowFunction interface. +func (w *percentRankWindowFunction) Dispose(ctx *sql.Context) {} diff --git a/server/functions/window/row_number.go b/server/functions/window/row_number.go index f3a379ca22..e975c6743e 100644 --- a/server/functions/window/row_number.go +++ b/server/functions/window/row_number.go @@ -27,6 +27,7 @@ func initRanking() { framework.RegisterWindowFunction(rowNumber) framework.RegisterWindowFunction(rank) framework.RegisterWindowFunction(denseRank) + framework.RegisterWindowFunction(percentRank) } // rowNumber represents the PostgreSQL row_number() window function, returning a bigint directly rather than diff --git a/testing/go/enginetest/doltgres_engine_test.go b/testing/go/enginetest/doltgres_engine_test.go index a152ec394d..2c345dd20a 100755 --- a/testing/go/enginetest/doltgres_engine_test.go +++ b/testing/go/enginetest/doltgres_engine_test.go @@ -836,20 +836,21 @@ func TestVersionedViews(t *testing.T) { } func TestWindowFunctions(t *testing.T) { - t.Skip() h := newDoltgresServerHarness(t) defer h.Close() enginetest.TestWindowFunctions(t, h) } func TestWindowRowFrames(t *testing.T) { - t.Skip() h := newDoltgresServerHarness(t) defer h.Close() enginetest.TestWindowRowFrames(t, h) } func TestWindowRangeFrames(t *testing.T) { + // TODO: RANGE frame offset boundaries (e.g. "RANGE 2 PRECEDING") are computed by GMS as float64, + // which panics against DoltgresType.Compare for an int32 column; named windows themselves work fine + // (see TestNamedWindows). t.Skip() h := newDoltgresServerHarness(t) defer h.Close() @@ -857,6 +858,8 @@ func TestWindowRangeFrames(t *testing.T) { } func TestNamedWindows(t *testing.T) { + // TODO: only remaining failures are GMS's own mergeWindowDefs error messages using empty window + // names ("window '' cannot inherit ''" instead of the actual names) - cosmetic, not a correctness bug. t.Skip() h := newDoltgresServerHarness(t) defer h.Close()