feat: stop receiving data from tablet when we know we don't need it

Signed-off-by: Andres Taylor <andres@planetscale.com>
vitessio · Dec 18, 2024 · 128b3d8 · 128b3d8
1 parent 880284d
commit 128b3d8
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 14 deletions.
diff --git a/go/vt/vtgate/engine/limit.go b/go/vt/vtgate/engine/limit.go
@@ -19,6 +19,7 @@ package engine
 import (
 	"context"
 	"fmt"
+	"io"
 	"strconv"
 	"sync"
 
@@ -36,9 +37,10 @@ var _ Primitive = (*Limit)(nil)
 
 // Limit is a primitive that performs the LIMIT operation.
 type Limit struct {
-	Count  evalengine.Expr
-	Offset evalengine.Expr
-	Input  Primitive
+	Count                evalengine.Expr
+	Offset               evalengine.Expr
+	RequireCompleteInput bool
+	Input                Primitive
 }
 
 var UpperLimitStr = "__upper_limit"
@@ -141,9 +143,16 @@ func (l *Limit) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars
 			return err
 		}
 
-		return nil
+		if l.RequireCompleteInput || vcursor.Session().InTransaction() {
+			return nil
+		}
+
+		return io.EOF
 	})
 
+	if err == io.EOF {
+		return nil
+	}
 	if err != nil {
 		return err
 	}

diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go
@@ -22,12 +22,10 @@ import (
 	"strconv"
 	"strings"
 
-	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/slice"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/sysvars"
-	"vitess.io/vitess/go/vt/vtenv"
 	"vitess.io/vitess/go/vt/vterrors"
 	"vitess.io/vitess/go/vt/vtgate/engine"
 	"vitess.io/vitess/go/vt/vtgate/engine/opcode"
@@ -887,18 +885,18 @@ func transformUnionPlan(ctx *plancontext.PlanningContext, op *operators.Union) (
 }
 
 func transformLimit(ctx *plancontext.PlanningContext, op *operators.Limit) (engine.Primitive, error) {
-	plan, err := transformToPrimitive(ctx, op.Source)
+	input, err := transformToPrimitive(ctx, op.Source)
 	if err != nil {
 		return nil, err
 	}
 
-	return createLimit(plan, op.AST, ctx.VSchema.Environment(), ctx.VSchema.ConnCollation())
+	return createLimit(ctx, input, op.AST)
 }
 
-func createLimit(input engine.Primitive, limit *sqlparser.Limit, env *vtenv.Environment, coll collations.ID) (engine.Primitive, error) {
+func createLimit(ctx *plancontext.PlanningContext, input engine.Primitive, limit *sqlparser.Limit) (engine.Primitive, error) {
 	cfg := &evalengine.Config{
-		Collation:   coll,
-		Environment: env,
+		Collation:   ctx.VSchema.ConnCollation(),
+		Environment: ctx.VSchema.Environment(),
 	}
 	count, err := evalengine.Translate(limit.Rowcount, cfg)
 	if err != nil {
@@ -913,9 +911,10 @@ func createLimit(input engine.Primitive, limit *sqlparser.Limit, env *vtenv.Envi
 	}
 
 	return &engine.Limit{
-		Input:  input,
-		Count:  count,
-		Offset: offset,
+		Count:                count,
+		Offset:               offset,
+		RequireCompleteInput: ctx.SemTable.ShouldFetchLastInsertID(),
+		Input:                input,
 	}, nil
 }