From a6de864816146400537b0dc06b5763e451ada216 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Wed, 27 Jul 2016 15:47:39 +0900
Subject: [PATCH 6/8] Introduce a PartitionTreeNode data structure.

It encapsulates the tree structure of a partition hierarchy which can be
arbitrarily deeply nested.  Every node in the tree represents a partitioned
table.  The only currently envisioned application is for tuple-routing.
---
 src/backend/catalog/partition.c |  213 +++++++++++++++++++++++++++++++++++++++
 src/include/catalog/partition.h |    5 +
 2 files changed, 218 insertions(+), 0 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index e21e7ad..7bc4b15 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -114,6 +114,64 @@ typedef struct PartitionListValue
 	int		index;
 } PartitionListValue;
 
+/*
+ * PartitionKeyExecInfo
+ *
+ *		This struct holds the information needed to extract partition
+ *		column values from a heap tuple.
+ *
+ *		Key					copy of the rd_partkey of rel
+ *		ExpressionState		exec state for expressions, or NIL if none
+ */
+typedef struct PartitionKeyExecInfo
+{
+	NodeTag			type;
+	PartitionKey	pi_Key;
+	List		   *pi_ExpressionState;	/* list of ExprState */
+} PartitionKeyExecInfo;
+
+/* ----------------
+ * Partition tree node (corresponding to one partitioned table in the
+ * partition tree)
+ *
+ *	pkinfo				PartitionKey executor state
+ *
+ *	pdesc				Info about immediate partitions (see
+ *						PartitionDescData)
+ *
+ *	relid				Table OID
+ *
+ *	index				If the table is partition itself, index in the
+ *						parent's partition array
+ *
+ *	num_leaf_parts		Number of leaf partitions in the partition
+ *						tree rooted at this node
+ *
+ *	offset				Index across the whole partition tree, of the first
+ *						leaf partition in the sub-tree rooted at this node
+ *
+ *	downlink			Link to the leftmost child node (corresponding to
+ *						the first of our partitions that are partitioned
+ *						themselves)
+ *
+ *	next				Link to the right sibling node (corresponding to the
+ *						next partition on the same level that is partitioned
+ *						itself)
+ * ----------------
+ */
+typedef struct PartitionTreeNodeData
+{
+	PartitionKeyExecInfo *pkinfo;
+	PartitionDesc		pdesc;
+	Oid					relid;
+	int					index;
+	int					offset;
+	int					num_leaf_parts;
+
+	struct PartitionTreeNodeData *downlink;
+	struct PartitionTreeNodeData *next;
+} PartitionTreeNodeData;
+
 static int32 qsort_partition_list_value_cmp(const void *a, const void *b, void *arg);
 static int32 qsort_partition_rbound_cmp(const void *a, const void *b, void *arg);
 
@@ -124,6 +182,9 @@ static Oid get_partition_operator(PartitionKey key, int col, StrategyNumber stra
 
 static List *generate_partition_qual(Relation rel, bool recurse);
 
+static PartitionTreeNode GetPartitionTreeNodeRecurse(Relation rel, int offset, int lockmode);
+static int get_leaf_partition_count(PartitionTreeNode parent);
+
 /* List partition related support functions */
 static bool equal_list_info(PartitionKey key,
 				PartitionListInfo *l1, PartitionListInfo *l2);
@@ -874,6 +935,55 @@ RelationGetPartitionQual(Relation rel, bool recurse)
 	return generate_partition_qual(rel, recurse);
 }
 
+/*
+ * RelationGetPartitionTreeNode
+ *		Recursively form the partition node tree with rel at root
+ *
+ * All the partitions will be locked with lockmode unless it is NoLock.
+ */
+PartitionTreeNode
+RelationGetPartitionTreeNode(Relation rel, int lockmode)
+{
+	PartitionTreeNode	root;
+
+	/*
+	 * We recurse to build the PartitionTreeNodes for any partitions in the
+	 * partition tree that are themselves partitioned tables.
+	 */
+	root = GetPartitionTreeNodeRecurse(rel, 0, lockmode);
+	root->index = 0;	/* Root table has no parent */
+	root->num_leaf_parts = get_leaf_partition_count(root);
+
+	return root;
+}
+
+/*
+ * get_leaf_partition_oids
+ * 		Recursively compute the list of OIDs of leaf partitions in the
+ *		partition tree rooted at parent
+ */
+List *
+get_leaf_partition_oids(PartitionTreeNode parent)
+{
+	int		i;
+	List   *result = NIL;
+	PartitionTreeNode node = parent->downlink;
+
+	for (i = 0; i < parent->pdesc->nparts; i++)
+	{
+		/* Indexes 0..(node->index - 1) are leaf partitions */
+		if (node && i == node->index)
+		{
+			result = list_concat(result, get_leaf_partition_oids(node));
+			node = node->next;
+		}
+		else
+			result = lappend_oid(result, parent->pdesc->oids[i]);
+	}
+
+	return result;
+}
+
 /* Module-local functions */
 
 /*
@@ -1314,6 +1424,109 @@ generate_partition_qual(Relation rel, bool recurse)
 	return result;
 }
 
+/*
+ * GetPartitionTreeNodeRecurse
+ *		Workhorse of RelationGetPartitionTreeNode
+ *
+ * offset is the index across the whole partition tree of the first leaf node
+ * in this subtree.
+ */
+static PartitionTreeNode
+GetPartitionTreeNodeRecurse(Relation rel, int offset, int lockmode)
+{
+	PartitionTreeNode	parent,
+						prev;
+	int					i;
+
+	/* Guard against stack overflow due to overly deep partition tree */
+	check_stack_depth();
+
+	/* First build our own node */
+	parent = (PartitionTreeNode) palloc0(sizeof(PartitionTreeNodeData));
+	parent->pkinfo = NULL;
+	parent->pdesc = RelationGetPartitionDesc(rel);
+	parent->relid = RelationGetRelid(rel);
+	parent->offset = offset;
+	parent->downlink = NULL;
+	parent->next = NULL;
+
+	/*
+	 * Go through rel's partitions and recursively add nodes for partitions
+	 * that are themselves partitioned.  Link parent to the first child node
+	 * using 'downlink'.  Each new child node is linked to its right sibling
+	 * using 'next'.  Offset value passed when creating a child node is
+	 * determined by looking at the left node if one exists or the parent
+	 * node if it is the first child node of this level.
+	 */
+	prev = NULL;
+	for (i = 0; i < parent->pdesc->nparts; i++)
+	{
+		Oid			relid = parent->pdesc->oids[i];
+		int			offset;
+		Relation	rel;
+		PartitionTreeNode child;
+
+		rel = heap_open(relid, lockmode);
+
+		/* Skip if a leaf partition */
+		if (rel->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
+		{
+			heap_close(rel, NoLock);
+			continue;
+		}
+
+		if (prev)
+			offset = prev->offset + prev->num_leaf_parts +
+												(i - prev->index - 1);
+		else
+			offset = parent->offset + i;
+
+		child = GetPartitionTreeNodeRecurse(rel, offset, lockmode);
+		child->index = i;
+		child->num_leaf_parts = get_leaf_partition_count(child);
+
+		heap_close(rel, NoLock);
+
+		/* Found first child; link to it. */
+		if (parent->downlink == NULL)
+			parent->downlink = child;
+
+		/* Link new node to the left sibling, if any  */
+		if (prev)
+			prev->next = child;
+		prev = child;
+	}
+
+	return parent;
+}
+
+/*
+ * get_leaf_partition_count
+ * 		Recursively count the number of leaf partitions in the partition
+ *		tree rooted at parent
+ */
+static int
+get_leaf_partition_count(PartitionTreeNode parent)
+{
+	int		i;
+	int 	result = 0;
+	PartitionTreeNode node = parent->downlink;
+
+	for (i = 0; i < parent->pdesc->nparts; i++)
+	{
+		/* Indexes 0..(node->index - 1) are of leaf partitions */
+		if (node && i == node->index)
+		{
+			result += get_leaf_partition_count(node);
+			node = node->next;
+		}
+		else
+			result += 1;
+	}
+
+	return result;
+}
+
 /* List partition related support functions */
 
 /*
diff --git a/src/include/catalog/partition.h b/src/include/catalog/partition.h
index 062de88..672183b 100644
--- a/src/include/catalog/partition.h
+++ b/src/include/catalog/partition.h
@@ -36,6 +36,7 @@ typedef struct PartitionDescData
 } PartitionDescData;
 
 typedef struct PartitionDescData *PartitionDesc;
+typedef struct PartitionTreeNodeData *PartitionTreeNode;
 
 extern void RelationBuildPartitionDesc(Relation relation);
 extern bool partition_bounds_equal(PartitionKey key,
@@ -45,4 +46,8 @@ extern void check_new_partition_bound(char *relname, Oid parentId, Node *bound);
 extern Oid get_partition_parent(Oid relid);
 extern List *get_qual_from_partbound(Relation rel, Relation parent, Node *bound);
 extern List *RelationGetPartitionQual(Relation rel, bool recurse);
+
+/* For tuple routing */
+extern PartitionTreeNode RelationGetPartitionTreeNode(Relation rel, int lockmode);
+extern List *get_leaf_partition_oids(PartitionTreeNode parent);
 #endif   /* PARTITION_H */
-- 
1.7.1

