From 47ef2a746e6f17a86dc468565db3a825cf0526fc Mon Sep 17 00:00:00 2001
From: jian he <jian.universality@gmail.com>
Date: Fri, 16 Aug 2024 18:08:39 +0800
Subject: [PATCH v5 1/1] minor refactor pg_stats_vacuum and sub-routines

---
 src/backend/utils/adt/pgstatfuncs.c | 45 +++++++++--------------------
 1 file changed, 13 insertions(+), 32 deletions(-)

diff --git a/src/backend/utils/adt/pgstatfuncs.c b/src/backend/utils/adt/pgstatfuncs.c
index 0c490ba5f1..adb04795ff 100644
--- a/src/backend/utils/adt/pgstatfuncs.c
+++ b/src/backend/utils/adt/pgstatfuncs.c
@@ -2078,8 +2078,8 @@ fetch_dbstat_tabentry(Oid dbid, Oid relid)
 }
 
 static void
-tuplestore_put_for_database(Oid dbid, Tuplestorestate *tupstore,
-			   TupleDesc tupdesc, PgStatShared_Database *dbentry, int ncolumns)
+tuplestore_put_for_database(Oid dbid, ReturnSetInfo *rsinfo,
+							PgStatShared_Database *dbentry)
 {
 	Datum		values[EXTVACDBSTAT_COLUMNS];
 	bool		nulls[EXTVACDBSTAT_COLUMNS];
@@ -2113,15 +2113,14 @@ tuplestore_put_for_database(Oid dbid, Tuplestorestate *tupstore,
 	values[i++] = Float8GetDatum(dbentry->stats.vacuum_ext.total_time);
 	values[i++] = Int32GetDatum(dbentry->stats.vacuum_ext.interrupts);
 
+	Assert(i == rsinfo->setDesc->natts);
 
-	Assert(i == ncolumns);
-
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
 }
 
 static void
-tuplestore_put_for_relation(Oid relid, Tuplestorestate *tupstore,
-			   TupleDesc tupdesc, PgStat_StatTabEntry *tabentry, int ncolumns)
+tuplestore_put_for_relation(Oid relid, ReturnSetInfo *rsinfo,
+							PgStat_StatTabEntry *tabentry)
 {
 	Datum		values[EXTVACSTAT_COLUMNS];
 	bool		nulls[EXTVACSTAT_COLUMNS];
@@ -2179,9 +2178,8 @@ tuplestore_put_for_relation(Oid relid, Tuplestorestate *tupstore,
 	values[i++] = Float8GetDatum(tabentry->vacuum_ext.total_time);
 	values[i++] = Int32GetDatum(tabentry->vacuum_ext.interrupts);
 
-	Assert(i == ncolumns);
-
-	tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+	Assert(i == rsinfo->setDesc->natts);
+	tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
 }
 
 /*
@@ -2191,10 +2189,6 @@ static Datum
 pg_stats_vacuum(FunctionCallInfo fcinfo, ExtVacReportType type, int ncolumns)
 {
 	ReturnSetInfo		   *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
-	MemoryContext			per_query_ctx;
-	MemoryContext			oldcontext;
-	Tuplestorestate		   *tupstore;
-	TupleDesc				tupdesc;
 	Oid						dbid = PG_GETARG_OID(0);
 	PgStat_StatTabEntry    *tabentry;
 
@@ -2205,22 +2199,9 @@ pg_stats_vacuum(FunctionCallInfo fcinfo, ExtVacReportType type, int ncolumns)
 		ereport(ERROR,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("set-valued function called in context that cannot accept a set")));
-	/* Switch to long-lived context to create the returned data structures */
-	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
-	oldcontext = MemoryContextSwitchTo(per_query_ctx);
 
-	/* Build a tuple descriptor for our result type */
-	if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
-		elog(ERROR, "return type must be a row type");
-
-	Assert(tupdesc->natts == ncolumns);
-
-	tupstore = tuplestore_begin_heap(true, false, work_mem);
-	Assert (tupstore != NULL);
-	rsinfo->setResult = tupstore;
-	rsinfo->setDesc = tupdesc;
-
-	MemoryContextSwitchTo(oldcontext);
+	Assert(rsinfo->setDesc->natts == ncolumns);
+	Assert(rsinfo->setResult != NULL);
 
 	if (type == PGSTAT_EXTVAC_INDEX || type == PGSTAT_EXTVAC_HEAP)
 	{
@@ -2234,7 +2215,7 @@ pg_stats_vacuum(FunctionCallInfo fcinfo, ExtVacReportType type, int ncolumns)
 				/* Table don't exists or isn't an heap relation. */
 				PG_RETURN_NULL();
 
-			tuplestore_put_for_relation(relid, tupstore, tupdesc, tabentry, ncolumns);
+			tuplestore_put_for_relation(relid, rsinfo, tabentry);
 		}
 		else
 		{
@@ -2259,7 +2240,7 @@ pg_stats_vacuum(FunctionCallInfo fcinfo, ExtVacReportType type, int ncolumns)
 				reloid = entry->key.objoid;
 
 				if (tabentry != NULL && tabentry->vacuum_ext.type == type)
-					tuplestore_put_for_relation(reloid, tupstore, tupdesc, tabentry, ncolumns);
+					tuplestore_put_for_relation(reloid, rsinfo, tabentry, ncolumns);
 			}
 		}
 	}
@@ -2282,7 +2263,7 @@ pg_stats_vacuum(FunctionCallInfo fcinfo, ExtVacReportType type, int ncolumns)
 				/* Table doesn't exist or isn't a heap relation */
 				PG_RETURN_NULL();
 
-			tuplestore_put_for_database(dbid, tupstore, tupdesc, dbentry, ncolumns);
+			tuplestore_put_for_database(dbid, rsinfo, dbentry, ncolumns);
 			pgstat_unlock_entry(entry_ref);
 		}
 		else
-- 
2.34.1