v8-0006-Parallel-Copy-For-Binary-Format-Files.patch

From 93042e3cab8e9a4f067e5f74b0dc67b0a6df2886 Mon Sep 17 00:00:00 2001
From: Vignesh C <vignesh21@gmail.com>
Date: Wed, 21 Oct 2020 11:33:36 +0530
Subject: [PATCH v8 6/6] Parallel Copy For Binary Format Files

Leader reads data from the file into the DSM data blocks each of 64K size.
It also identifies each tuple data block id, start offset, end offset,
tuple size and updates this information in the ring data structure.
Workers parallelly read the tuple information from the ring data structure,
the actual tuple data from the data blocks and parallelly insert the tuples
into the table.
---
 src/backend/commands/copy.c | 126 ++++++++++++++++++++------------------------
 src/include/commands/copy.h | 126 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 184 insertions(+), 68 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index e613fc0..5894b28 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -222,19 +222,14 @@ static void CopySendData(CopyState cstate, const void *databuf, int datasize);
 static void CopySendString(CopyState cstate, const char *str);
 static void CopySendChar(CopyState cstate, char c);
 static void CopySendEndOfRow(CopyState cstate);
-static int	CopyGetData(CopyState cstate, void *databuf,
-						int minread, int maxread);
 static void CopySendInt32(CopyState cstate, int32 val);
 static bool CopyGetInt32(CopyState cstate, int32 *val);
 static void CopySendInt16(CopyState cstate, int16 val);
 static bool CopyGetInt16(CopyState cstate, int16 *val);
 static bool CopyLoadRawBuf(CopyState cstate);
-static int	CopyReadBinaryData(CopyState cstate, char *dest, int nbytes);
-
 static void ClearEOLFromCopiedData(CopyState cstate, char *copy_line_data,
 								   int copy_line_pos, int *copy_line_size);
 
-
 /*
  * Send copy start/stop messages for frontend copies.  These have changed
  * in past protocol redesigns.
@@ -448,7 +443,7 @@ CopySendEndOfRow(CopyState cstate)
  *
  * NB: no data conversion is applied here.
  */
-static int
+int
 CopyGetData(CopyState cstate, void *databuf, int minread, int maxread)
 {
 	int			bytesread = 0;
@@ -581,10 +576,25 @@ CopyGetInt32(CopyState cstate, int32 *val)
 {
 	uint32		buf;
 
-	if (CopyReadBinaryData(cstate, (char *) &buf, sizeof(buf)) != sizeof(buf))
+	/*
+	 * For parallel copy, avoid reading data to raw buf, read directly from
+	 * file, later the data will be read to parallel copy data buffers.
+	 */
+	if (cstate->nworkers > 0)
 	{
-		*val = 0;				/* suppress compiler warning */
-		return false;
+		if (CopyGetData(cstate, &buf, sizeof(buf), sizeof(buf)) != sizeof(buf))
+		{
+			*val = 0;			/* suppress compiler warning */
+			return false;
+		}
+	}
+	else
+	{
+		if (CopyReadBinaryData(cstate, (char *) &buf, sizeof(buf)) != sizeof(buf))
+		{
+			*val = 0;			/* suppress compiler warning */
+			return false;
+		}
 	}
 	*val = (int32) pg_ntoh32(buf);
 	return true;
@@ -658,7 +668,7 @@ CopyLoadRawBuf(CopyState cstate)
  * and writes them to 'dest'.  Returns the number of bytes read (which
  * would be less than 'nbytes' only if we reach EOF).
  */
-static int
+int
 CopyReadBinaryData(CopyState cstate, char *dest, int nbytes)
 {
 	int			copied_bytes = 0;
@@ -3539,7 +3549,7 @@ BeginCopyFrom(ParseState *pstate,
 		int32		tmp;
 
 		/* Signature */
-		if (CopyReadBinaryData(cstate, readSig, 11) != 11 ||
+		if (CopyGetData(cstate, readSig, 11, 11) != 11 ||
 			memcmp(readSig, BinarySignature, 11) != 0)
 			ereport(ERROR,
 					(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
@@ -3567,7 +3577,7 @@ BeginCopyFrom(ParseState *pstate,
 		/* Skip extension header, if present */
 		while (tmp-- > 0)
 		{
-			if (CopyReadBinaryData(cstate, readSig, 1) != 1)
+			if (CopyGetData(cstate, readSig, 1, 1) != 1)
 				ereport(ERROR,
 						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
 						 errmsg("invalid COPY file header (wrong length)")));
@@ -3764,60 +3774,45 @@ NextCopyFrom(CopyState cstate, ExprContext *econtext,
 	else
 	{
 		/* binary */
-		int16		fld_count;
-		ListCell   *cur;
-
 		cstate->cur_lineno++;
+		cstate->max_fields = list_length(cstate->attnumlist);
 
-		if (!CopyGetInt16(cstate, &fld_count))
+		if (!IsParallelCopy())
 		{
-			/* EOF detected (end of file, or protocol-level EOF) */
-			return false;
-		}
-
-		if (fld_count == -1)
-		{
-			/*
-			 * Received EOF marker.  In a V3-protocol copy, wait for the
-			 * protocol-level EOF, and complain if it doesn't come
-			 * immediately.  This ensures that we correctly handle CopyFail,
-			 * if client chooses to send that now.
-			 *
-			 * Note that we MUST NOT try to read more data in an old-protocol
-			 * copy, since there is no protocol-level EOF marker then.  We
-			 * could go either way for copy from file, but choose to throw
-			 * error if there's data after the EOF marker, for consistency
-			 * with the new-protocol case.
-			 */
-			char		dummy;
+			int16		fld_count;
+			ListCell   *cur;
 
-			if (cstate->copy_dest != COPY_OLD_FE &&
-				CopyReadBinaryData(cstate, &dummy, 1) > 0)
-				ereport(ERROR,
-						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
-						 errmsg("received copy data after EOF marker")));
-			return false;
-		}
+			if (!CopyGetInt16(cstate, &fld_count))
+			{
+				/* EOF detected (end of file, or protocol-level EOF) */
+				return false;
+			}
 
-		if (fld_count != attr_count)
-			ereport(ERROR,
-					(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
-					 errmsg("row field count is %d, expected %d",
-							(int) fld_count, attr_count)));
+			CHECK_FIELD_COUNT;
 
-		foreach(cur, cstate->attnumlist)
+			foreach(cur, cstate->attnumlist)
+			{
+				int			attnum = lfirst_int(cur);
+				int			m = attnum - 1;
+				Form_pg_attribute att = TupleDescAttr(tupDesc, m);
+
+				cstate->cur_attname = NameStr(att->attname);
+				values[m] = CopyReadBinaryAttribute(cstate,
+													&in_functions[m],
+													typioparams[m],
+													att->atttypmod,
+													&nulls[m]);
+				cstate->cur_attname = NULL;
+			}
+		}
+		else
 		{
-			int			attnum = lfirst_int(cur);
-			int			m = attnum - 1;
-			Form_pg_attribute att = TupleDescAttr(tupDesc, m);
+			bool		eof = false;
 
-			cstate->cur_attname = NameStr(att->attname);
-			values[m] = CopyReadBinaryAttribute(cstate,
-												&in_functions[m],
-												typioparams[m],
-												att->atttypmod,
-												&nulls[m]);
-			cstate->cur_attname = NULL;
+			eof = CopyReadBinaryTupleWorker(cstate, values, nulls);
+
+			if (eof)
+				return false;
 		}
 	}
 
@@ -4828,18 +4823,15 @@ CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo,
 	Datum		result;
 
 	if (!CopyGetInt32(cstate, &fld_size))
-		ereport(ERROR,
-				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
-				 errmsg("unexpected EOF in COPY data")));
+		EOF_ERROR;
+
 	if (fld_size == -1)
 	{
 		*isnull = true;
 		return ReceiveFunctionCall(flinfo, NULL, typioparam, typmod);
 	}
-	if (fld_size < 0)
-		ereport(ERROR,
-				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
-				 errmsg("invalid field size")));
+
+	CHECK_FIELD_SIZE(fld_size);
 
 	/* reset attribute_buf to empty, and load raw data in it */
 	resetStringInfo(&cstate->attribute_buf);
@@ -4847,9 +4839,7 @@ CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo,
 	enlargeStringInfo(&cstate->attribute_buf, fld_size);
 	if (CopyReadBinaryData(cstate, cstate->attribute_buf.data,
 						   fld_size) != fld_size)
-		ereport(ERROR,
-				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT),
-				 errmsg("unexpected EOF in COPY data")));
+		EOF_ERROR;
 
 	cstate->attribute_buf.len = fld_size;
 	cstate->attribute_buf.data[fld_size] = '\0';
diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h
index 9b19dcb..49f438f 100644
--- a/src/include/commands/copy.h
+++ b/src/include/commands/copy.h
@@ -59,6 +59,109 @@
 
 
 /*
+ * CHECK_FIELD_COUNT - Handles the error cases for field count
+ * for binary format files.
+ */
+#define CHECK_FIELD_COUNT \
+{\
+	if (fld_count == -1) \
+	{ \
+		if (IsParallelCopy() && \
+			!IsLeader()) \
+			return true; \
+		else if (IsParallelCopy() && \
+			IsLeader()) \
+		{ \
+			if (cstate->pcdata->curr_data_block->data[cstate->raw_buf_index + sizeof(fld_count)] != 0) \
+				ereport(ERROR, \
+						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT), \
+						errmsg("received copy data after EOF marker"))); \
+			return true; \
+		} \
+		else \
+		{ \
+			/* \
+			 * Received EOF marker.  In a V3-protocol copy, wait for the \
+			 * protocol-level EOF, and complain if it doesn't come \
+			 * immediately.  This ensures that we correctly handle CopyFail, \
+			 * if client chooses to send that now. \
+			 * \
+			 * Note that we MUST NOT try to read more data in an old-protocol \
+			 * copy, since there is no protocol-level EOF marker then.  We \
+			 * could go either way for copy from file, but choose to throw \
+			 * error if there's data after the EOF marker, for consistency \
+			 * with the new-protocol case. \
+			 */ \
+			char		dummy; \
+			if (cstate->copy_dest != COPY_OLD_FE && \
+				CopyReadBinaryData(cstate, &dummy, 1) > 0) \
+				ereport(ERROR, \
+						(errcode(ERRCODE_BAD_COPY_FILE_FORMAT), \
+						errmsg("received copy data after EOF marker"))); \
+				return false; \
+		} \
+	} \
+	if (fld_count != cstate->max_fields) \
+		ereport(ERROR, \
+				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT), \
+				errmsg("row field count is %d, expected %d", \
+				(int) fld_count, cstate->max_fields))); \
+}
+
+/*
+ * CHECK_FIELD_SIZE - Handles the error case for field size
+ * for binary format files.
+ */
+#define CHECK_FIELD_SIZE(fld_size) \
+{ \
+	if (fld_size < -1) \
+		ereport(ERROR, \
+				(errcode(ERRCODE_BAD_COPY_FILE_FORMAT), \
+				errmsg("invalid field size")));\
+}
+
+/*
+ * EOF_ERROR - Error statement for EOF for binary format
+ * files.
+ */
+#define EOF_ERROR \
+{ \
+	ereport(ERROR, \
+		(errcode(ERRCODE_BAD_COPY_FILE_FORMAT), \
+		errmsg("unexpected EOF in COPY data")));\
+}
+
+/*
+ * GET_RAW_BUF_INDEX - Calculates the raw buf index for the cases
+ * where the data spread is across multiple data blocks.
+ */
+#define GET_RAW_BUF_INDEX(raw_buf_index, fld_size, required_blks, curr_blk_bytes) \
+{ \
+	raw_buf_index = fld_size - (((required_blks - 1) * DATA_BLOCK_SIZE) + curr_blk_bytes); \
+}
+
+/*
+ * GET_REQUIRED_BLOCKS - Calculates the number of data
+ * blocks required for the cases where the data spread
+ * is across multiple data blocks.
+ */
+#define GET_REQUIRED_BLOCKS(required_blks, fld_size, curr_blk_bytes) \
+{ \
+	/* \
+	 * field size can spread across multiple data blocks, \
+	 * calculate the number of required data blocks and try to get \
+	 * those many data blocks. \
+	 */ \
+	required_blks = (int32)(fld_size - curr_blk_bytes)/(int32)DATA_BLOCK_SIZE; \
+	/* \
+	 * check if we need the data block for the field data \
+	 * bytes that are not modulus of data block size. \
+	 */ \
+	if ((fld_size - curr_blk_bytes)%DATA_BLOCK_SIZE != 0) \
+		required_blks++; \
+}
+
+/*
  * Represents the different source/dest cases we need to worry about at
  * the bottom level
  */
@@ -236,6 +339,17 @@ typedef struct ParallelCopyLineBuf
 } ParallelCopyLineBuf;
 
 /*
+ * Represents the usage mode for CopyReadBinaryGetDataBlock.
+ */
+typedef enum FieldInfoType
+{
+	FIELD_NONE = 0,
+	FIELD_COUNT,
+	FIELD_SIZE,
+	FIELD_DATA
+}			FieldInfoType;
+
+/*
  * This structure helps in storing the common data from CopyStateData that are
  * required by the workers. This information will then be allocated and stored
  * into the DSM for the worker to retrieve and copy it to CopyStateData.
@@ -258,6 +372,7 @@ typedef struct SerializedParallelCopyState
 	/* Working state for COPY FROM */
 	AttrNumber	num_defaults;
 	Oid			relid;
+	bool		binary;
 } SerializedParallelCopyState;
 
 /*
@@ -284,6 +399,9 @@ typedef struct ParallelCopyData
 
 	/* Current position in worker_line_buf */
 	uint32		worker_line_buf_pos;
+
+	/* For binary formatted files */
+	ParallelCopyDataBlock *curr_data_block;
 } ParallelCopyData;
 
 typedef int (*copy_data_source_cb) (void *outbuf, int minread, int maxread);
@@ -470,4 +588,12 @@ extern uint32 UpdateSharedLineInfo(CopyState cstate, uint32 blk_pos, uint32 offs
 								   uint32 line_size, uint32 line_state, uint32 blk_line_pos);
 extern void EndLineParallelCopy(CopyState cstate, uint32 line_pos, uint32 line_size,
 								uint32 raw_buf_ptr);
+extern int	CopyGetData(CopyState cstate, void *databuf, int minread, int maxread);
+extern int	CopyReadBinaryData(CopyState cstate, char *dest, int nbytes);
+extern bool CopyReadBinaryTupleLeader(CopyState cstate);
+extern bool CopyReadBinaryTupleWorker(CopyState cstate, Datum *values, bool *nulls);
+extern void CopyReadBinaryFindTupleSize(CopyState cstate, uint32 *line_size);
+extern Datum CopyReadBinaryAttributeWorker(CopyState cstate, FmgrInfo *flinfo,
+										   Oid typioparam, int32 typmod, bool *isnull);
+extern void CopyReadBinaryGetDataBlock(CopyState cstate, FieldInfoType field_info);
 #endif							/* COPY_H */
-- 
1.8.3.1