Fix tuple routing for hypertables

For tuple routing we set up a ChunkDispatch subplan to the ModifyTable
plan that reads the tuples to insert and locates the correct chunk to
insert them into and then passes a tuple table slot matching that chunk
to ModifyTable.

`ExecModifyTable()` expects a tuple table slot with the same number of
attributes as the parent table. If an attribute is dropped in the
parent table and a chunk is subsequently created, the dropped attribute
will be present in the tuple table slot for the parent table but not
for the new chunk. As a result, the slot passed back to ModifyTable
will have a different number of attributes compared to what is expected
and if `ExecCopySlot()` is called, it will assert with a failure because
the two slots are expected to have the same number of attributes.

Partition tables deal with this by passing back a virtual tuple to
`ExecModifyTable()` and then locate the correct partition and
corresponding slot in `ExecInsert()` by calling the function
`ExecPrepareTupleRouting()` and mapping the attributes between the
virtual slot and the actual partition slot.

This commit solves this in a similar manner by moving part of the
re-mapping done by `chunk_dispatch_exec()` to `ExecInsert()`, similar
to how it is done for partition tables, and returning the slot from the
subplan, allowing `ExecModifyTable()` to keep the same logic as before.

src/import/ht_hypertable_modify.c

@@ -380,7 +380,8 @@ ExecProcessReturning(ResultRelInfo *resultRelInfo, TupleTableSlot *tupleSlot,
 #if PG15_GE
 
 TupleTableSlot *ExecInsert(ModifyTableContext * context, ResultRelInfo * resultRelInfo,
-			   TupleTableSlot * slot, bool canSetTag);
+						   ChunkDispatchState* cds,
+						   TupleTableSlot * slot, bool canSetTag);
 
 static TupleTableSlot * mergeGetUpdateNewTuple(ResultRelInfo * relinfo, TupleTableSlot * planSlot,
 		    TupleTableSlot * oldSlot, MergeActionState * relaction);
@@ -943,14 +944,15 @@ ht_ExecMergeNotMatched(ModifyTableContext * context, ResultRelInfo * resultRelIn
 												MakeSingleTupleTableSlot(chunktupdesc,
 																		&TTSOpsVirtual));
 				rslot =  ExecInsert(context,
-									cds->rri,
+									resultRelInfo,
+									cds,
 									(chunk_slot ? chunk_slot : newslot),
 									canSetTag);
 				if (chunk_slot)
 					ExecDropSingleTupleTableSlot(chunk_slot);
 			}
 			else
-				rslot = ExecInsert(context, cds->rri, newslot, canSetTag);
+				rslot = ExecInsert(context, resultRelInfo, cds, newslot, canSetTag);
 			mtstate->mt_merge_inserted = 1;
 			break;
 		case CMD_NOTHING:

src/nodes/chunk_dispatch/chunk_dispatch.c

@@ -438,6 +438,32 @@ chunk_dispatch_exec(CustomScanState *node)
 
 	MemoryContextSwitchTo(old);
 
+	/*
+	 * Save away the insert state for using it in ExecInsert().
+	 *
+	 * We need to return the original slot from the subplan since otherwise
+	 * the slot might not match what is expected. The expected slot should
+	 * contain a tuple with the same definition as the parent hypertable while
+	 * the slot in the chunk insert state is a slot matching the chunk
+	 * definition.
+	 *
+	 * If columns have been dropped from the parent hypertable, new chunks
+	 * will not have these dropped attributes, which will cause problems later
+	 * in the insert execution (see ExecGetInsertNewTuple()).
+	 *
+	 * We can probably improve this code by removing ChunkDispatch. See
+	 * hypertable_modify.c for more information.
+	 */
+	state->cis = cis;
+
+	return slot;
+}
+
+TupleTableSlot *
+ts_chunk_dispatch_prepare_tuple_routing(ChunkDispatchState *state, TupleTableSlot *slot)
+{
+	ChunkInsertState *cis = state->cis;
+
 	/* Convert the tuple to the chunk's rowtype, if necessary */
 	if (cis->hyper_to_chunk_map != NULL && state->is_dropped_attr_exists == false)
 		slot = execute_attr_map_slot(cis->hyper_to_chunk_map->attrMap, slot, cis->slot);

src/nodes/chunk_dispatch/chunk_dispatch.h

@@ -72,6 +72,14 @@ typedef struct ChunkDispatchState
 	 */
 	ChunkDispatch *dispatch;
 	ResultRelInfo *rri;
+
+	/*
+	 * Keep the chunk insert state available to pass it from
+	 * ExecGetInsertNewTuple() to ExecInsert(), where the actual slot to
+	 * use is decided.
+	 */
+	ChunkInsertState *cis;
+
 	/* flag to represent dropped attributes */
 	bool is_dropped_attr_exists;
 	int64 batches_deleted;
@@ -98,6 +106,8 @@ ts_chunk_dispatch_get_chunk_insert_state(ChunkDispatch *dispatch, Point *p,
 extern void ts_chunk_dispatch_decompress_batches_for_insert(ChunkDispatch *dispatch,
 															ChunkInsertState *cis,
 															TupleTableSlot *slot);
+extern TupleTableSlot *ts_chunk_dispatch_prepare_tuple_routing(ChunkDispatchState *state,
+															   TupleTableSlot *slot);
 
 extern TSDLLEXPORT Path *ts_chunk_dispatch_path_create(PlannerInfo *root, ModifyTablePath *mtpath,
 													   Index hypertable_rti, int subpath_index);

src/nodes/hypertable_modify.c

@@ -929,7 +929,7 @@ ExecModifyTable(CustomScanState *cs_node, PlanState *pstate)
 				if (unlikely(!resultRelInfo->ri_projectNewInfoValid))
 					ExecInitInsertProjection(node, resultRelInfo);
 				slot = ExecGetInsertNewTuple(resultRelInfo, context.planSlot);
-				slot = ExecInsert(&context, cds->rri, slot, node->canSetTag);
+				slot = ExecInsert(&context, resultRelInfo, cds, slot, node->canSetTag);
 				break;
 			case CMD_UPDATE:
 				/* Initialize projection info if first time for this table */
@@ -1431,8 +1431,8 @@ ExecGetUpdateNewTuple(ResultRelInfo *relinfo, TupleTableSlot *planSlot, TupleTab
  * copied and modified version of ExecInsert from executor/nodeModifyTable.c
  */
 TupleTableSlot *
-ExecInsert(ModifyTableContext *context, ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
-		   bool canSetTag)
+ExecInsert(ModifyTableContext *context, ResultRelInfo *resultRelInfo, ChunkDispatchState *cds,
+		   TupleTableSlot *slot, bool canSetTag)
 {
 	ModifyTableState *mtstate = context->mtstate;
 	EState *estate = context->estate;
@@ -1447,6 +1447,36 @@ ExecInsert(ModifyTableContext *context, ResultRelInfo *resultRelInfo, TupleTable
 
 	Assert(!mtstate->mt_partition_tuple_routing);
 
+	/*
+	 * Fetch the chunk dispatch state similar to how it is done in
+	 * nodeModifyTable.c. For us, this is stored in the chunk insert state,
+	 * which we add in to the chunk dispatch state in chunk_dispatch_exec().
+	 *
+	 * We can probably improve this code by removing ChunkDispatch. It is
+	 * currently the immediate child of ModifyTable and placed between
+	 * ModifyTable and the original subplan of ModifyTable. This was
+	 * previously necessary because we didn't have our own version of
+	 * ModifyTable, but since PG14 we have our own version of
+	 * ModifyTable. This means that we can move the logic to make a
+	 * partition/chunk lookup into a separate function similar to how
+	 * ExecPrepareTupleRouting() does it in nodeModifyTable.c.
+	 *
+	 *    if (proute)
+	 *    {
+	 *        ResultRelInfo *partRelInfo;
+	 *
+	 *        slot = ExecPrepareTupleRouting(mtstate, estate, proute,
+	 *                                       resultRelInfo, slot,
+	 *                                       &partRelInfo);
+	 *        resultRelInfo = partRelInfo;
+	 *  }
+	 *
+	 * The current approach is a quick fix to avoid changing too much code at
+	 * the same time and risk introducing a bug.
+	 */
+	slot = ts_chunk_dispatch_prepare_tuple_routing(cds, slot);
+	resultRelInfo = cds->rri;
+
 	ExecMaterializeSlot(slot);
 
 	resultRelationDesc = resultRelInfo->ri_RelationDesc;

src/nodes/hypertable_modify.h

@@ -12,6 +12,9 @@
 #include "hypertable.h"
 #include "import/ht_hypertable_modify.h"
 
+/* Forward declarations */
+struct ChunkDispatchState;
+
 typedef struct HypertableModifyPath
 {
 	CustomPath cpath;
@@ -41,4 +44,5 @@ extern Path *ts_hypertable_modify_path_create(PlannerInfo *root, ModifyTablePath
 extern List *ts_replace_rowid_vars(PlannerInfo *root, List *tlist, int varno);
 
 extern TupleTableSlot *ExecInsert(ModifyTableContext *context, ResultRelInfo *resultRelInfo,
-								  TupleTableSlot *slot, bool canSetTag);
+								  struct ChunkDispatchState *cds, TupleTableSlot *slot,
+								  bool canSetTag);