Add get_relation_info hook to avoid crash adjusted partitioning

If we allow to adjust partitioning, we get a crash when accessing
amcostestimate of partitioned indexes, because amcostestimate is NULL
for them. The following PG commit is the culprit
3c569049b7
3c569049b7b502bb4952483d19ce622ff0af5fd6

Previously, partitioned indexes would just be ignored.

Now, they are added in the list. However get_relation_info expects
the tables which have partitioned indexes to have the inh flag set
properly. AdjustPartitioningForDistributedPlanning plays with that
flag, hence we don't get the desired behaviour.

The hook is simply removing all partitioned indexes from the list.

src/backend/distributed/planner/distributed_planner.c

@@ -16,6 +16,7 @@
 #include <float.h>
 #include <limits.h>
 
+#include "access/genam.h"
 #include "access/htup_details.h"
 #include "access/xact.h"
 #include "catalog/pg_class.h"
@@ -2010,6 +2011,71 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 }
 
 
+/*
+ * 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.
+ *
+ * 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, causing high planning times for
+ * such queries.
+ */
+void
+multi_partitioned_index_hook(PlannerInfo *root, Oid relationObjectId, bool inhparent,
+							 RelOptInfo *rel)
+{
+	bool partitioningAdjusted = false;
+	List *rangeTableList = ExtractRangeTableEntryList(root->parse);
+
+	ListCell *rangeTableCell = NULL;
+	foreach(rangeTableCell, rangeTableList)
+	{
+		RangeTblEntry *rangeTableEntry = (RangeTblEntry *) lfirst(rangeTableCell);
+
+		if (rangeTableEntry->rtekind == RTE_RELATION &&
+			PartitionedTable(rangeTableEntry->relid) &&
+			rangeTableEntry->inh == false)
+		{
+			partitioningAdjusted = true;
+			break;
+		}
+	}
+
+	if (partitioningAdjusted && PartitionedTable(relationObjectId))
+	{
+		Index varno = rel->relid;
+		LOCKMODE lmode = root->simple_rte_array[varno]->rellockmode;
+		List *indexinfos = NIL;
+		IndexOptInfo *indexOptInfo = NULL;
+		foreach_ptr(indexOptInfo, rel->indexlist)
+		{
+			Relation indexRelation = index_open(indexOptInfo->indexoid, lmode);
+			if (indexRelation->rd_rel->relkind != RELKIND_PARTITIONED_INDEX)
+			{
+				indexinfos = lappend(indexinfos, indexOptInfo);
+			}
+			index_close(indexRelation, NoLock);
+		}
+		rel->indexlist = indexinfos;
+	}
+}
+
+
 /*
  * 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_partitioned_index_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_partitioned_index_hook(PlannerInfo *root, Oid relationObjectId, bool
+										 inhparent, RelOptInfo *rel);
 extern void multi_join_restriction_hook(PlannerInfo *root,
 										RelOptInfo *joinrel,
 										RelOptInfo *outerrel,