diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index ac86f3d5be..63b3430f2b 100644 --- a/src/backend/commands/copy.c +++ b/src/backend/commands/copy.c @@ -2648,6 +2648,18 @@ CopyMultiInsertInfoStore(CopyMultiInsertInfo *miinfo, ResultRelInfo *rri, miinfo->bufferedBytes += tuplen; } +/* + * callback function for ExecCustomTupleRoutingAction to flush the bulk + * insert state of each partition at the end of a COPY FROM. + */ +static void +finish_partition_bulk_insert(ResultRelInfo *rri, void *state) +{ + int ti_options = *((int *) state); + + table_finish_bulk_insert(rri->ri_RelationDesc, ti_options); +} + /* * Copy FROM file to relation. */ @@ -3359,7 +3371,15 @@ CopyFrom(CopyState cstate) /* Close all the partitioned tables, leaf partitions, and their indices */ if (proute) + { + /* Perform finish_partition_bulk_insert on each ResultRelInfo used */ + if (insertMethod != CIM_SINGLE) + ExecCustomTupleRoutingAction(proute, + finish_partition_bulk_insert, + &ti_options); + ExecCleanupTupleRouting(mtstate, proute); + } /* Close any trigger target relations */ ExecCleanUpTriggerState(estate); diff --git a/src/backend/executor/execPartition.c b/src/backend/executor/execPartition.c index 6f2b4d62b4..6b5fc5bb56 100644 --- a/src/backend/executor/execPartition.c +++ b/src/backend/executor/execPartition.c @@ -1097,6 +1097,45 @@ ExecInitPartitionDispatchInfo(EState *estate, return pd; } +/* + * ExecCustomTupleRoutingAction + * Invokes callback function on each partition's ResultRelInfo that has + * been allocated for tuple routing. Ignores any ResultRelInfo that's + * found in proute->subplan_resultrel_htab. + */ +void +ExecCustomTupleRoutingAction(PartitionTupleRouting *proute, + TupleRoutingActionCallback callback, + void *callback_state) +{ + HTAB *htab = proute->subplan_resultrel_htab; + int i; + + for (i = 0; i < proute->num_partitions; i++) + { + ResultRelInfo *resultRelInfo = proute->partitions[i]; + + /* + * Check if this result rel is one belonging to the node's subplans, + * if so, we don't invoke the callback function. + */ + if (htab) + { + Oid partoid; + bool found; + + partoid = RelationGetRelid(resultRelInfo->ri_RelationDesc); + + (void) hash_search(htab, &partoid, HASH_FIND, &found); + if (found) + continue; + } + + /* Invoke callback function */ + callback(resultRelInfo, callback_state); + } +} + /* * ExecCleanupTupleRouting -- Clean up objects allocated for partition tuple * routing. diff --git a/src/include/executor/execPartition.h b/src/include/executor/execPartition.h index 580734e9c9..f5234402b9 100644 --- a/src/include/executor/execPartition.h +++ b/src/include/executor/execPartition.h @@ -49,6 +49,12 @@ typedef struct PartitionRoutingInfo TupleTableSlot *pi_PartitionTupleSlot; } PartitionRoutingInfo; +/* + * Typedef for callback function for performing custom actions on each + * partition's ResultRelInfo in tuple routing + */ +typedef void(*TupleRoutingActionCallback) (ResultRelInfo *rri, void *state); + /* * PartitionedRelPruningData - Per-partitioned-table data for run-time pruning * of partitions. For a multilevel partitioned table, we have one of these @@ -145,6 +151,9 @@ extern ResultRelInfo *ExecFindPartition(ModifyTableState *mtstate, PartitionTupleRouting *proute, TupleTableSlot *slot, EState *estate); +extern void ExecCustomTupleRoutingAction(PartitionTupleRouting *proute, + TupleRoutingActionCallback callback, + void *callback_state); extern void ExecCleanupTupleRouting(ModifyTableState *mtstate, PartitionTupleRouting *proute); extern PartitionPruneState *ExecCreatePartitionPruneState(PlanState *planstate,