Merge pull request duckdb#13752 from Mytherin/describeunion

Clean-up code around transform select nodes and statements
lnkuiper · Sep 5, 2024 · a6e32b1 · a6e32b1
2 parents 1f01ef8 + b5156e0
commit a6e32b1
Show file tree

Hide file tree

Showing 19 changed files with 75 additions and 58 deletions.
diff --git a/src/include/duckdb/parser/transformer.hpp b/src/include/duckdb/parser/transformer.hpp
@@ -112,9 +112,8 @@ class Transformer {
 	// Statement transformation
 	//===--------------------------------------------------------------------===//
 	//! Transform a Postgres duckdb_libpgquery::T_PGSelectStmt node into a SelectStatement
-	unique_ptr<SelectStatement> TransformSelect(optional_ptr<duckdb_libpgquery::PGNode> node, bool is_select = true);
-	//! Transform a Postgres duckdb_libpgquery::T_PGSelectStmt node into a SelectStatement
-	unique_ptr<SelectStatement> TransformSelect(duckdb_libpgquery::PGSelectStmt &select, bool is_select = true);
+	unique_ptr<SelectStatement> TransformSelectStmt(duckdb_libpgquery::PGSelectStmt &select, bool is_select = true);
+	unique_ptr<SelectStatement> TransformSelectStmt(duckdb_libpgquery::PGNode &node, bool is_select = true);
 	//! Transform a Postgres T_AlterStmt node into a AlterStatement
 	unique_ptr<AlterStatement> TransformAlter(duckdb_libpgquery::PGAlterTableStmt &stmt);
 	//! Transform a Postgres duckdb_libpgquery::T_PGRenameStmt node into a RenameStatement
@@ -170,8 +169,10 @@ class Transformer {
 	unique_ptr<PragmaStatement> TransformImport(duckdb_libpgquery::PGImportStmt &stmt);
 	unique_ptr<ExplainStatement> TransformExplain(duckdb_libpgquery::PGExplainStmt &stmt);
 	unique_ptr<SQLStatement> TransformVacuum(duckdb_libpgquery::PGVacuumStmt &stmt);
-	unique_ptr<SelectStatement> TransformShow(duckdb_libpgquery::PGVariableShowStmt &stmt);
-	unique_ptr<SelectStatement> TransformShowSelect(duckdb_libpgquery::PGVariableShowSelectStmt &stmt);
+	unique_ptr<QueryNode> TransformShow(duckdb_libpgquery::PGVariableShowStmt &stmt);
+	unique_ptr<SelectStatement> TransformShowStmt(duckdb_libpgquery::PGVariableShowStmt &stmt);
+	unique_ptr<QueryNode> TransformShowSelect(duckdb_libpgquery::PGVariableShowSelectStmt &stmt);
+	unique_ptr<SelectStatement> TransformShowSelectStmt(duckdb_libpgquery::PGVariableShowSelectStmt &stmt);
 	unique_ptr<AttachStatement> TransformAttach(duckdb_libpgquery::PGAttachStmt &stmt);
 	unique_ptr<DetachStatement> TransformDetach(duckdb_libpgquery::PGDetachStmt &stmt);
 	unique_ptr<SetStatement> TransformUse(duckdb_libpgquery::PGUseStmt &stmt);
@@ -204,7 +205,8 @@ class Transformer {
 	// Query Node Transform
 	//===--------------------------------------------------------------------===//
 	//! Transform a Postgres duckdb_libpgquery::T_PGSelectStmt node into a QueryNode
-	unique_ptr<QueryNode> TransformSelectNode(duckdb_libpgquery::PGSelectStmt &select);
+	unique_ptr<QueryNode> TransformSelectNode(duckdb_libpgquery::PGNode &select, bool is_select = true);
+	unique_ptr<QueryNode> TransformSelectNodeInternal(duckdb_libpgquery::PGSelectStmt &select, bool is_select = true);
 	unique_ptr<QueryNode> TransformSelectInternal(duckdb_libpgquery::PGSelectStmt &select);
 	void TransformModifiers(duckdb_libpgquery::PGSelectStmt &stmt, QueryNode &node);
 

diff --git a/src/parser/transform/expression/transform_subquery.cpp b/src/parser/transform/expression/transform_subquery.cpp
@@ -24,7 +24,7 @@ void RemoveOrderQualificationRecursive(unique_ptr<ParsedExpression> &expr) {
 unique_ptr<ParsedExpression> Transformer::TransformSubquery(duckdb_libpgquery::PGSubLink &root) {
 	auto subquery_expr = make_uniq<SubqueryExpression>();
 
-	subquery_expr->subquery = TransformSelect(root.subselect);
+	subquery_expr->subquery = TransformSelectStmt(*root.subselect);
 	SetQueryLocation(*subquery_expr, root.location);
 	D_ASSERT(subquery_expr->subquery);
 	D_ASSERT(!subquery_expr->subquery->node->GetSelectList().empty());

diff --git a/src/parser/transform/helpers/transform_cte.cpp b/src/parser/transform/helpers/transform_cte.cpp
@@ -70,8 +70,7 @@ void Transformer::TransformCTE(duckdb_libpgquery::PGWithClause &de_with_clause,
 			info->query = TransformRecursiveCTE(cte, *info);
 		} else {
 			Transformer cte_transformer(*this);
-			info->query =
-			    cte_transformer.TransformSelect(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(cte.ctequery));
+			info->query = cte_transformer.TransformSelectStmt(*cte.ctequery);
 		}
 		D_ASSERT(info->query);
 		auto cte_name = string(cte.ctename);
@@ -114,16 +113,20 @@ unique_ptr<SelectStatement> Transformer::TransformRecursiveCTE(duckdb_libpgquery
 			auto with_clause = PGPointerCast<duckdb_libpgquery::PGWithClause>(stmt.withClause);
 			TransformCTE(*with_clause, result.cte_map);
 		}
-		result.left = TransformSelectNode(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(stmt.larg));
-		result.right = TransformSelectNode(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(stmt.rarg));
+		result.left = TransformSelectNode(*stmt.larg);
+		result.right = TransformSelectNode(*stmt.rarg);
 		result.aliases = info.aliases;
 		break;
 	}
 	case duckdb_libpgquery::PG_SETOP_EXCEPT:
 	case duckdb_libpgquery::PG_SETOP_INTERSECT:
-	default:
+	default: {
 		// This CTE is not recursive. Fallback to regular query transformation.
-		return TransformSelect(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(cte.ctequery));
+		auto node = TransformSelectNode(*cte.ctequery);
+		auto result = make_uniq<SelectStatement>();
+		result->node = std::move(node);
+		return result;
+	}
 	}
 
 	if (stmt.limitCount || stmt.limitOffset) {

diff --git a/src/parser/transform/statement/transform_copy.cpp b/src/parser/transform/statement/transform_copy.cpp
@@ -117,7 +117,7 @@ unique_ptr<CopyStatement> Transformer::TransformCopy(duckdb_libpgquery::PGCopySt
 		info.schema = table.schema_name;
 		info.catalog = table.catalog_name;
 	} else {
-		info.select_statement = TransformSelectNode(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(stmt.query));
+		info.select_statement = TransformSelectNode(*stmt.query);
 	}
 
 	// handle the different options of the COPY statement

diff --git a/src/parser/transform/statement/transform_create_function.cpp b/src/parser/transform/statement/transform_create_function.cpp
@@ -13,8 +13,7 @@ unique_ptr<MacroFunction> Transformer::TransformMacroFunction(duckdb_libpgquery:
 		auto expression = TransformExpression(def.function);
 		macro_func = make_uniq<ScalarMacroFunction>(std::move(expression));
 	} else if (def.query) {
-		auto query_node =
-		    TransformSelect(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(def.query), true)->node->Copy();
+		auto query_node = TransformSelectNode(*def.query);
 		macro_func = make_uniq<TableMacroFunction>(std::move(query_node));
 	}
 

diff --git a/src/parser/transform/statement/transform_create_table_as.cpp b/src/parser/transform/statement/transform_create_table_as.cpp
@@ -15,7 +15,7 @@ unique_ptr<CreateStatement> Transformer::TransformCreateTableAs(duckdb_libpgquer
 	if (stmt.query->type != duckdb_libpgquery::T_PGSelectStmt) {
 		throw ParserException("CREATE TABLE AS requires a SELECT clause");
 	}
-	auto query = TransformSelect(stmt.query, false);
+	auto query = TransformSelectStmt(*stmt.query, false);
 
 	auto result = make_uniq<CreateStatement>();
 	auto info = make_uniq<CreateTableInfo>();

diff --git a/src/parser/transform/statement/transform_create_type.cpp b/src/parser/transform/statement/transform_create_type.cpp
@@ -49,7 +49,7 @@ unique_ptr<CreateStatement> Transformer::TransformCreateType(duckdb_libpgquery::
 		if (stmt.query) {
 			// CREATE TYPE mood AS ENUM (SELECT ...)
 			D_ASSERT(stmt.vals == nullptr);
-			auto query = TransformSelect(stmt.query, false);
+			auto query = TransformSelectStmt(*stmt.query, false);
 			info->query = std::move(query);
 			info->type = LogicalType::INVALID;
 		} else {

diff --git a/src/parser/transform/statement/transform_create_view.cpp b/src/parser/transform/statement/transform_create_view.cpp
@@ -21,7 +21,7 @@ unique_ptr<CreateStatement> Transformer::TransformCreateView(duckdb_libpgquery::
 	}
 	info->on_conflict = TransformOnConflict(stmt.onconflict);
 
-	info->query = TransformSelect(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(stmt.query), false);
+	info->query = TransformSelectStmt(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(stmt.query), false);
 
 	PivotEntryCheck("view");
 

diff --git a/src/parser/transform/statement/transform_insert.cpp b/src/parser/transform/statement/transform_insert.cpp
@@ -41,7 +41,7 @@ unique_ptr<InsertStatement> Transformer::TransformInsert(duckdb_libpgquery::PGIn
 		TransformExpressionList(*stmt.returningList, result->returning_list);
 	}
 	if (stmt.selectStmt) {
-		result->select_statement = TransformSelect(stmt.selectStmt, false);
+		result->select_statement = TransformSelectStmt(*stmt.selectStmt, false);
 	} else {
 		result->default_values = true;
 	}

diff --git a/src/parser/transform/statement/transform_select.cpp b/src/parser/transform/statement/transform_select.cpp
@@ -4,20 +4,19 @@
 
 namespace duckdb {
 
-unique_ptr<QueryNode> Transformer::TransformSelectNode(duckdb_libpgquery::PGSelectStmt &select) {
-	unique_ptr<QueryNode> stmt = nullptr;
-	if (select.pivot) {
-		stmt = TransformPivotStatement(select);
-	} else {
-		stmt = TransformSelectInternal(select);
+unique_ptr<QueryNode> Transformer::TransformSelectNode(duckdb_libpgquery::PGNode &node, bool is_select) {
+	switch (node.type) {
+	case duckdb_libpgquery::T_PGVariableShowSelectStmt:
+		return TransformShowSelect(PGCast<duckdb_libpgquery::PGVariableShowSelectStmt>(node));
+	case duckdb_libpgquery::T_PGVariableShowStmt:
+		return TransformShow(PGCast<duckdb_libpgquery::PGVariableShowStmt>(node));
+	default:
+		return TransformSelectNodeInternal(PGCast<duckdb_libpgquery::PGSelectStmt>(node), is_select);
 	}
-
-	return TransformMaterializedCTE(std::move(stmt));
 }
 
-unique_ptr<SelectStatement> Transformer::TransformSelect(duckdb_libpgquery::PGSelectStmt &select, bool is_select) {
-	auto result = make_uniq<SelectStatement>();
-
+unique_ptr<QueryNode> Transformer::TransformSelectNodeInternal(duckdb_libpgquery::PGSelectStmt &select,
+                                                               bool is_select) {
 	// Both Insert/Create Table As uses this.
 	if (is_select) {
 		if (select.intoClause) {
@@ -27,20 +26,26 @@ unique_ptr<SelectStatement> Transformer::TransformSelect(duckdb_libpgquery::PGSe
 			throw ParserException("SELECT locking clause is not supported!");
 		}
 	}
+	unique_ptr<QueryNode> stmt = nullptr;
+	if (select.pivot) {
+		stmt = TransformPivotStatement(select);
+	} else {
+		stmt = TransformSelectInternal(select);
+	}
+	return TransformMaterializedCTE(std::move(stmt));
+}
 
-	result->node = TransformSelectNode(select);
+unique_ptr<SelectStatement> Transformer::TransformSelectStmt(duckdb_libpgquery::PGSelectStmt &select, bool is_select) {
+	auto result = make_uniq<SelectStatement>();
+	result->node = TransformSelectNodeInternal(select, is_select);
 	return result;
 }
 
-unique_ptr<SelectStatement> Transformer::TransformSelect(optional_ptr<duckdb_libpgquery::PGNode> node, bool is_select) {
-	switch (node->type) {
-	case duckdb_libpgquery::T_PGVariableShowSelectStmt:
-		return TransformShowSelect(PGCast<duckdb_libpgquery::PGVariableShowSelectStmt>(*node));
-	case duckdb_libpgquery::T_PGVariableShowStmt:
-		return TransformShow(PGCast<duckdb_libpgquery::PGVariableShowStmt>(*node));
-	default:
-		return TransformSelect(PGCast<duckdb_libpgquery::PGSelectStmt>(*node), is_select);
-	}
+unique_ptr<SelectStatement> Transformer::TransformSelectStmt(duckdb_libpgquery::PGNode &node, bool is_select) {
+	auto select_node = TransformSelectNode(node, is_select);
+	auto select_statement = make_uniq<SelectStatement>();
+	select_statement->node = std::move(select_node);
+	return select_statement;
 }
 
 } // namespace duckdb
diff --git a/src/parser/transform/statement/transform_show.cpp b/src/parser/transform/statement/transform_show.cpp
@@ -9,7 +9,7 @@
 
 namespace duckdb {
 
-unique_ptr<SelectStatement> Transformer::TransformShow(duckdb_libpgquery::PGVariableShowStmt &stmt) {
+unique_ptr<QueryNode> Transformer::TransformShow(duckdb_libpgquery::PGVariableShowStmt &stmt) {
 	string name = stmt.name;
 
 	auto select_node = make_uniq<SelectNode>();
@@ -18,9 +18,12 @@ unique_ptr<SelectStatement> Transformer::TransformShow(duckdb_libpgquery::PGVari
 	showref->table_name = std::move(name);
 	showref->show_type = stmt.is_summary ? ShowType::SUMMARY : ShowType::DESCRIBE;
 	select_node->from_table = std::move(showref);
+	return std::move(select_node);
+}
 
+unique_ptr<SelectStatement> Transformer::TransformShowStmt(duckdb_libpgquery::PGVariableShowStmt &stmt) {
 	auto result = make_uniq<SelectStatement>();
-	result->node = std::move(select_node);
+	result->node = TransformShow(stmt);
 	return result;
 }
 

diff --git a/src/parser/transform/statement/transform_show_select.cpp b/src/parser/transform/statement/transform_show_select.cpp
@@ -7,19 +7,21 @@
 
 namespace duckdb {
 
-unique_ptr<SelectStatement> Transformer::TransformShowSelect(duckdb_libpgquery::PGVariableShowSelectStmt &stmt) {
+unique_ptr<QueryNode> Transformer::TransformShowSelect(duckdb_libpgquery::PGVariableShowSelectStmt &stmt) {
 	// we capture the select statement of SHOW
 	auto select_node = make_uniq<SelectNode>();
 	select_node->select_list.push_back(make_uniq<StarExpression>());
 
 	auto show_ref = make_uniq<ShowRef>();
 	show_ref->show_type = stmt.is_summary ? ShowType::SUMMARY : ShowType::DESCRIBE;
-	auto select = TransformSelect(stmt.stmt);
-	show_ref->query = std::move(select->node);
+	show_ref->query = TransformSelectNode(*stmt.stmt);
 	select_node->from_table = std::move(show_ref);
+	return std::move(select_node);
+}
 
+unique_ptr<SelectStatement> Transformer::TransformShowSelectStmt(duckdb_libpgquery::PGVariableShowSelectStmt &stmt) {
 	auto result = make_uniq<SelectStatement>();
-	result->node = std::move(select_node);
+	result->node = TransformShowSelect(stmt);
 	return result;
 }
 

diff --git a/src/parser/transform/tableref/transform_pivot.cpp b/src/parser/transform/tableref/transform_pivot.cpp
@@ -79,7 +79,7 @@ PivotColumn Transformer::TransformPivotColumn(duckdb_libpgquery::PGPivot &pivot,
 		}
 	}
 	if (pivot.subquery) {
-		col.subquery = TransformSelectNode(*PGPointerCast<duckdb_libpgquery::PGSelectStmt>(pivot.subquery));
+		col.subquery = TransformSelectNode(*pivot.subquery);
 	}
 	if (pivot.pivot_enum) {
 		col.pivot_enum = pivot.pivot_enum;

diff --git a/src/parser/transform/tableref/transform_subquery.cpp b/src/parser/transform/tableref/transform_subquery.cpp
@@ -5,7 +5,7 @@ namespace duckdb {
 
 unique_ptr<TableRef> Transformer::TransformRangeSubselect(duckdb_libpgquery::PGRangeSubselect &root) {
 	Transformer subquery_transformer(*this);
-	auto subquery = subquery_transformer.TransformSelect(root.subquery);
+	auto subquery = subquery_transformer.TransformSelectStmt(*root.subquery);
 	if (!subquery) {
 		return nullptr;
 	}

diff --git a/src/parser/transformer.cpp b/src/parser/transformer.cpp
@@ -144,7 +144,7 @@ unique_ptr<SQLStatement> Transformer::TransformStatementInternal(duckdb_libpgque
 		return result;
 	}
 	case duckdb_libpgquery::T_PGSelectStmt:
-		return TransformSelect(PGCast<duckdb_libpgquery::PGSelectStmt>(stmt));
+		return TransformSelectStmt(PGCast<duckdb_libpgquery::PGSelectStmt>(stmt));
 	case duckdb_libpgquery::T_PGCreateStmt:
 		return TransformCreateTable(PGCast<duckdb_libpgquery::PGCreateStmt>(stmt));
 	case duckdb_libpgquery::T_PGCreateSchemaStmt:
@@ -194,9 +194,9 @@ unique_ptr<SQLStatement> Transformer::TransformStatementInternal(duckdb_libpgque
 	case duckdb_libpgquery::T_PGVacuumStmt:
 		return TransformVacuum(PGCast<duckdb_libpgquery::PGVacuumStmt>(stmt));
 	case duckdb_libpgquery::T_PGVariableShowStmt:
-		return TransformShow(PGCast<duckdb_libpgquery::PGVariableShowStmt>(stmt));
+		return TransformShowStmt(PGCast<duckdb_libpgquery::PGVariableShowStmt>(stmt));
 	case duckdb_libpgquery::T_PGVariableShowSelectStmt:
-		return TransformShowSelect(PGCast<duckdb_libpgquery::PGVariableShowSelectStmt>(stmt));
+		return TransformShowSelectStmt(PGCast<duckdb_libpgquery::PGVariableShowSelectStmt>(stmt));
 	case duckdb_libpgquery::T_PGCallStmt:
 		return TransformCall(PGCast<duckdb_libpgquery::PGCallStmt>(stmt));
 	case duckdb_libpgquery::T_PGVariableSetStmt:

diff --git a/test/sql/show_select/describe_subquery.test b/test/sql/show_select/describe_subquery.test
@@ -38,3 +38,6 @@ null
 key
 default
 extra
+
+statement ok
+(DESCRIBE (values(42))) UNION ALL (DESCRIBE (values(42)));
diff --git a/third_party/libpg_query/grammar/grammar.cpp b/third_party/libpg_query/grammar/grammar.cpp
@@ -429,8 +429,8 @@ makeSetOp(PGSetOperation op, bool all, PGNode *larg, PGNode *rarg)
 
 	n->op = op;
 	n->all = all;
-	n->larg = (PGSelectStmt *) larg;
-	n->rarg = (PGSelectStmt *) rarg;
+	n->larg = larg;
+	n->rarg = rarg;
 	return (PGNode *) n;
 }
 

diff --git a/third_party/libpg_query/include/nodes/parsenodes.hpp b/third_party/libpg_query/include/nodes/parsenodes.hpp
@@ -1287,8 +1287,8 @@ typedef struct PGSelectStmt {
 	 */
 	PGSetOperation op;         /* type of set op */
 	bool all;                  /* ALL specified? */
-	struct PGSelectStmt *larg; /* left child */
-	struct PGSelectStmt *rarg; /* right child */
+	struct PGNode *larg; /* left child */
+	struct PGNode *rarg; /* right child */
 	                           /* Eventually add fields for CORRESPONDING spec here */
 } PGSelectStmt;
 

diff --git a/third_party/libpg_query/src_backend_parser_gram.cpp b/third_party/libpg_query/src_backend_parser_gram.cpp
@@ -31729,8 +31729,8 @@ makeSetOp(PGSetOperation op, bool all, PGNode *larg, PGNode *rarg)
 
 	n->op = op;
 	n->all = all;
-	n->larg = (PGSelectStmt *) larg;
-	n->rarg = (PGSelectStmt *) rarg;
+	n->larg = larg;
+	n->rarg = rarg;
 	return (PGNode *) n;
 }