Merge pull request #3060 from dolthub/angela/union

Merge scopeColumn types in SetOp scope

Author: angelamayxie

enginetest/queries/integration_plans.go

@@ -8728,6 +8728,27 @@ where
 			},
 		},
 	},
+	{
+		// https://github.com/dolthub/dolt/issues/9024
+		Name:    "subqueries should coerce union types",
+		Dialect: "mysql",
+		SetUpScript: []string{
+			"create table enum_table (i int primary key, e enum('a','b') not null)",
+			"insert into enum_table values (1,'a'),(2,'b')",
+			"create table uv (u int primary key, v varchar(10))",
+			"insert into uv values (0, 'bug'),(1,'ant'),(3, null)",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "select * from (select e from enum_table union select v from uv) sq",
+				Expected: []sql.Row{{"a"}, {"b"}, {"bug"}, {"ant"}, {nil}},
+			},
+			{
+				Query:    "with a as (select e from enum_table union select v from uv) select * from a",
+				Expected: []sql.Row{{"a"}, {"b"}, {"bug"}, {"ant"}, {nil}},
+			},
+		},
+	},
 
 	// Enum tests
 	{

enginetest/queries/script_queries.go

@@ -24,6 +24,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"
 )
 
 func hasRecursiveCte(node sql.Node) bool {
@@ -144,10 +145,28 @@ func (b *Builder) buildSetOp(inScope *scope, u *ast.SetOp) (outScope *scope) {
 	tabId := b.tabId
 	ret := plan.NewSetOp(setOpType, leftScope.node, rightScope.node, distinct, limit, offset, sortFields).WithId(tabId).WithColumns(cols)
 	outScope = leftScope
+	outScope.cols = b.mergeSetOpScopeColumns(leftScope.cols, rightScope.cols, tabId)
 	outScope.node = b.mergeSetOpSchemas(ret.(*plan.SetOp))
 	return
 }
 
+func (b *Builder) mergeSetOpScopeColumns(left, right []scopeColumn, tabId sql.TableId) []scopeColumn {
+	merged := make([]scopeColumn, len(left))
+	for i := range left {
+		merged[i] = scopeColumn{
+			tableId:     tabId,
+			db:          left[i].db,
+			table:       left[i].table,
+			col:         left[i].col,
+			originalCol: left[i].originalCol,
+			id:          left[i].id,
+			typ:         types.GeneralizeTypes(left[i].typ, right[i].typ),
+			nullable:    left[i].nullable || right[i].nullable,
+		}
+	}
+	return merged
+}
+
 func (b *Builder) mergeSetOpSchemas(u *plan.SetOp) sql.Node {
 	ls, rs := u.Left().Schema(), u.Right().Schema()
 	if len(ls) != len(rs) {

sql/planbuilder/set_op.go

@@ -16,6 +16,7 @@ package types
 
 import (
 	"fmt"
+	"reflect"
 	"strconv"
 	"strings"
 	"time"
@@ -635,7 +636,12 @@ func generalizeNumberTypes(a, b sql.Type) sql.Type {
 // GeneralizeTypes returns the more "general" of two types as defined by
 // https://dev.mysql.com/doc/refman/8.4/en/flow-control-functions.html
 // TODO: Create and handle "Illegal mix of collations" error
+// TODO: Handle extended types, like DoltgresType
 func GeneralizeTypes(a, b sql.Type) sql.Type {
+	if reflect.DeepEqual(a, b) {
+		return a
+	}
+
 	if a == Null {
 		return b
 	}

sql/types/conversion.go

@@ -198,15 +198,16 @@ func TestGeneralizeTypes(t *testing.T) {
 		{Int8, Int8, Int8},
 		{Boolean, Int64, Int64},
 		{Boolean, Boolean, Boolean},
-		{Text, Text, LongText},
+		{Text, Text, Text},
+		{Text, LongText, LongText},
 		{Text, Float64, LongText},
 		{Int64, Text, LongText},
 		{Int8, Null, Int8},
 		{Time, Time, Time},
 		{Time, Date, DatetimeMaxPrecision},
 		{Date, Date, Date},
 		{Date, Timestamp, DatetimeMaxPrecision},
-		{Timestamp, Timestamp, TimestampMaxPrecision},
+		{Timestamp, Timestamp, Timestamp},
 		{Timestamp, Datetime, DatetimeMaxPrecision},
 		{Null, Int64, Int64},
 		{Null, Null, Null},