From ddffccd68513bb0e68d6cf75810cf64cf9a4d757 Mon Sep 17 00:00:00 2001
From: Amit Langote <amitlan@postgresql.org>
Date: Tue, 11 Nov 2025 22:30:52 +0900
Subject: [PATCH v1 4/4] Use pruning-aware locking in cached plans

Extend GetCachedPlan() to perform ExecutorPrep() on each planned
statement, capturing unpruned relids and initial pruning results.
Use this data to acquire execution locks only on surviving partitions,
avoiding unnecessary locking of pruned tables even when using cached
plans.

Introduce CachedPlanPrepData to carry ExecutorPrep results
through the plan caching layer. Adjust call sites in SPI,
functions, portals, and EXPLAIN to propagate this data.

This ensures pruning decisions made during initial pruning are
consistently reused without redoing pruning logic in executor paths
like parallel workers. It also lays the groundwork for
pruning-dependent lock behavior during plan reuse.

To maintain correctness when all target partitions are pruned, also
reinstate the firstResultRel locking behavior lost in commit
28317de72. That commit required the first ModifyTable target to
remain initialized for executor assumptions to hold. We now
explicitly track these relids in PlannerGlobal and PlannedStmt so they
are locked even if pruned, preserving that rule across cached plan
reuse.
---
 src/backend/commands/prepare.c         |  15 +-
 src/backend/executor/functions.c       |  14 +-
 src/backend/executor/nodeModifyTable.c |   4 +-
 src/backend/executor/spi.c             |  22 ++-
 src/backend/optimizer/plan/planner.c   |   1 +
 src/backend/optimizer/plan/setrefs.c   |   3 +
 src/backend/tcop/postgres.c            |   7 +-
 src/backend/utils/cache/plancache.c    | 223 ++++++++++++++++++++++++-
 src/include/nodes/pathnodes.h          |   3 +
 src/include/nodes/plannodes.h          |   7 +
 src/include/utils/plancache.h          |  23 ++-
 11 files changed, 299 insertions(+), 23 deletions(-)

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index afd449c73ba..10fdff403b9 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -154,6 +154,7 @@ ExecuteQuery(ParseState *pstate,
 {
 	PreparedStatement *entry;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *plan_list;
 	ParamListInfo paramLI = NULL;
 	EState	   *estate = NULL;
@@ -193,7 +194,9 @@ ExecuteQuery(ParseState *pstate,
 									   entry->plansource->query_string);
 
 	/* Replan if needed, and increment plan refcount for portal */
-	cplan = GetCachedPlan(entry->plansource, paramLI, NULL, NULL);
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(entry->plansource, paramLI, NULL, NULL, &cprep);
 	plan_list = cplan->stmt_list;
 
 	/*
@@ -205,7 +208,7 @@ ExecuteQuery(ParseState *pstate,
 					  query_string,
 					  entry->plansource->commandTag,
 					  plan_list,
-					  NIL,
+					  cprep.prep_list,
 					  cplan);
 
 	/*
@@ -575,6 +578,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	PreparedStatement *entry;
 	const char *query_string;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *plan_list;
 	List	   *prep_list;
 	ListCell   *p;
@@ -633,8 +637,11 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	}
 
 	/* Replan if needed, and acquire a transient refcount */
+	cprep.context = CurrentMemoryContext;
+	cprep.owner = CurrentResourceOwner;
 	cplan = GetCachedPlan(entry->plansource, paramLI,
-						  CurrentResourceOwner, pstate->p_queryEnv);
+						  CurrentResourceOwner, pstate->p_queryEnv,
+						  &cprep);
 
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
@@ -653,7 +660,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	}
 
 	plan_list = cplan->stmt_list;
-	prep_list = NIL;
+	prep_list = cprep.prep_list;
 
 	/* Explain each query */
 	i = 0;
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 633310c5f5b..8fc22fbd283 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -72,6 +72,7 @@ typedef struct execution_state
 	bool		setsResult;		/* true if this query produces func's result */
 	bool		lazyEval;		/* true if should fetch one row at a time */
 	PlannedStmt *stmt;			/* plan for this query */
+	ExecPrep   *prep;			/* ExecutorPrep() output for this plan */
 	QueryDesc  *qd;				/* null unless status == RUN */
 } execution_state;
 
@@ -657,6 +658,8 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	execution_state *lasttages = NULL;
 	int			nstmts;
 	ListCell   *lc;
+	CachedPlanPrepData cprep = {0};
+	int			i;
 
 	/*
 	 * Clean up after previous query, if there was one.
@@ -695,10 +698,13 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	 * CurrentResourceOwner will be the same when ShutdownSQLFunction runs.)
 	 */
 	fcache->cowner = CurrentResourceOwner;
+	cprep.context = fcache->fcontext;
+	cprep.owner = fcache->cowner;
 	fcache->cplan = GetCachedPlan(plansource,
 								  fcache->paramLI,
 								  fcache->cowner,
-								  NULL);
+								  NULL,
+								  &cprep);
 
 	/*
 	 * If necessary, make esarray[] bigger to hold the needed state.
@@ -719,9 +725,12 @@ init_execution_state(SQLFunctionCachePtr fcache)
 	/*
 	 * Build execution_state list to match the number of contained plans.
 	 */
+	i = 0;
 	foreach(lc, fcache->cplan->stmt_list)
 	{
 		PlannedStmt *stmt = lfirst_node(PlannedStmt, lc);
+		ExecPrep *prep = cprep.prep_list ? list_nth(cprep.prep_list, i) :
+			NULL;
 		execution_state *newes;
 
 		/*
@@ -763,6 +772,7 @@ init_execution_state(SQLFunctionCachePtr fcache)
 		newes->setsResult = false;	/* might change below */
 		newes->lazyEval = false;	/* might change below */
 		newes->stmt = stmt;
+		newes->prep = prep;
 		newes->qd = NULL;
 
 		if (stmt->canSetTag)
@@ -1362,7 +1372,7 @@ postquel_start(execution_state *es, SQLFunctionCachePtr fcache)
 		dest = None_Receiver;
 
 	es->qd = CreateQueryDesc(es->stmt,
-							 NULL,
+							 es->prep,
 							 fcache->func->src,
 							 GetActiveSnapshot(),
 							 InvalidSnapshot,
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 4c5647ac38a..c5812612f8d 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -4648,8 +4648,8 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	 * as a reference for building the ResultRelInfo of the target partition.
 	 * In either case, it doesn't matter which result relation is kept, so we
 	 * just keep the first one, if all others have been pruned.  See also,
-	 * ExecDoInitialPruning(), which ensures that this first result relation
-	 * has been locked.
+	 * AcquireExecutorLocksUnpruned(), which ensures that this first result
+	 * relation has been locked.
 	 */
 	i = 0;
 	foreach(l, node->resultRelations)
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 7a3cb944d6f..72d52baff4b 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -1579,6 +1579,7 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 {
 	CachedPlanSource *plansource;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	List	   *stmt_list;
 	char	   *query_string;
 	Snapshot	snapshot;
@@ -1659,7 +1660,10 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 	 */
 
 	/* Replan if needed, and increment plan refcount for portal */
-	cplan = GetCachedPlan(plansource, paramLI, NULL, _SPI_current->queryEnv);
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(plansource, paramLI, NULL, _SPI_current->queryEnv,
+						  &cprep);
 	stmt_list = cplan->stmt_list;
 
 	if (!plan->saved)
@@ -1685,7 +1689,7 @@ SPI_cursor_open_internal(const char *name, SPIPlanPtr plan,
 					  query_string,
 					  plansource->commandTag,
 					  stmt_list,
-					  NIL,
+					  cprep.prep_list,	/* XXX - need copy? */
 					  cplan);
 
 	/*
@@ -2078,6 +2082,7 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 {
 	CachedPlanSource *plansource;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	SPICallbackArg spicallbackarg;
 	ErrorContextCallback spierrcontext;
 
@@ -2101,9 +2106,12 @@ SPI_plan_get_cached_plan(SPIPlanPtr plan)
 	error_context_stack = &spierrcontext;
 
 	/* Get the generic plan for the query */
+	cprep.context = CurrentMemoryContext;
+	cprep.owner = CurrentResourceOwner;
 	cplan = GetCachedPlan(plansource, NULL,
 						  plan->saved ? CurrentResourceOwner : NULL,
-						  _SPI_current->queryEnv);
+						  _SPI_current->queryEnv,
+						  &cprep);
 	Assert(cplan == plansource->gplan);
 
 	/* Pop the error context stack */
@@ -2501,6 +2509,7 @@ _SPI_execute_plan(SPIPlanPtr plan, const SPIExecuteOptions *options,
 		CachedPlanSource *plansource = (CachedPlanSource *) lfirst(lc1);
 		List	   *stmt_list;
 		ListCell   *lc2;
+		CachedPlanPrepData cprep = {0};
 		List	   *prep_list;
 		int			i;
 
@@ -2577,11 +2586,14 @@ _SPI_execute_plan(SPIPlanPtr plan, const SPIExecuteOptions *options,
 		 * Replan if needed, and increment plan refcount.  If it's a saved
 		 * plan, the refcount must be backed by the plan_owner.
 		 */
+		cprep.context = CurrentMemoryContext;
+		cprep.owner = CurrentResourceOwner;
 		cplan = GetCachedPlan(plansource, options->params,
-							  plan_owner, _SPI_current->queryEnv);
+							  plan_owner, _SPI_current->queryEnv,
+							  &cprep);
 
 		stmt_list = cplan->stmt_list;
-		prep_list = NIL;
+		prep_list = cprep.prep_list;
 
 		/*
 		 * If we weren't given a specific snapshot to use, and the statement
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index c4fd646b999..4c76e78c1da 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -608,6 +608,7 @@ standard_planner(Query *parse, const char *query_string, int cursorOptions,
 											  glob->prunableRelids);
 	result->permInfos = glob->finalrteperminfos;
 	result->resultRelations = glob->resultRelations;
+	result->firstResultRels = glob->firstResultRels;
 	result->appendRelations = glob->appendRelations;
 	result->subplans = glob->subplans;
 	result->rewindPlanIDs = glob->rewindPlanIDs;
diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index ccdc9bc264a..229b39060ae 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -1274,6 +1274,9 @@ set_plan_refs(PlannerInfo *root, Plan *plan, int rtoffset)
 						lappend_int(root->glob->resultRelations,
 									splan->rootRelation);
 				}
+				root->glob->firstResultRels =
+					lappend_int(root->glob->firstResultRels,
+								linitial_int(splan->resultRelations));
 			}
 			break;
 		case T_Append:
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index d3964a12a14..82972beee70 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -1639,6 +1639,7 @@ exec_bind_message(StringInfo input_message)
 	int16	   *rformats = NULL;
 	CachedPlanSource *psrc;
 	CachedPlan *cplan;
+	CachedPlanPrepData cprep = {0};
 	Portal		portal;
 	char	   *query_string;
 	char	   *saved_stmt_name;
@@ -2021,7 +2022,9 @@ exec_bind_message(StringInfo input_message)
 	 * will be generated in MessageContext.  The plan refcount will be
 	 * assigned to the Portal, so it will be released at portal destruction.
 	 */
-	cplan = GetCachedPlan(psrc, params, NULL, NULL);
+	cprep.context = portal->portalContext;
+	cprep.owner = portal->resowner;
+	cplan = GetCachedPlan(psrc, params, NULL, NULL, &cprep);
 
 	/*
 	 * Now we can define the portal.
@@ -2034,7 +2037,7 @@ exec_bind_message(StringInfo input_message)
 					  query_string,
 					  psrc->commandTag,
 					  cplan->stmt_list,
-					  NIL,
+					  cprep.prep_list,
 					  cplan);
 
 	/* Portal is defined, set the plan ID based on its contents. */
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index 6661d2c6b73..ebcf601fce7 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -93,7 +93,7 @@ static bool StmtPlanRequiresRevalidation(CachedPlanSource *plansource);
 static bool BuildingPlanRequiresSnapshot(CachedPlanSource *plansource);
 static List *RevalidateCachedQuery(CachedPlanSource *plansource,
 								   QueryEnvironment *queryEnv);
-static bool CheckCachedPlan(CachedPlanSource *plansource);
+static bool PrepAndCheckCachedPlan(CachedPlanSource *plansource, CachedPlanPrepData *cprep);
 static CachedPlan *BuildCachedPlan(CachedPlanSource *plansource, List *qlist,
 								   ParamListInfo boundParams, QueryEnvironment *queryEnv);
 static bool choose_custom_plan(CachedPlanSource *plansource,
@@ -101,6 +101,8 @@ static bool choose_custom_plan(CachedPlanSource *plansource,
 static double cached_plan_cost(CachedPlan *plan, bool include_planner);
 static Query *QueryListGetPrimaryStmt(List *stmts);
 static void AcquireExecutorLocks(List *stmt_list, bool acquire);
+static void AcquireExecutorLocksUnpruned(List *stmt_list, bool acquire,
+										 CachedPlanPrepData *cprep);
 static void AcquirePlannerLocks(List *stmt_list, bool acquire);
 static void ScanQueryForLocks(Query *parsetree, bool acquire);
 static bool ScanQueryWalker(Node *node, bool *acquire);
@@ -137,6 +139,26 @@ ResourceOwnerForgetPlanCacheRef(ResourceOwner owner, CachedPlan *plan)
 /* GUC parameter */
 int			plan_cache_mode = PLAN_CACHE_MODE_AUTO;
 
+/*
+ * Lock acquisition policy for execution locks.
+ *
+ * LOCK_ALL acquires locks on all relations mentioned in the plan,
+ * reproducing the behavior of AcquireExecutorLocks().
+ *
+ * LOCK_UNPRUNED restricts locking to only the unpruned relations. That
+ * includes those mentioned in PlannedStmt.unprunableRelids and the leaf
+ * partitions remaining after performing initial pruning.
+ */
+typedef enum LockPolicy
+{
+	LOCK_ALL,
+	LOCK_UNPRUNED,
+} LockPolicy;
+
+static void AcquireExecutorLocksWithPolicy(List *stmt_list,
+										   LockPolicy policy, bool acquire,
+										   CachedPlanPrepData *cprep);
+
 /*
  * InitPlanCache: initialize module during InitPostgres.
  *
@@ -938,7 +960,12 @@ RevalidateCachedQuery(CachedPlanSource *plansource,
 }
 
 /*
- * CheckCachedPlan: see if the CachedPlanSource's generic plan is valid.
+ * PrepAndCheckCachedPlan: see if the CachedPlanSource's generic plan is valid.
+ *
+ * If 'cprep' is not NULL, ExecutorPrep() is applied to each PlannedStmt to
+ * compute the set of partitions that survive initial runtime pruning in order
+ * to only lock them.  The resulting ExecPrep structures are saved in cprep for
+ * later reuse by ExecutorStart().
  *
  * Caller must have already called RevalidateCachedQuery to verify that the
  * querytree is up to date.
@@ -947,7 +974,7 @@ RevalidateCachedQuery(CachedPlanSource *plansource,
  * (We must do this for the "true" result to be race-condition-free.)
  */
 static bool
-CheckCachedPlan(CachedPlanSource *plansource)
+PrepAndCheckCachedPlan(CachedPlanSource *plansource, CachedPlanPrepData *cprep)
 {
 	CachedPlan *plan = plansource->gplan;
 
@@ -975,13 +1002,15 @@ CheckCachedPlan(CachedPlanSource *plansource)
 	 */
 	if (plan->is_valid)
 	{
+		LockPolicy policy = !cprep ? LOCK_ALL : LOCK_UNPRUNED;
+
 		/*
 		 * Plan must have positive refcount because it is referenced by
 		 * plansource; so no need to fear it disappears under us here.
 		 */
 		Assert(plan->refcount > 0);
 
-		AcquireExecutorLocks(plan->stmt_list, true);
+		AcquireExecutorLocksWithPolicy(plan->stmt_list, policy, true, cprep);
 
 		/*
 		 * If plan was transient, check to see if TransactionXmin has
@@ -1003,7 +1032,7 @@ CheckCachedPlan(CachedPlanSource *plansource)
 		}
 
 		/* Oops, the race case happened.  Release useless locks. */
-		AcquireExecutorLocks(plan->stmt_list, false);
+		AcquireExecutorLocksWithPolicy(plan->stmt_list, policy, false, cprep);
 	}
 
 	/*
@@ -1283,6 +1312,10 @@ cached_plan_cost(CachedPlan *plan, bool include_planner)
  * On return, the plan is valid and we have sufficient locks to begin
  * execution.
  *
+ * If 'cprep' is not NULL and a generic plan is reused, the function prepares
+ * each PlannedStmt via ExecutorPrep() and stores the results in
+ * cprep->prep_list.  These are intended to be passed later to ExecutorStart().
+ *
  * On return, the refcount of the plan has been incremented; a later
  * ReleaseCachedPlan() call is expected.  If "owner" is not NULL then
  * the refcount has been reported to that ResourceOwner (note that this
@@ -1293,7 +1326,8 @@ cached_plan_cost(CachedPlan *plan, bool include_planner)
  */
 CachedPlan *
 GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
-			  ResourceOwner owner, QueryEnvironment *queryEnv)
+			  ResourceOwner owner, QueryEnvironment *queryEnv,
+			  CachedPlanPrepData *cprep)
 {
 	CachedPlan *plan = NULL;
 	List	   *qlist;
@@ -1315,7 +1349,9 @@ GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 
 	if (!customplan)
 	{
-		if (CheckCachedPlan(plansource))
+		if (cprep)
+			cprep->params = boundParams;
+		if (PrepAndCheckCachedPlan(plansource, cprep))
 		{
 			/* We want a generic plan, and we already have a valid one */
 			plan = plansource->gplan;
@@ -1902,6 +1938,32 @@ QueryListGetPrimaryStmt(List *stmts)
 	return NULL;
 }
 
+/*
+ * AcquireExecutorLocksWithPolicy
+ *      Acquire or release execution locks for a plan according to
+ *      the specified policy.
+ *
+ * The policy determines whether all relations or only unpruned ones are locked.
+ * For LOCK_UNPRUNED, ExecutorPrep is invoked to identify surviving partitions
+ * and its result is populated in cprep.
+ */
+static void
+AcquireExecutorLocksWithPolicy(List *stmt_list, LockPolicy policy, bool acquire,
+							   CachedPlanPrepData *cprep)
+{
+	switch (policy)
+	{
+		case LOCK_ALL:
+			AcquireExecutorLocks(stmt_list, acquire);
+			break;
+		case LOCK_UNPRUNED:
+			AcquireExecutorLocksUnpruned(stmt_list, acquire, cprep);
+			break;
+		default:
+			elog(ERROR, "invalid LockPolicy");
+	}
+}
+
 /*
  * AcquireExecutorLocks: acquire locks needed for execution of a cached plan;
  * or release them if acquire is false.
@@ -1954,6 +2016,153 @@ AcquireExecutorLocks(List *stmt_list, bool acquire)
 	}
 }
 
+/*
+ * LockRelids
+ * 		Acquire or release locks on the specified relids, which reference
+ * 		entries in the provided range table.
+ *
+ * Helper for AcquireExecutorLocksUnpruned().
+ */
+static void
+LockRelids(List *rtable, Bitmapset *relids, bool acquire)
+{
+	int	rtindex = -1;
+
+	while ((rtindex = bms_next_member(relids, rtindex)) >= 0)
+	{
+		RangeTblEntry *rte = list_nth_node(RangeTblEntry, rtable, rtindex - 1);
+
+		Assert(rte->rtekind == RTE_RELATION ||
+			   (rte->rtekind == RTE_SUBQUERY && OidIsValid(rte->relid)));
+
+		/*
+		 * Acquire the appropriate type of lock on each relation OID. Note
+		 * that we don't actually try to open the rel, and hence will not
+		 * fail if it's been dropped entirely --- we'll just transiently
+		 * acquire a non-conflicting lock.
+		 */
+		if (acquire)
+			LockRelationOid(rte->relid, rte->rellockmode);
+		else
+			UnlockRelationOid(rte->relid, rte->rellockmode);
+	}
+}
+
+/*
+ * AcquireExecutorLocksUnpruned
+ *		Acquire or release execution locks for only unpruned relations
+ *		referenced by the given PlannedStmts.
+ *
+ * This function uses ExecutorPrep to identify which partitions survive
+ * initial runtime pruning and locks only those, along with any unprunable
+ * base relations. During acquire, the resulting ExecPrep objects are stored
+ * in cprep->prep_list for later reuse. During release, those same ExecPrep
+ * objects are used to identify what to unlock.
+ *
+ * Unlike AcquireExecutorLocks(), which locks all relations listed in the
+ * PlannedStmt's rtable (LOCK_ALL policy), this function selectively locks
+ * only those rels that may be referenced during execution.
+ *
+ * prep_list is extended during acquire and must match stmt_list during
+ * release. Memory allocation happens in cprep->context.
+ */
+static void
+AcquireExecutorLocksUnpruned(List *stmt_list, bool acquire,
+							 CachedPlanPrepData *cprep)
+{
+	MemoryContext oldcontext = MemoryContextSwitchTo(cprep->context);
+	ListCell   *lc1;
+	List	   *prep_list;
+	int			i;
+
+	Assert(cprep);
+
+	/*
+	 * When releasing locks, use the ExecPrep list (if any) created during
+	 * acquisition to determine which relids to unlock. The list must match
+	 * the PlannedStmt list one-to-one.
+	 */
+	prep_list = cprep->prep_list;
+	Assert(acquire || list_length(prep_list) == list_length(stmt_list));
+
+	i = 0;
+	foreach(lc1, stmt_list)
+	{
+		PlannedStmt *plannedstmt = lfirst_node(PlannedStmt, lc1);
+		ExecPrep *prep;
+
+		if (plannedstmt->commandType == CMD_UTILITY)
+		{
+			/* Same as AcquireExecutorLocks(). */
+			Query	   *query = UtilityContainsQuery(plannedstmt->utilityStmt);
+
+			if (query)
+				ScanQueryForLocks(query, acquire);
+
+			/* Keep the list one-to-one with stmt_list. */
+			if (acquire)
+				cprep->prep_list = lappend(cprep->prep_list, NULL);
+			continue;
+		}
+
+		/*
+		 * Lock tables mentioned in the original query and other unprunable
+		 * relations that were added to the plan via inheritance expansion.
+		 */
+		LockRelids(plannedstmt->rtable, plannedstmt->unprunableRelids, acquire);
+
+		/* Lock partitions surviving runtime initial pruning. */
+		if (acquire)
+		{
+			prep = ExecutorPrep(plannedstmt, cprep->params, cprep->owner, true);
+			Assert(prep || plannedstmt->partPruneInfos == NULL);
+			cprep->prep_list = lappend(cprep->prep_list, prep);
+		}
+		else
+			prep = list_nth(prep_list, i++);
+
+		Assert(prep == NULL || prep->prep_estate);
+		if (prep)
+		{
+			EState *prep_estate = prep->prep_estate;
+
+			/*
+			 * es_unpruned_relids includes plannedstmt->unprunableRelids,
+			 * which we've already locked. Filter them out to avoid double-locking.
+			 */
+			Bitmapset *lock_relids = bms_difference(prep_estate->es_unpruned_relids,
+													plannedstmt->unprunableRelids);
+
+			/*
+			 * firstResultRels may contain pruned partitions that must still be
+			 * locked to satisfy executor assumptions (see comments in
+			 * ExecInitModifyTable(). Ensure they’re included here.
+			 */
+			if (plannedstmt->resultRelations)
+			{
+				ListCell *lc2;
+
+				foreach(lc2, plannedstmt->firstResultRels)
+				{
+					Index       firstResultRel = lfirst_int(lc2);
+
+					if (!bms_is_member(firstResultRel, lock_relids))
+						lock_relids = bms_add_member(lock_relids, firstResultRel);
+				}
+			}
+
+			LockRelids(plannedstmt->rtable, lock_relids, acquire);
+			bms_free(lock_relids);
+		}
+
+		/* Clean up prep if releasing locks. */
+		if (!acquire)
+			ExecPrepCleanup(prep);
+	}
+
+	MemoryContextSwitchTo(oldcontext);
+}
+
 /*
  * AcquirePlannerLocks: acquire locks needed for planning of a querytree list;
  * or release them if acquire is false.
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 30d889b54c5..6fb86dc05f6 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -141,6 +141,9 @@ typedef struct PlannerGlobal
 	/* "flat" list of integer RT indexes */
 	List	   *resultRelations;
 
+	/* "flat" list of integer RT indexes (one per ModifyTable node) */
+	List	   *firstResultRels;
+
 	/* "flat" list of AppendRelInfos */
 	List	   *appendRelations;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index c4393a94321..42b51299ece 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -123,6 +123,13 @@ typedef struct PlannedStmt
 	/* integer list of RT indexes, or NIL */
 	List	   *resultRelations;
 
+	/*
+	 * rtable indexes of first target relation in each ModifyTable node in the
+	 * plan for INSERT/UPDATE/DELETE/MERGE
+	 */
+	/* integer list of RT indexes, or NIL */
+	List	   *firstResultRels;
+
 	/* list of AppendRelInfo nodes */
 	List	   *appendRelations;
 
diff --git a/src/include/utils/plancache.h b/src/include/utils/plancache.h
index a82b66d4bc2..59f0b0fc4a4 100644
--- a/src/include/utils/plancache.h
+++ b/src/include/utils/plancache.h
@@ -197,6 +197,26 @@ typedef struct CachedExpression
 } CachedExpression;
 
 
+/*
+ * CachedPlanPrepData
+ *	Carries ExecutorPrep results for each PlannedStmt in a CachedPlan,
+ *	along with context and owner information needed to allocate them.
+ *
+ * Populated by GetCachedPlan() when ExecutorPrep is run on a generic plan.
+ *
+ * prep_list: results from ExecutorPrep(), one per PlannedStmt
+ * params: parameters that may be used during ExecutorPrep (e.g., pruning)
+ * context: memory context to allocate ExecutorPrep results in
+ * owner: resource owner to associate ExecutorPrep resources with
+ */
+typedef struct CachedPlanPrepData
+{
+	List   *prep_list;		/* List of ExecPrep */
+	ParamListInfo params;
+	MemoryContext context;
+	ResourceOwner owner;
+} CachedPlanPrepData;
+
 extern void InitPlanCache(void);
 extern void ResetPlanCache(void);
 
@@ -240,7 +260,8 @@ extern List *CachedPlanGetTargetList(CachedPlanSource *plansource,
 extern CachedPlan *GetCachedPlan(CachedPlanSource *plansource,
 								 ParamListInfo boundParams,
 								 ResourceOwner owner,
-								 QueryEnvironment *queryEnv);
+								 QueryEnvironment *queryEnv,
+								 CachedPlanPrepData *cprep);
 extern void ReleaseCachedPlan(CachedPlan *plan, ResourceOwner owner);
 
 extern bool CachedPlanAllowsSimpleValidityCheck(CachedPlanSource *plansource,
-- 
2.47.3

