From 17652ea80aad319bd505519e0ab7768e3652f3ab Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Wed, 27 Jul 2016 15:47:39 +0900
Subject: [PATCH 7/9] 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 |  206 +++++++++++++++++++++++++++++++++++++++
 src/include/catalog/partition.h |    5 +
 2 files changed, 211 insertions(+), 0 deletions(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 05e63c3..402eb9e 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -131,6 +131,61 @@ typedef struct RangePartition
 	PartitionRange *range;
 } RangePartition;
 
+/*
+ * 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)
+ *
+ *	index				If a partition ourselves, index in the parent's
+ *						partition array
+ *
+ *	num_leaf_parts		Number of leaf partitions in the partition
+ *						tree rooted at this node
+ *
+ *	offset				0-based index of the first leaf partition
+ *						in the partition tree rooted at this node
+ *
+ *	downlink			Link to our leftmost child node (ie, corresponding
+ *						to first of our partitions that is itself
+ *						partitioned)
+ *
+ *	next				Link to the right sibling node on a given level
+ *						(ie, corresponding to the next partition on the same
+ *						level that is itself partitioned)
+ */
+typedef struct PartitionTreeNodeData
+{
+	PartitionKeyExecInfo *pkinfo;
+	PartitionDesc		pdesc;
+	Oid					relid;
+	int					index;
+	int					offset;
+	int					num_leaf_parts;
+
+	struct PartitionTreeNodeData *downlink;
+	struct PartitionTreeNodeData *next;
+} PartitionTreeNodeData;
+
 /* Support RelationBuildPartitionDesc() */
 static int32 list_value_cmp(const void *a, const void *b, void *arg);
 static int32 range_partition_cmp(const void *a, const void *b, void *arg);
@@ -167,6 +222,10 @@ static Oid get_partition_operator(PartitionKey key, int col, StrategyNumber stra
 /* Support RelationGetPartitionQual() */
 static List *generate_partition_qual(Relation rel, bool recurse);
 
+/* Support RelationGetPartitionTreeNode() */
+static PartitionTreeNode GetPartitionTreeNodeRecurse(Relation rel, int offset);
+static int get_leaf_partition_count(PartitionTreeNode ptnode);
+
 /* List partition related support functions */
 static PartitionList *make_list_from_spec(PartitionKey key,
 							PartitionBoundList *list_spec);
@@ -774,6 +833,53 @@ RelationGetPartitionQual(Relation rel, bool recurse)
 	return generate_partition_qual(rel, recurse);
 }
 
+/*
+ * RelationGetPartitionTreeNode
+ *		Recursively form partition tree rooted at this rel's node
+ */
+PartitionTreeNode
+RelationGetPartitionTreeNode(Relation rel)
+{
+	PartitionTreeNode	root;
+
+	/*
+	 * We recurse to build the PartitionTreeNodes for any partitions in the
+	 * partition hierarchy that are themselves partitioned.
+	 */
+	root = GetPartitionTreeNodeRecurse(rel, 0);
+	root->index = 0;	/* Root table has no parent */
+	root->num_leaf_parts = get_leaf_partition_count(root);
+
+	return root;
+}
+
+/*
+ * get_leaf_partition_oids_v2
+ * 		Recursively compute the list of OIDs of leaf partitions in the
+ *		partition tree rooted at ptnode
+ */
+List *
+get_leaf_partition_oids_v2(PartitionTreeNode ptnode)
+{
+	int		i;
+	List   *result = NIL;
+	PartitionTreeNode node = ptnode->downlink;
+
+	for (i = 0; i < ptnode->pdesc->nparts; i++)
+	{
+		/* Indexes 0..(node->index - 1) are leaf partitions */
+		if (node && i == node->index)
+		{
+			result = list_concat(result, get_leaf_partition_oids_v2(node));
+			node = node->next;
+		}
+		else
+			result = lappend_oid(result, ptnode->pdesc->oids[i]);
+	}
+
+	return result;
+}
+
 /* Module-local functions */
 
 /*
@@ -1423,6 +1529,106 @@ generate_partition_qual(Relation rel, bool recurse)
 	return result;
 }
 
+/*
+ * GetPartitionTreeNodeRecurse
+ *		Workhorse of RelationGetPartitionTreeNode
+ *
+ * 'offset' is 0-based index of the first leaf node in this subtree. During
+ * the first invocation, a 0 will be pass
+ */
+static PartitionTreeNode
+GetPartitionTreeNodeRecurse(Relation rel, int offset)
+{
+	PartitionTreeNode	parent,
+						prev;
+	int					i;
+
+	/* 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, AccessShareLock);
+
+		/* Skip if a leaf partition */
+		if (rel->rd_rel->relkind != RELKIND_PARTITIONED_TABLE)
+		{
+			heap_close(rel, AccessShareLock);
+			continue;
+		}
+
+		if (prev)
+			offset = prev->offset + prev->num_leaf_parts +
+												(i - prev->index - 1);
+		else
+			offset = parent->offset + i;
+
+		child = GetPartitionTreeNodeRecurse(rel, offset);
+		child->index = i;
+		child->num_leaf_parts = get_leaf_partition_count(child);
+
+		heap_close(rel, AccessShareLock);
+
+		/* Found our 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 ptnode
+ */
+static int
+get_leaf_partition_count(PartitionTreeNode ptnode)
+{
+	int		i;
+	int 	result = 0;
+	PartitionTreeNode node = ptnode->downlink;
+
+	for (i = 0; i < ptnode->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 ea7806e..1ecd5d6 100644
--- a/src/include/catalog/partition.h
+++ b/src/include/catalog/partition.h
@@ -43,6 +43,7 @@ typedef struct PartitionDescData
 } PartitionDescData;
 
 typedef struct PartitionDescData *PartitionDesc;
+typedef struct PartitionTreeNodeData *PartitionTreeNode;
 
 /* relcache support functions for partition descriptor */
 extern void RelationBuildPartitionDesc(Relation relation);
@@ -56,4 +57,8 @@ extern List *get_partition_ancestors(Oid relid);
 extern List *get_leaf_partition_oids(Oid relid, int lockmode);
 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);
+extern List *get_leaf_partition_oids_v2(PartitionTreeNode ptnode);
 #endif   /* PARTITION_H */
-- 
1.7.1

