Skip to content

Commit bc3a4f9

Browse files
committed
sql: make getPlanDistribution a method on planner
This will simplify the following commit a bit. Release note: None
1 parent e5f821c commit bc3a4f9

File tree

7 files changed

+12
-37
lines changed

7 files changed

+12
-37
lines changed

pkg/sql/apply_join.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -315,10 +315,7 @@ func runPlanInsidePlan(
315315
}
316316
}
317317

318-
distributePlan, distSQLProhibitedErr := getPlanDistribution(
319-
ctx, plannerCopy.Descriptors().HasUncommittedTypes(),
320-
plannerCopy.SessionData(), plan.main, &plannerCopy.distSQLVisitor,
321-
)
318+
distributePlan, distSQLProhibitedErr := plannerCopy.getPlanDistribution(ctx, plan.main)
322319
distributeType := DistributionType(LocalDistribution)
323320
if distributePlan.WillDistribute() {
324321
distributeType = FullDistribution

pkg/sql/conn_executor_exec.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2847,10 +2847,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
28472847
}
28482848
}
28492849
}
2850-
distributePlan, distSQLProhibitedErr := getPlanDistribution(
2851-
ctx, planner.Descriptors().HasUncommittedTypes(),
2852-
ex.sessionData(), planner.curPlan.main, &planner.distSQLVisitor,
2853-
)
2850+
distributePlan, distSQLProhibitedErr := planner.getPlanDistribution(ctx, planner.curPlan.main)
28542851
if afterGetPlanDistribution != nil {
28552852
afterGetPlanDistribution()
28562853
}

pkg/sql/distsql_running.go

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1901,10 +1901,7 @@ func (dsp *DistSQLPlanner) planAndRunSubquery(
19011901
skipDistSQLDiagramGeneration bool,
19021902
mustUseLeafTxn bool,
19031903
) error {
1904-
subqueryDistribution, distSQLProhibitedErr := getPlanDistribution(
1905-
ctx, planner.Descriptors().HasUncommittedTypes(),
1906-
planner.SessionData(), subqueryPlan.plan, &planner.distSQLVisitor,
1907-
)
1904+
subqueryDistribution, distSQLProhibitedErr := planner.getPlanDistribution(ctx, subqueryPlan.plan)
19081905
distribute := DistributionType(LocalDistribution)
19091906
if subqueryDistribution.WillDistribute() {
19101907
distribute = FullDistribution
@@ -2512,10 +2509,7 @@ func (dsp *DistSQLPlanner) planAndRunPostquery(
25122509
associateNodeWithComponents func(exec.Node, execComponents),
25132510
addTopLevelQueryStats func(stats *topLevelQueryStats),
25142511
) error {
2515-
postqueryDistribution, distSQLProhibitedErr := getPlanDistribution(
2516-
ctx, planner.Descriptors().HasUncommittedTypes(),
2517-
planner.SessionData(), postqueryPlan, &planner.distSQLVisitor,
2518-
)
2512+
postqueryDistribution, distSQLProhibitedErr := planner.getPlanDistribution(ctx, postqueryPlan)
25192513
distribute := DistributionType(LocalDistribution)
25202514
if postqueryDistribution.WillDistribute() {
25212515
distribute = FullDistribution

pkg/sql/exec_util.go

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2220,12 +2220,8 @@ func shouldDistributeGivenRecAndMode(
22202220
// remote node to the gateway.
22212221
// TODO(yuzefovich): this will be easy to solve once the DistSQL spec factory is
22222222
// completed but is quite annoying to do at the moment.
2223-
func getPlanDistribution(
2224-
ctx context.Context,
2225-
txnHasUncommittedTypes bool,
2226-
sd *sessiondata.SessionData,
2227-
plan planMaybePhysical,
2228-
distSQLVisitor *distSQLExprCheckVisitor,
2223+
func (p *planner) getPlanDistribution(
2224+
ctx context.Context, plan planMaybePhysical,
22292225
) (_ physicalplan.PlanDistribution, distSQLProhibitedErr error) {
22302226
if plan.isPhysicalPlan() {
22312227
// TODO(#47473): store the distSQLProhibitedErr for DistSQL spec factory
@@ -2236,10 +2232,11 @@ func getPlanDistribution(
22362232
// If this transaction has modified or created any types, it is not safe to
22372233
// distribute due to limitations around leasing descriptors modified in the
22382234
// current transaction.
2239-
if txnHasUncommittedTypes {
2235+
if p.Descriptors().HasUncommittedDescriptors() {
22402236
return physicalplan.LocalPlan, nil
22412237
}
22422238

2239+
sd := p.SessionData()
22432240
if sd.DistSQLMode == sessiondatapb.DistSQLOff {
22442241
return physicalplan.LocalPlan, nil
22452242
}
@@ -2249,7 +2246,7 @@ func getPlanDistribution(
22492246
return physicalplan.LocalPlan, nil
22502247
}
22512248

2252-
rec, err := checkSupportForPlanNode(ctx, plan.planNode, distSQLVisitor, sd)
2249+
rec, err := checkSupportForPlanNode(ctx, plan.planNode, &p.distSQLVisitor, sd)
22532250
if err != nil {
22542251
// Don't use distSQL for this request.
22552252
log.VEventf(ctx, 1, "query not supported for distSQL: %s", err)

pkg/sql/explain_plan.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,10 +58,7 @@ func (e *explainPlanNode) startExec(params runParams) error {
5858
// Note that we delay adding the annotation about the distribution until
5959
// after the plan is finalized (when the physical plan is successfully
6060
// created).
61-
distribution, _ := getPlanDistribution(
62-
params.ctx, params.p.Descriptors().HasUncommittedTypes(),
63-
params.extendedEvalCtx.SessionData(), plan.main, &params.p.distSQLVisitor,
64-
)
61+
distribution, _ := params.p.getPlanDistribution(params.ctx, plan.main)
6562

6663
outerSubqueries := params.p.curPlan.subqueryPlans
6764
distSQLPlanner := params.extendedEvalCtx.DistSQLPlanner

pkg/sql/explain_vec.go

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,7 @@ type explainVecNode struct {
3636
func (n *explainVecNode) startExec(params runParams) error {
3737
n.run.values = make(tree.Datums, 1)
3838
distSQLPlanner := params.extendedEvalCtx.DistSQLPlanner
39-
distribution, _ := getPlanDistribution(
40-
params.ctx, params.p.Descriptors().HasUncommittedTypes(),
41-
params.extendedEvalCtx.SessionData(), n.plan.main, &params.p.distSQLVisitor,
42-
)
39+
distribution, _ := params.p.getPlanDistribution(params.ctx, n.plan.main)
4340
outerSubqueries := params.p.curPlan.subqueryPlans
4441
planCtx := newPlanningCtxForExplainPurposes(distSQLPlanner, params, n.plan.subqueryPlans, distribution)
4542
defer func() {

pkg/sql/schema_changer.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -528,11 +528,7 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
528528
}
529529
}
530530
}
531-
planDistribution, _ := getPlanDistribution(
532-
ctx, localPlanner.Descriptors().HasUncommittedTypes(),
533-
localPlanner.extendedEvalCtx.SessionData(),
534-
localPlanner.curPlan.main, &localPlanner.distSQLVisitor,
535-
)
531+
planDistribution, _ := localPlanner.getPlanDistribution(ctx, localPlanner.curPlan.main)
536532
isLocal := !planDistribution.WillDistribute()
537533
out := execinfrapb.ProcessorCoreUnion{BulkRowWriter: &execinfrapb.BulkRowWriterSpec{
538534
Table: *table.TableDesc(),

0 commit comments

Comments
 (0)