Skip to content

Commit

Permalink
sql: new traversal method of planNodes in wrapPlan
Browse files Browse the repository at this point in the history
`wrapPlan` now traverses the `planNode` tree with a recursive function
that uses the `Input` and `InputCount` methods instead of the
`planVisitor`. This reduces the overhead of traversal.

Below are some benchmarks that show the improvement.

```
name                                        old time/op    new time/op    delta
EndToEnd/kv-read/vectorize=on/Simple-16        213µs ± 1%     205µs ± 1%  -4.02%  (p=0.000 n=20+20)
EndToEnd/kv-read/vectorize=off/Simple-16       207µs ± 1%     199µs ± 1%  -3.85%  (p=0.000 n=19+20)
EndToEnd/kv-read/vectorize=on/Prepared-16      146µs ± 1%     141µs ± 1%  -2.98%  (p=0.000 n=19+19)
EndToEnd/kv-read/vectorize=off/Prepared-16     140µs ± 1%     136µs ± 1%  -2.94%  (p=0.000 n=20+18)

name                                        old alloc/op   new alloc/op   delta
EndToEnd/kv-read/vectorize=on/Prepared-16     20.7kB ± 0%    20.6kB ± 0%  -0.23%  (p=0.000 n=19+19)
EndToEnd/kv-read/vectorize=on/Simple-16       33.4kB ± 0%    33.4kB ± 0%    ~     (p=0.670 n=19+19)
EndToEnd/kv-read/vectorize=off/Simple-16      33.8kB ± 0%    33.7kB ± 0%    ~     (p=0.760 n=18+18)
EndToEnd/kv-read/vectorize=off/Prepared-16    21.9kB ± 0%    21.9kB ± 0%    ~     (p=0.365 n=18+20)

name                                        old allocs/op  new allocs/op  delta
EndToEnd/kv-read/vectorize=on/Simple-16          283 ± 0%       282 ± 0%    ~     (p=0.105 n=20+20)
EndToEnd/kv-read/vectorize=on/Prepared-16        190 ± 0%       189 ± 0%    ~     (p=0.150 n=19+20)
EndToEnd/kv-read/vectorize=off/Simple-16         270 ± 0%       270 ± 0%    ~     (all equal)
EndToEnd/kv-read/vectorize=off/Prepared-16       182 ± 0%       182 ± 0%    ~     (p=0.447 n=18+20)
```

Release note: None
  • Loading branch information
mgartner committed Dec 20, 2024
1 parent d64d0ae commit 79bb063
Showing 1 changed file with 34 additions and 33 deletions.
67 changes: 34 additions & 33 deletions pkg/sql/distsql_physical_planner.go
Original file line number Diff line number Diff line change
Expand Up @@ -4148,26 +4148,27 @@ func (dsp *DistSQLPlanner) wrapPlan(
// DistSQL-enabled planNode in the tree. If we find one, we ask the planner to
// continue the DistSQL planning recursion on that planNode.
seenTop := false
nParents := uint32(0)
p := planCtx.NewPhysicalPlan()
// This will be set to first DistSQL-enabled planNode we find, if any. We'll
// modify its parent later to connect its source to the DistSQL-planned
// subtree.
var firstNotWrapped planNode
if err := walkPlan(ctx, n, planObserver{
enterNode: func(ctx context.Context, nodeName string, plan planNode) (bool, error) {
switch plan.(type) {
case *explainVecNode, *explainPlanNode, *explainDDLNode:
// Don't continue recursing into explain nodes - they need to be left
// alone since they handle their own planning later.
return false, nil
}
if !seenTop {
// We know we're wrapping the first node, so ignore it.
seenTop = true
return true, nil
var planFirstDistSQLNode func(plan planNode) (planNode, *PhysicalPlan, error)
planFirstDistSQLNode = func(plan planNode) (planNode, *PhysicalPlan, error) {
switch plan.(type) {
case *explainVecNode, *explainPlanNode, *explainDDLNode:
// Don't continue recursing into explain nodes - they need to be left
// alone since they handle their own planning later.
return nil, nil, nil
}
if !seenTop {
seenTop = true
} else if !dsp.mustWrapNode(planCtx, plan) || shouldWrapPlanNodeForExecStats(planCtx, plan) {
p, err := dsp.createPhysPlanForPlanNode(ctx, planCtx, plan)
if err != nil {
return nil, nil, err
}
var err error
return plan, p, nil
}
switch plan.InputCount() {
case 0:
return nil, nil, nil
case 1:
// Continue walking until we find a node that has a DistSQL
// representation - that's when we'll quit the wrapping process and
// hand control of planning back to the DistSQL physical planner.
Expand All @@ -4177,22 +4178,22 @@ func (dsp *DistSQLPlanner) wrapPlan(
// rowSourceToPlanNode adapters so that the execution statistics
// are collected for each planNode independently. This should have
// low enough overhead.
if !dsp.mustWrapNode(planCtx, plan) || shouldWrapPlanNodeForExecStats(planCtx, plan) {
firstNotWrapped = plan
p, err = dsp.createPhysPlanForPlanNode(ctx, planCtx, plan)
if err != nil {
return false, err
}
nParents++
return false, nil
input, err := plan.Input(0)
if err != nil {
return nil, nil, err
}
return true, nil
},
}); err != nil {
return planFirstDistSQLNode(input)
default:
// Can't wrap plans with more than 1 input.
return nil, nil, nil
}
}
firstNotWrapped, p, err := planFirstDistSQLNode(n)
if err != nil {
return nil, err
}
if nParents > 1 {
return nil, errors.Errorf("can't wrap plan %v %T with more than one input", n, n)
if p == nil {
p = planCtx.NewPhysicalPlan()
}

// Copy the evalCtx.
Expand Down Expand Up @@ -4228,7 +4229,7 @@ func (dsp *DistSQLPlanner) wrapPlan(
Input: input,
Core: execinfrapb.ProcessorCoreUnion{LocalPlanNode: &execinfrapb.LocalPlanNodeSpec{
RowSourceIdx: uint32(localProcIdx),
NumInputs: nParents,
NumInputs: uint32(len(input)),
Name: name,
}},
Post: execinfrapb.PostProcessSpec{},
Expand Down

0 comments on commit 79bb063

Please sign in to comment.