v10-0001-refactor-ExecGetExtraUpdatedCols.no-cfbot

From ea27806ed7ce524cff66d2f250e9625ae96bdc26 Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Tue, 3 Dec 2024 17:04:55 +0800
Subject: [PATCH v10 1/1] refactor ExecGetExtraUpdatedCols

---
 src/backend/executor/execUtils.c       | 68 ++++++++++++++++++++++++--
 src/backend/executor/nodeModifyTable.c | 21 +++-----
 src/include/executor/nodeModifyTable.h |  2 +-
 src/include/nodes/execnodes.h          |  3 --
 4 files changed, 73 insertions(+), 21 deletions(-)

diff --git a/src/backend/executor/execUtils.c b/src/backend/executor/execUtils.c
index 5e1bac04ee..81afeabdf9 100644
--- a/src/backend/executor/execUtils.c
+++ b/src/backend/executor/execUtils.c
@@ -55,6 +55,8 @@
 #include "miscadmin.h"
 #include "parser/parse_relation.h"
 #include "partitioning/partdesc.h"
+#include "rewrite/rewriteHandler.h"
+#include "optimizer/optimizer.h"
 #include "storage/lmgr.h"
 #include "utils/builtins.h"
 #include "utils/memutils.h"
@@ -1306,13 +1308,73 @@ ExecGetUpdatedCols(ResultRelInfo *relinfo, EState *estate)
 	return perminfo->updatedCols;
 }
 
-/* Return a bitmap representing generated columns being updated */
+/* Return a bitmap representing (viirtual|stored) generated columns being
+ * updated first, through ExecInitStoredGenerated colect all stored generated
+ * column.  then collect the virtual one.
+*/
 Bitmapset *
 ExecGetExtraUpdatedCols(ResultRelInfo *relinfo, EState *estate)
 {
+	Relation	rel = relinfo->ri_RelationDesc;
+	TupleDesc	tupdesc = RelationGetDescr(rel);
+	int			natts = tupdesc->natts;
+	Bitmapset  *updatedCols;
+	MemoryContext oldContext;
+
 	/* Compute the info if we didn't already */
-	if (!relinfo->ri_Generated_valid)
-		ExecInitGenerated(relinfo, estate, CMD_UPDATE);
+	if (relinfo->ri_GeneratedExprsU == NULL)
+		ExecInitStoredGenerated(relinfo, estate, CMD_UPDATE);
+
+	/* Nothing else to do if no virtual generated columns */
+	if (!(tupdesc->constr && tupdesc->constr->has_generated_virtual))
+		return relinfo->ri_extraUpdatedCols;
+
+	/*
+	 * In an UPDATE, we can skip computing any generated columns that do not
+	 * depend on any UPDATE target column.  But if there is a BEFORE ROW
+	 * UPDATE trigger, we cannot skip because the trigger might change more
+	 * columns.
+	 */
+	if (!(rel->trigdesc && rel->trigdesc->trig_update_before_row))
+		updatedCols = ExecGetUpdatedCols(relinfo, estate);
+	else
+		updatedCols = NULL;
+
+	oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
+	for (int i = 0; i < natts; i++)
+	{
+		char		attgenerated = TupleDescAttr(tupdesc, i)->attgenerated;
+
+		if (attgenerated == ATTRIBUTE_GENERATED_VIRTUAL)
+		{
+			Expr	   *expr;
+
+			/* Fetch the GENERATED AS expression tree */
+			expr = (Expr *) build_column_default(rel, i + 1);
+			if (expr == NULL)
+				elog(ERROR, "no generation expression found for column number %d of table \"%s\"",
+					 i + 1, RelationGetRelationName(rel));
+
+			/*
+			 * If it's an update with a known set of update target columns,
+			 * see if we can skip the computation.
+			 */
+			if (updatedCols)
+			{
+				Bitmapset  *attrs_used = NULL;
+
+				pull_varattnos((Node *) expr, 1, &attrs_used);
+
+				if (!bms_overlap(updatedCols, attrs_used))
+					continue;	/* need not update this column */
+			}
+
+			relinfo->ri_extraUpdatedCols =
+				bms_add_member(relinfo->ri_extraUpdatedCols,
+								i + 1 - FirstLowInvalidHeapAttributeNumber);
+		}
+	}
+	MemoryContextSwitchTo(oldContext);
 	return relinfo->ri_extraUpdatedCols;
 }
 
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index fc488231f3..3a8dbe6dd7 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -349,7 +349,7 @@ ExecCheckTIDVisible(EState *estate,
  * UPDATE and INSERT actions.
  */
 void
-ExecInitGenerated(ResultRelInfo *resultRelInfo,
+ExecInitStoredGenerated(ResultRelInfo *resultRelInfo,
 				  EState *estate,
 				  CmdType cmdtype)
 {
@@ -362,7 +362,7 @@ ExecInitGenerated(ResultRelInfo *resultRelInfo,
 	MemoryContext oldContext;
 
 	/* Nothing to do if no generated columns */
-	if (!(tupdesc->constr && (tupdesc->constr->has_generated_stored || tupdesc->constr->has_generated_virtual)))
+	if (!(tupdesc->constr && tupdesc->constr->has_generated_stored))
 		return;
 
 	/*
@@ -388,9 +388,7 @@ ExecInitGenerated(ResultRelInfo *resultRelInfo,
 
 	for (int i = 0; i < natts; i++)
 	{
-		char		attgenerated = TupleDescAttr(tupdesc, i)->attgenerated;
-
-		if (attgenerated)
+		if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_STORED)
 		{
 			Expr	   *expr;
 
@@ -415,11 +413,8 @@ ExecInitGenerated(ResultRelInfo *resultRelInfo,
 			}
 
 			/* No luck, so prepare the expression for execution */
-			if (attgenerated == ATTRIBUTE_GENERATED_STORED)
-			{
-				ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
-				ri_NumGeneratedNeeded++;
-			}
+			ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
+			ri_NumGeneratedNeeded++;
 
 			/* If UPDATE, mark column in resultRelInfo->ri_extraUpdatedCols */
 			if (cmdtype == CMD_UPDATE)
@@ -447,8 +442,6 @@ ExecInitGenerated(ResultRelInfo *resultRelInfo,
 		resultRelInfo->ri_NumGeneratedNeededI = ri_NumGeneratedNeeded;
 	}
 
-	resultRelInfo->ri_Generated_valid = true;
-
 	MemoryContextSwitchTo(oldContext);
 }
 
@@ -479,7 +472,7 @@ ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
 	if (cmdtype == CMD_UPDATE)
 	{
 		if (resultRelInfo->ri_GeneratedExprsU == NULL)
-			ExecInitGenerated(resultRelInfo, estate, cmdtype);
+			ExecInitStoredGenerated(resultRelInfo, estate, cmdtype);
 		if (resultRelInfo->ri_NumGeneratedNeededU == 0)
 			return;
 		ri_GeneratedExprs = resultRelInfo->ri_GeneratedExprsU;
@@ -487,7 +480,7 @@ ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
 	else
 	{
 		if (resultRelInfo->ri_GeneratedExprsI == NULL)
-			ExecInitGenerated(resultRelInfo, estate, cmdtype);
+			ExecInitStoredGenerated(resultRelInfo, estate, cmdtype);
 		/* Early exit is impossible given the prior Assert */
 		Assert(resultRelInfo->ri_NumGeneratedNeededI > 0);
 		ri_GeneratedExprs = resultRelInfo->ri_GeneratedExprsI;
diff --git a/src/include/executor/nodeModifyTable.h b/src/include/executor/nodeModifyTable.h
index 0d8a5f6613..2e6a2bab20 100644
--- a/src/include/executor/nodeModifyTable.h
+++ b/src/include/executor/nodeModifyTable.h
@@ -15,7 +15,7 @@
 
 #include "nodes/execnodes.h"
 
-extern void ExecInitGenerated(ResultRelInfo *resultRelInfo,
+extern void ExecInitStoredGenerated(ResultRelInfo *resultRelInfo,
 							  EState *estate,
 							  CmdType cmdtype);
 
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 5acf3c29de..182a6956bb 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -540,9 +540,6 @@ typedef struct ResultRelInfo
 	int			ri_NumGeneratedNeededI;
 	int			ri_NumGeneratedNeededU;
 
-	/* true if the above have been computed */
-	bool		ri_Generated_valid;
-
 	/* list of RETURNING expressions */
 	List	   *ri_returningList;
 
-- 
2.34.1