From 51779387b107e80ba598fe26f4ecec71c54f916b Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 4 Mar 2024 13:52:34 +0900 Subject: [PATCH v20 1/5] Add CopyToRoutine/CopyFromRountine They are for implementing custom COPY TO/FROM format. But this is not enough to implement custom COPY TO/FROM format yet. We'll export some APIs to receive/send data and add "format" option to COPY TO/FROM later. Existing text/csv/binary format implementations don't use CopyToRoutine/CopyFromRoutine for now. We have a patch for it but we defer it. Because there are some mysterious profile results in spite of we get faster runtimes. See [1] for details. [1] https://www.postgresql.org/message-id/ZdbtQJ-p5H1_EDwE%40paquier.xyz Note that this doesn't change existing text/csv/binary format implementations. --- src/backend/commands/copyfrom.c | 24 +++++- src/backend/commands/copyfromparse.c | 5 ++ src/backend/commands/copyto.c | 31 ++++++- src/include/commands/copyapi.h | 101 +++++++++++++++++++++++ src/include/commands/copyfrom_internal.h | 4 + src/tools/pgindent/typedefs.list | 2 + 6 files changed, 159 insertions(+), 8 deletions(-) create mode 100644 src/include/commands/copyapi.h diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c index 2d3462913e1..c42485bd9cb 100644 --- a/src/backend/commands/copyfrom.c +++ b/src/backend/commands/copyfrom.c @@ -1632,12 +1632,22 @@ BeginCopyFrom(ParseState *pstate, /* Fetch the input function and typioparam info */ if (cstate->opts.binary) + { getTypeBinaryInputInfo(att->atttypid, &in_func_oid, &typioparams[attnum - 1]); + fmgr_info(in_func_oid, &in_functions[attnum - 1]); + } + else if (cstate->routine) + cstate->routine->CopyFromInFunc(cstate, att->atttypid, + &in_functions[attnum - 1], + &typioparams[attnum - 1]); + else + { getTypeInputInfo(att->atttypid, &in_func_oid, &typioparams[attnum - 1]); - fmgr_info(in_func_oid, &in_functions[attnum - 1]); + fmgr_info(in_func_oid, &in_functions[attnum - 1]); + } /* Get default info if available */ defexprs[attnum - 1] = NULL; @@ -1777,10 +1787,13 @@ BeginCopyFrom(ParseState *pstate, /* Read and verify binary header */ ReceiveCopyBinaryHeader(cstate); } - - /* create workspace for CopyReadAttributes results */ - if (!cstate->opts.binary) + else if (cstate->routine) { + cstate->routine->CopyFromStart(cstate, tupDesc); + } + else + { + /* create workspace for CopyReadAttributes results */ AttrNumber attr_count = list_length(cstate->attnumlist); cstate->max_fields = attr_count; @@ -1798,6 +1811,9 @@ BeginCopyFrom(ParseState *pstate, void EndCopyFrom(CopyFromState cstate) { + if (cstate->routine) + cstate->routine->CopyFromEnd(cstate); + /* No COPY FROM related resources except memory. */ if (cstate->is_program) { diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c index 97a4c387a30..2e126448019 100644 --- a/src/backend/commands/copyfromparse.c +++ b/src/backend/commands/copyfromparse.c @@ -1012,6 +1012,11 @@ NextCopyFrom(CopyFromState cstate, ExprContext *econtext, Assert(fieldno == attr_count); } + else if (cstate->routine) + { + if (!cstate->routine->CopyFromOneRow(cstate, econtext, values, nulls)) + return false; + } else { /* binary */ diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c index 91de442f434..3c5a97679aa 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -20,6 +20,7 @@ #include "access/tableam.h" #include "commands/copy.h" +#include "commands/copyapi.h" #include "commands/progress.h" #include "executor/execdesc.h" #include "executor/executor.h" @@ -64,6 +65,9 @@ typedef enum CopyDest */ typedef struct CopyToStateData { + /* format routine */ + const CopyToRoutine *routine; + /* low-level state data */ CopyDest copy_dest; /* type of copy source/destination */ FILE *copy_file; /* used if copy_dest == COPY_FILE */ @@ -772,14 +776,22 @@ DoCopyTo(CopyToState cstate) Form_pg_attribute attr = TupleDescAttr(tupDesc, attnum - 1); if (cstate->opts.binary) + { getTypeBinaryOutputInfo(attr->atttypid, &out_func_oid, &isvarlena); + fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); + } + else if (cstate->routine) + cstate->routine->CopyToOutFunc(cstate, attr->atttypid, + &cstate->out_functions[attnum - 1]); else + { getTypeOutputInfo(attr->atttypid, &out_func_oid, &isvarlena); - fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); + fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); + } } /* @@ -806,6 +818,8 @@ DoCopyTo(CopyToState cstate) tmp = 0; CopySendInt32(cstate, tmp); } + else if (cstate->routine) + cstate->routine->CopyToStart(cstate, tupDesc); else { /* @@ -887,6 +901,8 @@ DoCopyTo(CopyToState cstate) /* Need to flush out the trailer */ CopySendEndOfRow(cstate); } + else if (cstate->routine) + cstate->routine->CopyToEnd(cstate); MemoryContextDelete(cstate->rowcontext); @@ -908,15 +924,22 @@ CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot) MemoryContextReset(cstate->rowcontext); oldcontext = MemoryContextSwitchTo(cstate->rowcontext); + /* Make sure the tuple is fully deconstructed */ + slot_getallattrs(slot); + + if (cstate->routine) + { + cstate->routine->CopyToOneRow(cstate, slot); + MemoryContextSwitchTo(oldcontext); + return; + } + if (cstate->opts.binary) { /* Binary per-tuple header */ CopySendInt16(cstate, list_length(cstate->attnumlist)); } - /* Make sure the tuple is fully deconstructed */ - slot_getallattrs(slot); - if (!cstate->opts.binary) { bool need_delim = false; diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h new file mode 100644 index 00000000000..d1289424c67 --- /dev/null +++ b/src/include/commands/copyapi.h @@ -0,0 +1,101 @@ +/*------------------------------------------------------------------------- + * + * copyapi.h + * API for COPY TO/FROM handlers + * + * + * Portions Copyright (c) 1996-2024, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * src/include/commands/copyapi.h + * + *------------------------------------------------------------------------- + */ +#ifndef COPYAPI_H +#define COPYAPI_H + +#include "executor/tuptable.h" +#include "nodes/execnodes.h" + +/* These are private in commands/copy[from|to].c */ +typedef struct CopyFromStateData *CopyFromState; +typedef struct CopyToStateData *CopyToState; + +/* + * API structure for a COPY FROM format implementation. Note this must be + * allocated in a server-lifetime manner, typically as a static const struct. + */ +typedef struct CopyFromRoutine +{ + /* + * Called when COPY FROM is started to set up the input functions + * associated with the relation's attributes writing to. `finfo` can be + * optionally filled to provide the catalog information of the input + * function. `typioparam` can be optionally filled to define the OID of + * the type to pass to the input function. `atttypid` is the OID of data + * type used by the relation's attribute. + */ + void (*CopyFromInFunc) (CopyFromState cstate, Oid atttypid, + FmgrInfo *finfo, Oid *typioparam); + + /* + * Called when COPY FROM is started. + * + * `tupDesc` is the tuple descriptor of the relation where the data needs + * to be copied. This can be used for any initialization steps required + * by a format. + */ + void (*CopyFromStart) (CopyFromState cstate, TupleDesc tupDesc); + + /* + * Copy one row to a set of `values` and `nulls` of size tupDesc->natts. + * + * 'econtext' is used to evaluate default expression for each column that + * is either not read from the file or is using the DEFAULT option of COPY + * FROM. It is NULL if no default values are used. + * + * Returns false if there are no more tuples to copy. + */ + bool (*CopyFromOneRow) (CopyFromState cstate, ExprContext *econtext, + Datum *values, bool *nulls); + + /* Called when COPY FROM has ended. */ + void (*CopyFromEnd) (CopyFromState cstate); +} CopyFromRoutine; + +/* + * API structure for a COPY TO format implementation. Note this must be + * allocated in a server-lifetime manner, typically as a static const struct. + */ +typedef struct CopyToRoutine +{ + /* + * Called when COPY TO is started to set up the output functions + * associated with the relation's attributes reading from. `finfo` can be + * optionally filled to provide the catalog information of the output + * function. `atttypid` is the OID of data type used by the relation's + * attribute. + */ + void (*CopyToOutFunc) (CopyToState cstate, Oid atttypid, + FmgrInfo *finfo); + + /* + * Called when COPY TO is started. + * + * `tupDesc` is the tuple descriptor of the relation from where the data + * is read. + */ + void (*CopyToStart) (CopyToState cstate, TupleDesc tupDesc); + + /* + * Copy one row for COPY TO. + * + * `slot` is the tuple slot where the data is emitted. + */ + void (*CopyToOneRow) (CopyToState cstate, TupleTableSlot *slot); + + /* Called when COPY TO has ended */ + void (*CopyToEnd) (CopyToState cstate); +} CopyToRoutine; + +#endif /* COPYAPI_H */ diff --git a/src/include/commands/copyfrom_internal.h b/src/include/commands/copyfrom_internal.h index cad52fcc783..509b9e92a18 100644 --- a/src/include/commands/copyfrom_internal.h +++ b/src/include/commands/copyfrom_internal.h @@ -15,6 +15,7 @@ #define COPYFROM_INTERNAL_H #include "commands/copy.h" +#include "commands/copyapi.h" #include "commands/trigger.h" #include "nodes/miscnodes.h" @@ -58,6 +59,9 @@ typedef enum CopyInsertMethod */ typedef struct CopyFromStateData { + /* format routine */ + const CopyFromRoutine *routine; + /* low-level state data */ CopySource copy_src; /* type of copy source */ FILE *copy_file; /* used if copy_src == COPY_FILE */ diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list index 5fabb127d7e..4c4bf60d9e5 100644 --- a/src/tools/pgindent/typedefs.list +++ b/src/tools/pgindent/typedefs.list @@ -492,6 +492,7 @@ ConvertRowtypeExpr CookedConstraint CopyDest CopyFormatOptions +CopyFromRoutine CopyFromState CopyFromStateData CopyHeaderChoice @@ -503,6 +504,7 @@ CopyMultiInsertInfo CopyOnErrorChoice CopySource CopyStmt +CopyToRoutine CopyToState CopyToStateData Cost -- 2.45.2