array_agg support and general framework for postgres aggregate functions

go.mod

@@ -10,9 +10,9 @@ require (
 	github.com/dolthub/dolt/go/gen/proto/dolt/services/eventsapi v0.0.0-20241119094239-f4e529af734d
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
 	github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad
-	github.com/dolthub/go-mysql-server v0.20.1-0.20250521012141-b56c7c6eb677
+	github.com/dolthub/go-mysql-server v0.20.1-0.20250527234113-f38274720545
 	github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216
-	github.com/dolthub/vitess v0.0.0-20250512224608-8fb9c6ea092c
+	github.com/dolthub/vitess v0.0.0-20250523011542-0f6cf9472d1c
 	github.com/fatih/color v1.13.0
 	github.com/goccy/go-json v0.10.2
 	github.com/gogo/protobuf v1.3.2

go.sum

@@ -266,8 +266,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad h1:66ZPawHszNu37VPQckdhX1BPPVzREsGgNxQeefnlm3g=
 github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad/go.mod h1:ylU4XjUpsMcvl/BKeRRMXSH7e7WBrPXdSLvnRJYrxEA=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250521012141-b56c7c6eb677 h1:Wn0v7xBxkdzYqDN/4ksI34jstZOskMccT2SYFrvUw4c=
-github.com/dolthub/go-mysql-server v0.20.1-0.20250521012141-b56c7c6eb677/go.mod h1:5ZdrW0fHZbz+8CngT9gksqSX4H3y+7v1pns7tJCEpu0=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250527234113-f38274720545 h1:O+/sjRQJadYzyVr89Zh9yCnhZJ0NlHwiDYsXHnj3LsU=
+github.com/dolthub/go-mysql-server v0.20.1-0.20250527234113-f38274720545/go.mod h1:Zn9XK5KLYwPbyMpwfeUP+TgnhlgyID2vXf1WcF0M6Fk=
 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-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=
@@ -276,8 +276,8 @@ github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 h1:bMGS25NWAGTE
 github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71/go.mod h1:2/2zjLQ/JOOSbbSboojeg+cAwcRV0fDLzIiWch/lhqI=
 github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216 h1:JWkKRE4EHUcEVQCMRBej8DYxjYjRz/9MdF/NNQh0o70=
 github.com/dolthub/sqllogictest/go v0.0.0-20240618184124-ca47f9354216/go.mod h1:e/FIZVvT2IR53HBCAo41NjqgtEnjMJGKca3Y/dAmZaA=
-github.com/dolthub/vitess v0.0.0-20250512224608-8fb9c6ea092c h1:imdag6PPCHAO2rZNsFoQoR4I/vIVTmO/czoOl5rUnbk=
-github.com/dolthub/vitess v0.0.0-20250512224608-8fb9c6ea092c/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
+github.com/dolthub/vitess v0.0.0-20250523011542-0f6cf9472d1c h1:23KvsBtJk2GmHpXwQ/RkwIkdNpWL8tWdHRCiidhnaUA=
+github.com/dolthub/vitess v0.0.0-20250523011542-0f6cf9472d1c/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=

server/analyzer/init.go

@@ -17,6 +17,7 @@ package analyzer
 import (
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
 	"github.com/dolthub/go-mysql-server/sql/plan"
+	"github.com/dolthub/go-mysql-server/sql/planbuilder"
 )
 
 // IDs are basically arbitrary, we just need to ensure that they do not conflict with existing IDs
@@ -101,6 +102,23 @@ func initEngine() {
 	// place to put it. Our foreign key validation logic is different from MySQL's, and since it's not an analyzer rule
 	// we can't swap out a rule like the rest of the logic in this package, we have to do a function swap.
 	plan.ValidateForeignKeyDefinition = validateForeignKeyDefinition
+
+	planbuilder.IsAggregateFunc = IsAggregateFunc
+}
+
+// IsAggregateFunc checks if the given function name is an aggregate function. This is the entire set supported by
+// MySQL plus some postgres specific ones.
+func IsAggregateFunc(name string) bool {
+	if planbuilder.IsMySQLAggregateFuncName(name) {
+		return true
+	}
+
+	switch name {
+	case "array_agg":
+		return true
+	}
+
+	return false
 }
 
 // insertAnalyzerRules inserts the given rule(s) before or after the given analyzer.RuleId, returning an updated slice.

server/functions/aggregate/array_agg.go

@@ -0,0 +1,66 @@
+// 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 aggregate
+
+import (
+	"github.com/dolthub/go-mysql-server/sql"
+
+	"github.com/dolthub/doltgresql/server/functions/framework"
+	pgtypes "github.com/dolthub/doltgresql/server/types"
+)
+
+// initArrayAgg registers the functions to the catalog.
+func initArrayAgg() {
+	framework.RegisterAggregateFunction(array_agg)
+}
+
+// array_agg represents the PostgreSQL array_agg function.
+var array_agg = framework.Func1Aggregate{
+	Function1: framework.Function1{
+		Name:   "array_agg",
+		Return: pgtypes.AnyArray,
+		Parameters: [1]*pgtypes.DoltgresType{
+			pgtypes.AnyElement,
+		},
+		Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) {
+			return nil, nil
+		},
+	},
+	NewAggBuffer: newArrayAggBuffer,
+}
+
+type arrayAggBuffer struct {
+	elements []any
+}
+
+func newArrayAggBuffer() (sql.AggregationBuffer, error) {
+	return &arrayAggBuffer{
+		elements: make([]any, 0),
+	}, nil
+}
+
+func (a *arrayAggBuffer) Dispose() {}
+
+func (a *arrayAggBuffer) Eval(context *sql.Context) (interface{}, error) {
+	if len(a.elements) == 0 {
+		return nil, nil
+	}
+	return a.elements, nil
+}
+
+func (a *arrayAggBuffer) Update(ctx *sql.Context, row sql.Row) error {
+	a.elements = append(a.elements, row[0])
+	return nil
+}

server/functions/aggregate/init.go

@@ -0,0 +1,19 @@
+// 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 aggregate
+
+func Init() {
+	initArrayAgg()
+}

server/functions/framework/catalog.go

@@ -15,6 +15,7 @@
 package framework
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/cockroachdb/errors"
@@ -28,6 +29,9 @@ import (
 // Catalog contains all of the PostgreSQL functions.
 var Catalog = map[string][]FunctionInterface{}
 
+// AggregateCatalog contains all of the PostgreSQL aggregate functions.
+var AggregateCatalog = map[string][]AggregateFunctionInterface{}
+
 // initializedFunctions simply states whether Initialize has been called yet.
 var initializedFunctions = false
 
@@ -61,6 +65,21 @@ func RegisterFunction(f FunctionInterface) {
 	}
 }
 
+// RegisterAggregateFunction registers the given function, so that it will be usable from a running server. This should be called
+// from within an init().
+func RegisterAggregateFunction(f AggregateFunctionInterface) {
+	if initializedFunctions {
+		panic("attempted to register a function after the init() phase")
+	}
+	switch f := f.(type) {
+	case Func1Aggregate:
+		name := strings.ToLower(f.Name)
+		AggregateCatalog[name] = append(AggregateCatalog[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() {
@@ -76,6 +95,7 @@ func Initialize() {
 	replaceGmsBuiltIns()
 	validateFunctions()
 	compileFunctions()
+	compileAggs()
 }
 
 // replaceGmsBuiltIns replaces all GMS built-ins that have conflicting names with PostgreSQL functions.
@@ -151,7 +171,29 @@ func compileNonOperatorFunction(funcName string, overloads []FunctionInterface)
 		Fn:   createFunc,
 	})
 	compiledCatalog[funcName] = createFunc
-	namedCatalog[funcName] = overloads
+}
+
+// compileNonOperatorFunction creates a CompiledFunction for each overload of the given function.
+func compileAggFunction(funcName string, overloads []AggregateFunctionInterface) {
+	var newBuffer func() (sql.AggregationBuffer, error)
+	overloadTree := NewOverloads()
+	for _, functionOverload := range overloads {
+		newBuffer = functionOverload.NewBuffer
+		if err := overloadTree.Add(functionOverload); err != nil {
+			panic(err)
+		}
+	}
+
+	// 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(params ...sql.Expression) (sql.Expression, error) {
+		return NewCompiledAggregateFunction(funcName, params, overloadTree, newBuffer), nil
+	}
+	function.BuiltIns = append(function.BuiltIns, sql.FunctionN{
+		Name: funcName,
+		Fn:   createFunc,
+	})
+	compiledCatalog[funcName] = createFunc
 }
 
 // compileFunctions creates a CompiledFunction for each overload of each function in the catalog.
@@ -165,32 +207,38 @@ func compileFunctions() {
 	// special rules, so it's far more efficient to reuse it for operators. Operators are also a special case since they
 	// all have different names, while standard overload deducers work on a function-name basis.
 	for signature, functionOverload := range unaryFunctions {
-		overloads, ok := unaryAggregateOverloads[signature.Operator]
+		overloads, ok := unaryOperatorOverloads[signature.Operator]
 		if !ok {
 			overloads = NewOverloads()
-			unaryAggregateOverloads[signature.Operator] = overloads
+			unaryOperatorOverloads[signature.Operator] = overloads
 		}
 		if err := overloads.Add(functionOverload); err != nil {
 			panic(err)
 		}
 	}
 
 	for signature, functionOverload := range binaryFunctions {
-		overloads, ok := binaryAggregateOverloads[signature.Operator]
+		overloads, ok := binaryOperatorOverloads[signature.Operator]
 		if !ok {
 			overloads = NewOverloads()
-			binaryAggregateOverloads[signature.Operator] = overloads
+			binaryOperatorOverloads[signature.Operator] = overloads
 		}
 		if err := overloads.Add(functionOverload); err != nil {
 			panic(err)
 		}
 	}
 
 	// Add all permutations for the unary and binary operators
-	for operator, overload := range unaryAggregateOverloads {
-		unaryAggregatePermutations[operator] = overload.overloadsForParams(1)
+	for operator, overload := range unaryOperatorOverloads {
+		unaryOperatorPermutations[operator] = overload.overloadsForParams(1)
+	}
+	for operator, overload := range binaryOperatorOverloads {
+		binaryOperatorPermutations[operator] = overload.overloadsForParams(2)
 	}
-	for operator, overload := range binaryAggregateOverloads {
-		binaryAggregatePermutations[operator] = overload.overloadsForParams(2)
+}
+
+func compileAggs() {
+	for funcName, overloads := range AggregateCatalog {
+		compileAggFunction(funcName, overloads)
 	}
 }