Merge pull request #3054 from dolthub/zachmu/project-iters

Support for set-returning functions

Author: zachmu

sql/core.go

@@ -45,6 +45,16 @@ type Expression interface {
 	WithChildren(children ...Expression) (Expression, error)
 }
 
+// RowIterExpression is an Expression that returns a RowIter rather than a scalar, used to implement functions that
+// return sets.
+type RowIterExpression interface {
+	Expression
+	// EvalRowIter evaluates the expression, which must be a RowIter
+	EvalRowIter(ctx *Context, r Row) (RowIter, error)
+	// ReturnsRowIter returns whether this expression returns a RowIter
+	ReturnsRowIter() bool
+}
+
 // ExpressionWithNodes is an expression that contains nodes as children.
 type ExpressionWithNodes interface {
 	Expression

sql/expression/case.go

@@ -45,7 +45,8 @@ func NewCase(expr sql.Expression, branches []CaseBranch, elseExpr sql.Expression
 
 // Type implements the sql.Expression interface.
 func (c *Case) Type() sql.Type {
-	curr := types.Null
+	var curr sql.Type
+	curr = types.Null
 	for _, b := range c.Branches {
 		curr = types.GeneralizeTypes(curr, b.Value.Type())
 	}

sql/expression/function/coalesce.go

@@ -58,7 +58,9 @@ func (c *Coalesce) Type() sql.Type {
 	if c.typ != nil {
 		return c.typ
 	}
-	retType := types.Null
+
+	var retType sql.Type
+	retType = types.Null
 	for i, arg := range c.args {
 		if arg == nil {
 			continue

sql/plan/project.go

@@ -26,9 +26,15 @@ import (
 // Project is a projection of certain expression from the children node.
 type Project struct {
 	UnaryNode
+	// Projections are the expressions to be projected on the row returned by the child node
 	Projections []sql.Expression
-	CanDefer    bool
-	deps        sql.ColSet
+	// CanDefer is true when the projection evaluation can be deferred to row spooling, which allows us to avoid a
+	// separate iterator for the project node.
+	CanDefer bool
+	// IncludesNestedIters is true when the projection includes nested iterators because of expressions that return
+	// a RowIter.
+	IncludesNestedIters bool
+	deps                sql.ColSet
 }
 
 var _ sql.Expressioner = (*Project)(nil)
@@ -202,8 +208,17 @@ func (p *Project) WithExpressions(exprs ...sql.Expression) (sql.Node, error) {
 	return &np, nil
 }
 
+// WithCanDefer returns a new Project with the CanDefer field set to the given value.
 func (p *Project) WithCanDefer(canDefer bool) *Project {
 	np := *p
 	np.CanDefer = canDefer
 	return &np
 }
+
+// WithIncludesNestedIters returns a new Project with the IncludesNestedIters field set to the given value.
+func (p *Project) WithIncludesNestedIters(includesNestedIters bool) *Project {
+	np := *p
+	np.IncludesNestedIters = includesNestedIters
+	np.CanDefer = false
+	return &np
+}

sql/rowexec/rel.go

@@ -312,9 +312,10 @@ func (b *BaseBuilder) buildProject(ctx *sql.Context, n *plan.Project, row sql.Ro
 	}
 
 	return sql.NewSpanIter(span, &ProjectIter{
-		projs:     n.Projections,
-		canDefer:  n.CanDefer,
-		childIter: i,
+		projs:          n.Projections,
+		canDefer:       n.CanDefer,
+		hasNestedIters: n.IncludesNestedIters,
+		childIter:      i,
 	}), nil
 }
 

sql/rowexec/rel_iters.go

@@ -25,6 +25,7 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/hash"
 	"github.com/dolthub/go-mysql-server/sql/iters"
 	"github.com/dolthub/go-mysql-server/sql/plan"
+	"github.com/dolthub/go-mysql-server/sql/transform"
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
 
@@ -126,16 +127,29 @@ func (i *offsetIter) Close(ctx *sql.Context) error {
 var _ sql.RowIter = &iters.JsonTableRowIter{}
 
 type ProjectIter struct {
-	projs     []sql.Expression
-	canDefer  bool
-	childIter sql.RowIter
+	projs          []sql.Expression
+	canDefer       bool
+	hasNestedIters bool
+	nestedState    *nestedIterState
+	childIter      sql.RowIter
+}
+
+type nestedIterState struct {
+	projections    []sql.Expression
+	sourceRow      sql.Row
+	iterEvaluators []*RowIterEvaluator
 }
 
 func (i *ProjectIter) Next(ctx *sql.Context) (sql.Row, error) {
+	if i.hasNestedIters {
+		return i.ProjectRowWithNestedIters(ctx)
+	}
+
 	childRow, err := i.childIter.Next(ctx)
 	if err != nil {
 		return nil, err
 	}
+
 	return ProjectRow(ctx, i.projs, childRow)
 }
 
@@ -155,6 +169,136 @@ func (i *ProjectIter) GetChildIter() sql.RowIter {
 	return i.childIter
 }
 
+// ProjectRowWithNestedIters evaluates a set of projections, allowing for nested iterators in the expressions.
+func (i *ProjectIter) ProjectRowWithNestedIters(
+	ctx *sql.Context,
+) (sql.Row, error) {
+
+	// For the set of iterators, we return one row each element in the longest of the iterators provided.
+	// Other iterator values will be NULL after they are depleted. All non-iterator fields for the row are returned
+	// identically for each row in the result set.
+	if i.nestedState != nil {
+		row, err := ProjectRow(ctx, i.nestedState.projections, i.nestedState.sourceRow)
+		if err != nil {
+			return nil, err
+		}
+
+		nestedIterationFinished := true
+		for _, evaluator := range i.nestedState.iterEvaluators {
+			if !evaluator.finished && evaluator.iter != nil {
+				nestedIterationFinished = false
+				break
+			}
+		}
+
+		if nestedIterationFinished {
+			i.nestedState = nil
+			return i.ProjectRowWithNestedIters(ctx)
+		}
+
+		return row, nil
+	}
+
+	row, err := i.childIter.Next(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	i.nestedState = &nestedIterState{
+		sourceRow: row,
+	}
+
+	// We need a new set of projections, with any iterator-returning expressions replaced by new expressions that will
+	// return the result of the iteration on each call to Eval. We also need to keep a list of all such iterators, so
+	// that we can tell when they have all finished their iterations.
+	var rowIterEvaluators []*RowIterEvaluator
+	newProjs := make([]sql.Expression, len(i.projs))
+	for i, proj := range i.projs {
+		p, _, err := transform.Expr(proj, func(e sql.Expression) (sql.Expression, transform.TreeIdentity, error) {
+			if rie, ok := e.(sql.RowIterExpression); ok && rie.ReturnsRowIter() {
+				ri, err := rie.EvalRowIter(ctx, row)
+				if err != nil {
+					return nil, false, err
+				}
+
+				evaluator := &RowIterEvaluator{
+					iter: ri,
+					typ:  rie.Type(),
+				}
+				rowIterEvaluators = append(rowIterEvaluators, evaluator)
+				return evaluator, transform.NewTree, nil
+			}
+
+			return e, transform.SameTree, nil
+		})
+
+		if err != nil {
+			return nil, err
+		}
+
+		newProjs[i] = p
+	}
+
+	i.nestedState.projections = newProjs
+	i.nestedState.iterEvaluators = rowIterEvaluators
+
+	return i.ProjectRowWithNestedIters(ctx)
+}
+
+// RowIterEvaluator is an expression that returns the next value from a sql.RowIter each time Eval is called.
+type RowIterEvaluator struct {
+	iter     sql.RowIter
+	typ      sql.Type
+	finished bool
+}
+
+var _ sql.Expression = (*RowIterEvaluator)(nil)
+
+func (r RowIterEvaluator) Resolved() bool {
+	return true
+}
+
+func (r RowIterEvaluator) String() string {
+	return "RowIterEvaluator"
+}
+
+func (r RowIterEvaluator) Type() sql.Type {
+	return r.typ
+}
+
+func (r RowIterEvaluator) IsNullable() bool {
+	return true
+}
+
+func (r *RowIterEvaluator) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	if r.finished || r.iter == nil {
+		return nil, nil
+	}
+
+	nextRow, err := r.iter.Next(ctx)
+	if err != nil {
+		if errors.Is(err, io.EOF) {
+			r.finished = true
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	// All of the set-returning functions return a single value per column
+	return nextRow[0], nil
+}
+
+func (r RowIterEvaluator) Children() []sql.Expression {
+	return nil
+}
+
+func (r RowIterEvaluator) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 0 {
+		return nil, sql.ErrInvalidChildrenNumber.New(r, len(children), 0)
+	}
+	return &r, nil
+}
+
 // ProjectRow evaluates a set of projections.
 func ProjectRow(
 	ctx *sql.Context,

sql/type.go

@@ -104,6 +104,9 @@ type Type interface {
 // NullType represents the type of NULL values
 type NullType interface {
 	Type
+
+	// IsNullType is a marker interface for types that represent NULL values.
+	IsNullType() bool
 }
 
 // DeferredType is a placeholder for prepared statements

sql/types/conversion.go

@@ -642,10 +642,10 @@ func GeneralizeTypes(a, b sql.Type) sql.Type {
 		return a
 	}
 
-	if a == Null {
+	if IsNullType(a) {
 		return b
 	}
-	if b == Null {
+	if IsNullType(b) {
 		return a
 	}
 
@@ -722,6 +722,16 @@ func GeneralizeTypes(a, b sql.Type) sql.Type {
 	if IsNumber(a) && IsNumber(b) {
 		return generalizeNumberTypes(a, b)
 	}
+
+	if IsText(a) && IsText(b) {
+		sta := a.(sql.StringType)
+		stb := b.(sql.StringType)
+		if sta.Length() > stb.Length() {
+			return a
+		}
+		return b
+	}
+
 	// TODO: decide if we want to make this VarChar to match MySQL, match VarChar length to max of two types
 	return LongText
 }

sql/types/conversion_test.go

@@ -208,6 +208,7 @@ func TestGeneralizeTypes(t *testing.T) {
 		{Date, Date, Date},
 		{Date, Timestamp, DatetimeMaxPrecision},
 		{Timestamp, Timestamp, Timestamp},
+		{Timestamp, TimestampMaxPrecision, TimestampMaxPrecision},
 		{Timestamp, Datetime, DatetimeMaxPrecision},
 		{Null, Int64, Int64},
 		{Null, Null, Null},

sql/types/null.go

@@ -34,6 +34,10 @@ var (
 
 type nullType struct{}
 
+func (t nullType) IsNullType() bool {
+	return true
+}
+
 // Compare implements Type interface. Note that while this returns 0 (equals)
 // for ordering purposes, in SQL NULL != NULL.
 func (t nullType) Compare(s context.Context, a interface{}, b interface{}) (int, error) {