sql: dynamically disable vectorized execution

Author: mgartner

pkg/sql/distsql_running.go

@@ -459,25 +459,29 @@ func (dsp *DistSQLPlanner) setupFlows(
 		CollectStats:      planCtx.collectExecStats,
 		StatementSQL:      statementSQL,
 	}
-	if localState.IsLocal {
-		// VectorizeMode is the only field that the setup code expects to be set
-		// in the local flows.
-		setupReq.EvalContext.SessionData.VectorizeMode = evalCtx.SessionData().VectorizeMode
-	} else {
-		// In distributed plans populate some extra state.
+	if !localState.IsLocal {
+		// In distributed plans populate some extra state. VectorizeMode is the
+		// only field that the setup code expects to be set in the local flows.
 		setupReq.EvalContext = execinfrapb.MakeEvalContext(evalCtx)
 		if jobTag, ok := logtags.FromContext(ctx).GetTag("job"); ok {
 			setupReq.JobTag = jobTag.ValueStr()
 		}
 	}
+	vectorizeMode := evalCtx.SessionData().VectorizeMode
+	if planCtx.planner != nil && planCtx.planner.curPlan.avoidVectorization {
+		// Vectorization has been automatically disabled for this plan.
+		vectorizeMode = sessiondatapb.VectorizeOff
+	}
+	setupReq.EvalContext.SessionData.VectorizeMode = vectorizeMode
+	// In local and distributed plans set the vectorized mode.
 	if evalCtx.SessionData().PropagateAdmissionHeaderToLeafTransactions && localState.Txn != nil {
 		// Propagate the admission control header so that leaf transactions
 		// correctly inherit it.
 		setupReq.LeafTxnAdmissionHeader = localState.Txn.AdmissionHeader()
 	}
 
 	var isVectorized bool
-	if vectorizeMode := evalCtx.SessionData().VectorizeMode; vectorizeMode != sessiondatapb.VectorizeOff {
+	if vectorizeMode != sessiondatapb.VectorizeOff {
 		// Now we determine whether the vectorized engine supports the flow
 		// specs.
 		isVectorized = true

pkg/sql/explain_plan.go

@@ -93,7 +93,7 @@ func (e *explainPlanNode) startExec(params runParams) error {
 
 			ctxSessionData := planCtx.EvalContext().SessionData()
 			var willVectorize bool
-			if ctxSessionData.VectorizeMode == sessiondatapb.VectorizeOff {
+			if ctxSessionData.VectorizeMode == sessiondatapb.VectorizeOff || params.p.curPlan.avoidVectorization {
 				willVectorize = false
 			} else {
 				willVectorize = true

pkg/sql/opt/memo/cost.go

@@ -20,6 +20,9 @@ type Cost struct {
 		// fullScanCount is the number of full table or index scans in a
 		// sub-plan, up to 255.
 		fullScanCount uint8
+		// maxCardinality is the maximum cardinality in any expression in a
+		// sub-plan.
+		maxCardinality uint32
 	}
 }
 
@@ -63,6 +66,7 @@ func (c *Cost) Add(other Cost) {
 	} else {
 		c.aux.fullScanCount += other.aux.fullScanCount
 	}
+	c.aux.maxCardinality = max(c.aux.maxCardinality, other.aux.maxCardinality)
 }
 
 // FullScanCount returns the number of full scans in the cost.
@@ -79,6 +83,17 @@ func (c *Cost) IncrFullScanCount() {
 	c.aux.fullScanCount++
 }
 
+// MaxCardinality returns the maximum cardinality in the cost.
+func (c Cost) MaxCardinality() uint32 {
+	return c.aux.maxCardinality
+}
+
+// RatchetMaxCardinality ratchets the max cardinality in the cost to card. No-op
+// if the max cardinality is greater than or equal to card.
+func (c *Cost) RatchetMaxCardinality(card uint32) {
+	c.aux.maxCardinality = max(c.aux.maxCardinality, card)
+}
+
 // CostFlags contains flags that penalize the cost of an operator.
 type CostFlags struct {
 	// FullScanPenalty is true if the cost of a full table or index scan is

pkg/sql/opt/xform/coster.go

@@ -622,12 +622,17 @@ func (c *coster) ComputeCost(candidate memo.RelExpr, required *physical.Required
 	// preferable, all else being equal.
 	cost.C += cpuCostFactor
 
-	// Add a one-time cost for any operator with unbounded cardinality. This
-	// ensures we prefer plans that push limits as far down the tree as possible,
-	// all else being equal.
-	//
-	// Also add a cost flag for unbounded cardinality.
-	if candidate.Relational().Cardinality.IsUnbounded() {
+	// Ratchet the max cardinality and add a one-time cost for any operator with
+	// unbounded cardinality. This ensures we prefer plans that push limits as
+	// far down the tree as possible, all else being equal. Also add a cost flag
+	// for unbounded cardinality.
+	card := candidate.Relational().Cardinality
+	if candidate.Op() != opt.ExplainOp {
+		// Do not ratchet the max cardinality for EXPLAIN so that we retain the
+		// max cardinality of that we would get without EXPLAIN.
+		cost.RatchetMaxCardinality(card.Max)
+	}
+	if card.IsUnbounded() {
 		cost.C += cpuCostFactor
 		if c.evalCtx.SessionData().OptimizerPreferBoundedCardinality {
 			cost.Flags.UnboundedCardinality = true

pkg/sql/plan.go

@@ -356,6 +356,9 @@ type planTop struct {
 	// avoidBuffering, when set, causes the execution to avoid buffering
 	// results.
 	avoidBuffering bool
+	// avoidVectorization, when set, causes the execution to avoid using
+	// vectorized execution.
+	avoidVectorization bool
 
 	// If we are collecting query diagnostics, flow information, including
 	// diagrams, are saved here.

pkg/sql/plan_opt.go

@@ -961,7 +961,15 @@ func (opc *optPlanningCtx) runExecBuilder(
 	if opc.gf.Initialized() {
 		planTop.instrumentation.planGist = opc.gf.PlanGist()
 	}
-	planTop.instrumentation.costEstimate = mem.RootExpr().(memo.RelExpr).Cost().C
+
+	cost := mem.RootExpr().(memo.RelExpr).Cost()
+	if cost.MaxCardinality() <= 1 {
+		// Use the row-by-row engine if the entire plan is guaranteed to operate
+		// on one row or less.
+		planTop.avoidVectorization = true
+	}
+	planTop.instrumentation.costEstimate = cost.C
+
 	available := mem.RootExpr().(memo.RelExpr).Relational().Statistics().Available
 	planTop.instrumentation.statsAvailable = available
 	if available {