From 9d3d8b3c3d0e510ad4bdf85dd4d7994724c555fa Mon Sep 17 00:00:00 2001
From: Tommy Reilly <treilly@cockroachlabs.com>
Date: Thu, 23 Mar 2023 17:00:10 -0400
Subject: [PATCH] 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
a normalization step to do this so it happens on the assignplaceholder
path.

Fixes: #99389
Epic: none
Release note: none
---
 pkg/sql/opt/norm/factory.go          | 18 +++++++
 pkg/sql/opt/norm/testdata/rules/with | 78 ++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)

diff --git a/pkg/sql/opt/norm/factory.go b/pkg/sql/opt/norm/factory.go
index 8b5978f46240..e3bf46077eef 100644
--- a/pkg/sql/opt/norm/factory.go
+++ b/pkg/sql/opt/norm/factory.go
@@ -16,6 +16,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"
@@ -354,6 +355,23 @@ func (f *Factory) AssignPlaceholders(from *memo.Memo) (err error) {
 			}
 			return f.ConstructConstVal(d, placeholder.DataType())
 		}
+		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().Statistics().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)
diff --git a/pkg/sql/opt/norm/testdata/rules/with b/pkg/sql/opt/norm/testdata/rules/with
index 3f024cb2a237..b2fbeeda1d67 100644
--- a/pkg/sql/opt/norm/testdata/rules/with
+++ b/pkg/sql/opt/norm/testdata/rules/with
@@ -3005,3 +3005,81 @@ project
  │              └── i:3 + 1 [as="?column?":4, outer=(3), immutable]
  └── projections
       └── i:2 [as=i:5, outer=(2)]
+
+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)]