opt: fixup CTE stats on placeholder queries

During optbuilder phase we copy the initial expressions stats into the
fake-rel but this value can change when placeholders are assigned so add
code in AssignPlaceholders to rebuild the cte if the stats change.

Fixes: cockroachdb#99389
Epic: none
Release note (sql change): Prepared statements using placeholders in
recursive CTEs sometimes did not re-optimize correctly after plugging
in the parameters leading to poor plan choices, this has been fixed.

pkg/sql/opt/norm/factory.go

@@ -14,6 +14,7 @@ import (
  "github.com/cockroachdb/cockroach/pkg/sql/opt"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+ "github.com/cockroachdb/cockroach/pkg/sql/opt/props"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
  "github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
  "github.com/cockroachdb/cockroach/pkg/sql/types"
@@ -301,6 +302,28 @@ func (f *Factory) AssignPlaceholders(from *memo.Memo) (err error) {
  }
  return f.ConstructConstVal(d, placeholder.DataType())
  }
+ // A recursive CTE may have the stats change on its Initial expression
+ // after placeholder assignment, if that happens we need to
+ // propagate that change to the Binding expression and rebuild
+ // everything.
+ if rcte, ok := e.(*memo.RecursiveCTEExpr); ok {
+ newInitial := f.CopyAndReplaceDefault(rcte.Initial, replaceFn).(memo.RelExpr)
+ if newInitial != rcte.Initial {
+ newBinding := f.ConstructFakeRel(&memo.FakeRelPrivate{
+ Props: MakeBindingPropsForRecursiveCTE(
+ props.AnyCardinality, rcte.Binding.Relational().OutputCols,
+ newInitial.Relational().Stats.RowCount)})
+ if id := rcte.WithBindingID(); id != 0 {
+ f.Metadata().AddWithBinding(id, newBinding)
+ }
+ return f.ConstructRecursiveCTE(
+ newBinding,
+ newInitial,
+ f.invokeReplace(rcte.Recursive, replaceFn).(memo.RelExpr),
+ &rcte.RecursiveCTEPrivate,
+ )
+ }
+ }
  return f.CopyAndReplaceDefault(e, replaceFn)
  }
  f.CopyAndReplace(from.RootExpr().(memo.RelExpr), from.RootProps(), replaceFn)

pkg/sql/opt/norm/testdata/rules/with

@@ -867,3 +867,82 @@ scalar-group-by
  └── aggregations
  └── sum [as=sum:6, outer=(5)]
  └── n:5
+
+
+exec-ddl
+CREATE TABLE yz (y INT, z INT, INDEX (y) STORING (z));
+----
+
+exec-ddl
+ALTER TABLE yz INJECT STATISTICS '[
+ {
+ "columns": ["y"],
+ "created_at": "2018-01-01 1:00:00.00000+00:00",
+ "row_count": 100000,
+ "distinct_count": 100000
+ },
+ {
+ "columns": ["z"],
+ "created_at": "2018-01-01 1:00:00.00000+00:00",
+ "row_count": 100000,
+ "distinct_count": 1
+ }
+]';
+----
+
+# Regression test for #99389.
+assign-placeholders-norm query-args=(1) format=show-stats
+WITH RECURSIVE cte(a,b) AS (
+ (SELECT * FROM yz WHERE y = $1)
+ UNION ALL
+ (SELECT y+1, z FROM yz LEFT JOIN cte ON b = z)
+)
+SELECT * FROM cte;
+----
+project
+ ├── columns: a:16 b:17
+ ├── immutable
+ ├── stats: [rows=10]
+ ├── recursive-c-t-e
+ │ ├── columns: a:6 b:7
+ │ ├── working table binding: &1
+ │ ├── initial columns: y:1 z:2
+ │ ├── recursive columns: "?column?":15 z:9
+ │ ├── immutable
+ │ ├── stats: [rows=10]
+ │ ├── fake-rel
+ │ │ ├── columns: a:6 b:7
+ │ │ ├── cardinality: [1 - ]
+ │ │ └── stats: [rows=1.00001, distinct(7)=0.100001, null(7)=0.0100001, avgsize(7)=4]
+ │ ├── select
+ │ │ ├── columns: y:1!null z:2
+ │ │ ├── stats: [rows=1.00001, distinct(1)=1, null(1)=0, avgsize(1)=4]
+ │ │ ├── fd: ()-->(1)
+ │ │ ├── scan yz
+ │ │ │ ├── columns: y:1 z:2
+ │ │ │ └── stats: [rows=100000, distinct(1)=100000, null(1)=0, avgsize(1)=4]
+ │ │ └── filters
+ │ │ └── y:1 = 1 [outer=(1), constraints=(/1: [/1 - /1]; tight), fd=()-->(1)]
+ │ └── project
+ │ ├── columns: "?column?":15 z:9
+ │ ├── immutable
+ │ ├── stats: [rows=100001]
+ │ ├── left-join (hash)
+ │ │ ├── columns: y:8 z:9 b:14
+ │ │ ├── stats: [rows=100001, distinct(14)=1, null(14)=0, avgsize(14)=4]
+ │ │ ├── scan yz
+ │ │ │ ├── columns: y:8 z:9
+ │ │ │ └── stats: [rows=100000, distinct(9)=1, null(9)=0, avgsize(9)=4]
+ │ │ ├── with-scan &1 (cte)
+ │ │ │ ├── columns: b:14
+ │ │ │ ├── mapping:
+ │ │ │ │ └── b:7 => b:14
+ │ │ │ ├── cardinality: [1 - ]
+ │ │ │ └── stats: [rows=1.00001, distinct(14)=0.100001, null(14)=0.0100001, avgsize(14)=4]
+ │ │ └── filters
+ │ │ └── b:14 = z:9 [outer=(9,14), constraints=(/9: (/NULL - ]; /14: (/NULL - ]), fd=(9)==(14), (14)==(9)]
+ │ └── projections
+ │ └── y:8 + 1 [as="?column?":15, outer=(8), immutable]
+ └── projections
+ ├── a:6 [as=a:16, outer=(6)]
+ └── b:7 [as=b:17, outer=(7)]

pkg/sql/opt/norm/with_funcs.go

@@ -13,6 +13,7 @@ package norm
 import (
  "github.com/cockroachdb/cockroach/pkg/sql/opt"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+ "github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 )
 
 // CanInlineWith returns whether or not it's valid to inline binding in expr.
@@ -61,3 +62,24 @@ func (c *CustomFuncs) InlineWith(binding, input memo.RelExpr, priv *memo.WithPri
 
  return replace(input).(memo.RelExpr)
 }
+
+// MakeBindingPropsForRecursiveCTE makes a Relational struct that applies to all
+// iterations of a recursive CTE. The caller must verify that the supplied
+// cardinality applies to all iterations.
+func MakeBindingPropsForRecursiveCTE(
+ card props.Cardinality, outCols opt.ColSet, rowCount float64,
+) *props.Relational {
+ // The working table will always have at least one row, because any iteration
+ // that returns zero rows will end the loop.
+ card = card.AtLeast(props.OneCardinality)
+ bindingProps := &props.Relational{}
+ bindingProps.OutputCols = outCols
+ bindingProps.Cardinality = card
+ bindingProps.Stats.RowCount = rowCount
+ // Row count must be greater than 0 or the stats code will throw an error.
+ // Set it to 1 to match the cardinality.
+ if bindingProps.Stats.RowCount < 1 {
+ bindingProps.Stats.RowCount = 1
+ }
+ return bindingProps
+}

pkg/sql/opt/optbuilder/with.go

@@ -14,6 +14,7 @@ import (
  "github.com/cockroachdb/cockroach/pkg/server/telemetry"
  "github.com/cockroachdb/cockroach/pkg/sql/opt"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
+ "github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/props"
  "github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
  "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
@@ -285,19 +286,11 @@ func (b *Builder) buildCTE(
  // recursive query is executed. We can't really say too much about what the
  // working table contains, except that it has at least one row (the recursive
  // query is never invoked with an empty working table).
- bindingProps := &props.Relational{}
- bindingProps.OutputCols = outScope.colSet()
- bindingProps.Cardinality = props.AnyCardinality.AtLeast(props.OneCardinality)
- // We don't really know the input row count, except for the first time we run
- // the recursive query. We don't have anything better though.
- bindingProps.Stats.RowCount = initialScope.expr.Relational().Stats.RowCount
- // Row count must be greater than 0 or the stats code will throw an error.
- // Set it to 1 to match the cardinality.
- if bindingProps.Stats.RowCount < 1 {
- bindingProps.Stats.RowCount = 1
- }
+ initialRowCount := initialScope.expr.Relational().Stats.RowCount
  cteSrc.expr = b.factory.ConstructFakeRel(&memo.FakeRelPrivate{
- Props: bindingProps,
+ Props: norm.MakeBindingPropsForRecursiveCTE(
+ props.AnyCardinality, outScope.colSet(), initialRowCount,
+ ),
  })
  b.factory.Metadata().AddWithBinding(withID, cteSrc.expr)