From 9b493c782685a05ed17f3c2dd3ff9673d8f5f6d4 Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathan@postgresql.org>
Date: Mon, 31 Mar 2025 20:48:07 -0500
Subject: [PATCH v12n 3/3] pg_dump: Batch queries for retrieving attribute
 statistics.

Currently, pg_dump gathers attribute statistics with a query per
relation, which can cause pg_dump to take significantly longer,
especially when there are many tables.  This commit improves
matters by gathering attribute statistics for 64 relations at a
time.  Some simple testing showed this was the ideal batch size,
but performance may vary depending on workload.

To construct the next set of relations for the query, we scan
through the TOC list for relevant entries.  Ordinarily, we can stop
issuing queries once we reach the end of the list.  However,
custom-format dumps that include data run the statistics queries
twice (thanks to commit XXXXXXXXXX), so we allow a second pass in
that case.

This change increases the memory usage of pg_dump a bit, but that
isn't expected to be too egregious and is arguably well worth the
trade-off.

Author: Corey Huinker <corey.huinker@gmail.com>
Discussion: https://postgr.es/m/CADkLM%3Dc%2Br05srPy9w%2B-%2BnbmLEo15dKXYQ03Q_xyK%2BriJerigLQ%40mail.gmail.com
---
 src/bin/pg_dump/pg_dump.c        | 131 +++++++++++++++++++++++++++----
 src/tools/pgindent/typedefs.list |   1 +
 2 files changed, 115 insertions(+), 17 deletions(-)

diff --git a/src/bin/pg_dump/pg_dump.c b/src/bin/pg_dump/pg_dump.c
index 9fa2cb0672e..5506286f26f 100644
--- a/src/bin/pg_dump/pg_dump.c
+++ b/src/bin/pg_dump/pg_dump.c
@@ -143,6 +143,13 @@ typedef enum OidOptions
 	zeroAsNone = 4,
 } OidOptions;
 
+typedef struct
+{
+	PGresult   *res;			/* most recent fetchAttributeStats() result */
+	int			idx;			/* first un-consumed row of results */
+	TocEntry   *te;				/* next TOC entry to search */
+} AttStatsCache;
+
 /* global decls */
 static bool dosync = true;		/* Issue fsync() to make dump durable on disk. */
 
@@ -209,6 +216,11 @@ static int	nbinaryUpgradeClassOids = 0;
 static SequenceItem *sequences = NULL;
 static int	nsequences = 0;
 
+static AttStatsCache attStats;
+
+/* Maximum number of relations to fetch in a fetchAttributeStats() call. */
+#define MAX_ATTR_STATS_RELS 64
+
 /*
  * The default number of rows per INSERT when
  * --inserts is specified without --rows-per-insert
@@ -10559,6 +10571,81 @@ appendNamedArgument(PQExpBuffer out, Archive *fout, const char *argname,
 	appendPQExpBuffer(out, "::%s", argtype);
 }
 
+/*
+ * fetchAttributeStats --
+ *
+ * Fetch next batch of rows for getAttributeStats().
+ */
+static void
+fetchAttributeStats(Archive *fout)
+{
+	ArchiveHandle *AH = (ArchiveHandle *) fout;
+	PQExpBuffer nspnames = createPQExpBuffer();
+	PQExpBuffer relnames = createPQExpBuffer();
+	int			count = 0;
+	static bool restarted;
+
+	/* free last result set */
+	PQclear(attStats.res);
+	attStats.res = NULL;
+
+	/*
+	 * If we're just starting, set our TOC pointer.
+	 */
+	if (!attStats.te)
+		attStats.te = AH->toc->next;
+
+	/*
+	 * Restart the TOC scan once for custom-format dumps that include data.
+	 * This is necessary because we'll call WriteToc() twice in that case.
+	 */
+	if (!restarted && attStats.te == AH->toc &&
+		AH->format == archCustom && fout->dopt->dumpData)
+	{
+		attStats.te = AH->toc->next;
+		restarted = true;
+	}
+
+	/*
+	 * Walk ahead looking for stats entries that are active in this section.
+	 */
+	while (attStats.te != AH->toc && count < MAX_ATTR_STATS_RELS)
+	{
+		if (attStats.te->reqs &&
+			strcmp(attStats.te->desc, "STATISTICS DATA") == 0)
+		{
+			RelStatsInfo *rsinfo = (RelStatsInfo *) attStats.te->defnDumperArg;
+
+			appendPQExpBuffer(nspnames, "%s%s", count ? "," : "",
+							  fmtId(rsinfo->dobj.namespace->dobj.name));
+			appendPQExpBuffer(relnames, "%s%s", count ? "," : "",
+							  fmtId(rsinfo->dobj.name));
+			count++;
+		}
+
+		attStats.te = attStats.te->next;
+	}
+
+	/*
+	 * Execute the query for the next batch of relations.
+	 */
+	if (count > 0)
+	{
+		PQExpBuffer query = createPQExpBuffer();
+
+		appendPQExpBuffer(query, "EXECUTE getAttributeStats("
+						  "'{%s}'::pg_catalog.name[],"
+						  "'{%s}'::pg_catalog.name[])",
+						  nspnames->data, relnames->data);
+		attStats.res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
+		attStats.idx = 0;
+		destroyPQExpBuffer(query);
+	}
+
+	destroyPQExpBuffer(nspnames);
+	destroyPQExpBuffer(relnames);
+}
+
 /*
  * dumpRelationStats_dumper --
  *
@@ -10575,6 +10662,8 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 	PGresult   *res;
 	PQExpBuffer query;
 	PQExpBuffer out;
+	int			i_schemaname;
+	int			i_tablename;
 	int			i_attname;
 	int			i_inherited;
 	int			i_null_frac;
@@ -10596,8 +10685,8 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 	if (!fout->is_prepared[PREPQUERY_GETATTRIBUTESTATS])
 	{
 		appendPQExpBufferStr(query,
-							 "PREPARE getAttributeStats(pg_catalog.name, pg_catalog.name) AS\n"
-							 "SELECT s.attname, s.inherited, "
+							 "PREPARE getAttributeStats(pg_catalog.name[], pg_catalog.name[]) AS\n"
+							 "SELECT s.schemaname, s.tablename, s.attname, s.inherited, "
 							 "s.null_frac, s.avg_width, s.n_distinct, "
 							 "s.most_common_vals, s.most_common_freqs, "
 							 "s.histogram_bounds, s.correlation, "
@@ -10617,9 +10706,11 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 
 		appendPQExpBufferStr(query,
 							 "FROM pg_catalog.pg_stats s "
-							 "WHERE s.schemaname = $1 "
-							 "AND s.tablename = $2 "
-							 "ORDER BY s.attname, s.inherited");
+							 "JOIN unnest($1, $2) WITH ORDINALITY AS u (schemaname, tablename, ord) "
+							 "ON s.schemaname = u.schemaname "
+							 "AND s.tablename = u.tablename "
+							 "WHERE s.tablename = ANY($2) "
+							 "ORDER BY u.ord, s.attname, s.inherited");
 
 		ExecuteSqlStatement(fout, query->data);
 
@@ -10649,16 +10740,16 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 
 	appendPQExpBufferStr(out, "\n);\n");
 
+	/*
+	 * If the attribute stats cache is uninitialized or exhausted, fetch the
+	 * next batch of attributes statistics.
+	 */
+	if (attStats.idx >= PQntuples(attStats.res))
+		fetchAttributeStats(fout);
+	res = attStats.res;
 
-	/* fetch attribute stats */
-	appendPQExpBufferStr(query, "EXECUTE getAttributeStats(");
-	appendStringLiteralAH(query, dobj->namespace->dobj.name, fout);
-	appendPQExpBufferStr(query, ", ");
-	appendStringLiteralAH(query, dobj->name, fout);
-	appendPQExpBufferStr(query, ");");
-
-	res = ExecuteSqlQuery(fout, query->data, PGRES_TUPLES_OK);
-
+	i_schemaname = PQfnumber(res, "schemaname");
+	i_tablename = PQfnumber(res, "tablename");
 	i_attname = PQfnumber(res, "attname");
 	i_inherited = PQfnumber(res, "inherited");
 	i_null_frac = PQfnumber(res, "null_frac");
@@ -10676,9 +10767,17 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 	i_range_bounds_histogram = PQfnumber(res, "range_bounds_histogram");
 
 	/* restore attribute stats */
-	for (int rownum = 0; rownum < PQntuples(res); rownum++)
+	for (; attStats.idx < PQntuples(res); attStats.idx++)
 	{
 		const char *attname;
+		int			rownum = attStats.idx;
+
+		/*
+		 * Stop if the next stat row in our cache isn't for this relation.
+		 */
+		if (strcmp(dobj->name, PQgetvalue(res, rownum, i_tablename)) != 0 ||
+			strcmp(dobj->namespace->dobj.name, PQgetvalue(res, rownum, i_schemaname)) != 0)
+			break;
 
 		appendPQExpBufferStr(out, "SELECT * FROM pg_catalog.pg_restore_attribute_stats(\n");
 		appendPQExpBuffer(out, "\t'version', '%u'::integer,\n",
@@ -10769,8 +10868,6 @@ dumpRelationStats_dumper(Archive *fout, const void *userArg)
 		appendPQExpBufferStr(out, "\n);\n");
 	}
 
-	PQclear(res);
-
 	destroyPQExpBuffer(query);
 	ret = out->data;
 	pg_free(out);
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index b66cecd8799..2719a6642ad 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -162,6 +162,7 @@ AsyncQueueEntry
 AsyncRequest
 ATAlterConstraint
 AttInMetadata
+AttStatsCache
 AttStatsSlot
 AttoptCacheEntry
 AttoptCacheKey
-- 
2.39.5 (Apple Git-154)

