Merge #138692

138692: sql: refactor `checkScanParallelizationIfLocal` r=mgartner a=mgartner

#### sql: replace `walkPlan` in `checkScanParallelizationIfLocal`

`checkScanParallelizationIfLocal` now traverses a `planNode` tree using
the `Input` and `InputCount` methods instead of `walkPlan`.

Epic: None
Release note: None

#### sql: fix checkScanParallelizationIfLocal for groupNode

This commit fixes a bug in `checkScanParallelizationIfLocal` that caused
it to ignore all but the last aggregate function in a `groupNode`. This
has been fixed and a test has been added. In addition, an early return
has been added to the `renderNode` case.

Release note: None


Co-authored-by: Marcus Gartner <[email protected]>

Author: craig[bot]

pkg/sql/distsql_physical_planner.go

@@ -5289,7 +5289,7 @@ func logAndSanitizeExportDestination(ctx context.Context, dest string) error {
 // processors eagerly move into the draining state which will cancel the context
 // of parallel TableReaders which might "poison" the transaction.
 func checkScanParallelizationIfLocal(
-	ctx context.Context, plan *planComponents, c *localScanParallelizationChecker,
+	ctx context.Context, plan *planComponents,
 ) (prohibitParallelization, hasScanNodeToParallelize bool) {
 	if plan.main.planNode == nil || len(plan.cascades) != 0 ||
 		len(plan.checkPlans) != 0 || len(plan.triggers) != 0 {
@@ -5298,82 +5298,76 @@ func checkScanParallelizationIfLocal(
 		// the scan parallelization.
 		return true, false
 	}
-	*c = localScanParallelizationChecker{}
-	o := planObserver{enterNode: c.enterNode}
-	_ = walkPlan(ctx, plan.main.planNode, o)
-	for _, s := range plan.subqueryPlans {
-		_ = walkPlan(ctx, s.plan.planNode, o)
-	}
-	return c.prohibitParallelization, c.hasScanNodeToParallelize
-}
-
-type localScanParallelizationChecker struct {
-	prohibitParallelization  bool
-	hasScanNodeToParallelize bool
-}
 
-func (c *localScanParallelizationChecker) enterNode(
-	ctx context.Context, _ string, plan planNode,
-) (bool, error) {
-	if c.prohibitParallelization {
-		return false, nil
-	}
-	switch n := plan.(type) {
-	case *distinctNode:
-		return true, nil
-	case *explainPlanNode:
-		// walkPlan doesn't recurse into explainPlanNode, so we have to manually
-		// walk over the wrapped plan.
-		plan := n.plan.WrappedPlan.(*planComponents)
-		prohibit, has := checkScanParallelizationIfLocal(ctx, plan, c)
-		c.prohibitParallelization = c.prohibitParallelization || prohibit
-		c.hasScanNodeToParallelize = c.hasScanNodeToParallelize || has
-		return false, nil
-	case *explainVecNode:
-		return true, nil
-	case *filterNode:
-		// Some filter expressions might be handled by falling back to the
-		// wrapped processors, so we choose to be safe.
-		c.prohibitParallelization = true
-		return false, nil
-	case *groupNode:
-		for _, f := range n.funcs {
-			c.prohibitParallelization = f.hasFilter()
+	var checkRec func(p planNode)
+	checkRec = func(p planNode) {
+		if prohibitParallelization {
+			return
 		}
-		return true, nil
-	case *indexJoinNode:
-		return true, nil
-	case *joinNode:
-		c.prohibitParallelization = n.pred.onCond != nil
-		return true, nil
-	case *limitNode:
-		return true, nil
-	case *ordinalityNode:
-		return true, nil
-	case *renderNode:
-		// Only support projections since render expressions might be handled
-		// via a wrapped row-by-row processor.
-		for _, e := range n.render {
-			if _, isIVar := e.(*tree.IndexedVar); !isIVar {
-				c.prohibitParallelization = true
+		switch n := p.(type) {
+		case *explainPlanNode:
+			// explainPlanNode is a zeroInputPlanNode, so we have to manually
+			// recurse.
+			plan := n.plan.WrappedPlan.(*planComponents)
+			prohibit, has := checkScanParallelizationIfLocal(ctx, plan)
+			prohibitParallelization = prohibitParallelization || prohibit
+			hasScanNodeToParallelize = hasScanNodeToParallelize || has
+			// Do not recurse.
+			return
+		case *filterNode:
+			// Some filter expressions might be handled by falling back to the
+			// wrapped processors, so we choose to be safe.
+			prohibitParallelization = true
+			// Do not recurse.
+			return
+		case *groupNode:
+			for _, f := range n.funcs {
+				if f.hasFilter() {
+					prohibitParallelization = true
+					// Do not recurse.
+					return
+				}
+			}
+		case *joinNode:
+			prohibitParallelization = n.pred.onCond != nil
+		case *renderNode:
+			// Only support projections since render expressions might be
+			// handled via a wrapped row-by-row processor.
+			for _, e := range n.render {
+				if _, isIVar := e.(*tree.IndexedVar); !isIVar {
+					prohibitParallelization = true
+					// Do not recurse.
+					return
+				}
 			}
+		case *scanNode:
+			if len(n.reqOrdering) == 0 && n.parallelize {
+				hasScanNodeToParallelize = true
+			}
+		case *distinctNode, *explainVecNode, *indexJoinNode, *limitNode,
+			*ordinalityNode, *sortNode, *unionNode, *valuesNode:
+		default:
+			prohibitParallelization = true
+			// Do not recurse into any unmatched nodes.
+			return
 		}
-		return true, nil
-	case *scanNode:
-		if len(n.reqOrdering) == 0 && n.parallelize {
-			c.hasScanNodeToParallelize = true
+		// Recurse into input nodes.
+		for i, n := 0, p.InputCount(); i < n; i++ {
+			input, err := p.Input(i)
+			if err != nil {
+				// This error should never occur with correct usage of Input
+				// and InputCount. If it does, ignore it and prohibit
+				// parallelization.
+				prohibitParallelization = true
+			}
+			checkRec(input)
 		}
-		return true, nil
-	case *sortNode:
-		return true, nil
-	case *unionNode:
-		return true, nil
-	case *valuesNode:
-		return true, nil
-	default:
-		c.prohibitParallelization = true
-		return false, nil
 	}
+	checkRec(plan.main.planNode)
+	for _, s := range plan.subqueryPlans {
+		checkRec(s.plan.planNode)
+	}
+	return prohibitParallelization, hasScanNodeToParallelize
 }
 
 // NewPlanningCtx returns a new PlanningCtx. When distribute is false, a
@@ -5435,7 +5429,7 @@ func (dsp *DistSQLPlanner) NewPlanningCtxWithOracle(
 			return planCtx
 		}
 		prohibitParallelization, hasScanNodeToParallelize := checkScanParallelizationIfLocal(
-			ctx, &planner.curPlan.planComponents, &planner.parallelizationChecker,
+			ctx, &planner.curPlan.planComponents,
 		)
 		if prohibitParallelization || !hasScanNodeToParallelize {
 			return planCtx

pkg/sql/distsql_physical_planner_test.go

@@ -1994,6 +1994,17 @@ func TestCheckScanParallelizationIfLocal(t *testing.T) {
 			// Filtering aggregation is not natively supported.
 			prohibitParallelization: true,
 		},
+		{
+			plan: planComponents{main: planMaybePhysical{planNode: &groupNode{
+				singleInputPlanNode: singleInputPlanNode{scanToParallelize},
+				funcs: []*aggregateFuncHolder{
+					{filterRenderIdx: 0},
+					{filterRenderIdx: tree.NoColumnIdx},
+				}},
+			}},
+			// Filtering aggregation is not natively supported.
+			prohibitParallelization: true,
+		},
 		{
 			plan: planComponents{main: planMaybePhysical{planNode: &indexJoinNode{
 				singleInputPlanNode: singleInputPlanNode{scanToParallelize},
@@ -2028,6 +2039,16 @@ func TestCheckScanParallelizationIfLocal(t *testing.T) {
 			// prohibit the parallelization for all non-IndexedVar expressions).
 			prohibitParallelization: true,
 		},
+		{
+			plan: planComponents{main: planMaybePhysical{planNode: &renderNode{
+				singleInputPlanNode: singleInputPlanNode{scanToParallelize},
+				render:              []tree.TypedExpr{&tree.IndexedVar{Idx: 0}, &tree.IsNullExpr{}},
+			}}},
+			// Not a simple projection (some expressions might be handled by
+			// wrapping a row-execution processor, so we choose to be safe and
+			// prohibit the parallelization for all non-IndexedVar expressions).
+			prohibitParallelization: true,
+		},
 		{
 			plan:                     planComponents{main: planMaybePhysical{planNode: &sortNode{singleInputPlanNode: singleInputPlanNode{scanToParallelize}}}},
 			hasScanNodeToParallelize: true,
@@ -2064,8 +2085,7 @@ func TestCheckScanParallelizationIfLocal(t *testing.T) {
 			prohibitParallelization: true,
 		},
 	} {
-		var c localScanParallelizationChecker
-		prohibitParallelization, hasScanNodeToParallize := checkScanParallelizationIfLocal(context.Background(), &tc.plan, &c)
+		prohibitParallelization, hasScanNodeToParallize := checkScanParallelizationIfLocal(context.Background(), &tc.plan)
 		require.Equal(t, tc.prohibitParallelization, prohibitParallelization)
 		require.Equal(t, tc.hasScanNodeToParallelize, hasScanNodeToParallize)
 	}

pkg/sql/planner.go

@@ -285,9 +285,6 @@ type planner struct {
 	// This field is embedded into the planner to avoid an allocation in
 	// checkExprForDistSQL.
 	distSQLVisitor distSQLExprCheckVisitor
-	// This field is embedded into the planner to avoid an allocation in
-	// checkScanParallelizationIfLocal.
-	parallelizationChecker localScanParallelizationChecker
 
 	// datumAlloc is used when decoding datums and running subqueries.
 	datumAlloc *tree.DatumAlloc