1 files changed, 107 insertions, 47 deletions

diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 56398c399c9..687a5422eab 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -54,6 +54,7 @@
 #include "foreign/fdwapi.h"
 #include "miscadmin.h"
 #include "nodes/nodeFuncs.h"
+#include "optimizer/optimizer.h"
 #include "rewrite/rewriteHandler.h"
 #include "storage/bufmgr.h"
 #include "storage/lmgr.h"
@@ -353,69 +354,120 @@ ExecCheckTIDVisible(EState *estate,
 }
 
 /*
- * Compute stored generated columns for a tuple
+ * Initialize to compute stored generated columns for a tuple
+ *
+ * This fills the resultRelInfo's ri_GeneratedExprs and ri_extraUpdatedCols
+ * fields.  (Currently, ri_extraUpdatedCols is consulted only in UPDATE,
+ * but we might as well fill it for INSERT too.)
  */
-void
-ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
-						   EState *estate, TupleTableSlot *slot,
-						   CmdType cmdtype)
+static void
+ExecInitStoredGenerated(ResultRelInfo *resultRelInfo,
+						EState *estate,
+						CmdType cmdtype)
 {
 	Relation	rel = resultRelInfo->ri_RelationDesc;
 	TupleDesc	tupdesc = RelationGetDescr(rel);
 	int			natts = tupdesc->natts;
+	Bitmapset  *updatedCols;
 	MemoryContext oldContext;
-	Datum	   *values;
-	bool	   *nulls;
 
-	Assert(tupdesc->constr && tupdesc->constr->has_generated_stored);
+	/* Don't call twice */
+	Assert(resultRelInfo->ri_GeneratedExprs == NULL);
+
+	/* Nothing to do if no generated columns */
+	if (!(tupdesc->constr && tupdesc->constr->has_generated_stored))
+		return;
 
 	/*
-	 * If first time through for this result relation, build expression
-	 * nodetrees for rel's stored generation expressions.  Keep them in the
-	 * per-query memory context so they'll survive throughout the query.
+	 * 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 (resultRelInfo->ri_GeneratedExprs == NULL)
-	{
-		oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
+	if (cmdtype == CMD_UPDATE &&
+		!(rel->trigdesc && rel->trigdesc->trig_update_before_row))
+		updatedCols = ExecGetUpdatedCols(resultRelInfo, estate);
+	else
+		updatedCols = NULL;
 
-		resultRelInfo->ri_GeneratedExprs =
-			(ExprState **) palloc(natts * sizeof(ExprState *));
-		resultRelInfo->ri_NumGeneratedNeeded = 0;
+	/*
+	 * Make sure these data structures are built in the per-query memory
+	 * context so they'll survive throughout the query.
+	 */
+	oldContext = MemoryContextSwitchTo(estate->es_query_cxt);
 
-		for (int i = 0; i < natts; i++)
+	resultRelInfo->ri_GeneratedExprs =
+		(ExprState **) palloc0(natts * sizeof(ExprState *));
+	resultRelInfo->ri_NumGeneratedNeeded = 0;
+
+	for (int i = 0; i < natts; i++)
+	{
+		if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_STORED)
 		{
-			if (TupleDescAttr(tupdesc, i)->attgenerated == ATTRIBUTE_GENERATED_STORED)
-			{
-				Expr	   *expr;
+			Expr	   *expr;
 
-				/*
-				 * If it's an update and the current column was not marked as
-				 * being updated, then we can skip the computation.  But if
-				 * there is a BEFORE ROW UPDATE trigger, we cannot skip
-				 * because the trigger might affect additional columns.
-				 */
-				if (cmdtype == CMD_UPDATE &&
-					!(rel->trigdesc && rel->trigdesc->trig_update_before_row) &&
-					!bms_is_member(i + 1 - FirstLowInvalidHeapAttributeNumber,
-								   ExecGetExtraUpdatedCols(resultRelInfo, estate)))
-				{
-					resultRelInfo->ri_GeneratedExprs[i] = NULL;
-					continue;
-				}
+			/* 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));
 
-				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);
 
-				resultRelInfo->ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
-				resultRelInfo->ri_NumGeneratedNeeded++;
+				if (!bms_overlap(updatedCols, attrs_used))
+					continue;	/* need not update this column */
 			}
-		}
 
-		MemoryContextSwitchTo(oldContext);
+			/* No luck, so prepare the expression for execution */
+			resultRelInfo->ri_GeneratedExprs[i] = ExecPrepareExpr(expr, estate);
+			resultRelInfo->ri_NumGeneratedNeeded++;
+
+			/* And mark this column in resultRelInfo->ri_extraUpdatedCols */
+			resultRelInfo->ri_extraUpdatedCols =
+				bms_add_member(resultRelInfo->ri_extraUpdatedCols,
+							   i + 1 - FirstLowInvalidHeapAttributeNumber);
+		}
 	}
 
+	MemoryContextSwitchTo(oldContext);
+}
+
+/*
+ * Compute stored generated columns for a tuple
+ */
+void
+ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
+						   EState *estate, TupleTableSlot *slot,
+						   CmdType cmdtype)
+{
+	Relation	rel = resultRelInfo->ri_RelationDesc;
+	TupleDesc	tupdesc = RelationGetDescr(rel);
+	int			natts = tupdesc->natts;
+	ExprContext *econtext = GetPerTupleExprContext(estate);
+	MemoryContext oldContext;
+	Datum	   *values;
+	bool	   *nulls;
+
+	/* We should not be called unless this is true */
+	Assert(tupdesc->constr && tupdesc->constr->has_generated_stored);
+
+	/*
+	 * For relations named directly in the query, ExecInitStoredGenerated
+	 * should have been called already; but this might not have happened yet
+	 * for a partition child rel.  Also, it's convenient for outside callers
+	 * to not have to call ExecInitStoredGenerated explicitly.
+	 */
+	if (resultRelInfo->ri_GeneratedExprs == NULL)
+		ExecInitStoredGenerated(resultRelInfo, estate, cmdtype);
+
 	/*
 	 * If no generated columns have been affected by this change, then skip
 	 * the rest.
@@ -435,14 +487,13 @@ ExecComputeStoredGenerated(ResultRelInfo *resultRelInfo,
 	{
 		Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
 
-		if (attr->attgenerated == ATTRIBUTE_GENERATED_STORED &&
-			resultRelInfo->ri_GeneratedExprs[i])
+		if (resultRelInfo->ri_GeneratedExprs[i])
 		{
-			ExprContext *econtext;
 			Datum		val;
 			bool		isnull;
 
-			econtext = GetPerTupleExprContext(estate);
+			Assert(attr->attgenerated == ATTRIBUTE_GENERATED_STORED);
+
 			econtext->ecxt_scantuple = slot;
 
 			val = ExecEvalExpr(resultRelInfo->ri_GeneratedExprs[i], econtext, &isnull);
@@ -4088,6 +4139,15 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 					elog(ERROR, "could not find junk wholerow column");
 			}
 		}
+
+		/*
+		 * For INSERT and UPDATE, prepare to evaluate any generated columns.
+		 * We must do this now, even if we never insert or update any rows,
+		 * because we have to fill resultRelInfo->ri_extraUpdatedCols for
+		 * possible use by the trigger machinery.
+		 */
+		if (operation == CMD_INSERT || operation == CMD_UPDATE)
+			ExecInitStoredGenerated(resultRelInfo, estate, operation);
 	}
 
 	/*