v2-0001-Expand-virtual-generated-columns-in-planner.patch

From 96c822e322652a29d0eab791565c54c9d0cd73af Mon Sep 17 00:00:00 2001
From: Richard Guo <guofenglinux@gmail.com>
Date: Wed, 12 Feb 2025 17:55:35 +0900
Subject: [PATCH v2] Expand virtual generated columns in planner

---
 src/backend/optimizer/plan/planner.c      |   6 +
 src/backend/optimizer/prep/prepjointree.c | 180 ++++++++++++++++++++++
 src/backend/rewrite/rewriteHandler.c      |  14 +-
 src/backend/rewrite/rewriteManip.c        |   2 +-
 src/include/optimizer/prep.h              |   1 +
 src/include/rewrite/rewriteManip.h        |   2 +
 6 files changed, 193 insertions(+), 12 deletions(-)

diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 7b1a8a0a9f..c80108a06c 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -749,6 +749,12 @@ subquery_planner(PlannerGlobal *glob, Query *parse, PlannerInfo *parent_root,
 	if (parse->setOperations)
 		flatten_simple_union_all(root);
 
+	/*
+	 * Expand virtual generated columns.
+	 * XXX more comments
+	 */
+	parse = root->parse = expand_virtual_generated_columns(root);
+
 	/*
 	 * Survey the rangetable to see what kinds of entries are present.  We can
 	 * skip some later processing if relevant SQL features are not used; for
diff --git a/src/backend/optimizer/prep/prepjointree.c b/src/backend/optimizer/prep/prepjointree.c
index 5d9225e990..99ef55f7ac 100644
--- a/src/backend/optimizer/prep/prepjointree.c
+++ b/src/backend/optimizer/prep/prepjointree.c
@@ -9,6 +9,7 @@
  *		preprocess_function_rtes
  *		pull_up_subqueries
  *		flatten_simple_union_all
+ *		expand_virtual_generated_columns
  *		do expression preprocessing (including flattening JOIN alias vars)
  *		reduce_outer_joins
  *		remove_useless_result_rtes
@@ -25,6 +26,7 @@
  */
 #include "postgres.h"
 
+#include "access/table.h"
 #include "catalog/pg_type.h"
 #include "funcapi.h"
 #include "miscadmin.h"
@@ -39,7 +41,9 @@
 #include "optimizer/tlist.h"
 #include "parser/parse_relation.h"
 #include "parser/parsetree.h"
+#include "rewrite/rewriteHandler.h"
 #include "rewrite/rewriteManip.h"
+#include "utils/rel.h"
 
 
 typedef struct nullingrel_info
@@ -57,6 +61,7 @@ typedef struct pullup_replace_vars_context
 {
 	PlannerInfo *root;
 	List	   *targetlist;		/* tlist of subquery being pulled up */
+	int         result_relation;	/* the index of the result relation */
 	RangeTblEntry *target_rte;	/* RTE of subquery */
 	Relids		relids;			/* relids within subquery, as numbered after
 								 * pullup (set only if target_rte->lateral) */
@@ -1273,6 +1278,7 @@ pull_up_simple_subquery(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte,
 	 */
 	rvcontext.root = root;
 	rvcontext.targetlist = subquery->targetList;
+	rvcontext.result_relation = 0;
 	rvcontext.target_rte = rte;
 	if (rte->lateral)
 	{
@@ -1833,6 +1839,7 @@ pull_up_simple_values(PlannerInfo *root, Node *jtnode, RangeTblEntry *rte)
 	}
 	rvcontext.root = root;
 	rvcontext.targetlist = tlist;
+	rvcontext.result_relation = 0;
 	rvcontext.target_rte = rte;
 	rvcontext.relids = NULL;	/* can't be any lateral references here */
 	rvcontext.nullinfo = NULL;
@@ -1992,6 +1999,7 @@ pull_up_constant_function(PlannerInfo *root, Node *jtnode,
 													  1,	/* resno */
 													  NULL, /* resname */
 													  false));	/* resjunk */
+	rvcontext.result_relation = 0;
 	rvcontext.target_rte = rte;
 
 	/*
@@ -2499,6 +2507,10 @@ pullup_replace_vars_callback(Var *var,
 	 */
 	need_phv = (var->varnullingrels != NULL) || rcon->wrap_non_vars;
 
+	/* System columns are not replaced. */
+	if (varattno < InvalidAttrNumber)
+		return (Node *) copyObject(var);
+
 	/*
 	 * If PlaceHolderVars are needed, we cache the modified expressions in
 	 * rcon->rv_cache[].  This is not in hopes of any material speed gain
@@ -2559,6 +2571,22 @@ pullup_replace_vars_callback(Var *var,
 		rowexpr->location = var->location;
 		newnode = (Node *) rowexpr;
 
+		/*
+		 * Handle any OLD/NEW RETURNING list Vars
+		 *
+		 * XXX comments about wrapping it into ReturningExpr
+		 */
+		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
+		{
+			ReturningExpr *rexpr = makeNode(ReturningExpr);
+
+			rexpr->retlevelsup = var->varlevelsup;
+			rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
+			rexpr->retexpr = (Expr *) rowexpr;
+
+			newnode = (Node *) rexpr;
+		}
+
 		/*
 		 * Insert PlaceHolderVar if needed.  Notice that we are wrapping one
 		 * PlaceHolderVar around the whole RowExpr, rather than putting one
@@ -2588,6 +2616,37 @@ pullup_replace_vars_callback(Var *var,
 		/* Make a copy of the tlist item to return */
 		newnode = (Node *) copyObject(tle->expr);
 
+		/*
+		 * Handle any OLD/NEW RETURNING list Vars
+		 *
+		 * XXX comments about wrapping it into ReturningExpr
+		 */
+		if (var->varreturningtype != VAR_RETURNING_DEFAULT)
+		{
+			/*
+			 * Copy varreturningtype onto any Vars in the tlist item that
+			 * refer to result_relation (which had better be non-zero).
+			 */
+			if (rcon->result_relation == 0)
+				elog(ERROR, "variable returning old/new found outside RETURNING list");
+
+			SetVarReturningType((Node *) newnode, rcon->result_relation,
+								var->varlevelsup, var->varreturningtype);
+
+			if (!IsA(newnode, Var) ||
+				((Var *) newnode)->varno != rcon->result_relation ||
+				((Var *) newnode)->varlevelsup != var->varlevelsup)
+			{
+				ReturningExpr *rexpr = makeNode(ReturningExpr);
+
+				rexpr->retlevelsup = var->varlevelsup;
+				rexpr->retold = (var->varreturningtype == VAR_RETURNING_OLD);
+				rexpr->retexpr = (Expr *) newnode;
+
+				newnode = (Node *) rexpr;
+			}
+		}
+
 		/* Insert PlaceHolderVar if needed */
 		if (need_phv)
 		{
@@ -2947,6 +3006,127 @@ flatten_simple_union_all(PlannerInfo *root)
 }
 
 
+/*
+ * expand_virtual_generated_columns
+ *		Expand all virtual generated column references in a query.
+ *
+ * XXX more comments
+ */
+Query *
+expand_virtual_generated_columns(PlannerInfo *root)
+{
+	Query	   *parse = root->parse;
+	int			rt_index;
+	ListCell   *lc;
+
+	rt_index = 0;
+	foreach(lc, parse->rtable)
+	{
+		RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
+		Relation	rel;
+		TupleDesc	tupdesc;
+
+		++rt_index;
+
+		/*
+		 * Only normal relations can have virtual generated columns.
+		 */
+		if (rte->rtekind != RTE_RELATION)
+			continue;
+
+		rel = table_open(rte->relid, NoLock);
+
+		tupdesc = RelationGetDescr(rel);
+		if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
+		{
+			List	   *tlist = NIL;
+			pullup_replace_vars_context rvcontext;
+
+			for (int i = 0; i < tupdesc->natts; i++)
+			{
+				Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
+				int				attnum = i + 1;
+				TargetEntry	   *te;
+
+				if (attr->attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+				{
+					Node	   *defexpr;
+					Oid			attcollid;
+
+					defexpr = build_column_default(rel, attnum);
+					if (defexpr == NULL)
+						elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
+							 attnum, RelationGetRelationName(rel));
+
+					/*
+					 * If the column definition has a collation and it is
+					 * different from the collation of the generation expression,
+					 * put a COLLATE clause around the expression.
+					 */
+					attcollid = attr->attcollation;
+					if (attcollid && attcollid != exprCollation(defexpr))
+					{
+						CollateExpr *ce = makeNode(CollateExpr);
+
+						ce->arg = (Expr *) defexpr;
+						ce->collOid = attcollid;
+						ce->location = -1;
+
+						defexpr = (Node *) ce;
+					}
+
+					ChangeVarNodes(defexpr, 1, rt_index, 0);
+
+					te = makeTargetEntry((Expr *) defexpr, attnum, 0, false);
+					tlist = lappend(tlist, te);
+				}
+				else
+				{
+					Var		   *var;
+
+					var = makeVar(rt_index,
+								  attnum,
+								  attr->atttypid,
+								  attr->atttypmod,
+								  attr->attcollation,
+								  0);
+					te = makeTargetEntry((Expr *) var, attnum, 0, false);
+					tlist = lappend(tlist, te);
+				}
+			}
+
+			Assert(list_length(tlist) > 0);
+
+			rvcontext.root = root;
+			rvcontext.targetlist = tlist;
+			rvcontext.result_relation = parse->resultRelation;
+			rvcontext.target_rte = rte;
+			rvcontext.relids = NULL;
+			rvcontext.nullinfo = NULL;
+			rvcontext.outer_hasSubLinks = NULL;
+			rvcontext.varno = rt_index;
+			rvcontext.wrap_non_vars = false;
+			rvcontext.rv_cache = palloc0((list_length(tlist) + 1) *
+										 sizeof(Node *));
+			/*
+			 * If the query uses grouping sets, we need a PlaceHolderVar for
+			 * anything that's not a simple Var.  This ensures that expressions
+			 * retain their separate identity so that they will match grouping
+			 * set columns when appropriate.
+			 */
+			if (parse->groupingSets)
+				rvcontext.wrap_non_vars = true;
+
+			parse = (Query *) pullup_replace_vars((Node *) parse, &rvcontext);
+		}
+
+		table_close(rel, NoLock);
+	}
+
+	return parse;
+}
+
+
 /*
  * reduce_outer_joins
  *		Attempt to reduce outer joins to plain inner joins.
diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index e996bdc0d2..0b1fcb3758 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -2211,7 +2211,9 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		 * Only normal relations can have RLS policies or virtual generated
 		 * columns.
 		 */
-		if (rte->rtekind != RTE_RELATION)
+		if (rte->rtekind != RTE_RELATION ||
+			(rte->relkind != RELKIND_RELATION &&
+			 rte->relkind != RELKIND_PARTITIONED_TABLE))
 			continue;
 
 		rel = table_open(rte->relid, NoLock);
@@ -2300,16 +2302,6 @@ fireRIRrules(Query *parsetree, List *activeRIRs)
 		if (hasSubLinks)
 			parsetree->hasSubLinks = true;
 
-		/*
-		 * Expand any references to virtual generated columns of this table.
-		 * Note that subqueries in virtual generated column expressions are
-		 * not currently supported, so this cannot add any more sublinks.
-		 */
-		parsetree = (Query *)
-			expand_generated_columns_internal((Node *) parsetree,
-											  rel, rt_index, rte,
-											  parsetree->resultRelation);
-
 		table_close(rel, NoLock);
 	}
 
diff --git a/src/backend/rewrite/rewriteManip.c b/src/backend/rewrite/rewriteManip.c
index 9433548d27..6994b8c542 100644
--- a/src/backend/rewrite/rewriteManip.c
+++ b/src/backend/rewrite/rewriteManip.c
@@ -1010,7 +1010,7 @@ SetVarReturningType_walker(Node *node, SetVarReturningType_context *context)
 	return expression_tree_walker(node, SetVarReturningType_walker, context);
 }
 
-static void
+void
 SetVarReturningType(Node *node, int result_relation, int sublevels_up,
 					VarReturningType returning_type)
 {
diff --git a/src/include/optimizer/prep.h b/src/include/optimizer/prep.h
index 0ae57ec24a..6bc75ac9f3 100644
--- a/src/include/optimizer/prep.h
+++ b/src/include/optimizer/prep.h
@@ -27,6 +27,7 @@ extern void pull_up_sublinks(PlannerInfo *root);
 extern void preprocess_function_rtes(PlannerInfo *root);
 extern void pull_up_subqueries(PlannerInfo *root);
 extern void flatten_simple_union_all(PlannerInfo *root);
+extern Query *expand_virtual_generated_columns(PlannerInfo *root);
 extern void reduce_outer_joins(PlannerInfo *root);
 extern void remove_useless_result_rtes(PlannerInfo *root);
 extern Relids get_relids_in_jointree(Node *jtnode, bool include_outer_joins,
diff --git a/src/include/rewrite/rewriteManip.h b/src/include/rewrite/rewriteManip.h
index 5ec475c63e..2f09657ab2 100644
--- a/src/include/rewrite/rewriteManip.h
+++ b/src/include/rewrite/rewriteManip.h
@@ -97,5 +97,7 @@ extern Node *ReplaceVarsFromTargetList(Node *node,
 									   ReplaceVarsNoMatchOption nomatch_option,
 									   int nomatch_varno,
 									   bool *outer_hasSubLinks);
+extern void SetVarReturningType(Node *node, int result_relation, int sublevels_up,
+					VarReturningType returning_type);
 
 #endif							/* REWRITEMANIP_H */
-- 
2.43.0