From dfc41510ef3ebec38e7a56b639ffa41193109b43 Mon Sep 17 00:00:00 2001
From: amitlan <amitlangote09@gmail.com>
Date: Thu, 9 Mar 2023 11:26:06 +0900
Subject: [PATCH v37 1/3] Add field to store partitioned relids to
 Append/MergeAppend

A future commit would like to move the timing of locking relations
referenced in a cached plan to ExecInitNode() traversal of the plan
tree from the current loop-over-rangetable in AcquireExecutorLocks().
Given that partitioned tables (their RT indexes) would not be
accessible via the new way of finding the relations to lock, add a
field to Append/MergeAppend to track them separately.

This refactors the code to look up partitioned parent relids from a
given list of leaf partition subpaths of an Append/MergeAppend out
of make_partition_pruneinfo() into its own function called
add_append_subpath_partrelids().  Though, the code needs to be
generalized to the cases where child rels can be joinrels or
upper (grouping) rels.  Also, to make it easier to traverse the parent
chain of a child grouping rel, this makes its RelOptInfo.parent to be
set, which is already done for baserels and joinrels.
---
 src/backend/optimizer/plan/createplan.c |  36 +++++--
 src/backend/optimizer/plan/planner.c    |   3 +
 src/backend/optimizer/util/appendinfo.c | 134 ++++++++++++++++++++++++
 src/backend/partitioning/partprune.c    | 124 +++-------------------
 src/include/nodes/plannodes.h           |  14 +++
 src/include/optimizer/appendinfo.h      |   3 +
 src/include/partitioning/partprune.h    |   3 +-
 7 files changed, 194 insertions(+), 123 deletions(-)

diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 910ffbf1e1..794cdb5e3b 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -25,6 +25,7 @@
 #include "nodes/extensible.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
+#include "optimizer/appendinfo.h"
 #include "optimizer/clauses.h"
 #include "optimizer/cost.h"
 #include "optimizer/optimizer.h"
@@ -1209,6 +1210,7 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path, int flags)
 	Oid		   *nodeCollations = NULL;
 	bool	   *nodeNullsFirst = NULL;
 	bool		consider_async = false;
+	List	   *allpartrelids = NIL;
 
 	/*
 	 * The subpaths list could be empty, if every child was proven empty by
@@ -1350,18 +1352,24 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path, int flags)
 			++nasyncplans;
 		}
 
+		/* Populate partitioned parent relids. */
+		allpartrelids = add_append_subpath_partrelids(root, subpath, rel,
+													  allpartrelids);
+
 		subplans = lappend(subplans, subplan);
 	}
 
+	plan->allpartrelids = allpartrelids;
+
 	/* Set below if we find quals that we can use to run-time prune */
 	plan->part_prune_index = -1;
 
 	/*
-	 * If any quals exist, they may be useful to perform further partition
-	 * pruning during execution.  Gather information needed by the executor to
-	 * do partition pruning.
+	 * If scanning partitions, check if there are quals that may be useful to
+	 * perform further partition pruning during execution.  Gather information
+	 * needed by the executor to do partition pruning.
 	 */
-	if (enable_partition_pruning)
+	if (enable_partition_pruning && allpartrelids != NIL)
 	{
 		List	   *prunequal;
 
@@ -1381,7 +1389,8 @@ create_append_plan(PlannerInfo *root, AppendPath *best_path, int flags)
 		if (prunequal != NIL)
 			plan->part_prune_index = make_partition_pruneinfo(root, rel,
 															  best_path->subpaths,
-															  prunequal);
+															  prunequal,
+															  allpartrelids);
 	}
 
 	plan->appendplans = subplans;
@@ -1425,6 +1434,7 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path,
 	List	   *subplans = NIL;
 	ListCell   *subpaths;
 	RelOptInfo *rel = best_path->path.parent;
+	List	   *allpartrelids = NIL;
 
 	/*
 	 * We don't have the actual creation of the MergeAppend node split out
@@ -1514,18 +1524,23 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path,
 			subplan = (Plan *) sort;
 		}
 
+		allpartrelids = add_append_subpath_partrelids(root, subpath, rel,
+													  allpartrelids);
+
 		subplans = lappend(subplans, subplan);
 	}
 
+	node->allpartrelids = allpartrelids;
+
 	/* Set below if we find quals that we can use to run-time prune */
 	node->part_prune_index = -1;
 
 	/*
-	 * If any quals exist, they may be useful to perform further partition
-	 * pruning during execution.  Gather information needed by the executor to
-	 * do partition pruning.
+	 * If scanning partitions, check if there are quals that may be useful to
+	 * perform further partition pruning during execution.  Gather information
+	 * needed by the executor to do partition pruning.
 	 */
-	if (enable_partition_pruning)
+	if (enable_partition_pruning && allpartrelids != NIL)
 	{
 		List	   *prunequal;
 
@@ -1537,7 +1552,8 @@ create_merge_append_plan(PlannerInfo *root, MergeAppendPath *best_path,
 		if (prunequal != NIL)
 			node->part_prune_index = make_partition_pruneinfo(root, rel,
 															  best_path->subpaths,
-															  prunequal);
+															  prunequal,
+															  allpartrelids);
 	}
 
 	node->mergeplans = subplans;
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index a1873ce26d..62b3ec96cc 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -7801,8 +7801,11 @@ create_partitionwise_grouping_paths(PlannerInfo *root,
 									   agg_costs, gd, &child_extra,
 									   &child_partially_grouped_rel);
 
+		/* Mark as child of grouped_rel. */
+		child_grouped_rel->parent = grouped_rel;
 		if (child_partially_grouped_rel)
 		{
+			child_partially_grouped_rel->parent = grouped_rel;
 			partially_grouped_live_children =
 				lappend(partially_grouped_live_children,
 						child_partially_grouped_rel);
diff --git a/src/backend/optimizer/util/appendinfo.c b/src/backend/optimizer/util/appendinfo.c
index 9d377385f1..4876742ab2 100644
--- a/src/backend/optimizer/util/appendinfo.c
+++ b/src/backend/optimizer/util/appendinfo.c
@@ -40,6 +40,7 @@ static void make_inh_translation_list(Relation oldrelation,
 									  AppendRelInfo *appinfo);
 static Node *adjust_appendrel_attrs_mutator(Node *node,
 											adjust_appendrel_attrs_context *context);
+static List *add_part_relids(List *allpartrelids, Bitmapset *partrelids);
 
 
 /*
@@ -1031,3 +1032,136 @@ distribute_row_identity_vars(PlannerInfo *root)
 		}
 	}
 }
+
+/*
+ * add_append_subpath_partrelids
+ *		Look up a child subpath's rel's partitioned parent relids up to
+ *		parentrel and add the bitmapset containing those into
+ *		'allpartrelids'
+ */
+List *
+add_append_subpath_partrelids(PlannerInfo *root, Path *subpath,
+							  RelOptInfo *parentrel,
+							  List *allpartrelids)
+{
+	RelOptInfo *prel = subpath->parent;
+	Relids		partrelids = NULL;
+
+	/* Nothing to do if there's no parent to begin with. */
+	if (!IS_OTHER_REL(prel))
+		return allpartrelids;
+
+	/*
+	 * Traverse up to the pathrel's topmost partitioned parent, collecting
+	 * parent relids as we go; but stop if we reach parentrel.  (Normally, a
+	 * pathrel's topmost partitioned parent is either parentrel or a UNION ALL
+	 * appendrel child of parentrel.  But when handling partitionwise joins of
+	 * multi-level partitioning trees, we can see an append path whose
+	 * parentrel is an intermediate partitioned table.)
+	 */
+	do
+	{
+		Relids	parent_relids = NULL;
+
+		/*
+		 * For simple child rels, we can simply get the parent relid from
+		 * prel->parent.  But for partitionwise join and aggregate child rels,
+		 * while we can use prel->parent to move up the tree, parent relids to
+		 * add into 'partrelids' must be found the hard way through the
+		 * AppendInfoInfos, because 1) a joinrel's relids may point to RTE_JOIN
+		 * entries, 2) topmost parent grouping rel's relids field is left NULL.
+		 */
+		if (IS_SIMPLE_REL(prel))
+		{
+			prel = prel->parent;
+			/* Stop once we reach the root partitioned rel. */
+			if (!IS_PARTITIONED_REL(prel))
+				break;
+			parent_relids = bms_add_members(parent_relids, prel->relids);
+		}
+		else
+		{
+			AppendRelInfo **appinfos;
+			int		nappinfos,
+					i;
+
+			appinfos = find_appinfos_by_relids(root, prel->relids,
+											   &nappinfos);
+			for (i = 0; i < nappinfos; i++)
+			{
+				AppendRelInfo *appinfo = appinfos[i];
+
+				parent_relids = bms_add_member(parent_relids,
+											   appinfo->parent_relid);
+			}
+			pfree(appinfos);
+			prel = prel->parent;
+		}
+		/* accept this level as an interesting parent */
+		partrelids = bms_add_members(partrelids, parent_relids);
+		if (prel == parentrel)
+			break;		/* don't traverse above parentrel */
+	} while (IS_OTHER_REL(prel));
+
+	if (partrelids == NULL)
+		return allpartrelids;
+
+	return add_part_relids(allpartrelids, partrelids);
+}
+
+/*
+ * add_part_relids
+ *		Add new info to a list of Bitmapsets of partitioned relids.
+ *
+ * Within 'allpartrelids', there is one Bitmapset for each topmost parent
+ * partitioned rel.  Each Bitmapset contains the RT indexes of the topmost
+ * parent as well as its relevant non-leaf child partitions.  Since (by
+ * construction of the rangetable list) parent partitions must have lower
+ * RT indexes than their children, we can distinguish the topmost parent
+ * as being the lowest set bit in the Bitmapset.
+ *
+ * 'partrelids' contains the RT indexes of a parent partitioned rel, and
+ * possibly some non-leaf children, that are newly identified as parents of
+ * some subpath rel passed to make_partition_pruneinfo().  These are added
+ * to an appropriate member of 'allpartrelids'.
+ *
+ * Note that the list contains only RT indexes of partitioned tables that
+ * are parents of some scan-level relation appearing in the 'subpaths' that
+ * make_partition_pruneinfo() is dealing with.  Also, "topmost" parents are
+ * not allowed to be higher than the 'parentrel' associated with the append
+ * path.  In this way, we avoid expending cycles on partitioned rels that
+ * can't contribute useful pruning information for the problem at hand.
+ * (It is possible for 'parentrel' to be a child partitioned table, and it
+ * is also possible for scan-level relations to be child partitioned tables
+ * rather than leaf partitions.  Hence we must construct this relation set
+ * with reference to the particular append path we're dealing with, rather
+ * than looking at the full partitioning structure represented in the
+ * RelOptInfos.)
+ */
+static List *
+add_part_relids(List *allpartrelids, Bitmapset *partrelids)
+{
+	Index		targetpart;
+	ListCell   *lc;
+
+	/* We can easily get the lowest set bit this way: */
+	targetpart = bms_next_member(partrelids, -1);
+	Assert(targetpart > 0);
+
+	/* Look for a matching topmost parent */
+	foreach(lc, allpartrelids)
+	{
+		Bitmapset  *currpartrelids = (Bitmapset *) lfirst(lc);
+		Index		currtarget = bms_next_member(currpartrelids, -1);
+
+		if (targetpart == currtarget)
+		{
+			/* Found a match, so add any new RT indexes to this hierarchy */
+			currpartrelids = bms_add_members(currpartrelids, partrelids);
+			lfirst(lc) = currpartrelids;
+			return allpartrelids;
+		}
+	}
+	/* No match, so add the new partition hierarchy to the list */
+	return lappend(allpartrelids, partrelids);
+}
diff --git a/src/backend/partitioning/partprune.c b/src/backend/partitioning/partprune.c
index 510145e3c0..3557e07082 100644
--- a/src/backend/partitioning/partprune.c
+++ b/src/backend/partitioning/partprune.c
@@ -138,7 +138,6 @@ typedef struct PruneStepResult
 } PruneStepResult;
 
 
-static List *add_part_relids(List *allpartrelids, Bitmapset *partrelids);
 static List *make_partitionedrel_pruneinfo(PlannerInfo *root,
 										   RelOptInfo *parentrel,
 										   List *prunequal,
@@ -221,33 +220,32 @@ static void partkey_datum_from_expr(PartitionPruneContext *context,
  * of scan paths for its child rels.
  * 'prunequal' is a list of potential pruning quals (i.e., restriction
  * clauses that are applicable to the appendrel).
+ * 'allpartrelids' contains Bitmapsets of RT indexes of partitioned parents
+ * whose partitions' Paths are in 'subpaths'; there's one Bitmapset for every
+ * partition tree involved.
  */
 int
 make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 						 List *subpaths,
-						 List *prunequal)
+						 List *prunequal,
+						 List *allpartrelids)
 {
 	PartitionPruneInfo *pruneinfo;
 	Bitmapset  *allmatchedsubplans = NULL;
-	List	   *allpartrelids;
 	List	   *prunerelinfos;
 	int		   *relid_subplan_map;
 	ListCell   *lc;
 	int			i;
 
+	Assert(list_length(allpartrelids) > 0);
+
 	/*
-	 * Scan the subpaths to see which ones are scans of partition child
-	 * relations, and identify their parent partitioned rels.  (Note: we must
-	 * restrict the parent partitioned rels to be parentrel or children of
-	 * parentrel, otherwise we couldn't translate prunequal to match.)
-	 *
-	 * Also construct a temporary array to map from partition-child-relation
-	 * relid to the index in 'subpaths' of the scan plan for that partition.
+	 * Construct a temporary array to map from partition-child-relation relid
+	 * to the index in 'subpaths' of the scan plan for that partition.
 	 * (Use of "subplan" rather than "subpath" is a bit of a misnomer, but
 	 * we'll let it stand.)  For convenience, we use 1-based indexes here, so
 	 * that zero can represent an un-filled array entry.
 	 */
-	allpartrelids = NIL;
 	relid_subplan_map = palloc0(sizeof(int) * root->simple_rel_array_size);
 
 	i = 1;
@@ -256,50 +254,9 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 		Path	   *path = (Path *) lfirst(lc);
 		RelOptInfo *pathrel = path->parent;
 
-		/* We don't consider partitioned joins here */
-		if (pathrel->reloptkind == RELOPT_OTHER_MEMBER_REL)
-		{
-			RelOptInfo *prel = pathrel;
-			Bitmapset  *partrelids = NULL;
-
-			/*
-			 * Traverse up to the pathrel's topmost partitioned parent,
-			 * collecting parent relids as we go; but stop if we reach
-			 * parentrel.  (Normally, a pathrel's topmost partitioned parent
-			 * is either parentrel or a UNION ALL appendrel child of
-			 * parentrel.  But when handling partitionwise joins of
-			 * multi-level partitioning trees, we can see an append path whose
-			 * parentrel is an intermediate partitioned table.)
-			 */
-			do
-			{
-				AppendRelInfo *appinfo;
-
-				Assert(prel->relid < root->simple_rel_array_size);
-				appinfo = root->append_rel_array[prel->relid];
-				prel = find_base_rel(root, appinfo->parent_relid);
-				if (!IS_PARTITIONED_REL(prel))
-					break;		/* reached a non-partitioned parent */
-				/* accept this level as an interesting parent */
-				partrelids = bms_add_member(partrelids, prel->relid);
-				if (prel == parentrel)
-					break;		/* don't traverse above parentrel */
-			} while (prel->reloptkind == RELOPT_OTHER_MEMBER_REL);
-
-			if (partrelids)
-			{
-				/*
-				 * Found some relevant parent partitions, which may or may not
-				 * overlap with partition trees we already found.  Add new
-				 * information to the allpartrelids list.
-				 */
-				allpartrelids = add_part_relids(allpartrelids, partrelids);
-				/* Also record the subplan in relid_subplan_map[] */
-				/* No duplicates please */
-				Assert(relid_subplan_map[pathrel->relid] == 0);
-				relid_subplan_map[pathrel->relid] = i;
-			}
-		}
+		/* No duplicates please */
+		Assert(relid_subplan_map[pathrel->relid] == 0);
+		relid_subplan_map[pathrel->relid] = i;
 		i++;
 	}
 
@@ -368,63 +325,6 @@ make_partition_pruneinfo(PlannerInfo *root, RelOptInfo *parentrel,
 	return list_length(root->partPruneInfos) - 1;
 }
 
-/*
- * add_part_relids
- *		Add new info to a list of Bitmapsets of partitioned relids.
- *
- * Within 'allpartrelids', there is one Bitmapset for each topmost parent
- * partitioned rel.  Each Bitmapset contains the RT indexes of the topmost
- * parent as well as its relevant non-leaf child partitions.  Since (by
- * construction of the rangetable list) parent partitions must have lower
- * RT indexes than their children, we can distinguish the topmost parent
- * as being the lowest set bit in the Bitmapset.
- *
- * 'partrelids' contains the RT indexes of a parent partitioned rel, and
- * possibly some non-leaf children, that are newly identified as parents of
- * some subpath rel passed to make_partition_pruneinfo().  These are added
- * to an appropriate member of 'allpartrelids'.
- *
- * Note that the list contains only RT indexes of partitioned tables that
- * are parents of some scan-level relation appearing in the 'subpaths' that
- * make_partition_pruneinfo() is dealing with.  Also, "topmost" parents are
- * not allowed to be higher than the 'parentrel' associated with the append
- * path.  In this way, we avoid expending cycles on partitioned rels that
- * can't contribute useful pruning information for the problem at hand.
- * (It is possible for 'parentrel' to be a child partitioned table, and it
- * is also possible for scan-level relations to be child partitioned tables
- * rather than leaf partitions.  Hence we must construct this relation set
- * with reference to the particular append path we're dealing with, rather
- * than looking at the full partitioning structure represented in the
- * RelOptInfos.)
- */
-static List *
-add_part_relids(List *allpartrelids, Bitmapset *partrelids)
-{
-	Index		targetpart;
-	ListCell   *lc;
-
-	/* We can easily get the lowest set bit this way: */
-	targetpart = bms_next_member(partrelids, -1);
-	Assert(targetpart > 0);
-
-	/* Look for a matching topmost parent */
-	foreach(lc, allpartrelids)
-	{
-		Bitmapset  *currpartrelids = (Bitmapset *) lfirst(lc);
-		Index		currtarget = bms_next_member(currpartrelids, -1);
-
-		if (targetpart == currtarget)
-		{
-			/* Found a match, so add any new RT indexes to this hierarchy */
-			currpartrelids = bms_add_members(currpartrelids, partrelids);
-			lfirst(lc) = currpartrelids;
-			return allpartrelids;
-		}
-	}
-	/* No match, so add the new partition hierarchy to the list */
-	return lappend(allpartrelids, partrelids);
-}
-
 /*
  * make_partitionedrel_pruneinfo
  *		Build a List of PartitionedRelPruneInfos, one for each interesting
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 659bd05c0c..a0bb16cff4 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -270,6 +270,13 @@ typedef struct Append
 	List	   *appendplans;
 	int			nasyncplans;	/* # of asynchronous plans */
 
+	/*
+	 * RTIs of all partitioned tables whose children are scanned by
+	 * appendplans. The list contains a bitmapset for every partition tree
+	 * covered by this Append.
+	 */
+	List	   *allpartrelids;
+
 	/*
 	 * All 'appendplans' preceding this index are non-partial plans. All
 	 * 'appendplans' from this index onwards are partial plans.
@@ -294,6 +301,13 @@ typedef struct MergeAppend
 
 	List	   *mergeplans;
 
+	/*
+	 * RTIs of all partitioned tables whose children are scanned by
+	 * mergeplans. The list contains a bitmapset for every partition tree
+	 * covered by this MergeAppend.
+	 */
+	List	   *allpartrelids;
+
 	/* these fields are just like the sort-key info in struct Sort: */
 
 	/* number of sort-key columns */
diff --git a/src/include/optimizer/appendinfo.h b/src/include/optimizer/appendinfo.h
index a05f91f77d..1621a7319a 100644
--- a/src/include/optimizer/appendinfo.h
+++ b/src/include/optimizer/appendinfo.h
@@ -46,5 +46,8 @@ extern void add_row_identity_columns(PlannerInfo *root, Index rtindex,
 									 RangeTblEntry *target_rte,
 									 Relation target_relation);
 extern void distribute_row_identity_vars(PlannerInfo *root);
+extern List *add_append_subpath_partrelids(PlannerInfo *root, Path *subpath,
+							  RelOptInfo *parentrel,
+							  List *allpartrelids);
 
 #endif							/* APPENDINFO_H */
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index c0d6889d47..2d907d31d4 100644
--- a/src/include/partitioning/partprune.h
+++ b/src/include/partitioning/partprune.h
@@ -73,7 +73,8 @@ typedef struct PartitionPruneContext
 extern int make_partition_pruneinfo(struct PlannerInfo *root,
 									struct RelOptInfo *parentrel,
 									List *subpaths,
-									List *prunequal);
+									List *prunequal,
+									List *allpartrelids);
 extern Bitmapset *prune_append_rel_partitions(struct RelOptInfo *rel);
 extern Bitmapset *get_matching_partitions(PartitionPruneContext *context,
 										  List *pruning_steps);
-- 
2.35.3

