Refactor partition tuple routing code to reduce duplication.

Amit Langote

src/backend/commands/copy.c

@@ -1406,64 +1406,22 @@ BeginCopy(ParseState *pstate,
 		/* Initialize state for CopyFrom tuple routing. */
 		if (is_from && rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 		{
-			List	   *leaf_parts;
-			ListCell   *cell;
-			int			i,
-						num_parted;
-			ResultRelInfo *leaf_part_rri;
-
-			/* Get the tuple-routing information and lock partitions */
-			cstate->partition_dispatch_info =
-				RelationGetPartitionDispatchInfo(rel, RowExclusiveLock,
-												 &num_parted,
-												 &leaf_parts);
+			PartitionDispatch  *partition_dispatch_info;
+			ResultRelInfo	   *partitions;
+			TupleConversionMap **partition_tupconv_maps;
+			int					num_parted,
+								num_partitions;
+
+			ExecSetupPartitionTupleRouting(rel,
+										   &partition_dispatch_info,
+										   &partitions,
+										   &partition_tupconv_maps,
+										   &num_parted, &num_partitions);
+			cstate->partition_dispatch_info = partition_dispatch_info;
 			cstate->num_dispatch = num_parted;
-			cstate->num_partitions = list_length(leaf_parts);
-			cstate->partitions = (ResultRelInfo *)
-				palloc(cstate->num_partitions *
-					   sizeof(ResultRelInfo));
-			cstate->partition_tupconv_maps = (TupleConversionMap **)
-				palloc0(cstate->num_partitions *
-						sizeof(TupleConversionMap *));
-
-			leaf_part_rri = cstate->partitions;
-			i = 0;
-			foreach(cell, leaf_parts)
-			{
-				Relation	partrel;
-
-				/*
-				 * We locked all the partitions above including the leaf
-				 * partitions.  Note that each of the relations in
-				 * cstate->partitions will be closed by CopyFrom() after it's
-				 * finished with its processing.
-				 */
-				partrel = heap_open(lfirst_oid(cell), NoLock);
-
-				/*
-				 * Verify result relation is a valid target for the current
-				 * operation.
-				 */
-				CheckValidResultRel(partrel, CMD_INSERT);
-
-				InitResultRelInfo(leaf_part_rri,
-								  partrel,
-								  1,	/* dummy */
-								  false,		/* no partition constraint
-												 * check */
-								  0);
-
-				/* Open partition indices */
-				ExecOpenIndices(leaf_part_rri, false);
-
-				if (!equalTupleDescs(tupDesc, RelationGetDescr(partrel)))
-					cstate->partition_tupconv_maps[i] =
-						convert_tuples_by_name(tupDesc,
-											   RelationGetDescr(partrel),
-								 gettext_noop("could not convert row type"));
-				leaf_part_rri++;
-				i++;
-			}
+			cstate->partitions = partitions;
+			cstate->num_partitions = num_partitions;
+			cstate->partition_tupconv_maps = partition_tupconv_maps;
 		}
 	}
 	else

src/backend/executor/execMain.c

@@ -51,6 +51,7 @@
 #include "miscadmin.h"
 #include "optimizer/clauses.h"
 #include "parser/parsetree.h"
+#include "rewrite/rewriteManip.h"
 #include "storage/bufmgr.h"
 #include "storage/lmgr.h"
 #include "tcop/utility.h"
@@ -2998,6 +2999,97 @@ EvalPlanQualEnd(EPQState *epqstate)
 	epqstate->origslot = NULL;
 }
 
+/*
+ * ExecSetupPartitionTupleRouting - set up information needed during
+ * tuple routing for partitioned tables
+ *
+ * Output arguments:
+ * 'pd' receives an array of PartitionDispatch objects with one entry for
+ *		every partitioned table in the partition tree
+ * 'partitions' receives an array of ResultRelInfo objects with one entry for
+ *		every leaf partition in the partition tree
+ * 'tup_conv_maps' receives an array of TupleConversionMap objects with one
+ *		entry for every leaf partition (required to convert input tuple based
+ *		on the root table's rowtype to a leaf partition's rowtype after tuple
+ *		routing is done
+ * 'num_parted' receives the number of partitioned tables in the partition
+ *		tree (= the number of entries in the 'pd' output array)
+ * 'num_partitions' receives the number of leaf partitions in the partition
+ *		tree (= the number of entries in the 'partitions' and 'tup_conv_maps'
+ *		output arrays
+ *
+ * Note that all the relations in the partition tree are locked using the
+ * RowExclusiveLock mode upon return from this function.
+ */
+void
+ExecSetupPartitionTupleRouting(Relation rel,
+							   PartitionDispatch **pd,
+							   ResultRelInfo **partitions,
+							   TupleConversionMap ***tup_conv_maps,
+							   int *num_parted, int *num_partitions)
+{
+	TupleDesc	tupDesc = RelationGetDescr(rel);
+	List	   *leaf_parts;
+	ListCell   *cell;
+	int			i;
+	ResultRelInfo *leaf_part_rri;
+
+	/* Get the tuple-routing information and lock partitions */
+	*pd = RelationGetPartitionDispatchInfo(rel, RowExclusiveLock, num_parted,
+										   &leaf_parts);
+	*num_partitions = list_length(leaf_parts);
+	*partitions = (ResultRelInfo *) palloc(*num_partitions *
+										   sizeof(ResultRelInfo));
+	*tup_conv_maps = (TupleConversionMap **) palloc0(*num_partitions *
+										   sizeof(TupleConversionMap *));
+
+	leaf_part_rri = *partitions;
+	i = 0;
+	foreach(cell, leaf_parts)
+	{
+		Relation	partrel;
+		TupleDesc	part_tupdesc;
+
+		/*
+		 * We locked all the partitions above including the leaf partitions.
+		 * Note that each of the relations in *partitions are eventually
+		 * closed by the caller.
+		 */
+		partrel = heap_open(lfirst_oid(cell), NoLock);
+		part_tupdesc = RelationGetDescr(partrel);
+
+		/*
+		 * Verify result relation is a valid target for the current operation.
+		 */
+		CheckValidResultRel(partrel, CMD_INSERT);
+
+		/*
+		 * Save a tuple conversion map to convert a tuple routed to this
+		 * partition from the parent's type to the partition's.
+		 */
+		(*tup_conv_maps)[i] = convert_tuples_by_name(tupDesc, part_tupdesc,
+								 gettext_noop("could not convert row type"));
+
+		InitResultRelInfo(leaf_part_rri,
+						  partrel,
+						  1,	 /* dummy */
+						  false,
+						  0);
+
+		/*
+		 * Open partition indices (remember we do not support ON CONFLICT in
+		 * case of partitioned tables, so we do not need support information
+		 * for speculative insertion)
+		 */
+		if (leaf_part_rri->ri_RelationDesc->rd_rel->relhasindex &&
+			leaf_part_rri->ri_IndexRelationDescs == NULL)
+			ExecOpenIndices(leaf_part_rri, false);
+
+		leaf_part_rri++;
+		i++;
+	}
+}
+
 /*
  * ExecFindPartition -- Find a leaf partition in the partition tree rooted
  * at parent, for the heap tuple contained in *slot

src/backend/executor/nodeModifyTable.c

@@ -1718,68 +1718,22 @@ ExecInitModifyTable(ModifyTable *node, EState *estate, int eflags)
 	if (operation == CMD_INSERT &&
 		rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 	{
-		int					i,
-							j,
-							num_parted;
-		List			   *leaf_parts;
-		ListCell		   *cell;
-		ResultRelInfo	   *leaf_part_rri;
-
-		/* Get the tuple-routing information and lock partitions */
-		mtstate->mt_partition_dispatch_info =
-					RelationGetPartitionDispatchInfo(rel, RowExclusiveLock,
-													 &num_parted,
-													 &leaf_parts);
+		PartitionDispatch  *partition_dispatch_info;
+		ResultRelInfo	   *partitions;
+		TupleConversionMap **partition_tupconv_maps;
+		int					num_parted,
+							num_partitions;
+
+		ExecSetupPartitionTupleRouting(rel,
+									   &partition_dispatch_info,
+									   &partitions,
+									   &partition_tupconv_maps,
+									   &num_parted, &num_partitions);
+		mtstate->mt_partition_dispatch_info = partition_dispatch_info;
 		mtstate->mt_num_dispatch = num_parted;
-		mtstate->mt_num_partitions = list_length(leaf_parts);
-		mtstate->mt_partitions = (ResultRelInfo *)
-						palloc0(mtstate->mt_num_partitions *
-												sizeof(ResultRelInfo));
-		mtstate->mt_partition_tupconv_maps = (TupleConversionMap **)
-						palloc0(mtstate->mt_num_partitions *
-												sizeof(TupleConversionMap *));
-
-		leaf_part_rri = mtstate->mt_partitions;
-		i = j = 0;
-		foreach(cell, leaf_parts)
-		{
-			Oid			partrelid = lfirst_oid(cell);
-			Relation	partrel;
-
-			/*
-			 * We locked all the partitions above including the leaf
-			 * partitions.  Note that each of the relations in
-			 * mtstate->mt_partitions will be closed by ExecEndModifyTable().
-			 */
-			partrel = heap_open(partrelid, NoLock);
-
-			/*
-			 * Verify result relation is a valid target for the current
-			 * operation
-			 */
-			CheckValidResultRel(partrel, CMD_INSERT);
-
-			InitResultRelInfo(leaf_part_rri,
-							  partrel,
-							  1,		/* dummy */
-							  false,	/* no partition constraint checks */
-							  eflags);
-
-			/* Open partition indices (note: ON CONFLICT unsupported)*/
-			if (partrel->rd_rel->relhasindex && operation != CMD_DELETE &&
-				leaf_part_rri->ri_IndexRelationDescs == NULL)
-				ExecOpenIndices(leaf_part_rri, false);
-
-			if (!equalTupleDescs(RelationGetDescr(rel),
-								 RelationGetDescr(partrel)))
-				mtstate->mt_partition_tupconv_maps[i] =
-							convert_tuples_by_name(RelationGetDescr(rel),
-												   RelationGetDescr(partrel),
-								  gettext_noop("could not convert row type"));
-
-			leaf_part_rri++;
-			i++;
-		}
+		mtstate->mt_partitions = partitions;
+		mtstate->mt_num_partitions = num_partitions;
+		mtstate->mt_partition_tupconv_maps = partition_tupconv_maps;
 	}
 
 	/*

src/include/executor/executor.h

@@ -213,6 +213,11 @@ extern void EvalPlanQualSetPlan(EPQState *epqstate,
 extern void EvalPlanQualSetTuple(EPQState *epqstate, Index rti,
 					 HeapTuple tuple);
 extern HeapTuple EvalPlanQualGetTuple(EPQState *epqstate, Index rti);
+extern void ExecSetupPartitionTupleRouting(Relation rel,
+							   PartitionDispatch **pd,
+							   ResultRelInfo **partitions,
+							   TupleConversionMap ***tup_conv_maps,
+							   int *num_parted, int *num_partitions);
 extern int ExecFindPartition(ResultRelInfo *resultRelInfo,
 				  PartitionDispatch *pd,
 				  TupleTableSlot *slot,