From 814fa7bdafba2c004b923ce577c0c6f53cfc5387 Mon Sep 17 00:00:00 2001
From: Tatsuo Ishii <ishii@postgresql.org>
Date: Thu, 19 Sep 2024 19:14:07 +0900
Subject: [PATCH v3] Add memory/disk usage for more executor nodes.

This commit is similar to 95d6e9af07, expanding the idea to CTE scan,
table function scan and recursive union scan nodes so that the maximum
tuplestore memory or disk usage is shown with EXPLAIN ANALYZE command.
---
 src/backend/commands/explain.c        | 85 +++++++++++++++++++++++++++
 src/test/regress/expected/explain.out | 38 ++++++++++++
 src/test/regress/sql/explain.sql      |  6 ++
 3 files changed, 129 insertions(+)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index aaec439892..725c3e88db 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -129,6 +129,9 @@ static void show_incremental_sort_info(IncrementalSortState *incrsortstate,
 static void show_hash_info(HashState *hashstate, ExplainState *es);
 static void show_material_info(MaterialState *mstate, ExplainState *es);
 static void show_windowagg_info(WindowAggState *winstate, ExplainState *es);
+static void show_ctescan_info(CteScanState *ctescanstate, ExplainState *es);
+static void show_table_func_scan_info(TableFuncScanState *tscanstate, ExplainState *es);
+static void show_recursive_union_info(RecursiveUnionState *rstate, ExplainState *es);
 static void show_memoize_info(MemoizeState *mstate, List *ancestors,
 							  ExplainState *es);
 static void show_hashagg_info(AggState *aggstate, ExplainState *es);
@@ -2046,6 +2049,8 @@ ExplainNode(PlanState *planstate, List *ancestors,
 			if (plan->qual)
 				show_instrumentation_count("Rows Removed by Filter", 1,
 										   planstate, es);
+			if (IsA(plan, CteScan))
+				show_ctescan_info(castNode(CteScanState, planstate), es);
 			break;
 		case T_Gather:
 			{
@@ -2127,6 +2132,7 @@ ExplainNode(PlanState *planstate, List *ancestors,
 			if (plan->qual)
 				show_instrumentation_count("Rows Removed by Filter", 1,
 										   planstate, es);
+			show_table_func_scan_info(castNode(TableFuncScanState, planstate), es);
 			break;
 		case T_TidScan:
 			{
@@ -2278,6 +2284,10 @@ ExplainNode(PlanState *planstate, List *ancestors,
 			show_memoize_info(castNode(MemoizeState, planstate), ancestors,
 							  es);
 			break;
+		case T_RecursiveUnion:
+			show_recursive_union_info(castNode(RecursiveUnionState, planstate),
+									  es);
+			break;
 		default:
 			break;
 	}
@@ -3411,6 +3421,81 @@ show_windowagg_info(WindowAggState *winstate, ExplainState *es)
 	show_storage_info(tupstore, es);
 }
 
+/*
+ * Show information on CTE Scan node, storage method and maximum memory/disk
+ * space used.
+ */
+static void
+show_ctescan_info(CteScanState *ctescanstate, ExplainState *es)
+{
+	Tuplestorestate *tupstore = ctescanstate->leader->cte_table;
+
+	if (!es->analyze || tupstore == NULL)
+		return;
+
+	show_storage_info(tupstore, es);
+}
+
+/*
+ * Show information on Table Function Scan node, storage method and maximum
+ * memory/disk space used.
+ */
+static void
+show_table_func_scan_info(TableFuncScanState *tscanstate, ExplainState *es)
+{
+	Tuplestorestate *tupstore = tscanstate->tupstore;
+
+	if (!es->analyze || tupstore == NULL)
+		return;
+
+	show_storage_info(tupstore, es);
+}
+
+/*
+ * Show information on Recursive Union node, storage method and maximum
+ * memory/disk space used.
+ */
+static void
+show_recursive_union_info(RecursiveUnionState *rstate, ExplainState *es)
+{
+	char	   *maxStorageType,
+			   *tempStorageType;
+	int64		maxSpaceUsed,
+				maxSpaceUsedKB,
+				tempSpaceUsed;
+
+	if (!es->analyze)
+		return;
+
+	/*
+	 * Recursive union node uses two tuplestores.  We employ the storage type
+	 * from one of them which consumed more memory/disk than the other.  The
+	 * storage size is sum of the two.
+	 */
+	tuplestore_get_stats(rstate->working_table, &tempStorageType, &tempSpaceUsed);
+	tuplestore_get_stats(rstate->intermediate_table, &maxStorageType, &maxSpaceUsed);
+
+	if (tempSpaceUsed > maxSpaceUsed)
+		maxStorageType = tempStorageType;
+
+	maxSpaceUsed += tempSpaceUsed;
+	maxSpaceUsedKB = BYTES_TO_KILOBYTES(maxSpaceUsed);
+
+	if (es->format != EXPLAIN_FORMAT_TEXT)
+	{
+		ExplainPropertyText("Storage", maxStorageType, es);
+		ExplainPropertyInteger("Maximum Storage", "kB", maxSpaceUsedKB, es);
+	}
+	else
+	{
+		ExplainIndentText(es);
+		appendStringInfo(es->str,
+						 "Storage: %s  Maximum Storage: " INT64_FORMAT "kB\n",
+						 maxStorageType,
+						 maxSpaceUsedKB);
+	}
+}
+
 /*
  * Show information on memoize hits/misses/evictions and memory usage.
  */
diff --git a/src/test/regress/expected/explain.out b/src/test/regress/expected/explain.out
index d01c304c24..87bfaa6d79 100644
--- a/src/test/regress/expected/explain.out
+++ b/src/test/regress/expected/explain.out
@@ -729,3 +729,41 @@ select explain_filter('explain (analyze,costs off) select sum(n) over(partition
 (8 rows)
 
 reset work_mem;
+-- Test tuplestore storage usage in CTE scan
+select explain_filter('explain (analyze,costs off) with w(n) as (select n from generate_series(1,10) a(n)) select sum(n) from w');
+                                 explain_filter                                 
+--------------------------------------------------------------------------------
+ Aggregate (actual time=N.N..N.N rows=N loops=N)
+   ->  Function Scan on generate_series a (actual time=N.N..N.N rows=N loops=N)
+ Planning Time: N.N ms
+ Execution Time: N.N ms
+(4 rows)
+
+-- Test tuplestore storage usage in table function scan
+select explain_filter('explain (analyze,costs off) select * from json_table(''[]'', ''strict $.a'' columns (i int path ''$''))');
+                              explain_filter                               
+---------------------------------------------------------------------------
+ Table Function Scan on "json_table" (actual time=N.N..N.N rows=N loops=N)
+   Storage: Memory  Maximum Storage: NkB
+ Planning Time: N.N ms
+ Execution Time: N.N ms
+(4 rows)
+
+-- Test tuplestore storage usage in recurive union scan
+select explain_filter('explain (analyze,costs off) with recursive t(n) as (values(1) union all select n+1 from t where n < 10) select sum(n) from t');
+                               explain_filter                                
+-----------------------------------------------------------------------------
+ Aggregate (actual time=N.N..N.N rows=N loops=N)
+   CTE t
+     ->  Recursive Union (actual time=N.N..N.N rows=N loops=N)
+           Storage: Memory  Maximum Storage: NkB
+           ->  Result (actual time=N.N..N.N rows=N loops=N)
+           ->  WorkTable Scan on t t_1 (actual time=N.N..N.N rows=N loops=N)
+                 Filter: (n < N)
+                 Rows Removed by Filter: N
+   ->  CTE Scan on t (actual time=N.N..N.N rows=N loops=N)
+         Storage: Memory  Maximum Storage: NkB
+ Planning Time: N.N ms
+ Execution Time: N.N ms
+(12 rows)
+
diff --git a/src/test/regress/sql/explain.sql b/src/test/regress/sql/explain.sql
index b861e2b53d..5d5c68443e 100644
--- a/src/test/regress/sql/explain.sql
+++ b/src/test/regress/sql/explain.sql
@@ -178,3 +178,9 @@ select explain_filter('explain (analyze,costs off) select sum(n) over() from gen
 -- Test tuplestore storage usage in Window aggregate (memory and disk case, final result is disk)
 select explain_filter('explain (analyze,costs off) select sum(n) over(partition by m) from (SELECT n < 3 as m, n from generate_series(1,2000) a(n))');
 reset work_mem;
+-- Test tuplestore storage usage in CTE scan
+select explain_filter('explain (analyze,costs off) with w(n) as (select n from generate_series(1,10) a(n)) select sum(n) from w');
+-- Test tuplestore storage usage in table function scan
+select explain_filter('explain (analyze,costs off) select * from json_table(''[]'', ''strict $.a'' columns (i int path ''$''))');
+-- Test tuplestore storage usage in recurive union scan
+select explain_filter('explain (analyze,costs off) with recursive t(n) as (values(1) union all select n+1 from t where n < 10) select sum(n) from t');
-- 
2.25.1