Merge pull request cockroachdb#100327 from cucaroach/backport22.2-99433

release-22.2: opt: fixup CTE stats on placeholder queries

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/builtins" // register all builtins in builtins:init() for memo package
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/builtins/builtinsregistry"
@@ -350,6 +351,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

@@ -868,3 +868,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]
+ │    ├── select
+ │    │    ├── columns: y:1!null z:2
+ │    │    ├── stats: [rows=1.00001, distinct(1)=1, null(1)=0]
+ │    │    ├── fd: ()-->(1)
+ │    │    ├── scan yz
+ │    │    │    ├── columns: y:1 z:2
+ │    │    │    └── stats: [rows=100000, distinct(1)=100000, null(1)=0]
+ │    │    └── 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]
+ │         │    ├── scan yz
+ │         │    │    ├── columns: y:8 z:9
+ │         │    │    └── stats: [rows=100000, distinct(9)=1, null(9)=0]
+ │         │    ├── 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]
+ │         │    └── 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"
@@ -289,19 +290,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)