diff --git a/go.mod b/go.mod index 3a58e4a5a8..2f1107915b 100644 --- a/go.mod +++ b/go.mod @@ -9,7 +9,7 @@ require ( github.com/dolthub/dolt/go v0.40.5-0.20260710173237-7bdcf6ee8148 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.20260709233935-b8ef331dca43 + 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 1e16f71961..ff292cd169 100644 --- a/go.sum +++ b/go.sum @@ -256,8 +256,10 @@ github.com/dolthub/fslock v0.0.5 h1:QoXhBhgY1oumHE26qyE7tgmXUT8qjJwxsIzo54O/B/k= github.com/dolthub/fslock v0.0.5/go.mod h1:sdofYYqE0D79zNZyB4/kmlnsQOVap1C2yByjGKSirEM= github.com/dolthub/go-icu-regex v0.0.0-20260610153742-72563bc7ca83 h1:FEMjCGEroDnY/BXyAffVZxUpXhP2GpoUJyyq5KaLn8c= 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.20260709233935-b8ef331dca43 h1:BrDle6fyYfw4wxfzD4mD0R+PXT80e7Vl9UHWtKA58qw= -github.com/dolthub/go-mysql-server v0.20.1-0.20260709233935-b8ef331dca43/go.mod h1:mj5/QX3V8i92REbA1w6CzyknJAFdKtdE7l931405C/E= +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= 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/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..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": 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..458aa6e10d --- /dev/null +++ b/server/functions/aggregate/avg_aggregates.go @@ -0,0 +1,339 @@ +// 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/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// 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) + } + p += avgGuardDigits + result := new(apd.Decimal) + if _, err := sql.DecimalCtx.WithPrecision(uint32(p)).Quo(result, dividend, divisor); err != nil { + return nil, err + } + result.Reduce(result) + return result, 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: newBuffer, + NewAggWindowFunc: newWindowFunc, + } +} + +// 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 = (*intAvgBuffer[int32])(nil) + +func newIntAvgBuffer[T int16 | int32](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &intAvgBuffer[T]{expr: exprs[0]}, nil +} + +func (b *intAvgBuffer[T]) Dispose(ctx *sql.Context) {} + +func (b *intAvgBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return quoAvg(apd.New(b.sum, 0), apd.New(b.count, 0)) +} + +func (b *intAvgBuffer[T]) 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.(T) + if !ok { + return errors.Errorf("avg: expected %T, got %T", i, v) + } + b.sum += int64(i) + b.count++ + return nil +} + +// intAvgWindowFunction is the sql.WindowFunction used for avg(int2)/avg(int4) within an OVER(...) clause. +type intAvgWindowFunction[T int16 | int32] struct { + 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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil +} + +func (w *intAvgWindowFunction[T]) 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.(T) + if !ok { + return nil, errors.Errorf("avg: expected %T, got %T", iv, v) + } + sum += int64(iv) + count++ + } + if count == 0 { + return nil, nil + } + return quoAvg(apd.New(sum, 0), apd.New(count, 0)) +} + +// 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 +} + +var _ sql.AggregationBuffer = (*decimalAvgBuffer[int64])(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 + } +} + +func (b *decimalAvgBuffer[T]) Dispose(ctx *sql.Context) {} + +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)) +} + +func (b *decimalAvgBuffer[T]) 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 + } + typedV, ok := v.(T) + if !ok { + return errors.Errorf("avg: expected %T, got %T", typedV, v) + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, b.convert(typedV)) + b.count++ + return err +} + +// 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 + convert func(T) *apd.Decimal +} + +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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil + } +} + +func (w *decimalAvgWindowFunction[T]) 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 + } + typedV, ok := v.(T) + if !ok { + return nil, errors.Errorf("avg: expected %T, got %T", typedV, v) + } + if _, err = sql.DecimalCtx.Add(&sum, &sum, w.convert(typedV)); err != nil { + return nil, err + } + count++ + } + if count == 0 { + return nil, nil + } + return quoAvg(&sum, apd.New(count, 0)) +} + +// 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 = (*floatAvgBuffer[float64])(nil) + +func newFloatAvgBuffer[T float32 | float64](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &floatAvgBuffer[T]{expr: exprs[0]}, nil +} + +func (b *floatAvgBuffer[T]) Dispose(ctx *sql.Context) {} + +func (b *floatAvgBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { + if b.count == 0 { + return nil, nil + } + return b.sum / float64(b.count), nil +} + +func (b *floatAvgBuffer[T]) 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.(T) + if !ok { + return errors.Errorf("avg: expected %T, got %T", f, v) + } + b.sum += float64(f) + b.count++ + return nil +} + +// floatAvgWindowFunction is the sql.WindowFunction used for avg(float4)/avg(float8) within an OVER(...) +// clause. +type floatAvgWindowFunction[T float32 | float64] struct { + 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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil +} + +func (w *floatAvgWindowFunction[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 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.(T) + if !ok { + return nil, errors.Errorf("avg: expected %T, got %T", fv, v) + } + sum += float64(fv) + count++ + } + if count == 0 { + return nil, nil + } + return sum / float64(count), nil +} 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..dd254e6b8d --- /dev/null +++ b/server/functions/aggregate/numeric_aggregates.go @@ -0,0 +1,383 @@ +// 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" +) + +// 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). +func int64ToDecimal(v int64) *apd.Decimal { return apd.New(v, 0) } + +// 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. 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: newBuffer, + NewAggWindowFunc: newWindowFunc, + } +} + +// 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 = (*intSumBuffer[int32])(nil) + +func newIntSumBuffer[T int16 | int32](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &intSumBuffer[T]{expr: exprs[0]}, nil +} + +func (b *intSumBuffer[T]) Dispose(ctx *sql.Context) {} + +func (b *intSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +func (b *intSumBuffer[T]) 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.(T) + if !ok { + return errors.Errorf("sum: expected %T, got %T", i, v) + } + b.sum += int64(i) + b.sawOne = true + return nil +} + +// intSumWindowFunction is the sql.WindowFunction used for sum(int2)/sum(int4) within an OVER(...) clause. +type intSumWindowFunction[T int16 | int32] struct { + 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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil +} + +func (w *intSumWindowFunction[T]) 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.(T) + if !ok { + return nil, errors.Errorf("sum: expected %T, got %T", iv, v) + } + sum += int64(iv) + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, 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 +} + +var _ sql.AggregationBuffer = (*decimalSumBuffer[int64])(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 + } +} + +func (b *decimalSumBuffer[T]) Dispose(ctx *sql.Context) {} + +func (b *decimalSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + result := b.sum + return &result, nil +} + +func (b *decimalSumBuffer[T]) 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 + } + typedV, ok := v.(T) + if !ok { + return errors.Errorf("sum: expected %T, got %T", typedV, v) + } + d := b.convert(typedV) + if !b.sawOne { + b.sum.Set(d) + b.sawOne = true + return nil + } + _, err = sql.DecimalCtx.Add(&b.sum, &b.sum, d) + return err +} + +// 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 + convert func(T) *apd.Decimal +} + +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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil + } +} + +func (w *decimalSumWindowFunction[T]) 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 + } + typedV, ok := v.(T) + if !ok { + return nil, errors.Errorf("sum: expected %T, got %T", typedV, v) + } + d := w.convert(typedV) + 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 +} + +// 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 T + sawOne bool +} + +var _ sql.AggregationBuffer = (*floatSumBuffer[float64])(nil) + +func newFloatSumBuffer[T float32 | float64](exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &floatSumBuffer[T]{expr: exprs[0]}, nil +} + +func (b *floatSumBuffer[T]) Dispose(ctx *sql.Context) {} + +func (b *floatSumBuffer[T]) Eval(ctx *sql.Context) (interface{}, error) { + if !b.sawOne { + return nil, nil + } + return b.sum, nil +} + +func (b *floatSumBuffer[T]) 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.(T) + if !ok { + return errors.Errorf("sum: expected %T, got %T", f, v) + } + b.sum += f + b.sawOne = true + return nil +} + +// floatSumWindowFunction is the sql.WindowFunction used for sum(float4)/sum(float8) within an OVER(...) +// clause. +type floatSumWindowFunction[T float32 | float64] struct { + 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 err := wf.bindFramer(window); err != nil { + return nil, err + } + return wf, nil +} + +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 T + 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.(T) + if !ok { + return nil, errors.Errorf("sum: expected %T, got %T", fv, v) + } + sum += fv + sawOne = true + } + if !sawOne { + return nil, nil + } + return sum, nil +} 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..144c6dfd7d --- /dev/null +++ b/server/functions/window/rank.go @@ -0,0 +1,200 @@ +// 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, +} + +// 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 +// 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) {} + +// 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 new file mode 100644 index 0000000000..e975c6743e --- /dev/null +++ b/server/functions/window/row_number.go @@ -0,0 +1,78 @@ +// 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) + framework.RegisterWindowFunction(percentRank) +} + +// 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..2c345dd20a 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 @@ -834,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() @@ -855,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() 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")}, + }, + }, + }, + }, + }) +}