Fixing double-cast panic on window SUM/AVG in subqueries

Author: fulghum

server/analyzer/type_sanitizer.go

@@ -70,25 +70,12 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 					}
 					// After this point expr IS a DoltgresType.
 					doltType := expr.Type(ctx).(*pgtypes.DoltgresType)
-					// GMS window SUM propagates the child column's DoltgresType via Sum.Type(),
-					// but always accumulates float64 at runtime. Match by ColumnId in
-					// Window.SelectExprs; for integer aggregates apply Float64→Int64.
-					if winNode := findWindowNode(child); winNode != nil {
-						for _, selectExpr := range winNode.SelectExprs {
-							ide, ok := selectExpr.(sql.IdExpression)
-							if !ok || ide.Id() != expr.Id() {
-								continue
-							}
-							if _, isAgg := selectExpr.(sql.Aggregation); isAgg {
-								if doltType.Equals(pgtypes.Int32) || doltType.Equals(pgtypes.Int16) || doltType.Equals(pgtypes.Int64) {
-									return pgexprs.NewAssignmentCast(expr, pgtypes.Float64, pgtypes.Int64), transform.NewTree, nil
-								}
-							}
-							break
-						}
-					}
-					// An outer Project may carry a stale declared type; child.Schema() holds
-					// the corrected type after bottom-up transform. Re-annotate on mismatch.
+					// 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
@@ -148,8 +135,15 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 					//   SUM(real)             -> real (needs narrowing back from the float64 accumulator)
 					//   AVG(smallint/integer/bigint) -> numeric
 					//   AVG(real)             -> double precision
-					// In GroupBy context, wrap with AggCast to convert the float64 result to the
-					// correct target type while preserving the sql.Aggregation interface.
+					// 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() {
@@ -172,7 +166,8 @@ func TypeSanitizer(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 						}
 					}
 					if aggTargetType != nil {
-						if _, isGroupBy := n.(*plan.GroupBy); isGroupBy {
+						switch n.(type) {
+						case *plan.GroupBy, *plan.Window:
 							return pgexprs.NewAggCast(expr.(sql.Aggregation), aggTargetType), transform.NewTree, nil
 						}
 					}
@@ -361,7 +356,7 @@ func findWindowNode(n sql.Node) *plan.Window {
 		return w
 	}
 	switch n.(type) {
-	case *plan.Sort, *plan.Limit, *plan.Offset, *plan.Distinct, *plan.Filter, *plan.Project:
+	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])

server/expression/agg_cast.go

@@ -25,21 +25,23 @@ import (
 	pgtypes "github.com/dolthub/doltgresql/server/types"
 )
 
-// AggCast wraps a sql.Aggregation to override its declared return type and
-// post-convert the buffer's Eval result. It preserves the sql.Aggregation
-// interface so GroupBy aggregation machinery works correctly.
+// AggCast wraps a sql.Aggregation to override its declared return type and post-convert
+// its result, whether reached through the GroupBy buffer path (NewBuffer/Eval) or the
+// window function path (NewWindowFunction/Compute). It preserves both the sql.Aggregation
+// and sql.WindowAdaptableExpression interfaces so GroupBy and window execution both work.
 //
 // GMS SUM and AVG over integer columns always accumulate as float64 internally, but
 // Postgres specifies SUM(int2/int4/int8) → bigint and AVG(int2/int4/int8) → numeric.
-// AggCast intercepts the float64 buffer result and converts it to the target type so
-// the correct wire type is used.
+// AggCast intercepts the float64 result and converts it to the target type so the
+// correct wire type is used.
 type AggCast struct {
 	inner      sql.Aggregation
 	targetType *pgtypes.DoltgresType
 }
 
 var _ sql.Expression = (*AggCast)(nil)
 var _ sql.Aggregation = (*AggCast)(nil)
+var _ sql.WindowFunction = (*aggCastWindowFunction)(nil)
 
 // NewAggCast wraps inner so that its declared type is targetType and its buffer
 // Eval result is converted from float64 to int64.
@@ -59,9 +61,14 @@ func (a *AggCast) NewBuffer(ctx *sql.Context) (sql.AggregationBuffer, error) {
 	return &aggCastBuffer{inner: buf, targetType: a.targetType}, nil
 }
 
-// NewWindowFunction delegates to inner.
+// NewWindowFunction delegates to inner but wraps the result so Compute's float64 output is
+// converted to match targetType, the same way NewBuffer wraps the AggregationBuffer path.
 func (a *AggCast) NewWindowFunction(ctx *sql.Context) (sql.WindowFunction, error) {
-	return a.inner.NewWindowFunction(ctx)
+	fn, err := a.inner.NewWindowFunction(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &aggCastWindowFunction{inner: fn, targetType: a.targetType}, nil
 }
 
 // WithWindow delegates to inner.
@@ -126,26 +133,60 @@ func (b *aggCastBuffer) Eval(ctx *sql.Context) (any, error) {
 	if err != nil || v == nil {
 		return v, err
 	}
+	return convertAggResult(v, b.targetType)
+}
+
+func (b *aggCastBuffer) Dispose(ctx *sql.Context) {
+	b.inner.Dispose(ctx)
+}
+
+// aggCastWindowFunction wraps a sql.WindowFunction and post-converts its float64 Compute
+// result the same way aggCastBuffer does for the sql.AggregationBuffer (GroupBy) path.
+type aggCastWindowFunction struct {
+	inner      sql.WindowFunction
+	targetType *pgtypes.DoltgresType
+}
+
+func (w *aggCastWindowFunction) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buffer sql.WindowBuffer) error {
+	return w.inner.StartPartition(ctx, interval, buffer)
+}
+
+func (w *aggCastWindowFunction) DefaultFramer() sql.WindowFramer {
+	return w.inner.DefaultFramer()
+}
+
+func (w *aggCastWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buffer sql.WindowBuffer) (any, error) {
+	v, err := w.inner.Compute(ctx, interval, buffer)
+	if err != nil || v == nil {
+		return v, err
+	}
+	return convertAggResult(v, w.targetType)
+}
+
+func (w *aggCastWindowFunction) Dispose(ctx *sql.Context) {
+	w.inner.Dispose(ctx)
+}
+
+// convertAggResult converts a raw aggregation result (always float64 from GMS's SUM/AVG
+// implementations, regardless of the input column's width) to match targetType: numeric,
+// float32, float64 (a no-op), or int64 for everything else (bigint).
+func convertAggResult(v any, targetType *pgtypes.DoltgresType) (any, error) {
 	f, ok := v.(float64)
 	if !ok {
 		return v, nil
 	}
 	switch {
-	case b.targetType.Equals(pgtypes.Numeric):
+	case targetType.Equals(pgtypes.Numeric):
 		d, _, err := apd.NewFromString(strconv.FormatFloat(f, 'f', -1, 64))
 		if err != nil {
 			return nil, err
 		}
 		return d, nil
-	case b.targetType.Equals(pgtypes.Float32):
+	case targetType.Equals(pgtypes.Float32):
 		return float32(f), nil
-	case b.targetType.Equals(pgtypes.Float64):
+	case targetType.Equals(pgtypes.Float64):
 		return f, nil
 	default:
 		return int64(math.RoundToEven(f)), nil
 	}
 }
-
-func (b *aggCastBuffer) Dispose(ctx *sql.Context) {
-	b.inner.Dispose(ctx)
-}

testing/go/window_test.go

@@ -139,5 +139,30 @@ func TestWindowFunctions(t *testing.T) {
 				},
 			},
 		},
+		{
+			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")},
+					},
+				},
+			},
+		},
 	})
 }