v6-0001-Implement-retrieval-of-results-in-chunks-with-lib.patch

From 8cfb82b1e36e96996637948a231ae35b9af1e074 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20V=C3=A9rit=C3=A9?= <daniel@manitou-mail.org>
Date: Tue, 30 Jan 2024 14:38:21 +0100
Subject: [PATCH v6 1/2] Implement retrieval of results in chunks with libpq.

This mode is similar to the single-row mode except that chunks
of results contain up to N rows instead of a single row.
It is meant to reduce the overhead of the row-by-row allocations
for large result sets.
The mode is selected with PQsetChunkedRowsMode(int maxRows) and results
have the new status code PGRES_TUPLES_CHUNK.
---
 doc/src/sgml/libpq.sgml                       |  96 ++++++++++----
 .../libpqwalreceiver/libpqwalreceiver.c       |   1 +
 src/bin/pg_amcheck/pg_amcheck.c               |   1 +
 src/interfaces/libpq/exports.txt              |   1 +
 src/interfaces/libpq/fe-exec.c                | 117 +++++++++++++++---
 src/interfaces/libpq/libpq-fe.h               |   4 +-
 src/interfaces/libpq/libpq-int.h              |   7 +-
 7 files changed, 184 insertions(+), 43 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index d0d5aefadc..f7f5a04df6 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -3537,7 +3537,20 @@ ExecStatusType PQresultStatus(const PGresult *res);
             The <structname>PGresult</structname> contains a single result tuple
             from the current command.  This status occurs only when
             single-row mode has been selected for the query
-            (see <xref linkend="libpq-single-row-mode"/>).
+            (see <xref linkend="libpq-chunked-results-modes"/>).
+           </para>
+          </listitem>
+         </varlistentry>
+
+         <varlistentry id="libpq-pgres-tuples-chunk">
+          <term><literal>PGRES_TUPLES_CHUNK</literal></term>
+          <listitem>
+           <para>
+            The <structname>PGresult</structname> contains several tuples
+            from the current command. The count of tuples cannot exceed
+            the maximum passed to <xref linkend="libpq-PQsetChunkedRowsMode"/>.
+            This status occurs only when the chunked mode has been selected
+            for the query (see <xref linkend="libpq-chunked-results-modes"/>).
            </para>
           </listitem>
          </varlistentry>
@@ -5189,8 +5202,8 @@ PGresult *PQgetResult(PGconn *conn);
   <para>
    Another frequently-desired feature that can be obtained with
    <xref linkend="libpq-PQsendQuery"/> and <xref linkend="libpq-PQgetResult"/>
-   is retrieving large query results a row at a time.  This is discussed
-   in <xref linkend="libpq-single-row-mode"/>.
+   is retrieving large query results a limited number of rows at a time.  This is discussed
+   in <xref linkend="libpq-chunked-results-modes"/>.
   </para>
 
   <para>
@@ -5554,12 +5567,13 @@ int PQflush(PGconn *conn);
     </para>
 
     <para>
-     To enter single-row mode, call <function>PQsetSingleRowMode</function>
-     before retrieving results with <function>PQgetResult</function>.
-     This mode selection is effective only for the query currently
-     being processed. For more information on the use of
-     <function>PQsetSingleRowMode</function>,
-     refer to <xref linkend="libpq-single-row-mode"/>.
+     To enter single-row or chunked modes, call
+     respectively <function>PQsetSingleRowMode</function>
+     or <function>PQsetChunkedRowsMode</function> before retrieving results
+     with <function>PQgetResult</function>.  This mode selection is effective
+     only for the query currently being processed. For more information on the
+     use of these functions refer
+     to <xref linkend="libpq-chunked-results-modes" />.
     </para>
 
     <para>
@@ -5926,10 +5940,10 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
   </sect2>
  </sect1>
 
- <sect1 id="libpq-single-row-mode">
-  <title>Retrieving Query Results Row-by-Row</title>
+ <sect1 id="libpq-chunked-results-modes">
+  <title>Retrieving Query Results by chunks</title>
 
-  <indexterm zone="libpq-single-row-mode">
+  <indexterm zone="libpq-chunked-results-modes">
    <primary>libpq</primary>
    <secondary>single-row mode</secondary>
   </indexterm>
@@ -5940,13 +5954,15 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
    <structname>PGresult</structname>.  This can be unworkable for commands
    that return a large number of rows.  For such cases, applications can use
    <xref linkend="libpq-PQsendQuery"/> and <xref linkend="libpq-PQgetResult"/> in
-   <firstterm>single-row mode</firstterm>.  In this mode, the result row(s) are
-   returned to the application one at a time, as they are received from the
-   server.
+   <firstterm>single-row mode</firstterm> or <firstterm>chunked mode</firstterm>.
+   In these modes, the result row(s) are returned to the application one at a
+   time for the single-row mode and by chunks for the chunked mode, as they
+   are received from the server.
   </para>
 
   <para>
-   To enter single-row mode, call <xref linkend="libpq-PQsetSingleRowMode"/>
+   To enter these modes, call <xref linkend="libpq-PQsetSingleRowMode"/>
+    or <xref linkend="libpq-PQsetChunkedRowsMode"/>
    immediately after a successful call of <xref linkend="libpq-PQsendQuery"/>
    (or a sibling function).  This mode selection is effective only for the
    currently executing query.  Then call <xref linkend="libpq-PQgetResult"/>
@@ -5954,7 +5970,8 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
    linkend="libpq-async"/>.  If the query returns any rows, they are returned
    as individual <structname>PGresult</structname> objects, which look like
    normal query results except for having status code
-   <literal>PGRES_SINGLE_TUPLE</literal> instead of
+   <literal>PGRES_SINGLE_TUPLE</literal> for the single-row mode and
+   <literal>PGRES_TUPLES_CHUNK</literal> for the chunked mode, instead of
    <literal>PGRES_TUPLES_OK</literal>.  After the last row, or immediately if
    the query returns zero rows, a zero-row object with status
    <literal>PGRES_TUPLES_OK</literal> is returned; this is the signal that no
@@ -5967,9 +5984,9 @@ UPDATE mytable SET x = x + 1 WHERE id = 42;
   </para>
 
   <para>
-   When using pipeline mode, single-row mode needs to be activated for each
-   query in the pipeline before retrieving results for that query
-   with <function>PQgetResult</function>.
+   When using pipeline mode, the single-row or chunked mode need to be
+   activated for each query in the pipeline before retrieving results for that
+   query with <function>PQgetResult</function>.
    See <xref linkend="libpq-pipeline-mode"/> for more information.
   </para>
 
@@ -6003,14 +6020,49 @@ int PQsetSingleRowMode(PGconn *conn);
    </variablelist>
   </para>
 
+    <para>
+   <variablelist>
+    <varlistentry id="libpq-PQsetChunkedRowsMode">
+      <term><function>PQsetChunkedRowsMode</function>
+        <indexterm><primary>PQsetChunkedRowsMode</primary></indexterm></term>
+     <listitem>
+      <para>
+       Select the mode retrieving results in chunks for the currently-executing query.
+
+<synopsis>
+  int PQsetChunkedRowsMode(PGconn *conn,
+                           int maxRows);
+</synopsis>
+      </para>
+
+      <para>
+       This function is similar to <xref linkend="libpq-PQsetSingleRowMode"/>,
+       except that it can retrieve a user-specified number of rows
+       per call to <xref linkend="libpq-PQgetResult"/>, instead of a single row.
+       This function can only be called immediately after
+       <xref linkend="libpq-PQsendQuery"/> or one of its sibling functions,
+       before any other operation on the connection such as
+       <xref linkend="libpq-PQconsumeInput"/> or
+       <xref linkend="libpq-PQgetResult"/>.  If called at the correct time,
+       the function activates the chunked mode for the current query and
+       returns 1.  Otherwise the mode stays unchanged and the function
+       returns 0.  In any case, the mode reverts to normal after
+       completion of the current query.
+      </para>
+     </listitem>
+    </varlistentry>
+   </variablelist>
+  </para>
+
   <caution>
    <para>
     While processing a query, the server may return some rows and then
     encounter an error, causing the query to be aborted.  Ordinarily,
     <application>libpq</application> discards any such rows and reports only the
-    error.  But in single-row mode, those rows will have already been
+    error.  But in single-row or chunked modes, those rows will have already been
     returned to the application.  Hence, the application will see some
-    <literal>PGRES_SINGLE_TUPLE</literal> <structname>PGresult</structname>
+    <literal>PGRES_SINGLE_TUPLE</literal> or <literal>PGRES_TUPLES_CHUNK</literal>
+    <structname>PGresult</structname>
     objects followed by a <literal>PGRES_FATAL_ERROR</literal> object.  For
     proper transactional behavior, the application must be designed to
     discard or undo whatever has been done with the previously-processed
diff --git a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
index 2439733b55..97111d966d 100644
--- a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
+++ b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
@@ -1183,6 +1183,7 @@ libpqrcv_exec(WalReceiverConn *conn, const char *query,
 	switch (PQresultStatus(pgres))
 	{
 		case PGRES_SINGLE_TUPLE:
+		case PGRES_TUPLES_CHUNK:
 		case PGRES_TUPLES_OK:
 			walres->status = WALRCV_OK_TUPLES;
 			libpqrcv_processTuples(pgres, walres, nRetTypes, retTypes);
diff --git a/src/bin/pg_amcheck/pg_amcheck.c b/src/bin/pg_amcheck/pg_amcheck.c
index e5f9eedc47..728305a7cf 100644
--- a/src/bin/pg_amcheck/pg_amcheck.c
+++ b/src/bin/pg_amcheck/pg_amcheck.c
@@ -989,6 +989,7 @@ should_processing_continue(PGresult *res)
 		case PGRES_COPY_IN:
 		case PGRES_COPY_BOTH:
 		case PGRES_SINGLE_TUPLE:
+		case PGRES_TUPLES_CHUNK:
 		case PGRES_PIPELINE_SYNC:
 		case PGRES_PIPELINE_ABORTED:
 			return false;
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index 088592deb1..20effc8337 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -193,3 +193,4 @@ PQsendClosePrepared       190
 PQsendClosePortal         191
 PQchangePassword          192
 PQsendPipelineSync        193
+PQsetChunkedRowsMode	  194
diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index c02a9180b2..b9a73b583e 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -41,7 +41,8 @@ char	   *const pgresStatus[] = {
 	"PGRES_COPY_BOTH",
 	"PGRES_SINGLE_TUPLE",
 	"PGRES_PIPELINE_SYNC",
-	"PGRES_PIPELINE_ABORTED"
+	"PGRES_PIPELINE_ABORTED",
+	"PGRES_TUPLES_CHUNK"
 };
 
 /* We return this if we're unable to make a PGresult at all */
@@ -83,7 +84,7 @@ static int	check_field_number(const PGresult *res, int field_num);
 static void pqPipelineProcessQueue(PGconn *conn);
 static int	pqPipelineSyncInternal(PGconn *conn, bool immediate_flush);
 static int	pqPipelineFlush(PGconn *conn);
-
+static bool canChangeRowMode(PGconn *conn);
 
 /* ----------------
  * Space management for PGresult.
@@ -200,6 +201,7 @@ PQmakeEmptyPGresult(PGconn *conn, ExecStatusType status)
 			case PGRES_COPY_IN:
 			case PGRES_COPY_BOTH:
 			case PGRES_SINGLE_TUPLE:
+			case PGRES_TUPLES_CHUNK:
 				/* non-error cases */
 				break;
 			default:
@@ -913,8 +915,9 @@ pqPrepareAsyncResult(PGconn *conn)
 	/*
 	 * Replace conn->result with next_result, if any.  In the normal case
 	 * there isn't a next result and we're just dropping ownership of the
-	 * current result.  In single-row mode this restores the situation to what
-	 * it was before we created the current single-row result.
+	 * current result.  In single-row and chunked modes this restores the
+	 * situation to what it was before we created the current single-row or
+	 * chunk-of-rows result.
 	 */
 	conn->result = conn->next_result;
 	conn->error_result = false; /* next_result is never an error */
@@ -1200,10 +1203,11 @@ pqSaveParameterStatus(PGconn *conn, const char *name, const char *value)
  * (Such a string should already be translated via libpq_gettext().)
  * If it is left NULL, the error is presumed to be "out of memory".
  *
- * In single-row mode, we create a new result holding just the current row,
- * stashing the previous result in conn->next_result so that it becomes
- * active again after pqPrepareAsyncResult().  This allows the result metadata
- * (column descriptions) to be carried forward to each result row.
+ * In single-row or chunked mode, we create a new result holding just the
+ * current set of rows, stashing the previous result in conn->next_result so
+ * that it becomes active again after pqPrepareAsyncResult().  This allows the
+ * result metadata (column descriptions) to be carried forward to each result
+ * row.
  */
 int
 pqRowProcessor(PGconn *conn, const char **errmsgp)
@@ -1228,6 +1232,28 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 		if (!res)
 			return 0;
 	}
+	else if (conn->rowsChunkSize > 0)
+	{
+		/*
+		 * In chunked mode, make a new PGresult that will hold N rows; the
+		 * original conn->result is left unchanged, as in the single-row mode.
+		 */
+		if (!conn->chunk_result)
+		{
+			/* Allocate and initialize the result to hold a chunk of rows */
+			res = PQcopyResult(res,
+							   PG_COPYRES_ATTRS | PG_COPYRES_EVENTS |
+							   PG_COPYRES_NOTICEHOOKS);
+			if (!res)
+				return 0;
+			/* Change result status to special chunk-of-rows value */
+			res->resultStatus = PGRES_TUPLES_CHUNK;
+			/* Keep this result to reuse for the next rows of the chunk */
+			conn->chunk_result = res;
+		}
+		else
+			res = conn->chunk_result;	/* Use the current chunk */
+	}
 
 	/*
 	 * Basically we just allocate space in the PGresult for each field and
@@ -1290,6 +1316,21 @@ pqRowProcessor(PGconn *conn, const char **errmsgp)
 		conn->asyncStatus = PGASYNC_READY_MORE;
 	}
 
+	/*
+	 * In chunked mode, if the count has reached the requested limit, make the
+	 * rows of the current chunk available immediately.
+	 */
+	else if (conn->rowsChunkSize > 0 && res->ntups >= conn->rowsChunkSize)
+	{
+		/* Stash old result for re-use later */
+		conn->next_result = conn->result;
+		conn->result = res;
+		/* Do not reuse that chunk of results */
+		conn->chunk_result = NULL;
+		/* And mark the result ready to return */
+		conn->asyncStatus = PGASYNC_READY_MORE;
+	}
+
 	return 1;
 
 fail:
@@ -1745,8 +1786,9 @@ PQsendQueryStart(PGconn *conn, bool newQuery)
 		 */
 		pqClearAsyncResult(conn);
 
-		/* reset single-row processing mode */
+		/* reset row-by-row and chunked processing modes */
 		conn->singleRowMode = false;
+		conn->rowsChunkSize = 0;
 	}
 
 	/* ready to send command message */
@@ -1930,25 +1972,51 @@ sendFailed:
  */
 int
 PQsetSingleRowMode(PGconn *conn)
+{
+	if (canChangeRowMode(conn))
+	{
+		conn->singleRowMode = true;
+		return 1;
+	}
+	else
+		return 0;
+}
+
+/*
+ * Select chunked results processing mode
+ */
+int
+PQsetChunkedRowsMode(PGconn *conn, int chunkSize)
+{
+	if (chunkSize >= 0 && canChangeRowMode(conn))
+	{
+		conn->rowsChunkSize = chunkSize;
+		return 1;
+	}
+	else
+		return 0;
+}
+
+static
+bool
+canChangeRowMode(PGconn *conn)
 {
 	/*
-	 * Only allow setting the flag when we have launched a query and not yet
-	 * received any results.
+	 * Only allow setting the row-by-row or by-chunks modes when we have
+	 * launched a query and not yet received any results.
 	 */
 	if (!conn)
-		return 0;
+		return false;
 	if (conn->asyncStatus != PGASYNC_BUSY)
-		return 0;
+		return false;
 	if (!conn->cmd_queue_head ||
 		(conn->cmd_queue_head->queryclass != PGQUERY_SIMPLE &&
 		 conn->cmd_queue_head->queryclass != PGQUERY_EXTENDED))
-		return 0;
+		return false;
 	if (pgHavePendingResult(conn))
-		return 0;
+		return false;
 
-	/* OK, set flag */
-	conn->singleRowMode = true;
-	return 1;
+	return true;
 }
 
 /*
@@ -2115,6 +2183,16 @@ PQgetResult(PGconn *conn)
 			break;
 
 		case PGASYNC_READY:
+			/*
+			 * If there is a pending chunk of results, return it
+			 */
+			if (conn->chunk_result != NULL)
+			{
+				res = conn->chunk_result;
+				conn->chunk_result = NULL;
+				break;
+			}
+
 			res = pqPrepareAsyncResult(conn);
 
 			/* Advance the queue as appropriate */
@@ -3173,10 +3251,11 @@ pqPipelineProcessQueue(PGconn *conn)
 	}
 
 	/*
-	 * Reset single-row processing mode.  (Client has to set it up for each
+	 * Reset to full result sets mode. (Client has to set it up for each
 	 * query, if desired.)
 	 */
 	conn->singleRowMode = false;
+	conn->rowsChunkSize = 0;
 
 	/*
 	 * If there are no further commands to process in the queue, get us in
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index defc415fa3..21de21fbe8 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -109,8 +109,9 @@ typedef enum
 	PGRES_COPY_BOTH,			/* Copy In/Out data transfer in progress */
 	PGRES_SINGLE_TUPLE,			/* single tuple from larger resultset */
 	PGRES_PIPELINE_SYNC,		/* pipeline synchronization point */
-	PGRES_PIPELINE_ABORTED		/* Command didn't run because of an abort
+	PGRES_PIPELINE_ABORTED,		/* Command didn't run because of an abort
 								 * earlier in a pipeline */
+	PGRES_TUPLES_CHUNK			/* set of tuples from larger resultset */
 } ExecStatusType;
 
 typedef enum
@@ -463,6 +464,7 @@ extern int	PQsendQueryPrepared(PGconn *conn,
 								const int *paramFormats,
 								int resultFormat);
 extern int	PQsetSingleRowMode(PGconn *conn);
+extern int	PQsetChunkedRowsMode(PGconn *conn, int chunkSize);
 extern PGresult *PQgetResult(PGconn *conn);
 
 /* Routines for managing an asynchronous query */
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index ff8e0dce77..76130e2912 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -431,6 +431,8 @@ struct pg_conn
 								 * sending semantics */
 	PGpipelineStatus pipelineStatus;	/* status of pipeline mode */
 	bool		singleRowMode;	/* return current query result row-by-row? */
+	int			rowsChunkSize;	/* non-zero to return query results by chunks
+								 * not exceeding that number of rows */
 	char		copy_is_binary; /* 1 = copy binary, 0 = copy text */
 	int			copy_already_done;	/* # bytes already returned in COPY OUT */
 	PGnotify   *notifyHead;		/* oldest unreported Notify msg */
@@ -536,7 +538,10 @@ struct pg_conn
 	 */
 	PGresult   *result;			/* result being constructed */
 	bool		error_result;	/* do we need to make an ERROR result? */
-	PGresult   *next_result;	/* next result (used in single-row mode) */
+	PGresult   *next_result;	/* next result (used in single-row and
+								 * by-chunks modes) */
+	PGresult   *chunk_result;	/* current chunk of results (limited to
+								 * rowsChunkSize) */
 
 	/* Assorted state for SASL, SSL, GSS, etc */
 	const pg_fe_sasl_mech *sasl;
-- 
2.34.1