Simplify and encapsulate tuple routing support code.

Instead of having ExecSetupPartitionTupleRouting return multiple out
parameters, have it return a pointer to a structure containing all of
those different things.  Also, provide and use a cleanup function,
ExecCleanupTupleRouting, instead of cleaning up all of the resources
allocated by ExecSetupPartitionTupleRouting individually.

Amit Khandekar, reviewed by Amit Langote, David Rowley, and me

Discussion: http://postgr.es/m/CAJ3gD9fWfxgKC+PfJZF3hkgAcNOy-LpfPxVYitDEXKHjeieWQQ@mail.gmail.com

Author: robertmhaas

src/backend/commands/copy.c

@@ -166,12 +166,9 @@ typedef struct CopyStateData
 	bool		volatile_defexprs;	/* is any of defexprs volatile? */
 	List	   *range_table;
 
-	PartitionDispatch *partition_dispatch_info;
-	int			num_dispatch;	/* Number of entries in the above array */
-	int			num_partitions; /* Number of members in the following arrays */
-	ResultRelInfo **partitions; /* Per partition result relation pointers */
-	TupleConversionMap **partition_tupconv_maps;
-	TupleTableSlot *partition_tuple_slot;
+	/* Tuple-routing support info */
+	PartitionTupleRouting *partition_tuple_routing;
+
 	TransitionCaptureState *transition_capture;
 	TupleConversionMap **transition_tupconv_maps;
 
@@ -2472,28 +2469,10 @@ CopyFrom(CopyState cstate)
 	 */
 	if (cstate->rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
 	{
-		PartitionDispatch *partition_dispatch_info;
-		ResultRelInfo **partitions;
-		TupleConversionMap **partition_tupconv_maps;
-		TupleTableSlot *partition_tuple_slot;
-		int			num_parted,
-					num_partitions;
-
-		ExecSetupPartitionTupleRouting(NULL,
-									   cstate->rel,
-									   1,
-									   estate,
-									   &partition_dispatch_info,
-									   &partitions,
-									   &partition_tupconv_maps,
-									   &partition_tuple_slot,
-									   &num_parted, &num_partitions);
-		cstate->partition_dispatch_info = partition_dispatch_info;
-		cstate->num_dispatch = num_parted;
-		cstate->partitions = partitions;
-		cstate->num_partitions = num_partitions;
-		cstate->partition_tupconv_maps = partition_tupconv_maps;
-		cstate->partition_tuple_slot = partition_tuple_slot;
+		PartitionTupleRouting *proute;
+
+		proute = cstate->partition_tuple_routing =
+			ExecSetupPartitionTupleRouting(NULL, cstate->rel, 1, estate);
 
 		/*
 		 * If we are capturing transition tuples, they may need to be
@@ -2506,11 +2485,11 @@ CopyFrom(CopyState cstate)
 			int			i;
 
 			cstate->transition_tupconv_maps = (TupleConversionMap **)
-				palloc0(sizeof(TupleConversionMap *) * cstate->num_partitions);
-			for (i = 0; i < cstate->num_partitions; ++i)
+				palloc0(sizeof(TupleConversionMap *) * proute->num_partitions);
+			for (i = 0; i < proute->num_partitions; ++i)
 			{
 				cstate->transition_tupconv_maps[i] =
-					convert_tuples_by_name(RelationGetDescr(cstate->partitions[i]->ri_RelationDesc),
+					convert_tuples_by_name(RelationGetDescr(proute->partitions[i]->ri_RelationDesc),
 										   RelationGetDescr(cstate->rel),
 										   gettext_noop("could not convert row type"));
 			}
@@ -2530,7 +2509,7 @@ CopyFrom(CopyState cstate)
 	if ((resultRelInfo->ri_TrigDesc != NULL &&
 		 (resultRelInfo->ri_TrigDesc->trig_insert_before_row ||
 		  resultRelInfo->ri_TrigDesc->trig_insert_instead_row)) ||
-		cstate->partition_dispatch_info != NULL ||
+		cstate->partition_tuple_routing != NULL ||
 		cstate->volatile_defexprs)
 	{
 		useHeapMultiInsert = false;
@@ -2605,10 +2584,11 @@ CopyFrom(CopyState cstate)
 		ExecStoreTuple(tuple, slot, InvalidBuffer, false);
 
 		/* Determine the partition to heap_insert the tuple into */
-		if (cstate->partition_dispatch_info)
+		if (cstate->partition_tuple_routing)
 		{
 			int			leaf_part_index;
 			TupleConversionMap *map;
+			PartitionTupleRouting *proute = cstate->partition_tuple_routing;
 
 			/*
 			 * Away we go ... If we end up not finding a partition after all,
@@ -2619,11 +2599,11 @@ CopyFrom(CopyState cstate)
 			 * partition, respectively.
 			 */
 			leaf_part_index = ExecFindPartition(resultRelInfo,
-												cstate->partition_dispatch_info,
+												proute->partition_dispatch_info,
 												slot,
 												estate);
 			Assert(leaf_part_index >= 0 &&
-				   leaf_part_index < cstate->num_partitions);
+				   leaf_part_index < proute->num_partitions);
 
 			/*
 			 * If this tuple is mapped to a partition that is not same as the
@@ -2641,7 +2621,7 @@ CopyFrom(CopyState cstate)
 			 * to the selected partition.
 			 */
 			saved_resultRelInfo = resultRelInfo;
-			resultRelInfo = cstate->partitions[leaf_part_index];
+			resultRelInfo = proute->partitions[leaf_part_index];
 
 			/* We do not yet have a way to insert into a foreign partition */
 			if (resultRelInfo->ri_FdwRoutine)
@@ -2688,7 +2668,7 @@ CopyFrom(CopyState cstate)
 			 * We might need to convert from the parent rowtype to the
 			 * partition rowtype.
 			 */
-			map = cstate->partition_tupconv_maps[leaf_part_index];
+			map = proute->partition_tupconv_maps[leaf_part_index];
 			if (map)
 			{
 				Relation	partrel = resultRelInfo->ri_RelationDesc;
@@ -2700,7 +2680,7 @@ CopyFrom(CopyState cstate)
 				 * point on.  Use a dedicated slot from this point on until
 				 * we're finished dealing with the partition.
 				 */
-				slot = cstate->partition_tuple_slot;
+				slot = proute->partition_tuple_slot;
 				Assert(slot != NULL);
 				ExecSetSlotDescriptor(slot, RelationGetDescr(partrel));
 				ExecStoreTuple(tuple, slot, InvalidBuffer, true);
@@ -2852,34 +2832,8 @@ CopyFrom(CopyState cstate)
 	ExecCloseIndices(resultRelInfo);
 
 	/* Close all the partitioned tables, leaf partitions, and their indices */
-	if (cstate->partition_dispatch_info)
-	{
-		int			i;
-
-		/*
-		 * Remember cstate->partition_dispatch_info[0] corresponds to the root
-		 * partitioned table, which we must not try to close, because it is
-		 * the main target table of COPY that will be closed eventually by
-		 * DoCopy().  Also, tupslot is NULL for the root partitioned table.
-		 */
-		for (i = 1; i < cstate->num_dispatch; i++)
-		{
-			PartitionDispatch pd = cstate->partition_dispatch_info[i];
-
-			heap_close(pd->reldesc, NoLock);
-			ExecDropSingleTupleTableSlot(pd->tupslot);
-		}
-		for (i = 0; i < cstate->num_partitions; i++)
-		{
-			ResultRelInfo *resultRelInfo = cstate->partitions[i];
-
-			ExecCloseIndices(resultRelInfo);
-			heap_close(resultRelInfo->ri_RelationDesc, NoLock);
-		}
-
-		/* Release the standalone partition tuple descriptor */
-		ExecDropSingleTupleTableSlot(cstate->partition_tuple_slot);
-	}
+	if (cstate->partition_tuple_routing)
+		ExecCleanupTupleRouting(cstate->partition_tuple_routing);
 
 	/* Close any trigger target relations */
 	ExecCleanUpTriggerState(estate);

src/backend/executor/execPartition.c

@@ -38,68 +38,50 @@ static char *ExecBuildSlotPartitionKeyDescription(Relation rel,
 									 int maxfieldlen);
 
 /*
- * 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)
- * 'partition_tuple_slot' receives a standalone TupleTableSlot to be used
- *		to manipulate any given leaf partition's rowtype after that partition
- *		is chosen by tuple-routing.
- * '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
+ * ExecSetupPartitionTupleRouting - sets up information needed during
+ * tuple routing for partitioned tables, encapsulates it in
+ * PartitionTupleRouting, and returns it.
  *
  * Note that all the relations in the partition tree are locked using the
  * RowExclusiveLock mode upon return from this function.
  */
-void
+PartitionTupleRouting *
 ExecSetupPartitionTupleRouting(ModifyTableState *mtstate,
-							   Relation rel,
-							   Index resultRTindex,
-							   EState *estate,
-							   PartitionDispatch **pd,
-							   ResultRelInfo ***partitions,
-							   TupleConversionMap ***tup_conv_maps,
-							   TupleTableSlot **partition_tuple_slot,
-							   int *num_parted, int *num_partitions)
+							   Relation rel, Index resultRTindex,
+							   EState *estate)
 {
 	TupleDesc	tupDesc = RelationGetDescr(rel);
 	List	   *leaf_parts;
 	ListCell   *cell;
 	int			i;
 	ResultRelInfo *leaf_part_rri;
+	PartitionTupleRouting *proute;
 
 	/*
 	 * Get the information about the partition tree after locking all the
 	 * partitions.
 	 */
 	(void) find_all_inheritors(RelationGetRelid(rel), RowExclusiveLock, NULL);
-	*pd = RelationGetPartitionDispatchInfo(rel, 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 *));
+	proute = (PartitionTupleRouting *) palloc0(sizeof(PartitionTupleRouting));
+	proute->partition_dispatch_info =
+		RelationGetPartitionDispatchInfo(rel, &proute->num_dispatch,
+										 &leaf_parts);
+	proute->num_partitions = list_length(leaf_parts);
+	proute->partitions = (ResultRelInfo **) palloc(proute->num_partitions *
+												   sizeof(ResultRelInfo *));
+	proute->partition_tupconv_maps =
+		(TupleConversionMap **) palloc0(proute->num_partitions *
+										sizeof(TupleConversionMap *));
 
 	/*
 	 * Initialize an empty slot that will be used to manipulate tuples of any
 	 * given partition's rowtype.  It is attached to the caller-specified node
 	 * (such as ModifyTableState) and released when the node finishes
 	 * processing.
 	 */
-	*partition_tuple_slot = MakeTupleTableSlot();
+	proute->partition_tuple_slot = MakeTupleTableSlot();
 
-	leaf_part_rri = (ResultRelInfo *) palloc0(*num_partitions *
+	leaf_part_rri = (ResultRelInfo *) palloc0(proute->num_partitions *
 											  sizeof(ResultRelInfo));
 	i = 0;
 	foreach(cell, leaf_parts)
@@ -109,8 +91,8 @@ ExecSetupPartitionTupleRouting(ModifyTableState *mtstate,
 
 		/*
 		 * We locked all the partitions above including the leaf partitions.
-		 * Note that each of the relations in *partitions are eventually
-		 * closed by the caller.
+		 * Note that each of the relations in proute->partitions are
+		 * eventually closed by the caller.
 		 */
 		partrel = heap_open(lfirst_oid(cell), NoLock);
 		part_tupdesc = RelationGetDescr(partrel);
@@ -119,8 +101,9 @@ ExecSetupPartitionTupleRouting(ModifyTableState *mtstate,
 		 * 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"));
+		proute->partition_tupconv_maps[i] =
+			convert_tuples_by_name(tupDesc, part_tupdesc,
+								   gettext_noop("could not convert row type"));
 
 		InitResultRelInfo(leaf_part_rri,
 						  partrel,
@@ -149,9 +132,11 @@ ExecSetupPartitionTupleRouting(ModifyTableState *mtstate,
 		estate->es_leaf_result_relations =
 			lappend(estate->es_leaf_result_relations, leaf_part_rri);
 
-		(*partitions)[i] = leaf_part_rri++;
+		proute->partitions[i] = leaf_part_rri++;
 		i++;
 	}
+
+	return proute;
 }
 
 /*
@@ -272,6 +257,45 @@ ExecFindPartition(ResultRelInfo *resultRelInfo, PartitionDispatch *pd,
 	return result;
 }
 
+/*
+ * ExecCleanupTupleRouting -- Clean up objects allocated for partition tuple
+ * routing.
+ *
+ * Close all the partitioned tables, leaf partitions, and their indices.
+ */
+void
+ExecCleanupTupleRouting(PartitionTupleRouting * proute)
+{
+	int			i;
+
+	/*
+	 * Remember, proute->partition_dispatch_info[0] corresponds to the root
+	 * partitioned table, which we must not try to close, because it is the
+	 * main target table of the query that will be closed by callers such as
+	 * ExecEndPlan() or DoCopy(). Also, tupslot is NULL for the root
+	 * partitioned table.
+	 */
+	for (i = 1; i < proute->num_dispatch; i++)
+	{
+		PartitionDispatch pd = proute->partition_dispatch_info[i];
+
+		heap_close(pd->reldesc, NoLock);
+		ExecDropSingleTupleTableSlot(pd->tupslot);
+	}
+
+	for (i = 0; i < proute->num_partitions; i++)
+	{
+		ResultRelInfo *resultRelInfo = proute->partitions[i];
+
+		ExecCloseIndices(resultRelInfo);
+		heap_close(resultRelInfo->ri_RelationDesc, NoLock);
+	}
+
+	/* Release the standalone partition tuple descriptor, if any */
+	if (proute->partition_tuple_slot)
+		ExecDropSingleTupleTableSlot(proute->partition_tuple_slot);
+}
+
 /*
  * RelationGetPartitionDispatchInfo
  *		Returns information necessary to route tuples down a partition tree