opt: introduce placeholder fast path

Currently running a query with placeholders happens as follows:
 - at perparation time, we build a memo with the normalized expression
   *with* placeholders.
 - at execution time, we copy the expression into a new memo,
   replacing placeholders with their value (AssignPlaceholders).
 - then we run exploration which performs the actual optimization.

For trivial queries like KV reads, we do too much work during
execution (profiles show it is 10-20% of the runtime for KV workloads
with high read rates).

This commit introduces the concept of a "placeholder fast path". The
idea is that we can make specific checks for simple expressions and
produce (at preparation time) a fully optimized expression (which
still depends on placeholders). For this, we introduce a new operator,
PlaceholderScan which is similar to a Scan except that it always scans
one span with the same start and end key, and the key values are child
scalar expressions (either constants or placeholders).

We use this new operator only for simple SELECTs where the filters
constrain a prefix of an index to constant values; in addition, the
index must be covering and there must be only one such index. With
these conditions, we know the optimal plan upfront.

For now, the placeholder fast path transformation is Go code; if it
gets more complicated, it should be switched to use optgen rules.

A benchmark on my machine shows a kv95 workload going from 34kiops to
38kiops with this change.

Release note: None

pkg/sql/opt/bench/bench_test.go

@@ -294,9 +294,9 @@ func BenchmarkPhases(b *testing.B) {
  })
  b.Run("Prepared", func(b *testing.B) {
  phases := PreparedPhases
- if !h.prepMemo.HasPlaceholders() {
- // If the query has no placeholders, the only phase which does
- // something is ExecBuild.
+ if h.prepMemo.IsOptimized() {
+ // If the query has no placeholders or the placeholder fast path
+ // succeeded, the only phase which does something is ExecBuild.
  phases = []Phase{ExecBuild}
  }
  for _, phase := range phases {
@@ -358,6 +358,10 @@ func newHarness(tb testing.TB, query benchQuery) *harness {
  if _, err := h.optimizer.Optimize(); err != nil {
  tb.Fatalf("%v", err)
  }
+ } else {
+ if _, _, err := h.optimizer.TryPlaceholderFastPath(); err != nil {
+ tb.Fatalf("%v", err)
+ }
  }
  h.prepMemo = h.optimizer.DetachMemo()
  h.optimizer = xform.Optimizer{}
@@ -447,7 +451,7 @@ func (h *harness) runSimple(tb testing.TB, query benchQuery, phase Phase) {
 func (h *harness) runPrepared(tb testing.TB, phase Phase) {
  h.optimizer.Init(&h.evalCtx, h.testCat)
 
- if h.prepMemo.HasPlaceholders() {
+ if !h.prepMemo.IsOptimized() {
  if phase == AssignPlaceholdersNoNorm {
  h.optimizer.DisableOptimizations()
  }
@@ -462,8 +466,9 @@ func (h *harness) runPrepared(tb testing.TB, phase Phase) {
  }
 
  var execMemo *memo.Memo
- if !h.prepMemo.HasPlaceholders() {
- // No placeholders, we already did the exploration at prepare time.
+ if h.prepMemo.IsOptimized() {
+ // No placeholders or the placeholder fast path succeeded; we already did
+ // the exploration at prepare time.
  execMemo = h.prepMemo
  } else {
  if _, err := h.optimizer.Optimize(); err != nil {

pkg/sql/opt/exec/execbuilder/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
  "//pkg/sql/mutations",
  "//pkg/sql/opt",
  "//pkg/sql/opt/cat",
+ "//pkg/sql/opt/constraint",
  "//pkg/sql/opt/exec",
  "//pkg/sql/opt/memo",
  "//pkg/sql/opt/norm",

pkg/sql/opt/exec/execbuilder/mutation.go

@@ -719,7 +719,7 @@ func (b *Builder) buildDeleteRange(
  // We can't calculate the maximum number of keys if there are interleaved
  // children, as we don't know how many children rows may be in range.
  if len(interleavedTables) == 0 {
- if maxRows, ok := b.indexConstraintMaxResults(scan); ok {
+ if maxRows, ok := b.indexConstraintMaxResults(&scan.ScanPrivate, scan); ok {
  if maxKeys := maxRows * uint64(tab.FamilyCount()); maxKeys <= row.TableTruncateChunkSize {
  // Other mutations only allow auto-commit if there are no FK checks or
  // cascades. In this case, we won't actually execute anything for the

pkg/sql/opt/exec/execbuilder/relational.go

@@ -21,6 +21,7 @@ import (
  "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
  "github.com/cockroachdb/cockroach/pkg/sql/opt"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
+ "github.com/cockroachdb/cockroach/pkg/sql/opt/constraint"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
@@ -192,6 +193,9 @@ func (b *Builder) buildRelational(e memo.RelExpr) (execPlan, error) {
  case *memo.ScanExpr:
  ep, err = b.buildScan(t)
 
+ case *memo.PlaceholderScanExpr:
+ ep, err = b.buildPlaceholderScan(t)
+
  case *memo.SelectExpr:
  ep, err = b.buildSelect(t)
 
@@ -468,7 +472,9 @@ func (b *Builder) getColumns(
 // indexConstraintMaxResults returns the maximum number of results for a scan;
 // if successful (ok=true), the scan is guaranteed never to return more results
 // than maxRows.
-func (b *Builder) indexConstraintMaxResults(scan *memo.ScanExpr) (maxRows uint64, ok bool) {
+func (b *Builder) indexConstraintMaxResults(
+ scan *memo.ScanPrivate, expr memo.RelExpr,
+) (maxRows uint64, ok bool) {
  c := scan.Constraint
  if c == nil || c.IsContradiction() || c.IsUnconstrained() {
  return 0, false
@@ -479,7 +485,7 @@ func (b *Builder) indexConstraintMaxResults(scan *memo.ScanExpr) (maxRows uint64
  for i := 0; i < numCols; i++ {
  indexCols.Add(c.Columns.Get(i).ID())
  }
- rel := scan.Relational()
+ rel := expr.Relational()
  if !rel.FuncDeps.ColsAreLaxKey(indexCols) {
  return 0, false
  }
@@ -488,7 +494,7 @@ func (b *Builder) indexConstraintMaxResults(scan *memo.ScanExpr) (maxRows uint64
 }
 
 func (b *Builder) scanParams(
- tab cat.Table, scan *memo.ScanExpr,
+ tab cat.Table, scan *memo.ScanPrivate, relExpr memo.RelExpr,
 ) (exec.ScanParams, opt.ColMap, error) {
  // Check if we tried to force a specific index but there was no Scan with that
  // index in the memo.
@@ -535,28 +541,29 @@ func (b *Builder) scanParams(
 
  needed, outputMap := b.getColumns(scan.Cols, scan.Table)
 
- // Get the estimated row count from the statistics.
+ // Get the estimated row count from the statistics. When there are no
+ // statistics available, we construct a scan node with
+ // the estimated row count of zero rows.
+ //
  // Note: if this memo was originally created as part of a PREPARE
  // statement or was stored in the query cache, the column stats would have
  // been removed by DetachMemo. Update that function if the column stats are
  // needed here in the future.
- rowCount := scan.Relational().Stats.RowCount
- if !scan.Relational().Stats.Available {
- // When there are no statistics available, we construct a scan node with
- // the estimated row count of zero rows.
- rowCount = 0
+ var rowCount float64
+ if relProps := relExpr.Relational(); relProps.Stats.Available {
+ rowCount = relProps.Stats.RowCount
  }
 
  if scan.PartitionConstrainedScan {
  sqltelemetry.IncrementPartitioningCounter(sqltelemetry.PartitionConstrainedScan)
  }
 
- softLimit := int64(math.Ceil(scan.RequiredPhysical().LimitHint))
+ softLimit := int64(math.Ceil(relExpr.RequiredPhysical().LimitHint))
  hardLimit := scan.HardLimit.RowCount()
 
  parallelize := false
  if hardLimit == 0 && softLimit == 0 {
- maxResults, ok := b.indexConstraintMaxResults(scan)
+ maxResults, ok := b.indexConstraintMaxResults(scan, relExpr)
  if ok && maxResults < ParallelScanResultThreshold {
  // Don't set the flag when we have a single span which returns a single
  // row: it does nothing in this case except litter EXPLAINs.
@@ -568,18 +575,28 @@ func (b *Builder) scanParams(
  }
  }
 
+ // Figure out if we need to scan in reverse (ScanPrivateCanProvide takes
+ // HardLimit.Reverse() into account).
+ ok, reverse := ordering.ScanPrivateCanProvide(
+ b.mem.Metadata(),
+ scan,
+ &relExpr.RequiredPhysical().Ordering,
+ )
+ if !ok {
+ return exec.ScanParams{}, opt.ColMap{}, errors.AssertionFailedf("scan can't provide required ordering")
+ }
+
  return exec.ScanParams{
  NeededCols: needed,
  IndexConstraint: scan.Constraint,
  InvertedConstraint: scan.InvertedConstraint,
  HardLimit: hardLimit,
  SoftLimit: softLimit,
- // HardLimit.Reverse() is taken into account by ScanIsReverse.
- Reverse: ordering.ScanIsReverse(scan, &scan.RequiredPhysical().Ordering),
- Parallelize: parallelize,
- Locking: locking,
- EstimatedRowCount: rowCount,
- LocalityOptimized: scan.LocalityOptimized,
+ Reverse: reverse,
+ Parallelize: parallelize,
+ Locking: locking,
+ EstimatedRowCount: rowCount,
+ LocalityOptimized: scan.LocalityOptimized,
  }, outputMap, nil
 }
 
@@ -591,7 +608,7 @@ func (b *Builder) buildScan(scan *memo.ScanExpr) (execPlan, error) {
  telemetry.Inc(sqltelemetry.PartialIndexScanUseCounter)
  }
 
- params, outputCols, err := b.scanParams(tab, scan)
+ params, outputCols, err := b.scanParams(tab, &scan.ScanPrivate, scan)
  if err != nil {
  return execPlan{}, err
  }
@@ -620,6 +637,72 @@ func (b *Builder) buildScan(scan *memo.ScanExpr) (execPlan, error) {
  return res, nil
 }
 
+func (b *Builder) buildPlaceholderScan(scan *memo.PlaceholderScanExpr) (execPlan, error) {
+ if scan.Constraint != nil || scan.InvertedConstraint != nil {
+ return execPlan{}, errors.AssertionFailedf("PlaceholderScan cannot have constraints")
+ }
+
+ md := b.mem.Metadata()
+ tab := md.Table(scan.Table)
+ idx := tab.Index(scan.Index)
+
+ // Build the index constraint.
+ spanColumns := make([]opt.OrderingColumn, len(scan.Span))
+ for i := range spanColumns {
+ col := idx.Column(i)
+ ordinal := col.Ordinal()
+ colID := scan.Table.ColumnID(ordinal)
+ spanColumns[i] = opt.MakeOrderingColumn(colID, col.Descending)
+ }
+ var columns constraint.Columns
+ columns.Init(spanColumns)
+ keyCtx := constraint.MakeKeyContext(&columns, b.evalCtx)
+
+ values := make([]tree.Datum, len(scan.Span))
+ for i, expr := range scan.Span {
+ // The expression is either a placeholder or a constant.
+ if p, ok := expr.(*memo.PlaceholderExpr); ok {
+ val, err := p.Value.(*tree.Placeholder).Eval(b.evalCtx)
+ if err != nil {
+ return execPlan{}, err
+ }
+ values[i] = val
+ } else {
+ values[i] = memo.ExtractConstDatum(expr)
+ }
+ }
+
+ key := constraint.MakeCompositeKey(values...)
+ var span constraint.Span
+ span.Init(key, constraint.IncludeBoundary, key, constraint.IncludeBoundary)
+ var spans constraint.Spans
+ spans.InitSingleSpan(&span)
+
+ var c constraint.Constraint
+ c.Init(&keyCtx, &spans)
+
+ private := scan.ScanPrivate
+ private.Constraint = &c
+
+ params, outputCols, err := b.scanParams(tab, &private, scan)
+ if err != nil {
+ return execPlan{}, err
+ }
+ res := execPlan{outputCols: outputCols}
+ root, err := b.factory.ConstructScan(
+ tab,
+ tab.Index(scan.Index),
+ params,
+ res.reqOrdering(scan),
+ )
+ if err != nil {
+ return execPlan{}, err
+ }
+
+ res.root = root
+ return res, nil
+}
+
 func (b *Builder) buildSelect(sel *memo.SelectExpr) (execPlan, error) {
  input, err := b.buildRelational(sel.Input)
  if err != nil {

pkg/sql/opt/exec/execbuilder/testdata/prepare

@@ -102,3 +102,47 @@ vectorized: true
  estimated row count: 1 (100% of the table; stats collected <hidden> ago)
  table: ab@primary
  spans: FULL SCAN
+
+# Verify the plan of a very simple query which should be using the placeholder
+# fast path.
+statement ok
+PREPARE pklookup AS SELECT b FROM ab WHERE a = $1
+
+query T
+EXPLAIN ANALYZE EXECUTE pklookup(1)
+----
+planning time: 10µs
+execution time: 100µs
+distribution: <hidden>
+vectorized: <hidden>
+rows read from KV: 1 (8 B)
+maximum memory usage: <hidden>
+network usage: <hidden>
+·
+• scan
+ cluster nodes: <hidden>
+ actual row count: 1
+ KV rows read: 1
+ KV bytes read: 8 B
+ estimated row count: 0
+ table: ab@primary
+ spans: [/1 - /1]
+
+query T
+EXPLAIN ANALYZE EXECUTE pklookup(2)
+----
+planning time: 10µs
+execution time: 100µs
+distribution: <hidden>
+vectorized: <hidden>
+maximum memory usage: <hidden>
+network usage: <hidden>
+·
+• scan
+ cluster nodes: <hidden>
+ actual row count: 0
+ KV rows read: 0
+ KV bytes read: 0 B
+ estimated row count: 0
+ table: ab@primary
+ spans: [/2 - /2]

pkg/sql/opt/memo/check_expr.go

@@ -40,7 +40,7 @@ func (m *Memo) CheckExpr(e opt.Expr) {
  // If the expression was added to an existing group, cross-check its
  // properties against the properties of the group. Skip this check if the
  // operator is known to not have code for building logical props.
- if t != t.FirstExpr() && t.Op() != opt.MergeJoinOp {
+ if t != t.FirstExpr() && t.Op() != opt.MergeJoinOp && t.Op() != opt.PlaceholderScanOp {
  var relProps props.Relational
  // Don't build stats when verifying logical props - unintentionally
  // building stats for non-normalized expressions could add extra colStats