Merge #47341

47341: sql: support MATERIALIZED syntax in CTEs r=RaduBerinde a=Anthuang

The MATERIALIZED/NOT MATERIALIZED syntax, introduced in PostgresSQL 12,
allows the user to override the optimizer's default decision to
not materialize/materialize a table, respectively.

For example, the following query is non-recursive and side-effect-free
so the optimizer will not materialize the table by default:
```
WITH cte AS (SELECT 1) SELECT * FROM cte
```

To force materialization, `MATERIALIZED` can be specified explcitly:
```
WITH cte AS MATERIALIZED (SELECT 1) SELECT * FROM cte
```

Similarly, the opposite can be done:
```
WITH cte AS NOT MATERIALIZED (SELECT 1/0) SELECT * FROM cte
```

Or even for multiple different CTEs:
```
WITH cte (x) AS MATERIALIZED (INSERT INTO abc VALUES (1, 2)), cte2 (y)
  AS NOT MATERIALIZED (SELECT x + 1 FROM cte) SELECT * FROM cte, cte2
```

Resolves: #45863 

Release note (sql change): CTEs were extended to support
the MATERIALIZED/NOT MATERIALIZED syntax added in
PostgreSQL 12.

Co-authored-by: Anthony Huang <anthony.huang@affirm.com>

docs/generated/sql/bnf/stmt_block.bnf

@@ -1813,6 +1813,7 @@ enum_val_list ::=
 
 common_table_expr ::=
 	table_alias_name opt_column_list 'AS' '(' preparable_stmt ')'
+	| table_alias_name opt_column_list 'AS' materialize_clause '(' preparable_stmt ')'
 
 iconst64 ::=
 	'ICONST'
@@ -2103,6 +2104,10 @@ create_as_col_qual_list ::=
 create_as_constraint_def ::=
 	create_as_constraint_elem
 
+materialize_clause ::=
+	'MATERIALIZED'
+	| 'NOT' 'MATERIALIZED'
+
 index_flags_param ::=
 	'FORCE_INDEX' '=' index_name
 	| 'NO_INDEX_JOIN'

docs/generated/sql/bnf/with_clause.bnf

@@ -1,3 +1,3 @@
 with_clause ::=
-	'WITH' ( ( ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' ) ) ( ( ',' ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' ) ) )* ) ( insert_stmt | update_stmt | delete_stmt | upsert_stmt | select_stmt )
-	| 'WITH' 'RECURSIVE' ( ( ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' ) ) ( ( ',' ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' ) ) )* ) ( insert_stmt | update_stmt | delete_stmt | upsert_stmt | select_stmt )
+	'WITH' ( ( ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' | table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' materialize_clause '(' preparable_stmt ')' ) ) ( ( ',' ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' | table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' materialize_clause '(' preparable_stmt ')' ) ) )* ) ( insert_stmt | update_stmt | delete_stmt | upsert_stmt | select_stmt )
+	| 'WITH' 'RECURSIVE' ( ( ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' | table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' materialize_clause '(' preparable_stmt ')' ) ) ( ( ',' ( table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' '(' preparable_stmt ')' | table_alias_name ( '(' ( ( name ) ( ( ',' name ) )* ) ')' |  ) 'AS' materialize_clause '(' preparable_stmt ')' ) ) )* ) ( insert_stmt | update_stmt | delete_stmt | upsert_stmt | select_stmt )

pkg/sql/opt/memo/expr_format.go

@@ -496,6 +496,15 @@ func (f *ExprFmtCtx) formatRelational(e RelExpr, tp treeprinter.Node) {
 			f.formatMutationCommon(tp, &t.MutationPrivate)
 		}
 
+	case *WithExpr:
+		if t.Mtr.Set {
+			if t.Mtr.Materialize {
+				tp.Child("materialized")
+			} else {
+				tp.Child("not-materialized")
+			}
+		}
+
 	case *WithScanExpr:
 		if !f.HasFlags(ExprFmtHideColumns) {
 			child := tp.Child("mapping:")

pkg/sql/opt/memo/interner.go

@@ -633,6 +633,11 @@ func (h *hasher) HashPointer(val unsafe.Pointer) {
 	h.HashUint64(uint64(uintptr(val)))
 }
 
+func (h *hasher) HashMaterializeClause(val tree.MaterializeClause) {
+	h.HashBool(val.Set)
+	h.HashBool(val.Materialize)
+}
+
 // ----------------------------------------------------------------------
 //
 // Equality functions
@@ -1006,6 +1011,10 @@ func (h *hasher) IsOpaqueMetadataEqual(l, r opt.OpaqueMetadata) bool {
 	return l == r
 }
 
+func (h *hasher) IsMaterializeClauseEqual(l, r tree.MaterializeClause) bool {
+	return l.Set == r.Set && l.Materialize == r.Materialize
+}
+
 // encodeDatum turns the given datum into an encoded string of bytes. If two
 // datums are equivalent, then their encoded bytes will be identical.
 // Conversely, if two datums are not equivalent, then their encoded bytes will

pkg/sql/opt/norm/custom_funcs.go

@@ -2076,11 +2076,15 @@ func (c *CustomFuncs) CanAddConstInts(first tree.Datum, second tree.Datum) bool
 // ----------------------------------------------------------------------
 
 // CanInlineWith returns whether or not it's valid to inline binding in expr.
-// This is the case when:
+// This is the case when materialize is explicitly set to false, or when:
 // 1. binding has no side-effects (because once it's inlined, there's no
 //    guarantee it will be executed fully), and
 // 2. binding is referenced at most once in expr.
 func (c *CustomFuncs) CanInlineWith(binding, expr memo.RelExpr, private *memo.WithPrivate) bool {
+	// If materialization is set, ignore the checks below.
+	if private.Mtr.Set {
+		return !private.Mtr.Materialize
+	}
 	if binding.Relational().CanHaveSideEffects {
 		return false
 	}

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

@@ -674,3 +674,212 @@ with &2 (cte)
                 └── plus [as="?column?":9, type=int, outer=(8)]
                      ├── variable: c:8 [type=int]
                      └── const: 1 [type=int]
+
+# Original CTE is inlined, adding "NOT MATERIALIZED" should not change the behavior.
+norm format=show-all expect=InlineWith
+WITH foo AS NOT MATERIALIZED (SELECT 1) SELECT * FROM foo
+----
+values
+ ├── columns: "?column?":2(int!null)
+ ├── cardinality: [1 - 1]
+ ├── stats: [rows=1]
+ ├── cost: 0.02
+ ├── key: ()
+ ├── fd: ()-->(2)
+ ├── prune: (2)
+ └── tuple [type=tuple{int}]
+      └── const: 1 [type=int]
+
+# Original CTE is inlined, adding "MATERIALIZED" should prevent inlining.
+norm format=show-all expect-not=InlineWith
+WITH foo AS MATERIALIZED (SELECT 1) SELECT * FROM foo
+----
+with &1 (foo)
+ ├── columns: "?column?":2(int!null)
+ ├── materialized
+ ├── cardinality: [1 - 1]
+ ├── stats: [rows=1]
+ ├── cost: 0.04
+ ├── key: ()
+ ├── fd: ()-->(2)
+ ├── prune: (2)
+ ├── values
+ │    ├── columns: "?column?":1(int!null)
+ │    ├── cardinality: [1 - 1]
+ │    ├── stats: [rows=1]
+ │    ├── cost: 0.02
+ │    ├── key: ()
+ │    ├── fd: ()-->(1)
+ │    ├── prune: (1)
+ │    └── tuple [type=tuple{int}]
+ │         └── const: 1 [type=int]
+ └── with-scan &1 (foo)
+      ├── columns: "?column?":2(int!null)
+      ├── mapping:
+      │    └──  "?column?":1(int) => "?column?":2(int)
+      ├── cardinality: [1 - 1]
+      ├── stats: [rows=1]
+      ├── cost: 0.01
+      ├── key: ()
+      ├── fd: ()-->(2)
+      └── prune: (2)
+
+# Original CTE is not inlined, adding "MATERIALIZED" should not change the behavior.
+norm expect-not=InlineWith
+WITH foo AS MATERIALIZED (SELECT 1/0) SELECT * FROM foo
+----
+with &1 (foo)
+ ├── columns: "?column?":2
+ ├── materialized
+ ├── cardinality: [1 - 1]
+ ├── side-effects
+ ├── key: ()
+ ├── fd: ()-->(2)
+ ├── values
+ │    ├── columns: "?column?":1
+ │    ├── cardinality: [1 - 1]
+ │    ├── side-effects
+ │    ├── key: ()
+ │    ├── fd: ()-->(1)
+ │    └── (1 / 0,)
+ └── with-scan &1 (foo)
+      ├── columns: "?column?":2
+      ├── mapping:
+      │    └──  "?column?":1 => "?column?":2
+      ├── cardinality: [1 - 1]
+      ├── key: ()
+      └── fd: ()-->(2)
+
+# Original CTE is not inlined, adding "NOT MATERIALIZED" should force the inline.
+norm expect=InlineWith
+WITH foo AS NOT MATERIALIZED (SELECT 1/0) SELECT * FROM foo
+----
+project
+ ├── columns: "?column?":2
+ ├── cardinality: [1 - 1]
+ ├── side-effects
+ ├── key: ()
+ ├── fd: ()-->(2)
+ ├── values
+ │    ├── columns: "?column?":1
+ │    ├── cardinality: [1 - 1]
+ │    ├── side-effects
+ │    ├── key: ()
+ │    ├── fd: ()-->(1)
+ │    └── (1 / 0,)
+ └── projections
+      └── "?column?":1 [as="?column?":2, outer=(1)]
+
+# Original CTE is not inlined, adding "NOT MATERIALIZED" should force the inline.
+norm expect=InlineWith
+WITH foo AS NOT MATERIALIZED (SELECT 1) SELECT * FROM foo UNION ALL SELECT * FROM foo;
+----
+union-all
+ ├── columns: "?column?":4!null
+ ├── left columns: "?column?":2
+ ├── right columns: "?column?":3
+ ├── cardinality: [2 - 2]
+ ├── values
+ │    ├── columns: "?column?":2!null
+ │    ├── cardinality: [1 - 1]
+ │    ├── key: ()
+ │    ├── fd: ()-->(2)
+ │    └── (1,)
+ └── values
+      ├── columns: "?column?":3!null
+      ├── cardinality: [1 - 1]
+      ├── key: ()
+      ├── fd: ()-->(3)
+      └── (1,)
+
+# Recursive CTEs should respect "MATERIALIZED".
+norm expect-not=InlineWith
+WITH RECURSIVE t(n) AS MATERIALIZED (VALUES (1) UNION ALL SELECT n+1 FROM t WHERE n < 100 ) SELECT sum(n) FROM t;
+----
+with &2 (t)
+ ├── columns: sum:6
+ ├── materialized
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(6)
+ ├── recursive-c-t-e
+ │    ├── columns: n:2
+ │    ├── working table binding: &1
+ │    ├── initial columns: column1:1
+ │    ├── recursive columns: "?column?":4
+ │    ├── cardinality: [1 - ]
+ │    ├── values
+ │    │    ├── columns: column1:1!null
+ │    │    ├── cardinality: [1 - 1]
+ │    │    ├── key: ()
+ │    │    ├── fd: ()-->(1)
+ │    │    └── (1,)
+ │    └── project
+ │         ├── columns: "?column?":4!null
+ │         ├── select
+ │         │    ├── columns: n:3!null
+ │         │    ├── with-scan &1 (t)
+ │         │    │    ├── columns: n:3
+ │         │    │    ├── mapping:
+ │         │    │    │    └──  n:2 => n:3
+ │         │    │    └── cardinality: [1 - ]
+ │         │    └── filters
+ │         │         └── n:3 < 100 [outer=(3), constraints=(/3: (/NULL - /99]; tight)]
+ │         └── projections
+ │              └── n:3 + 1 [as="?column?":4, outer=(3)]
+ └── scalar-group-by
+      ├── columns: sum:6
+      ├── cardinality: [1 - 1]
+      ├── key: ()
+      ├── fd: ()-->(6)
+      ├── with-scan &2 (t)
+      │    ├── columns: n:5
+      │    ├── mapping:
+      │    │    └──  n:2 => n:5
+      │    └── cardinality: [1 - ]
+      └── aggregations
+           └── sum [as=sum:6, outer=(5)]
+                └── n:5
+
+# Recursive CTEs should respect "NOT MATERIALIZED".
+norm expect=InlineWith
+WITH RECURSIVE t(n) AS NOT MATERIALIZED (VALUES (1) UNION ALL SELECT n+1 FROM t WHERE n < 100 ) SELECT sum(n) FROM t;
+----
+scalar-group-by
+ ├── columns: sum:6
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(6)
+ ├── project
+ │    ├── columns: n:5
+ │    ├── cardinality: [1 - ]
+ │    ├── recursive-c-t-e
+ │    │    ├── columns: n:2
+ │    │    ├── working table binding: &1
+ │    │    ├── initial columns: column1:1
+ │    │    ├── recursive columns: "?column?":4
+ │    │    ├── cardinality: [1 - ]
+ │    │    ├── values
+ │    │    │    ├── columns: column1:1!null
+ │    │    │    ├── cardinality: [1 - 1]
+ │    │    │    ├── key: ()
+ │    │    │    ├── fd: ()-->(1)
+ │    │    │    └── (1,)
+ │    │    └── project
+ │    │         ├── columns: "?column?":4!null
+ │    │         ├── select
+ │    │         │    ├── columns: n:3!null
+ │    │         │    ├── with-scan &1 (t)
+ │    │         │    │    ├── columns: n:3
+ │    │         │    │    ├── mapping:
+ │    │         │    │    │    └──  n:2 => n:3
+ │    │         │    │    └── cardinality: [1 - ]
+ │    │         │    └── filters
+ │    │         │         └── n:3 < 100 [outer=(3), constraints=(/3: (/NULL - /99]; tight)]
+ │    │         └── projections
+ │    │              └── n:3 + 1 [as="?column?":4, outer=(3)]
+ │    └── projections
+ │         └── n:2 [as=n:5, outer=(2)]
+ └── aggregations
+      └── sum [as=sum:6, outer=(5)]
+           └── n:5

pkg/sql/opt/ops/relational.opt

@@ -870,6 +870,10 @@ define WithPrivate {
     # it in the EXPLAIN plan).
     OriginalExpr Statement
 
+    # Mtr is used to specify whether or not to override the optimizer's
+    # default decision for materializing or not materializing tables.
+    Mtr MaterializeClause
+
     # Name is used to identify the with for debugging purposes.
     Name string
 }

pkg/sql/opt/optbuilder/scope.go

@@ -116,6 +116,7 @@ type cteSource struct {
 	originalExpr tree.Statement
 	bindingProps *props.Relational
 	expr         memo.RelExpr
+	mtr          tree.MaterializeClause
 	// If set, this function is called when a CTE is referenced. It can throw an
 	// error.
 	onRef func()

pkg/sql/opt/optbuilder/select.go

@@ -728,6 +728,7 @@ func (b *Builder) buildCTEs(with *tree.With, inScope *scope) (outScope *scope) {
 			expr:         cteExpr,
 			bindingProps: cteExpr.Relational(),
 			id:           id,
+			mtr:          cte.Mtr,
 		}
 		cte := &addedCTEs[i]
 		outScope.ctes[cte.name.Alias.String()] = cte
@@ -782,6 +783,7 @@ func (b *Builder) flushCTEs(expr memo.RelExpr) memo.RelExpr {
 			&memo.WithPrivate{
 				ID:           ctes[i].id,
 				Name:         string(ctes[i].name.Alias),
+				Mtr:          ctes[i].mtr,
 				OriginalExpr: ctes[i].originalExpr,
 			},
 		)

pkg/sql/opt/optbuilder/testdata/with

@@ -1445,3 +1445,59 @@ with &1 (a)
       ├── columns: testval:2!null
       └── mapping:
            └──  testval:1 => testval:2
+
+build
+WITH t AS MATERIALIZED (SELECT a FROM y WHERE a < 3)
+  SELECT * FROM x NATURAL JOIN t
+----
+with &1 (t)
+ ├── columns: a:3!null b:4
+ ├── materialized
+ ├── project
+ │    ├── columns: y.a:1!null
+ │    └── select
+ │         ├── columns: y.a:1!null y.rowid:2!null
+ │         ├── scan y
+ │         │    └── columns: y.a:1 y.rowid:2!null
+ │         └── filters
+ │              └── y.a:1 < 3
+ └── project
+      ├── columns: x.a:3!null b:4
+      └── inner-join (hash)
+           ├── columns: x.a:3!null b:4 x.rowid:5!null a:6!null
+           ├── scan x
+           │    └── columns: x.a:3 b:4 x.rowid:5!null
+           ├── with-scan &1 (t)
+           │    ├── columns: a:6!null
+           │    └── mapping:
+           │         └──  y.a:1 => a:6
+           └── filters
+                └── x.a:3 = a:6
+
+build
+WITH t AS NOT MATERIALIZED (SELECT a FROM y WHERE a < 3)
+  SELECT * FROM x NATURAL JOIN t
+----
+with &1 (t)
+ ├── columns: a:3!null b:4
+ ├── not-materialized
+ ├── project
+ │    ├── columns: y.a:1!null
+ │    └── select
+ │         ├── columns: y.a:1!null y.rowid:2!null
+ │         ├── scan y
+ │         │    └── columns: y.a:1 y.rowid:2!null
+ │         └── filters
+ │              └── y.a:1 < 3
+ └── project
+      ├── columns: x.a:3!null b:4
+      └── inner-join (hash)
+           ├── columns: x.a:3!null b:4 x.rowid:5!null a:6!null
+           ├── scan x
+           │    └── columns: x.a:3 b:4 x.rowid:5!null
+           ├── with-scan &1 (t)
+           │    ├── columns: a:6!null
+           │    └── mapping:
+           │         └──  y.a:1 => a:6
+           └── filters
+                └── x.a:3 = a:6