From 589bb816499e98b60aa5c406c409fb27b42b0e39 Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Thu, 30 Mar 2023 13:24:09 +0200
Subject: [PATCH] Fix setrefs.c code for adjusting partPruneInfos

We were transferring partPruneInfos from PlannerInfo into PlannerGlobal
wrong, essentially relying on all of them being transferred, and
adjusting their list indexes based on that.  But apparently it's
possible that some of them are skipped, so that strategy leads to a
corrupted execution tree.  Instead, adjust each Append/MergeAppend's
partpruneinfo index as we copy from one list to the other, which seems
safer anyway.  This requires adjusting the RT offset of the RTE
referenced in each partPruneInfo ahead of actually adjusting the RTE
itself, which seems a bit too ad-hoc.

This problem was introduced by commit ec386948948c.  However, it may be
that we no longer require the change introduced there, so perhaps we
should revert both the present commit and that one.

Problem noticed by sqlsmith.

Author: Amit Langote <amitlangote09@gmail.com
Discussion: https://postgr.es/m/CA+HiwqG6tbc2oadsbyyy24b2AL295XHQgyLRWghmA7u_SL1K8A@mail.gmail.com
---
 src/backend/optimizer/plan/setrefs.c | 83 +++++++++++++++++-----------
 1 file changed, 52 insertions(+), 31 deletions(-)

diff --git a/src/backend/optimizer/plan/setrefs.c b/src/backend/optimizer/plan/setrefs.c
index 5cc8366af6..b56666398e 100644
--- a/src/backend/optimizer/plan/setrefs.c
+++ b/src/backend/optimizer/plan/setrefs.c
@@ -350,6 +350,29 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 			palloc0(list_length(glob->subplans) * sizeof(bool));
 	}
 
+	/* Also fix up the information in PartitionPruneInfos. */
+	foreach(lc, root->partPruneInfos)
+	{
+		PartitionPruneInfo *pruneinfo = lfirst(lc);
+		ListCell   *l;
+
+		pruneinfo->root_parent_relids =
+			offset_relid_set(pruneinfo->root_parent_relids, rtoffset);
+		foreach(l, pruneinfo->prune_infos)
+		{
+			List	   *prune_infos = lfirst(l);
+			ListCell   *l2;
+
+			foreach(l2, prune_infos)
+			{
+				PartitionedRelPruneInfo *pinfo = lfirst(l2);
+
+				/* RT index of the table to which the pinfo belongs. */
+				pinfo->rtindex += rtoffset;
+			}
+		}
+	}
+
 	/* Now fix the Plan tree */
 	result = set_plan_refs(root, plan, rtoffset);
 
@@ -378,31 +401,6 @@ set_plan_references(PlannerInfo *root, Plan *plan)
 		}
 	}
 
-	/* Also fix up the information in PartitionPruneInfos. */
-	foreach(lc, root->partPruneInfos)
-	{
-		PartitionPruneInfo *pruneinfo = lfirst(lc);
-		ListCell   *l;
-
-		pruneinfo->root_parent_relids =
-			offset_relid_set(pruneinfo->root_parent_relids, rtoffset);
-		foreach(l, pruneinfo->prune_infos)
-		{
-			List	   *prune_infos = lfirst(l);
-			ListCell   *l2;
-
-			foreach(l2, prune_infos)
-			{
-				PartitionedRelPruneInfo *pinfo = lfirst(l2);
-
-				/* RT index of the table to which the pinfo belongs. */
-				pinfo->rtindex += rtoffset;
-			}
-		}
-
-		glob->partPruneInfos = lappend(glob->partPruneInfos, pruneinfo);
-	}
-
 	return result;
 }
 
@@ -1718,6 +1716,29 @@ set_customscan_references(PlannerInfo *root,
 	cscan->custom_relids = offset_relid_set(cscan->custom_relids, rtoffset);
 }
 
+/*
+ * register_partpruneinfo
+ *		Subroutine for set_append_references and set_mergeappend_references
+ *
+ * Add the PartitionPruneInfo from root->partPruneInfos at the given index
+ * into PlannerGlobal->partPruneInfos and return its index there.
+ */
+static int
+register_partpruneinfo(PlannerInfo *root, int part_prune_index)
+{
+	PlannerGlobal  *glob = root->glob;
+	PartitionPruneInfo *pruneinfo;
+
+	Assert(part_prune_index >= 0 &&
+		   part_prune_index < list_length(root->partPruneInfos));
+	pruneinfo = list_nth_node(PartitionPruneInfo, root->partPruneInfos,
+							  part_prune_index);
+
+	glob->partPruneInfos = lappend(glob->partPruneInfos, pruneinfo);
+
+	return list_length(glob->partPruneInfos) - 1;
+}
+
 /*
  * set_append_references
  *		Do set_plan_references processing on an Append
@@ -1771,11 +1792,11 @@ set_append_references(PlannerInfo *root,
 	aplan->apprelids = offset_relid_set(aplan->apprelids, rtoffset);
 
 	/*
-	 * PartitionPruneInfos will be added to a list in PlannerGlobal, so update
-	 * the index.
+	 * Add PartitionPruneInfo, if any, to PlannerGlobal and update the index.
 	 */
 	if (aplan->part_prune_index >= 0)
-		aplan->part_prune_index += list_length(root->glob->partPruneInfos);
+		aplan->part_prune_index =
+			register_partpruneinfo(root, aplan->part_prune_index);
 
 	/* We don't need to recurse to lefttree or righttree ... */
 	Assert(aplan->plan.lefttree == NULL);
@@ -1838,11 +1859,11 @@ set_mergeappend_references(PlannerInfo *root,
 	mplan->apprelids = offset_relid_set(mplan->apprelids, rtoffset);
 
 	/*
-	 * PartitionPruneInfos will be added to a list in PlannerGlobal, so update
-	 * the index.
+	 * Add PartitionPruneInfo, if any, to PlannerGlobal and update the index.
 	 */
 	if (mplan->part_prune_index >= 0)
-		mplan->part_prune_index += list_length(root->glob->partPruneInfos);
+		mplan->part_prune_index =
+			register_partpruneinfo(root, mplan->part_prune_index);
 
 	/* We don't need to recurse to lefttree or righttree ... */
 	Assert(mplan->plan.lefttree == NULL);