v4-0001-Virtual-generated-columns_minorchange.no-cfbot

From 8f0e74efdb4ae92f6b0a27d2446688461f7346f5 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Thu, 29 Aug 2024 21:22:06 +0800
Subject: [PATCH v4 1/1] Virtual generated columns

address the get_attgenerated overhead.

Discussion: https://www.postgresql.org/message-id/flat/a368248e-69e4-40be-9c07-6c3b5880b0a6@eisentraut.org
---
 src/backend/rewrite/rewriteHandler.c | 135 +++++++++++++++++++--------
 1 file changed, 98 insertions(+), 37 deletions(-)

diff --git a/src/backend/rewrite/rewriteHandler.c b/src/backend/rewrite/rewriteHandler.c
index f2713eaef2..d4da0178f2 100644
--- a/src/backend/rewrite/rewriteHandler.c
+++ b/src/backend/rewrite/rewriteHandler.c
@@ -4380,6 +4380,12 @@ struct expand_generated_context
 
 	/* incremented for every level where it's true */
 	int			ancestor_has_virtual;
+
+	/* list of relation oids that have virtual generated column */
+	List		*virtual_gen_rel;
+
+	/* list of attnums that have virtual generated column */
+	List		*virtual_attnums;
 };
 
 static Node *
@@ -4393,6 +4399,9 @@ expand_generated_columns_mutator(Node *node, struct expand_generated_context *co
 		Var		   *v = (Var *) node;
 		Oid			relid;
 		AttrNumber	attnum;
+		ListCell   *lc,
+				   *lc2;
+		Oid			attcollid;
 		List	   *rtable = list_nth_node(List,
 										   context->rtables,
 										   list_length(context->rtables) - v->varlevelsup - 1);
@@ -4403,39 +4412,45 @@ expand_generated_columns_mutator(Node *node, struct expand_generated_context *co
 		if (!relid || !attnum)
 			return node;
 
-		if (get_attgenerated(relid, attnum) == ATTRIBUTE_GENERATED_VIRTUAL)
+		forboth(lc, context->virtual_gen_rel, lc2, context->virtual_attnums)
 		{
-			Relation	rt_entry_relation = table_open(relid, NoLock);
-			Oid			attcollid;
+			Oid	 virtual_rel = lfirst_oid(lc);
+			AttrNumber	attno = lfirst_int(lc2);
 
-			node = build_column_default(rt_entry_relation, attnum);
-			if (node == NULL)
-				elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
-					 attnum, RelationGetRelationName(rt_entry_relation));
-
-			/*
-			 * 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 = GetSysCacheOid(ATTNUM, Anum_pg_attribute_attcollation, relid, attnum, 0, 0);
-			if (attcollid && attcollid != exprCollation(node))
+			if(attno == attnum && virtual_rel == relid)
 			{
-				CollateExpr *ce = makeNode(CollateExpr);
+				Relation	rt_entry_relation = table_open(relid, NoLock);
 
-				ce->arg = (Expr *) node;
-				ce->collOid = attcollid;
-				ce->location = -1;
+				Assert(get_attgenerated(relid, attnum) == ATTRIBUTE_GENERATED_VIRTUAL);
 
-				node = (Node *) ce;
+				node = build_column_default(rt_entry_relation, attnum);
+				if (node == NULL)
+					elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
+						attnum, RelationGetRelationName(rt_entry_relation));
+
+				/*
+				* 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 = GetSysCacheOid(ATTNUM, Anum_pg_attribute_attcollation, relid, attnum, 0, 0);
+				if (attcollid && attcollid != exprCollation(node))
+				{
+					CollateExpr *ce = makeNode(CollateExpr);
+
+					ce->arg = (Expr *) node;
+					ce->collOid = attcollid;
+					ce->location = -1;
+
+					node = (Node *) ce;
+				}
+
+				IncrementVarSublevelsUp(node, v->varlevelsup, 0);
+				ChangeVarNodes(node, 1, v->varno, v->varlevelsup);
+
+				table_close(rt_entry_relation, NoLock);
 			}
-
-			IncrementVarSublevelsUp(node, v->varlevelsup, 0);
-			ChangeVarNodes(node, 1, v->varno, v->varlevelsup);
-
-			table_close(rt_entry_relation, NoLock);
 		}
-
 		return node;
 	}
 	else if (IsA(node, Query))
@@ -4457,6 +4472,7 @@ expand_generated_columns_in_expr(Node *node, Relation rel)
 
 	if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
 	{
+		int	i;
 		RangeTblEntry *rte;
 		List	   *rtable;
 		struct expand_generated_context context;
@@ -4470,6 +4486,23 @@ expand_generated_columns_in_expr(Node *node, Relation rel)
 		rte->relid = RelationGetRelid(rel);
 		rtable = list_make2(rte, rte);
 		context.rtables = list_make1(rtable);
+		context.virtual_gen_rel = NIL;
+		context.virtual_attnums = NIL;
+
+		for (i = 0; i < tupdesc->natts; i++)
+		{
+			Form_pg_attribute att = TupleDescAttr(tupdesc, i);
+			if (att->attisdropped)
+				continue;
+
+			if (att->attgenerated == 'v')
+			{
+				context.virtual_gen_rel = lappend_oid(context.virtual_gen_rel, rte->relid);
+				context.virtual_attnums = lappend_int(context.virtual_attnums, att->attnum);
+			}
+		}
+		Assert(list_length(context.virtual_gen_rel) > 0);
+		Assert(list_length(context.virtual_gen_rel) == list_length(context.virtual_attnums));
 
 		return expression_tree_mutator(node, expand_generated_columns_mutator, &context);
 	}
@@ -4486,13 +4519,49 @@ expand_generated_columns_in_query(Query *query, struct expand_generated_context
 {
 	context->rtables = lappend(context->rtables, query->rtable);
 	if (query->hasGeneratedVirtual)
+	{
+		List		*rte_relids = NIL;
 		context->ancestor_has_virtual++;
+		context->ancestor_has_virtual++;
+
+		foreach_node(RangeTblEntry, rte, query->rtable)
+		{
+			if(OidIsValid(rte->relid))
+				rte_relids = list_append_unique_oid(rte_relids, rte->relid);
+		}
+
+		foreach_oid(relid, rte_relids)
+		{
+			int	i;
+			Relation	rt_entry_relation = table_open(relid, NoLock);
+			TupleDesc	tupdesc = RelationGetDescr(rt_entry_relation);
+
+			if (tupdesc->constr && tupdesc->constr->has_generated_virtual)
+			{
+				for (i = 0; i < tupdesc->natts; i++)
+				{
+					Form_pg_attribute att = TupleDescAttr(tupdesc, i);
+					if (att->attisdropped)
+						continue;
+
+					if (att->attgenerated == 'v')
+					{
+						context->virtual_gen_rel = lappend_oid(context->virtual_gen_rel, relid);
+						context->virtual_attnums = lappend_int(context->virtual_attnums, att->attnum);
+					}
+				}
+				Assert(list_length(context->virtual_gen_rel) > 0);
+				Assert(list_length(context->virtual_gen_rel) == list_length(context->virtual_attnums));
+			}
+			table_close(rt_entry_relation, NoLock);
+		}
 
+	}
 	/*
 	 * If any table in the query has a virtual column or there is a sublink,
 	 * then we need to do the whole walk.
 	 */
-	if (query->hasGeneratedVirtual || query->hasSubLinks || context->ancestor_has_virtual)
+	if (context->ancestor_has_virtual)
 	{
 		query = query_tree_mutator(query,
 								   expand_generated_columns_mutator,
@@ -4505,22 +4574,14 @@ expand_generated_columns_in_query(Query *query, struct expand_generated_context
 	 */
 	else
 	{
-		ListCell   *lc;
-
-		foreach(lc, query->rtable)
+		foreach_node(RangeTblEntry, rte, query->rtable)
 		{
-			RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
-
 			if (rte->rtekind == RTE_SUBQUERY)
 				rte->subquery = expand_generated_columns_in_query(rte->subquery, context);
 		}
 
-		foreach(lc, query->cteList)
-		{
-			CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc);
-
+		foreach_node(CommonTableExpr, cte, query->cteList)
 			cte->ctequery = (Node *) expand_generated_columns_in_query(castNode(Query, cte->ctequery), context);
-		}
 	}
 
 	if (query->hasGeneratedVirtual)
-- 
2.34.1