Map tuple to partition before INSERT / UPDATE.

This is needed because partition can have different order of columns if was
attached (as opposed to CREATE TABLE ... PARTITION OF ...).

Author: Antonin Houska

concurrent.c

@@ -278,12 +278,27 @@ apply_concurrent_changes(EState *estate, ModifyTableState *mtstate,
 			ExecStoreHeapTuple(tup, slot, false);
 			if (proute)
 			{
+				PartitionEntry	*entry;
+
 				rri = ExecFindPartition(mtstate, mtstate->rootResultRelInfo,
 										proute, slot, estate);
 				rel_ins = rri->ri_RelationDesc;
 
-				bistate = get_partition_insert_state(partitions,
-													 RelationGetRelid(rel_ins));
+				entry = get_partition_entry(partitions,
+											RelationGetRelid(rel_ins));
+				bistate = entry->bistate;
+
+				/*
+				 * Make sure the tuple matches the partition. The typical
+				 * problem we address here is that a partition was attached
+				 * that has a different order of columns.
+				 */
+				if (entry->conv_map)
+				{
+					tup = convert_tuple_for_dest_table(tup, entry->conv_map);
+					ExecClearTuple(slot);
+					ExecStoreHeapTuple(tup, slot, false);
+				}
 			}
 			else
 			{
@@ -361,8 +376,12 @@ apply_concurrent_changes(EState *estate, ModifyTableState *mtstate,
 			{
 				ItemPointerData ctid;
 				List	   *recheck;
+				PartitionEntry *entry;
 
-				/* Delete the old tuple from its partition. */
+				/*
+				 * Cross-partition update. Delete the old tuple from its
+				 * partition.
+				 */
 				find_tuple_in_partition(tup_old, rri_old->ri_RelationDesc,
 										partitions, key, nkeys, &ctid);
 				simple_heap_delete(rri_old->ri_RelationDesc, &ctid);
@@ -372,7 +391,14 @@ apply_concurrent_changes(EState *estate, ModifyTableState *mtstate,
 				MyWorkerTask->progress.del++;
 				SpinLockRelease(&MyWorkerTask->mutex);
 
-				/* Insert the new tuple into its partition. */
+				/*
+				 * Insert the new tuple into its partition. This might include
+				 * conversion to match the partition, see above.
+				 */
+				entry = get_partition_entry(partitions,
+											RelationGetRelid(rri->ri_RelationDesc));
+				if (entry->conv_map)
+					tup = convert_tuple_for_dest_table(tup, entry->conv_map);
 				ExecStoreHeapTuple(tup, slot, false);
 				table_tuple_insert(rri->ri_RelationDesc, slot,
 								   GetCurrentCommandId(true), 0, NULL);
@@ -438,6 +464,20 @@ apply_concurrent_changes(EState *estate, ModifyTableState *mtstate,
 					TU_UpdateIndexes	update_indexes;
 #endif
 
+					if (partitions)
+					{
+						PartitionEntry *entry;
+
+						/*
+						 * Make sure the tuple matches the partition.
+						 */
+						entry = get_partition_entry(partitions,
+													RelationGetRelid(rri->ri_RelationDesc));
+						if (entry->conv_map)
+							tup = convert_tuple_for_dest_table(tup,
+															   entry->conv_map);
+					}
+
 					simple_heap_update(rri->ri_RelationDesc, &ctid, tup
 #if PG_VERSION_NUM >= 160000
 									   , &update_indexes

pg_rewrite.c

@@ -1017,6 +1017,9 @@ rewrite_table_impl(char *relschema_src, char *relname_src,
 
 			entry->ident_index = get_identity_index(partition, rel_src);
 			entry->ind_slot = table_slot_create(partition, NULL);
+			entry->conv_map =
+				convert_tuples_by_name_ext(RelationGetDescr(rel_dst),
+										   RelationGetDescr(partition));
 			/* Expect many insertions. */
 			entry->bistate = GetBulkInsertState();
 
@@ -1284,11 +1287,11 @@ rewrite_table_impl(char *relschema_src, char *relname_src,
 
 	if (part_desc)
 	{
+		close_partitions(partitions);
+
 		for (i = 0; i < part_desc->nparts; i++)
 			table_close(parts_dst[i], AccessExclusiveLock);
-
 		pfree(parts_dst);
-		close_partitions(partitions);
 	}
 	else
 	{
@@ -2685,20 +2688,39 @@ perform_initial_load(EState *estate, ModifyTableState *mtstate,
 			if (tup_out == NULL)
 				break;
 
-			/* Convert the tuple if needed. */
+			/*
+			 * If needed, convert the tuple so it matches the destination
+			 * table.
+			 */
 			if (conv_map)
 				tup_out = convert_tuple_for_dest_table(tup_out, conv_map);
 			ExecStoreHeapTuple(tup_out, slot_dst, false);
 
 			if (proute)
 			{
+				PartitionEntry *entry;
+
 				/* Find out which partition the tuple belongs to. */
 				rri = ExecFindPartition(mtstate, mtstate->rootResultRelInfo,
 										proute, slot_dst, estate);
 
 				rel_ins = rri->ri_RelationDesc;
-				bistate = get_partition_insert_state(partitions,
-													 RelationGetRelid(rri->ri_RelationDesc));
+				entry = get_partition_entry(partitions,
+											RelationGetRelid(rri->ri_RelationDesc));
+				bistate = entry->bistate;
+
+				/*
+				 * Make sure the tuple matches the partition. The typical
+				 * problem we address here is that a partition was attached
+				 * that has a different order of columns.
+				 */
+				if (entry->conv_map)
+				{
+					tup_out = convert_tuple_for_dest_table(tup_out,
+														   entry->conv_map);
+					ExecClearTuple(slot_dst);
+					ExecStoreHeapTuple(tup_out, slot_dst, false);
+				}
 			}
 			else
 			{
@@ -3061,17 +3083,19 @@ close_partitions(partitions_hash *partitions)
 	{
 		index_close(entry->ident_index, AccessShareLock);
 		ExecDropSingleTupleTableSlot(entry->ind_slot);
+		if (entry->conv_map)
+			free_conversion_map_ext(entry->conv_map);
 		FreeBulkInsertState(entry->bistate);
 	}
 
 	partitions_destroy(partitions);
 }
 
 /*
- * Find BulkInsertState for given partition.
+ * Find hash entry for given partition.
  */
-BulkInsertState
-get_partition_insert_state(partitions_hash *partitions, Oid part_oid)
+PartitionEntry *
+get_partition_entry(partitions_hash *partitions, Oid part_oid)
 {
 	PartitionEntry *entry;
 
@@ -3080,7 +3104,7 @@ get_partition_insert_state(partitions_hash *partitions, Oid part_oid)
 		elog(ERROR, "bulk insert state not found for partition %u", part_oid);
 	Assert(entry->part_oid == part_oid);
 
-	return entry->bistate;
+	return entry;
 }
 
 /*

pg_rewrite.h

@@ -185,39 +185,6 @@ typedef struct CatalogState
 	bool		have_pk_index;
 } CatalogState;
 
-/*
- * Hash table to cache partition-specific information.
- */
-typedef struct PartitionEntry
-{
-	Oid			part_oid;		/* key */
-	Relation	ident_index;
-
-	/*
-	 * Slot to retrieve tuples from identity index. Since we only allow
-	 * partitions to have exactly the same attributes as the parent table, it
-	 * should work if we used the same slot for all partitions. However it
-	 * seems cleaner if separate slots are used.
-	 */
-	TupleTableSlot *ind_slot;
-
-	/* This should make insertions into partitions more efficient. */
-	BulkInsertState bistate;
-
-	char		status;			/* used by simplehash */
-} PartitionEntry;
-
-#define SH_PREFIX partitions
-#define SH_ELEMENT_TYPE PartitionEntry
-#define SH_KEY_TYPE Oid
-#define SH_KEY part_oid
-#define SH_HASH_KEY(tb, key) (key)
-#define SH_EQUAL(tb, a, b) ((a) == (b))
-#define SH_SCOPE static inline
-#define SH_DECLARE
-#define SH_DEFINE
-#include "lib/simplehash.h"
-
 /* Progress tracking. */
 typedef struct TaskProgress
 {
@@ -233,6 +200,7 @@ typedef struct TaskProgress
 	int64		del;
 } TaskProgress;
 
+/* TODO Remove */
 typedef enum WorkerTaskKind
 {
 	WORKER_TASK_PARTITION
@@ -325,6 +293,40 @@ typedef struct TupleConversionMapExt
 								 * coercion. */
 } TupleConversionMapExt;
 
+/*
+ * Hash table to cache partition-specific information.
+ */
+typedef struct PartitionEntry
+{
+	Oid			part_oid;		/* key */
+	Relation	ident_index;
+
+	/* Slot to retrieve tuples from identity index. */
+	TupleTableSlot *ind_slot;
+
+	/* This should make insertions into partitions more efficient. */
+	BulkInsertState bistate;
+
+	/*
+	 * Map to convert tuples that match the partitioned table so they match
+	 * this partition.
+	 */
+	TupleConversionMapExt	*conv_map;
+
+	char		status;			/* used by simplehash */
+} PartitionEntry;
+
+#define SH_PREFIX partitions
+#define SH_ELEMENT_TYPE PartitionEntry
+#define SH_KEY_TYPE Oid
+#define SH_KEY part_oid
+#define SH_HASH_KEY(tb, key) (key)
+#define SH_EQUAL(tb, a, b) ((a) == (b))
+#define SH_SCOPE static inline
+#define SH_DECLARE
+#define SH_DEFINE
+#include "lib/simplehash.h"
+
 extern PGDLLEXPORT void rewrite_worker_main(Datum main_arg);
 
 extern void pg_rewrite_exit_if_requested(void);
@@ -357,7 +359,7 @@ extern bool pg_rewrite_decode_concurrent_changes(LogicalDecodingContext *ctx,
 extern HeapTuple convert_tuple_for_dest_table(HeapTuple tuple,
 											  TupleConversionMapExt *conv_map);
 extern void _PG_output_plugin_init(OutputPluginCallbacks *cb);
-extern BulkInsertState get_partition_insert_state(partitions_hash *partitions,
-												  Oid part_oid);;
+extern PartitionEntry *get_partition_entry(partitions_hash *partitions,
+										   Oid part_oid);;
 extern HeapTuple pg_rewrite_execute_attr_map_tuple(HeapTuple tuple,
 												   TupleConversionMapExt *map);