dolthub--doltgresql/server/expression/array_agg.go

// Copyright 2025 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 expression

import (
	"context"
	"sort"
	"strings"

	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/expression"
	vitess "github.com/dolthub/vitess/go/vt/sqlparser"

	"github.com/dolthub/doltgresql/server/types"
)

type ArrayAgg struct {
	selectExprs []sql.Expression
	orderBy     sql.SortFields
	id          sql.ColumnId
	Distinct    bool
}

var _ sql.Aggregation = (*ArrayAgg)(nil)
var _ vitess.Injectable = (*ArrayAgg)(nil)
var _ sql.OrderedAggregation = (*ArrayAgg)(nil)

// WithResolvedChildren returns a new ArrayAgg with the provided children as its select expressions.
// The last child is expected to be the order by expressions.
func (a *ArrayAgg) WithResolvedChildren(ctx context.Context, children []any) (any, error) {
	a.selectExprs = make([]sql.Expression, len(children)-1)
	for i := 0; i < len(children)-1; i++ {
		a.selectExprs[i] = children[i].(sql.Expression)
	}

	a.orderBy = children[len(children)-1].(sql.SortFields)
	return a, nil
}

// Resolved implements sql.Expression
func (a *ArrayAgg) Resolved() bool {
	return expression.ExpressionsResolved(a.selectExprs...) && expression.ExpressionsResolved(a.orderBy.ToExpressions()...)
}

// String implements sql.Expression
func (a *ArrayAgg) String() string {
	sb := strings.Builder{}
	sb.WriteString("array_agg(")

	if a.selectExprs != nil {
		var exprs = make([]string, len(a.selectExprs))
		for i, expr := range a.selectExprs {
			exprs[i] = expr.String()
		}

		sb.WriteString(strings.Join(exprs, ", "))
	}

	if len(a.orderBy) > 0 {
		sb.WriteString(" order by ")
		for i, ob := range a.orderBy {
			if i > 0 {
				sb.WriteString(", ")
			}
			sb.WriteString(ob.String())
		}
	}

	sb.WriteString(")")
	return sb.String()
}

// Type implements sql.Expression
func (a *ArrayAgg) Type(ctx *sql.Context) sql.Type {
	dt := a.selectExprs[0].Type(ctx).(*types.DoltgresType)
	return dt.ToArrayType()
}

// IsNullable implements sql.Expression
func (a *ArrayAgg) IsNullable(ctx *sql.Context) bool {
	return true
}

// Eval implements sql.Expression
func (a *ArrayAgg) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
	panic("eval should never be called on an aggregation function")
}

// Children implements sql.Expression
func (a *ArrayAgg) Children() []sql.Expression {
	return append(a.selectExprs, a.orderBy.ToExpressions()...)
}

func (a *ArrayAgg) OutputExpressions() []sql.Expression {
	return a.selectExprs
}

// WithChildren implements sql.Expression
func (a ArrayAgg) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) {
	if len(children) != len(a.selectExprs)+len(a.orderBy) {
		return nil, sql.ErrInvalidChildrenNumber.New(a, len(children), len(a.selectExprs)+len(a.orderBy))
	}

	a.selectExprs = children[:len(a.selectExprs)]
	a.orderBy = a.orderBy.FromExpressions(ctx, children[len(a.selectExprs):]...)
	return &a, nil
}

// Id implements sql.IdExpression
func (a *ArrayAgg) Id() sql.ColumnId {
	return a.id
}

// WithId implements sql.IdExpression
func (a ArrayAgg) WithId(id sql.ColumnId) sql.IdExpression {
	a.id = id
	return &a
}

// NewWindowFunction implements sql.WindowAdaptableExpression
func (a *ArrayAgg) NewWindowFunction(ctx *sql.Context) (sql.WindowFunction, error) {
	panic("window functions not yet supported for array_agg")
}

// WithWindow implements sql.WindowAdaptableExpression
func (a *ArrayAgg) WithWindow(ctx *sql.Context, window *sql.WindowDefinition) sql.WindowAdaptableExpression {
	panic("window functions not yet supported for array_agg")
}

// Window implements sql.WindowAdaptableExpression
func (a *ArrayAgg) Window() *sql.WindowDefinition {
	return nil
}

// NewBuffer implements sql.Aggregation
func (a *ArrayAgg) NewBuffer(ctx *sql.Context) (sql.AggregationBuffer, error) {
	return &arrayAggBuffer{
		elements: make([]sql.Row, 0),
		a:        a,
	}, nil
}

// arrayAggBuffer is the buffer used to accumulate values for the array_agg aggregation function.
type arrayAggBuffer struct {
	elements []sql.Row
	seen     []interface{} // sorted, non-NULL distinct values for binary search
	seenNull bool
	a        *ArrayAgg
}

// Dispose implements sql.AggregationBuffer
func (a *arrayAggBuffer) Dispose(ctx *sql.Context) {}

// Eval implements sql.AggregationBuffer
func (a *arrayAggBuffer) Eval(ctx *sql.Context) (interface{}, error) {
	if len(a.elements) == 0 {
		return nil, nil
	}

	if a.a.orderBy != nil {
		sorter := &expression.Sorter{
			SortFields: a.a.orderBy,
			Rows:       a.elements,
			Ctx:        ctx,
		}

		sort.Stable(sorter)
		if sorter.LastError != nil {
			return nil, sorter.LastError
		}
	}

	// convert to []interface for return. The last element in each row is the one we want to return, the rest are sort fields.
	result := make([]interface{}, len(a.elements))
	for i, row := range a.elements {
		result[i] = row[(len(row) - 1)]
	}

	return result, nil
}

// Update implements sql.AggregationBuffer
func (a *arrayAggBuffer) Update(ctx *sql.Context, row sql.Row) error {
	evalRow, err := evalExprs(ctx, a.a.selectExprs, row)
	if err != nil {
		return err
	}

	if a.a.Distinct {
		val := evalRow[0]
		if val == nil {
			if a.seenNull {
				return nil
			}
			a.seenNull = true
		} else {
			exprType := a.a.selectExprs[0].Type(ctx).(*types.DoltgresType)
			lo, hi := 0, len(a.seen)
			for lo < hi {
				mid := (lo + hi) / 2
				cmp, err := exprType.Compare(ctx, val, a.seen[mid])
				if err != nil {
					return err
				}
				if cmp == 0 {
					return nil
				} else if cmp < 0 {
					hi = mid
				} else {
					lo = mid + 1
				}
			}
			a.seen = append(a.seen, nil)
			copy(a.seen[lo+1:], a.seen[lo:])
			a.seen[lo] = val
		}
	}

	// Append the current value to the end of the row. We want to preserve the row's original structure
	// for sort ordering in the final step.
	a.elements = append(a.elements, append(row, evalRow[0]))
	return nil
}

// evalExprs evaluates the provided expressions against the given row and returns the results as a new row.
func evalExprs(ctx *sql.Context, exprs []sql.Expression, row sql.Row) (sql.Row, error) {
	result := make(sql.Row, len(exprs))
	for i, expr := range exprs {
		var err error
		result[i], err = expr.Eval(ctx, row)
		if err != nil {
			return nil, err
		}
	}

	return result, nil
}