Merge branch 'main' into citus_pause_node

src/backend/distributed/planner/distributed_planner.c

@@ -108,6 +108,7 @@ static int AssignRTEIdentities(List *rangeTableList, int rteIdCounter);
 static void AssignRTEIdentity(RangeTblEntry *rangeTableEntry, int rteIdentifier);
 static void AdjustPartitioningForDistributedPlanning(List *rangeTableList,
 													 bool setPartitionedTablesInherited);
+static bool RTEWentThroughAdjustPartitioning(RangeTblEntry *rangeTableEntry);
 static PlannedStmt * FinalizeNonRouterPlan(PlannedStmt *localPlan,
 										   DistributedPlan *distributedPlan,
 										   CustomScan *customScan);
@@ -494,6 +495,20 @@ AdjustPartitioningForDistributedPlanning(List *rangeTableList,
 }
 
 
+/*
+ * RTEWentThroughAdjustPartitioning returns true if the given rangetableentry
+ * has been modified through AdjustPartitioningForDistributedPlanning
+ * function, false otherwise.
+ */
+static bool
+RTEWentThroughAdjustPartitioning(RangeTblEntry *rangeTableEntry)
+{
+	return (rangeTableEntry->rtekind == RTE_RELATION &&
+			PartitionedTable(rangeTableEntry->relid) &&
+			rangeTableEntry->inh == false);
+}
+
+
 /*
  * AssignRTEIdentity assigns the given rteIdentifier to the given range table
  * entry.
@@ -1976,6 +1991,62 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 }
 
 
+/*
+ * multi_get_relation_info_hook modifies the relation's indexlist
+ * if necessary, to avoid a crash in PG16 caused by our
+ * Citus function AdjustPartitioningForDistributedPlanning().
+ *
+ * AdjustPartitioningForDistributedPlanning() is a hack that we use
+ * to prevent Postgres' standard_planner() to expand all the partitions
+ * for the distributed planning when a distributed partitioned table
+ * is queried. It is required for both correctness and performance
+ * reasons. Although we can eliminate the use of the function for
+ * the correctness (e.g., make sure that rest of the planner can handle
+ * partitions), it's performance implication is hard to avoid. Certain
+ * planning logic of Citus (such as router or query pushdown) relies
+ * heavily on the relationRestrictionList. If
+ * AdjustPartitioningForDistributedPlanning() is removed, all the
+ * partitions show up in the relationRestrictionList, causing high
+ * planning times for such queries.
+ */
+void
+multi_get_relation_info_hook(PlannerInfo *root, Oid relationObjectId, bool inhparent,
+							 RelOptInfo *rel)
+{
+	if (!CitusHasBeenLoaded())
+	{
+		return;
+	}
+
+	Index varno = rel->relid;
+	RangeTblEntry *rangeTableEntry = planner_rt_fetch(varno, root);
+
+	if (RTEWentThroughAdjustPartitioning(rangeTableEntry))
+	{
+		ListCell *lc = NULL;
+		foreach(lc, rel->indexlist)
+		{
+			IndexOptInfo *indexOptInfo = (IndexOptInfo *) lfirst(lc);
+			if (get_rel_relkind(indexOptInfo->indexoid) == RELKIND_PARTITIONED_INDEX)
+			{
+				/*
+				 * Normally, we should not need this. However, the combination of
+				 * Postgres commit 3c569049b7b502bb4952483d19ce622ff0af5fd6 and
+				 * Citus function AdjustPartitioningForDistributedPlanning()
+				 * forces us to do this. The commit expects partitioned indexes
+				 * to belong to relations with "inh" flag set properly. Whereas, the
+				 * function overrides "inh" flag. To avoid a crash,
+				 * we go over the list of indexinfos and remove all partitioned indexes.
+				 * Partitioned indexes were ignored pre PG16 anyway, we are essentially
+				 * not breaking any logic.
+				 */
+				rel->indexlist = foreach_delete_current(rel->indexlist, lc);
+			}
+		}
+	}
+}
+
+
 /*
  * TranslatedVars deep copies the translated vars for the given relation index
  * if there is any append rel list.

src/backend/distributed/shared_library_init.c

@@ -104,6 +104,7 @@
 #include "replication/walsender.h"
 #include "storage/ipc.h"
 #include "optimizer/planner.h"
+#include "optimizer/plancat.h"
 #include "optimizer/paths.h"
 #include "tcop/tcopprot.h"
 #include "utils/guc.h"
@@ -452,6 +453,7 @@ _PG_init(void)
 
 	/* register for planner hook */
 	set_rel_pathlist_hook = multi_relation_restriction_hook;
+	get_relation_info_hook = multi_get_relation_info_hook;
 	set_join_pathlist_hook = multi_join_restriction_hook;
 	ExecutorStart_hook = CitusExecutorStart;
 	ExecutorRun_hook = CitusExecutorRun;

src/include/distributed/distributed_planner.h

@@ -234,6 +234,8 @@ extern List * TranslatedVarsForRteIdentity(int rteIdentity);
 extern struct DistributedPlan * GetDistributedPlan(CustomScan *node);
 extern void multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 											Index restrictionIndex, RangeTblEntry *rte);
+extern void multi_get_relation_info_hook(PlannerInfo *root, Oid relationObjectId, bool
+										 inhparent, RelOptInfo *rel);
 extern void multi_join_restriction_hook(PlannerInfo *root,
 										RelOptInfo *joinrel,
 										RelOptInfo *outerrel,