Optimization: Defer Projections for Server Queries (#2676)

Author: jycor

engine.go

@@ -210,7 +210,7 @@ func (e *Engine) AnalyzeQuery(
 	query string,
 ) (sql.Node, error) {
 	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
-	parsed, _, _, qFlags, err := binder.Parse(query, false)
+	parsed, _, _, qFlags, err := binder.Parse(query, nil, false)
 	if err != nil {
 		return nil, err
 	}
@@ -238,7 +238,7 @@ func (e *Engine) PrepareParsedQuery(
 	stmt sqlparser.Statement,
 ) (sql.Node, error) {
 	binder := planbuilder.New(ctx, e.Analyzer.Catalog, e.Parser)
-	node, _, err := binder.BindOnly(stmt, query)
+	node, _, err := binder.BindOnly(stmt, query, nil)
 
 	if err != nil {
 		return nil, err
@@ -586,7 +586,7 @@ func (e *Engine) bindQuery(ctx *sql.Context, query string, parsed sqlparser.Stat
 	var bound sql.Node
 	var err error
 	if parsed == nil {
-		bound, _, _, qFlags, err = binder.Parse(query, false)
+		bound, _, _, qFlags, err = binder.Parse(query, qFlags, false)
 		if err != nil {
 			clearAutocommitErr := clearAutocommitTransaction(ctx)
 			if clearAutocommitErr != nil {
@@ -595,7 +595,7 @@ func (e *Engine) bindQuery(ctx *sql.Context, query string, parsed sqlparser.Stat
 			return nil, nil, err
 		}
 	} else {
-		bound, qFlags, err = binder.BindOnly(parsed, query)
+		bound, qFlags, err = binder.BindOnly(parsed, query, qFlags)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -651,7 +651,7 @@ func (e *Engine) bindExecuteQueryNode(ctx *sql.Context, query string, eq *plan.E
 		binder.SetBindingsWithExpr(tempBindings)
 	}
 
-	bound, _, err := binder.BindOnly(prep, query)
+	bound, _, err := binder.BindOnly(prep, query, nil)
 	if err != nil {
 		clearAutocommitErr := clearAutocommitTransaction(ctx)
 		if clearAutocommitErr != nil {

enginetest/engine_only_test.go

@@ -511,7 +511,7 @@ func TestAnalyzer_Exp(t *testing.T) {
 
 			ctx := enginetest.NewContext(harness)
 			b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, sql.NewMysqlParser())
-			parsed, _, _, _, err := b.Parse(tt.query, false)
+			parsed, _, _, _, err := b.Parse(tt.query, nil, false)
 			require.NoError(t, err)
 
 			analyzed, err := e.EngineAnalyzer().Analyze(ctx, parsed, nil, nil)

enginetest/enginetests.go

@@ -5653,11 +5653,11 @@ func TestTypesOverWire(t *testing.T, harness ClientHarness, sessionBuilder serve
 							break
 						}
 						expectedEngineRow := make([]*string, len(engineRow))
-						for i := range engineRow {
-							sqlVal, err := sch[i].Type.SQL(ctx, nil, engineRow[i])
-							if !assert.NoError(t, err) {
-								break
-							}
+						row, err := server.RowToSQL(ctx, sch, engineRow, nil)
+						if !assert.NoError(t, err) {
+							break
+						}
+						for i, sqlVal := range row {
 							if !sqlVal.IsNull() {
 								str := sqlVal.ToString()
 								expectedEngineRow[i] = &str

enginetest/evaluation.go

@@ -528,7 +528,7 @@ func injectBindVarsAndPrepare(
 
 	b := planbuilder.New(ctx, e.EngineAnalyzer().Catalog, sql.NewMysqlParser())
 	b.SetParserOptions(sql.LoadSqlMode(ctx).ParserOptions())
-	resPlan, _, err := b.BindOnly(parsed, q)
+	resPlan, _, err := b.BindOnly(parsed, q, nil)
 	if err != nil {
 		return q, nil, err
 	}

enginetest/plangen/cmd/plangen/main.go

@@ -166,7 +166,7 @@ func generatePlansForSuite(spec PlanSpec, w *bytes.Buffer) error {
 		if !tt.Skip {
 			ctx := enginetest.NewContextWithEngine(harness, engine)
 			binder := planbuilder.New(ctx, engine.EngineAnalyzer().Catalog, sql.NewMysqlParser())
-			parsed, _, _, qFlags, err := binder.Parse(tt.Query, false)
+			parsed, _, _, qFlags, err := binder.Parse(tt.Query, nil, false)
 			if err != nil {
 				exit(fmt.Errorf("%w\nfailed to parse query: %s", err, tt.Query))
 			}

server/handler.go

@@ -40,7 +40,9 @@ import (
 	"github.com/dolthub/go-mysql-server/internal/sockstate"
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/analyzer"
+	"github.com/dolthub/go-mysql-server/sql/iters"
 	"github.com/dolthub/go-mysql-server/sql/plan"
+	"github.com/dolthub/go-mysql-server/sql/rowexec"
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
 
@@ -218,7 +220,7 @@ func (h *Handler) ComExecuteBound(ctx context.Context, conn *mysql.Conn, query s
 func (h *Handler) ComStmtExecute(ctx context.Context, c *mysql.Conn, prepare *mysql.PrepareData, callback func(*sqltypes.Result) error) error {
 	_, err := h.errorWrappedDoQuery(ctx, c, prepare.PrepareStmt, nil, MultiStmtModeOff, prepare.BindVars, func(res *sqltypes.Result, more bool) error {
 		return callback(res)
-	}, nil)
+	}, &sql.QueryFlags{})
 	return err
 }
 
@@ -295,7 +297,7 @@ func (h *Handler) ComMultiQuery(
 	query string,
 	callback mysql.ResultSpoolFn,
 ) (string, error) {
-	return h.errorWrappedDoQuery(ctx, c, query, nil, MultiStmtModeOn, nil, callback, nil)
+	return h.errorWrappedDoQuery(ctx, c, query, nil, MultiStmtModeOn, nil, callback, &sql.QueryFlags{})
 }
 
 // ComQuery executes a SQL query on the SQLe engine.
@@ -305,7 +307,7 @@ func (h *Handler) ComQuery(
 	query string,
 	callback mysql.ResultSpoolFn,
 ) error {
-	_, err := h.errorWrappedDoQuery(ctx, c, query, nil, MultiStmtModeOff, nil, callback, nil)
+	_, err := h.errorWrappedDoQuery(ctx, c, query, nil, MultiStmtModeOff, nil, callback, &sql.QueryFlags{})
 	return err
 }
 
@@ -317,7 +319,7 @@ func (h *Handler) ComParsedQuery(
 	parsed sqlparser.Statement,
 	callback mysql.ResultSpoolFn,
 ) error {
-	_, err := h.errorWrappedDoQuery(ctx, c, query, parsed, MultiStmtModeOff, nil, callback, nil)
+	_, err := h.errorWrappedDoQuery(ctx, c, query, parsed, MultiStmtModeOff, nil, callback, &sql.QueryFlags{})
 	return err
 }
 
@@ -424,6 +426,7 @@ func (h *Handler) doQuery(
 		}
 	}()
 
+	qFlags.Set(sql.QFlagDeferProjections)
 	schema, rowIter, qFlags, err := queryExec(sqlCtx, query, parsed, analyzedPlan, bindings, qFlags)
 	if err != nil {
 		sqlCtx.GetLogger().WithError(err).Warn("error running query")
@@ -511,6 +514,37 @@ func resultForEmptyIter(ctx *sql.Context, iter sql.RowIter, resultFields []*quer
 	return &sqltypes.Result{Fields: resultFields}, nil
 }
 
+// GetDeferredProjections looks for a top-level deferred projection, retrieves its projections, and removes it from the
+// iterator tree.
+func GetDeferredProjections(iter sql.RowIter) (sql.RowIter, []sql.Expression) {
+	switch i := iter.(type) {
+	case *rowexec.ExprCloserIter:
+		_, projs := GetDeferredProjections(i.GetIter())
+		return i, projs
+	case *plan.TrackedRowIter:
+		_, projs := GetDeferredProjections(i.GetIter())
+		return i, projs
+	case *rowexec.TransactionCommittingIter:
+		newChild, projs := GetDeferredProjections(i.GetIter())
+		if projs != nil {
+			i.WithChildIter(newChild)
+		}
+		return i, projs
+	case *iters.LimitIter:
+		newChild, projs := GetDeferredProjections(i.ChildIter)
+		if projs != nil {
+			i.ChildIter = newChild
+		}
+		return i, projs
+	case *rowexec.ProjectIter:
+		if i.CanDefer() {
+			return i.GetChildIter(), i.GetProjections()
+		}
+		return i, nil
+	}
+	return iter, nil
+}
+
 // resultForMax1RowIter ensures that an empty iterator returns at most one row
 func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter, resultFields []*querypb.Field) (*sqltypes.Result, error) {
 	defer trace.StartRegion(ctx, "Handler.resultForMax1RowIter").End()
@@ -527,8 +561,7 @@ func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter,
 	if err := iter.Close(ctx); err != nil {
 		return nil, err
 	}
-
-	outputRow, err := rowToSQL(ctx, schema, row)
+	outputRow, err := RowToSQL(ctx, schema, row, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -558,16 +591,11 @@ func (h *Handler) resultForDefaultIter(
 		}
 	}
 
-	pollCtx, cancelF := ctx.NewSubContext()
-	eg.Go(func() error {
-		defer pan2err()
-		return h.pollForClosedConnection(pollCtx, c)
-	})
-
 	wg := sync.WaitGroup{}
 	wg.Add(2)
 
 	// Read rows off the row iterator and send them to the row channel.
+	iter, projs := GetDeferredProjections(iter)
 	var rowChan = make(chan sql.Row, 512)
 	eg.Go(func() error {
 		defer pan2err()
@@ -594,6 +622,12 @@ func (h *Handler) resultForDefaultIter(
 		}
 	})
 
+	pollCtx, cancelF := ctx.NewSubContext()
+	eg.Go(func() error {
+		defer pan2err()
+		return h.pollForClosedConnection(pollCtx, c)
+	})
+
 	// Default waitTime is one minute if there is no timeout configured, in which case
 	// it will loop to iterate again unless the socket died by the OS timeout or other problems.
 	// If there is a timeout, it will be enforced to ensure that Vitess has a chance to
@@ -639,7 +673,7 @@ func (h *Handler) resultForDefaultIter(
 					continue
 				}
 
-				outputRow, err := rowToSQL(ctx, schema, row)
+				outputRow, err := RowToSQL(ctx, schema, row, projs)
 				if err != nil {
 					return err
 				}
@@ -648,6 +682,7 @@ func (h *Handler) resultForDefaultIter(
 				r.Rows = append(r.Rows, outputRow)
 				r.RowsAffected++
 			case <-timer.C:
+				// TODO: timer should probably go in its own thread, as rowChan is blocking
 				if h.readTimeout != 0 {
 					// Cancel and return so Vitess can call the CloseConnection callback
 					ctx.GetLogger().Tracef("connection timeout")
@@ -901,25 +936,43 @@ func updateMaxUsedConnectionsStatusVariable() {
 	}()
 }
 
-func rowToSQL(ctx *sql.Context, s sql.Schema, row sql.Row) ([]sqltypes.Value, error) {
-	o := make([]sqltypes.Value, len(row))
+func RowToSQL(ctx *sql.Context, sch sql.Schema, row sql.Row, projs []sql.Expression) ([]sqltypes.Value, error) {
 	// need to make sure the schema is not null as some plan schema is defined as null (e.g. IfElseBlock)
-	if len(s) == 0 {
-		return o, nil
+	if len(sch) == 0 {
+		return []sqltypes.Value{}, nil
 	}
-	var err error
-	for i, v := range row {
-		if v == nil {
-			o[i] = sqltypes.NULL
+
+	outVals := make([]sqltypes.Value, len(sch))
+	if len(projs) == 0 {
+		for i, col := range sch {
+			if row[i] == nil {
+				outVals[i] = sqltypes.NULL
+				continue
+			}
+			var err error
+			outVals[i], err = col.Type.SQL(ctx, nil, row[i])
+			if err != nil {
+				return nil, err
+			}
+		}
+		return outVals, nil
+	}
+
+	for i, col := range sch {
+		field, err := projs[i].Eval(ctx, row)
+		if err != nil {
+			return nil, err
+		}
+		if field == nil {
+			outVals[i] = sqltypes.NULL
 			continue
 		}
-		o[i], err = s[i].Type.SQL(ctx, nil, v)
+		outVals[i], err = col.Type.SQL(ctx, nil, field)
 		if err != nil {
 			return nil, err
 		}
 	}
-
-	return o, nil
+	return outVals, nil
 }
 
 func schemaToFields(ctx *sql.Context, s sql.Schema) []*querypb.Field {

sql/analyzer/optimization_rules_test.go

@@ -223,7 +223,7 @@ func TestPushNotFilters(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.in, func(t *testing.T) {
 			q := fmt.Sprintf("SELECT 1 from xy WHERE %s", tt.in)
-			node, _, _, _, err := b.Parse(q, false)
+			node, _, _, _, err := b.Parse(q, nil, false)
 			require.NoError(t, err)
 
 			cmp, _, err := pushNotFilters(ctx, nil, node, nil, nil, nil)

sql/analyzer/rules.go

@@ -39,7 +39,7 @@ var OnceBeforeDefault = []Rule{
 	{applyDefaultSelectLimitId, applyDefaultSelectLimit},
 	{replaceCountStarId, replaceCountStar},
 	{applyEventSchedulerId, applyEventScheduler},
-	{validateOffsetAndLimitId, validateLimitAndOffset},
+	{validateOffsetAndLimitId, validateOffsetAndLimit},
 	{validateCreateTableId, validateCreateTable},
 	{validateAlterTableId, validateAlterTable},
 	{validateExprSemId, validateExprSem},

sql/analyzer/stored_procedures.go

@@ -56,7 +56,7 @@ func loadStoredProcedures(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan
 				var parsedProcedure sql.Node
 				b := planbuilder.New(ctx, a.Catalog, sql.NewMysqlParser())
 				b.SetParserOptions(sql.NewSqlModeFromString(procedure.SqlMode).ParserOptions())
-				parsedProcedure, _, _, _, err = b.Parse(procedure.CreateStatement, false)
+				parsedProcedure, _, _, _, err = b.Parse(procedure.CreateStatement, nil, false)
 				if err != nil {
 					procToRegister = &plan.Procedure{
 						CreateProcedureString: procedure.CreateStatement,
@@ -300,7 +300,7 @@ func applyProcedures(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scop
 				b.ProcCtx().AsOf = asOf
 			}
 			b.ProcCtx().DbName = call.Database().Name()
-			parsedProcedure, _, _, _, err = b.Parse(procedure.CreateStatement, false)
+			parsedProcedure, _, _, _, err = b.Parse(procedure.CreateStatement, nil, false)
 			if err != nil {
 				return nil, transform.SameTree, err
 			}

sql/analyzer/triggers.go

@@ -204,7 +204,7 @@ func applyTriggers(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope,
 			var parsedTrigger sql.Node
 			sqlMode := sql.NewSqlModeFromString(trigger.SqlMode)
 			b.SetParserOptions(sqlMode.ParserOptions())
-			parsedTrigger, _, _, _, err = b.Parse(trigger.CreateStatement, false)
+			parsedTrigger, _, _, _, err = b.Parse(trigger.CreateStatement, nil, false)
 			b.Reset()
 			if err != nil {
 				return nil, transform.SameTree, err
@@ -225,7 +225,7 @@ func applyTriggers(ctx *sql.Context, a *Analyzer, n sql.Node, scope *plan.Scope,
 				// first pass allows unresolved before we know whether trigger is relevant
 				// TODO store destination table name with trigger, so we don't have to do parse twice
 				b.TriggerCtx().Call = true
-				parsedTrigger, _, _, _, err = b.Parse(trigger.CreateStatement, false)
+				parsedTrigger, _, _, _, err = b.Parse(trigger.CreateStatement, nil, false)
 				b.TriggerCtx().Call = false
 				b.Reset()
 				if err != nil {