Window function fixes

server/analyzer/init.go

@@ -68,8 +68,11 @@ func Init() {
 		// because of custom batch set optimization in GMS skipping OnceBeforeDefault batch for some nodes.
 		analyzer.Rule{Id: ruleId_ResolveType, Apply: ResolveType},
 		analyzer.Rule{Id: ruleId_SetRunner, Apply: SetRunner},
-		analyzer.Rule{Id: ruleId_TypeSanitizer, Apply: TypeSanitizer},
+		// ResolveValuesTypes must run before TypeSanitizer: it unifies VALUES-clause column
+		// types (e.g. VALUES(1),(2.01) -> numeric) which TypeSanitizer's aggregate-type
+		// corrections (e.g. SUM(int) -> bigint) depend on being already resolved.
 		analyzer.Rule{Id: ruleId_ResolveValuesTypes, Apply: ResolveValuesTypes},
+		analyzer.Rule{Id: ruleId_TypeSanitizer, Apply: TypeSanitizer},
 		analyzer.Rule{Id: ruleId_GenerateForeignKeyName, Apply: generateForeignKeyName},
 		analyzer.Rule{Id: ruleId_AddDomainConstraints, Apply: AddDomainConstraints},
 		analyzer.Rule{Id: ruleId_ValidateColumnDefaults, Apply: ValidateColumnDefaults},

server/analyzer/resolve_values_types.go

@@ -23,6 +23,7 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/expression"
 	"github.com/dolthub/go-mysql-server/sql/plan"
 	"github.com/dolthub/go-mysql-server/sql/transform"
+	"github.com/dolthub/go-mysql-server/sql/types"
 
 	pgexprs "github.com/dolthub/doltgresql/server/expression"
 	"github.com/dolthub/doltgresql/server/functions/framework"
@@ -189,7 +190,9 @@ func transformValuesNode(ctx *sql.Context, n sql.Node) (sql.Node, transform.Tree
 		columnTypes[colIdx] = make([]*pgtypes.DoltgresType, len(values.ExpressionTuples))
 		for rowIdx, row := range values.ExpressionTuples {
 			exprType := row[colIdx].Type(ctx)
-			if exprType == nil {
+			if exprType == nil || exprType == types.Null {
+				// A raw NULL literal reports the GMS sentinel Null type until TypeSanitizer
+				// converts it to a pgtypes.Unknown-typed literal.
 				columnTypes[colIdx][rowIdx] = pgtypes.Unknown
 			} else if pgType, ok := exprType.(*pgtypes.DoltgresType); ok {
 				columnTypes[colIdx][rowIdx] = pgType

server/analyzer/type_sanitizer.go

@@ -42,13 +42,18 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 	// TODO: this probably should not be opaque, we should let the analyzer dig into subqueries and analyze them when
 	//  it chooses. Doing all type transformations upfront like this masks bugs where certain tyupe conversion errors
 	//  only manifest in a subquery
+	// groupByTypeCache memoizes each GroupBy's ColumnId->DoltgresType map (see
+	// groupBySchemaTypesById) for the duration of this single TypeSanitizer call, so a
+	// GroupBy with M output columns pays the cost of scanning its SelectDeps once, not once
+	// per column reference to it.
+	groupByTypeCache := make(map[*plan.GroupBy]map[sql.ColumnId]*pgtypes.DoltgresType)
 	return pgtransform.NodeExprsWithNodeWithOpaque(ctx, node, func(ctx *sql.Context, n sql.Node, expr sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
 		// This can be updated if we find more expressions that return GMS types.
 		// These should eventually be replaced with Doltgres-equivalents over time, rendering this function unnecessary.
 		switch expr := expr.(type) {
 		case *expression.GetField:
 			switch n := n.(type) {
-			case *plan.Project, *plan.Filter, *plan.GroupBy:
+			case *plan.Project, *plan.Filter, *plan.GroupBy, *plan.Having:
 				child := n.Children()[0]
 				// Some dolt_ tables do not have doltgres types for their columns, so we convert them here
 				if rt, ok := child.(*plan.ResolvedTable); ok && strings.HasPrefix(rt.Name(), "dolt_") {
@@ -57,6 +62,56 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 						return pgexprs.NewGMSCast(expr), transform.NewTree, nil
 					}
 				}
+				// Window function outputs carry GMS types in the Window schema; wrap with GMSCast
+				// so the type is properly converted to a Doltgres type for the wire protocol.
+				if isWindowOrWrapsWindow(child) {
+					if _, ok := expr.Type(ctx).(*pgtypes.DoltgresType); !ok {
+						// GMS ranking functions (rank, dense_rank, ntile) use Uint64 internally,
+						// but Postgres returns bigint. ExplicitCast overrides the Numeric mapping.
+						if expr.Type(ctx).Type() == query.Type_UINT64 {
+							return pgexprs.NewExplicitCast(expr, pgtypes.Int64), transform.NewTree, nil
+						}
+						return pgexprs.NewGMSCast(expr), transform.NewTree, nil
+					}
+					// After this point expr IS a DoltgresType.
+					doltType := expr.Type(ctx).(*pgtypes.DoltgresType)
+					// Window SUM/AVG over integer or real columns are fixed at the source (see the
+					// AggCast wrap for *plan.Window below), so the Window's own schema already
+					// reports the corrected type. An outer Project (possibly reached through a
+					// SubqueryAlias) may still carry a stale declared type from before that fix, or
+					// from GMS's own type propagation; child.Schema() holds the corrected type after
+					// the bottom-up transform. Re-annotate on mismatch.
+					if innerType := windowSchemaTypeByName(child, ctx, expr.Name()); innerType != nil {
+						if !doltType.Equals(innerType) {
+							return pgexprs.NewAssignmentCast(expr, innerType, innerType), transform.NewTree, nil
+						}
+					}
+				}
+				// If a GetField references a SubqueryAlias and has a GMS type but the SubqueryAlias
+				// schema already has a DoltgresType (e.g. from AggCast propagation through a CTE),
+				// rebuild the GetField with the correct type so downstream operators compile for the
+				// actual runtime type rather than the planbuilder-assigned GMS type.
+				if _, ok := child.(*plan.SubqueryAlias); ok {
+					if _, ok := expr.Type(ctx).(*pgtypes.DoltgresType); !ok {
+						if actualType := windowSchemaTypeByName(child, ctx, expr.Name()); actualType != nil {
+							return expression.NewGetField(expr.Index(), actualType, expr.Name(), expr.IsNullable(ctx)), transform.NewTree, nil
+						}
+					}
+				}
+				// When a GroupBy child (possibly through a Having wrapper) has AggCast applied
+				// (e.g. SUM over integers), its schema reports the corrected DoltgresType, but the
+				// GetField may have been constructed earlier with a stale type: either a GMS type
+				// (float64 from GMS's internal Sum type promotion), or the pre-AggCast DoltgresType
+				// (e.g. Int32, from Sum.Type() before AggCast overrode it to Int64). Re-annotate so
+				// rowToBytes uses the correct Doltgres type for wire serialization. The runtime value
+				// is already correct from AggCast.
+				if gb := findGroupByChild(child); gb != nil {
+					if doltType := groupBySchemaTypesById(groupByTypeCache, gb, ctx)[expr.Id()]; doltType != nil {
+						if currentType, ok := expr.Type(ctx).(*pgtypes.DoltgresType); !ok || !currentType.Equals(doltType) {
+							return pgexprs.NewAssignmentCast(expr, doltType, doltType), transform.NewTree, nil
+						}
+					}
+				}
 			}
 			return expr, transform.SameTree, nil
 		case *expression.Literal:
@@ -75,6 +130,54 @@ 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 {
+				// Window functions (and GMS Sum) implement sql.WindowAdaptableExpression and must
+				// NOT be wrapped in GMSCast — that hides the interface from windowToIter.
+				if _, ok := expr.(sql.WindowAdaptableExpression); ok {
+					// GMS SUM and AVG always accumulate as float64 internally, regardless of the
+					// input column's width, but Postgres promotes their result types:
+					//   SUM(smallint/integer) -> bigint
+					//   SUM(bigint)           -> numeric (a running sum of bigints can itself overflow bigint)
+					//   SUM(real)             -> real (needs narrowing back from the float64 accumulator)
+					//   AVG(smallint/integer/bigint) -> numeric
+					//   AVG(real)             -> double precision
+					// In GroupBy or Window context, wrap with AggCast to convert the float64 result
+					// to the correct target type while preserving the sql.Aggregation and
+					// sql.WindowAdaptableExpression interfaces. Fixing this at the source (the
+					// aggregate itself) rather than at each outer reference to it is what makes this
+					// idempotent: AggCast doesn't implement sql.FunctionExpression, so once applied it
+					// no longer matches this case on a later revisit (e.g. via the opaque traversal
+					// that re-walks an already-analyzed subquery), unlike wrapping an outer GetField
+					// reference, whose declared type never changes and so looks like it "still needs
+					// fixing" on every subsequent visit.
+					var aggTargetType *pgtypes.DoltgresType
+					if doltType, ok := expr.Type(ctx).(*pgtypes.DoltgresType); ok {
+						switch expr.FunctionName() {
+						case "Sum":
+							switch {
+							case doltType.Equals(pgtypes.Int16), doltType.Equals(pgtypes.Int32):
+								aggTargetType = pgtypes.Int64
+							case doltType.Equals(pgtypes.Int64):
+								aggTargetType = pgtypes.Numeric
+							case doltType.Equals(pgtypes.Float32):
+								aggTargetType = pgtypes.Float32
+							}
+						case "Avg":
+							switch {
+							case doltType.Equals(pgtypes.Int16), doltType.Equals(pgtypes.Int32), doltType.Equals(pgtypes.Int64):
+								aggTargetType = pgtypes.Numeric
+							case doltType.Equals(pgtypes.Float32):
+								aggTargetType = pgtypes.Float64
+							}
+						}
+					}
+					if aggTargetType != nil {
+						switch n.(type) {
+						case *plan.GroupBy, *plan.Window:
+							return pgexprs.NewAggCast(expr.(sql.Aggregation), aggTargetType), transform.NewTree, nil
+						}
+					}
+					return expr, transform.SameTree, nil
+				}
 				// 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":
@@ -245,3 +348,81 @@ func typeSanitizerLiterals(ctx *sql.Context, gmsLiteral *expression.Literal) (sq
 		return nil, transform.NewTree, errors.Errorf("SANITIZER: encountered a GMS type that cannot be handled: %s", gmsLiteral.Type(ctx).String())
 	}
 }
+
+// isWindowOrWrapsWindow reports whether n is, or transitively wraps, a *plan.Window.
+func isWindowOrWrapsWindow(n sql.Node) bool {
+	return findWindowNode(n) != nil
+}
+
+// findWindowNode traverses Sort/Limit/Offset/Distinct/Filter/Project wrappers to find
+// the underlying *plan.Window, or nil if n does not wrap a Window.
+func findWindowNode(n sql.Node) *plan.Window {
+	if w, ok := n.(*plan.Window); ok {
+		return w
+	}
+	switch n.(type) {
+	case *plan.Sort, *plan.Limit, *plan.Offset, *plan.Distinct, *plan.Filter, *plan.Project, *plan.SubqueryAlias:
+		children := n.Children()
+		if len(children) == 1 {
+			return findWindowNode(children[0])
+		}
+	}
+	return nil
+}
+
+// findGroupByChild returns the GroupBy if n is a GroupBy or is a transparent wrapper
+// (Having, Project, Sort, Limit, Offset, Distinct, Filter) whose sole child transitively
+// wraps one. A HAVING clause's own condition, in particular, sits under a Having node whose
+// child is a Project (not the GroupBy directly) — e.g. Having -> Project -> GroupBy — so
+// this must traverse through Project to find it. Returns nil if n does not expose a GroupBy
+// this way.
+func findGroupByChild(n sql.Node) *plan.GroupBy {
+	if gb, ok := n.(*plan.GroupBy); ok {
+		return gb
+	}
+	switch n.(type) {
+	case *plan.Having, *plan.Project, *plan.Sort, *plan.Limit, *plan.Offset, *plan.Distinct, *plan.Filter:
+		children := n.Children()
+		if len(children) == 1 {
+			return findGroupByChild(children[0])
+		}
+	}
+	return nil
+}
+
+// windowSchemaTypeByName returns the DoltgresType for a named column in n's schema,
+// or nil if no match is found.
+func windowSchemaTypeByName(n sql.Node, ctx *sql.Context, name string) *pgtypes.DoltgresType {
+	for _, col := range n.Schema(ctx) {
+		if strings.EqualFold(col.Name, name) {
+			if t, ok := col.Type.(*pgtypes.DoltgresType); ok {
+				return t
+			}
+		}
+	}
+	return nil
+}
+
+// groupBySchemaTypesById returns a map from ColumnId to DoltgresType for every SelectDeps
+// expression in gb that has one, building it once per gb and caching the result in cache for
+// the remainder of the current TypeSanitizer call. Without this, resolving a single GetField
+// against gb required a linear scan of all of gb's SelectDeps; since this lookup runs once per
+// column reference to gb, a GroupBy with M output columns paid O(M) per column — O(M^2) total
+// — even when nothing needed correcting. Caching drops that to O(M) total per GroupBy.
+func groupBySchemaTypesById(cache map[*plan.GroupBy]map[sql.ColumnId]*pgtypes.DoltgresType, gb *plan.GroupBy, ctx *sql.Context) map[sql.ColumnId]*pgtypes.DoltgresType {
+	if m, ok := cache[gb]; ok {
+		return m
+	}
+	m := make(map[sql.ColumnId]*pgtypes.DoltgresType, len(gb.SelectDeps))
+	for _, e := range gb.SelectDeps {
+		ide, ok := e.(sql.IdExpression)
+		if !ok {
+			continue
+		}
+		if t, ok := e.Type(ctx).(*pgtypes.DoltgresType); ok {
+			m[ide.Id()] = t
+		}
+	}
+	cache[gb] = m
+	return m
+}