Enable materialized CTEs in UPDATE and DELETE statements

src/common/enums/statement_type.cpp

@@ -61,6 +61,8 @@ string StatementTypeToString(StatementType type) {
 		return "DETACH";
 	case StatementType::MULTI_STATEMENT:
 		return "MULTI";
+	case StatementType::MATERIALIZED_CTE_STATEMENT:
+		return "WITH";
 	case StatementType::INVALID_STATEMENT:
 		break;
 	}

src/include/duckdb.h

@@ -180,6 +180,7 @@ typedef enum {
 	DUCKDB_STATEMENT_TYPE_ATTACH = 25,
 	DUCKDB_STATEMENT_TYPE_DETACH = 26,
 	DUCKDB_STATEMENT_TYPE_MULTI = 27,
+	DUCKDB_STATEMENT_TYPE_MATERIALIZED_CTE = 28,
 } duckdb_statement_type;
 
 //===--------------------------------------------------------------------===//

src/include/duckdb/common/enums/statement_type.hpp

@@ -46,6 +46,7 @@ enum class StatementType : uint8_t {
 	DETACH_STATEMENT,
 	MULTI_STATEMENT,
 	COPY_DATABASE_STATEMENT,
+	MATERIALIZED_CTE_STATEMENT,
 };
 
 DUCKDB_API string StatementTypeToString(StatementType type);

src/include/duckdb/parser/statement/list.hpp

@@ -1,19 +1,20 @@
 #include "duckdb/parser/statement/alter_statement.hpp"
 #include "duckdb/parser/statement/attach_statement.hpp"
 #include "duckdb/parser/statement/call_statement.hpp"
-#include "duckdb/parser/statement/copy_statement.hpp"
 #include "duckdb/parser/statement/copy_database_statement.hpp"
+#include "duckdb/parser/statement/copy_statement.hpp"
 #include "duckdb/parser/statement/create_statement.hpp"
 #include "duckdb/parser/statement/delete_statement.hpp"
 #include "duckdb/parser/statement/detach_statement.hpp"
 #include "duckdb/parser/statement/drop_statement.hpp"
 #include "duckdb/parser/statement/execute_statement.hpp"
 #include "duckdb/parser/statement/explain_statement.hpp"
-#include "duckdb/parser/statement/extension_statement.hpp"
 #include "duckdb/parser/statement/export_statement.hpp"
+#include "duckdb/parser/statement/extension_statement.hpp"
 #include "duckdb/parser/statement/insert_statement.hpp"
 #include "duckdb/parser/statement/load_statement.hpp"
 #include "duckdb/parser/statement/logical_plan_statement.hpp"
+#include "duckdb/parser/statement/materialized_cte_statement.hpp"
 #include "duckdb/parser/statement/pragma_statement.hpp"
 #include "duckdb/parser/statement/prepare_statement.hpp"
 #include "duckdb/parser/statement/relation_statement.hpp"

src/include/duckdb/parser/statement/materialized_cte_statement.hpp

@@ -0,0 +1,42 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// duckdb/parser/statement/materialized_cte_statement.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb/parser/sql_statement.hpp"
+#include "duckdb/parser/query_node.hpp"
+
+namespace duckdb {
+
+class MaterializedCTEStatement : public SQLStatement {
+public:
+	static constexpr const StatementType TYPE = StatementType::MATERIALIZED_CTE_STATEMENT;
+
+public:
+	MaterializedCTEStatement() : SQLStatement(StatementType::MATERIALIZED_CTE_STATEMENT) {
+	}
+
+	string ctename;
+	//! The query of the CTE
+	unique_ptr<QueryNode> query;
+	//! Child
+	unique_ptr<SQLStatement> child;
+	//! Aliases of the CTE node
+	vector<string> aliases;
+
+	//! CTEs
+	CommonTableExpressionMap cte_map;
+
+protected:
+	MaterializedCTEStatement(const MaterializedCTEStatement &other);
+
+public:
+	string ToString() const override;
+	unique_ptr<SQLStatement> Copy() const override;
+};
+} // namespace duckdb

src/include/duckdb/parser/tokens.hpp

@@ -42,6 +42,7 @@ class LoadStatement;
 class LogicalPlanStatement;
 class MultiStatement;
 class CopyDatabaseStatement;
+class MaterializedCTEStatement;
 
 //===--------------------------------------------------------------------===//
 // Query Node

src/include/duckdb/parser/transformer.hpp

@@ -155,9 +155,9 @@ class Transformer {
 	//! Transform a Postgres duckdb_libpgquery::T_PGTransactionStmt node into a TransactionStatement
 	unique_ptr<TransactionStatement> TransformTransaction(duckdb_libpgquery::PGTransactionStmt &stmt);
 	//! Transform a Postgres T_DeleteStatement node into a DeleteStatement
-	unique_ptr<DeleteStatement> TransformDelete(duckdb_libpgquery::PGDeleteStmt &stmt);
+	unique_ptr<SQLStatement> TransformDelete(duckdb_libpgquery::PGDeleteStmt &stmt);
 	//! Transform a Postgres duckdb_libpgquery::T_PGUpdateStmt node into a UpdateStatement
-	unique_ptr<UpdateStatement> TransformUpdate(duckdb_libpgquery::PGUpdateStmt &stmt);
+	unique_ptr<SQLStatement> TransformUpdate(duckdb_libpgquery::PGUpdateStmt &stmt);
 	//! Transform a Postgres duckdb_libpgquery::T_PGPragmaStmt node into a PragmaStatement
 	unique_ptr<SQLStatement> TransformPragma(duckdb_libpgquery::PGPragmaStmt &stmt);
 	//! Transform a Postgres duckdb_libpgquery::T_PGExportStmt node into a ExportStatement
@@ -290,6 +290,9 @@ class Transformer {
 	                  vector<unique_ptr<CTENode>> &materialized_ctes);
 	static unique_ptr<QueryNode> TransformMaterializedCTE(unique_ptr<QueryNode> root,
 	                                                      vector<unique_ptr<CTENode>> &materialized_ctes);
+	static unique_ptr<SQLStatement> TransformMaterializedCTEStatement(unique_ptr<SQLStatement> root,
+	                                                                  vector<unique_ptr<CTENode>> &materialized_ctes,
+	                                                                  CommonTableExpressionMap &cte_map);
 	unique_ptr<SelectStatement> TransformRecursiveCTE(duckdb_libpgquery::PGCommonTableExpr &node,
 	                                                  CommonTableExpressionInfo &info);
 

src/include/duckdb/planner/binder.hpp

@@ -251,6 +251,7 @@ class Binder : public std::enable_shared_from_this<Binder> {
 	BoundStatement Bind(AttachStatement &stmt);
 	BoundStatement Bind(DetachStatement &stmt);
 	BoundStatement Bind(CopyDatabaseStatement &stmt);
+	BoundStatement Bind(MaterializedCTEStatement &stmt);
 
 	BoundStatement BindReturning(vector<unique_ptr<ParsedExpression>> returning_list, TableCatalogEntry &table,
 	                             const string &alias, idx_t update_table_index,

src/main/capi/helper-c.cpp

@@ -242,6 +242,8 @@ duckdb_statement_type StatementTypeToC(duckdb::StatementType statement_type) {
 		return DUCKDB_STATEMENT_TYPE_DETACH;
 	case duckdb::StatementType::MULTI_STATEMENT:
 		return DUCKDB_STATEMENT_TYPE_MULTI;
+	case duckdb::StatementType::MATERIALIZED_CTE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_MATERIALIZED_CTE;
 	default:
 		return DUCKDB_STATEMENT_TYPE_INVALID;
 	}

src/parser/statement/CMakeLists.txt

@@ -24,7 +24,8 @@ add_library_unity(
   set_statement.cpp
   transaction_statement.cpp
   update_statement.cpp
-  vacuum_statement.cpp)
+  vacuum_statement.cpp
+  materialized_cte_statement.cpp)
 set(ALL_OBJECT_FILES
     ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:duckdb_statement>
     PARENT_SCOPE)

src/parser/statement/materialized_cte_statement.cpp

@@ -0,0 +1,21 @@
+#include "duckdb/parser/statement/materialized_cte_statement.hpp"
+
+namespace duckdb {
+
+MaterializedCTEStatement::MaterializedCTEStatement(const MaterializedCTEStatement &other)
+    : SQLStatement(other), ctename(other.ctename), query(other.query->Copy()), child(other.child->Copy()),
+      cte_map(other.cte_map.Copy()) {
+}
+
+string MaterializedCTEStatement::ToString() const {
+	string result;
+	result = cte_map.ToString();
+	result += child->ToString();
+	return result;
+}
+
+unique_ptr<SQLStatement> MaterializedCTEStatement::Copy() const {
+	return unique_ptr<MaterializedCTEStatement>(new MaterializedCTEStatement(*this));
+}
+
+} // namespace duckdb

src/parser/transform/statement/transform_delete.cpp

@@ -1,17 +1,16 @@
 #include "duckdb/parser/statement/delete_statement.hpp"
+#include "duckdb/parser/query_node/select_node.hpp"
 #include "duckdb/parser/transformer.hpp"
 
 namespace duckdb {
 
-unique_ptr<DeleteStatement> Transformer::TransformDelete(duckdb_libpgquery::PGDeleteStmt &stmt) {
+unique_ptr<SQLStatement> Transformer::TransformDelete(duckdb_libpgquery::PGDeleteStmt &stmt) {
 	auto result = make_uniq<DeleteStatement>();
 	vector<unique_ptr<CTENode>> materialized_ctes;
+	CommonTableExpressionMap cte_map;
 	if (stmt.withClause) {
-		TransformCTE(*PGPointerCast<duckdb_libpgquery::PGWithClause>(stmt.withClause), result->cte_map,
-		             materialized_ctes);
-		if (!materialized_ctes.empty()) {
-			throw NotImplementedException("Materialized CTEs are not implemented for delete.");
-		}
+		TransformCTE(*PGPointerCast<duckdb_libpgquery::PGWithClause>(stmt.withClause), cte_map, materialized_ctes);
+		result->cte_map = cte_map.Copy();
 	}
 
 	result->condition = TransformExpression(stmt.whereClause);
@@ -30,7 +29,11 @@ unique_ptr<DeleteStatement> Transformer::TransformDelete(duckdb_libpgquery::PGDe
 	if (stmt.returningList) {
 		TransformExpressionList(*stmt.returningList, result->returning_list);
 	}
-	return result;
+
+	// Handle materialized CTEs
+	auto cte_result = Transformer::TransformMaterializedCTEStatement(std::move(result), materialized_ctes, cte_map);
+
+	return cte_result;
 }
 
 } // namespace duckdb

src/parser/transform/statement/transform_update.cpp

@@ -17,15 +17,13 @@ unique_ptr<UpdateSetInfo> Transformer::TransformUpdateSetInfo(duckdb_libpgquery:
 	return result;
 }
 
-unique_ptr<UpdateStatement> Transformer::TransformUpdate(duckdb_libpgquery::PGUpdateStmt &stmt) {
+unique_ptr<SQLStatement> Transformer::TransformUpdate(duckdb_libpgquery::PGUpdateStmt &stmt) {
 	auto result = make_uniq<UpdateStatement>();
 	vector<unique_ptr<CTENode>> materialized_ctes;
+	CommonTableExpressionMap cte_map;
 	if (stmt.withClause) {
-		TransformCTE(*PGPointerCast<duckdb_libpgquery::PGWithClause>(stmt.withClause), result->cte_map,
-		             materialized_ctes);
-		if (!materialized_ctes.empty()) {
-			throw NotImplementedException("Materialized CTEs are not implemented for update.");
-		}
+		TransformCTE(*PGPointerCast<duckdb_libpgquery::PGWithClause>(stmt.withClause), cte_map, materialized_ctes);
+		result->cte_map = cte_map.Copy();
 	}
 
 	result->table = TransformRangeVar(*stmt.relation);
@@ -40,7 +38,10 @@ unique_ptr<UpdateStatement> Transformer::TransformUpdate(duckdb_libpgquery::PGUp
 		TransformExpressionList(*stmt.returningList, result->returning_list);
 	}
 
-	return result;
+	// Handle materialized CTEs
+	auto cte_result = Transformer::TransformMaterializedCTEStatement(std::move(result), materialized_ctes, cte_map);
+
+	return cte_result;
 }
 
 } // namespace duckdb

src/parser/transformer.cpp

@@ -236,6 +236,25 @@ unique_ptr<QueryNode> Transformer::TransformMaterializedCTE(unique_ptr<QueryNode
 	return root;
 }
 
+unique_ptr<SQLStatement> Transformer::TransformMaterializedCTEStatement(unique_ptr<SQLStatement> root,
+                                                                        vector<unique_ptr<CTENode>> &materialized_ctes,
+                                                                        CommonTableExpressionMap &cte_map) {
+	while (!materialized_ctes.empty()) {
+		unique_ptr<CTENode> node_result;
+		node_result = std::move(materialized_ctes.back());
+		unique_ptr<MaterializedCTEStatement> stmt_result = make_uniq<MaterializedCTEStatement>();
+		stmt_result->ctename = node_result->ctename;
+		stmt_result->query = std::move(node_result->query);
+		stmt_result->aliases = node_result->aliases;
+		stmt_result->cte_map = cte_map.Copy();
+		stmt_result->child = std::move(root);
+		root = std::move(stmt_result);
+		materialized_ctes.pop_back();
+	}
+
+	return root;
+}
+
 void Transformer::SetQueryLocation(ParsedExpression &expr, int query_location) {
 	if (query_location < 0) {
 		return;

src/planner/binder.cpp

@@ -118,6 +118,8 @@ BoundStatement Binder::Bind(SQLStatement &statement) {
 		return Bind(statement.Cast<DetachStatement>());
 	case StatementType::COPY_DATABASE_STATEMENT:
 		return Bind(statement.Cast<CopyDatabaseStatement>());
+	case StatementType::MATERIALIZED_CTE_STATEMENT:
+		return Bind(statement.Cast<MaterializedCTEStatement>());
 	default: // LCOV_EXCL_START
 		throw NotImplementedException("Unimplemented statement type \"%s\" for Bind",
 		                              StatementTypeToString(statement.type));

src/planner/binder/statement/CMakeLists.txt

@@ -25,7 +25,8 @@ add_library_unity(
   bind_simple.cpp
   bind_summarize.cpp
   bind_update.cpp
-  bind_vacuum.cpp)
+  bind_vacuum.cpp
+  bind_materialized_cte.cpp)
 set(ALL_OBJECT_FILES
     ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:duckdb_bind_statement>
     PARENT_SCOPE)

src/planner/binder/statement/bind_materialized_cte.cpp

@@ -0,0 +1,87 @@
+#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp"
+#include "duckdb/common/string_util.hpp"
+#include "duckdb/parser/expression/constant_expression.hpp"
+#include "duckdb/parser/expression_map.hpp"
+#include "duckdb/parser/query_node/cte_node.hpp"
+#include "duckdb/parser/query_node/select_node.hpp"
+#include "duckdb/parser/statement/materialized_cte_statement.hpp"
+#include "duckdb/planner/binder.hpp"
+#include "duckdb/planner/operator/logical_get.hpp"
+#include "duckdb/planner/operator/logical_materialized_cte.hpp"
+#include "duckdb/planner/operator/logical_projection.hpp"
+#include "duckdb/planner/query_node/bound_cte_node.hpp"
+#include "duckdb/planner/query_node/bound_select_node.hpp"
+#include "duckdb/planner/tableref/bound_basetableref.hpp"
+
+namespace duckdb {
+
+BoundStatement Binder::Bind(MaterializedCTEStatement &statement) {
+	auto result = make_uniq<BoundCTENode>();
+
+	// first recursively visit the materialized CTE operations
+	// the left side is visited first and is added to the BindContext of the right side
+	D_ASSERT(statement.query);
+	D_ASSERT(statement.child);
+
+	// Add CTEs as bindable
+	AddCTEMap(statement.cte_map);
+
+	result->ctename = statement.ctename;
+	result->setop_index = GenerateTableIndex();
+
+	result->query_binder = Binder::CreateBinder(context, this);
+	result->query = result->query_binder->BindNode(*statement.query);
+
+	// the result types of the CTE are the types of the LHS
+	result->types = result->query->types;
+	// names are picked from the LHS, unless aliases are explicitly specified
+	result->names = result->query->names;
+	for (idx_t i = 0; i < statement.aliases.size() && i < result->names.size(); i++) {
+		result->names[i] = statement.aliases[i];
+	}
+
+	// Rename columns if duplicate names are detected
+	idx_t index = 1;
+	vector<string> names;
+	for (auto &n : result->names) {
+		string name = n;
+		while (find(names.begin(), names.end(), name) != names.end()) {
+			name = n + "_" + std::to_string(index++);
+		}
+		names.push_back(name);
+	}
+
+	// This allows the right side to reference the CTE
+	bind_context.AddGenericBinding(result->setop_index, statement.ctename, names, result->types);
+
+	result->child_binder = Binder::CreateBinder(context, this);
+
+	// Add bindings of left side to temporary CTE bindings context
+	result->child_binder->bind_context.AddCTEBinding(result->setop_index, statement.ctename, names, result->types);
+
+	auto child_result = result->child_binder->Bind(*statement.child);
+
+	for (auto &c : result->query_binder->correlated_columns) {
+		result->child_binder->AddCorrelatedColumn(c);
+	}
+
+	// Create plan of sub-statement
+	auto cte_query = CreatePlan(*result->query);
+
+	D_ASSERT(child_result.plan->children.size() == 1);
+
+	// extract operator below root operation
+	auto plan = std::move(child_result.plan->children[0]);
+	child_result.plan->children.clear();
+
+	// add logical plan for materialized CTE in with children as right side operation
+	auto root = make_uniq<LogicalMaterializedCTE>(result->ctename, result->setop_index, result->types.size(),
+	                                              std::move(cte_query), std::move(plan));
+
+	// re-construct statement plan with materialized CTE
+	child_result.plan->children.push_back(std::move(root));
+
+	return child_result;
+}
+
+} // namespace duckdb

src/planner/planner.cpp

@@ -137,6 +137,7 @@ void Planner::CreatePlan(unique_ptr<SQLStatement> statement) {
 	case StatementType::ATTACH_STATEMENT:
 	case StatementType::DETACH_STATEMENT:
 	case StatementType::COPY_DATABASE_STATEMENT:
+	case StatementType::MATERIALIZED_CTE_STATEMENT:
 		CreatePlan(*statement);
 		break;
 	default: