From 00c7546659e305be45dbeb13a14bcde5066be76f Mon Sep 17 00:00:00 2001
From: Tatsuo Ishii <ishii@postgresql.org>
Date: Sat, 6 Jul 2024 19:48:30 +0900
Subject: [PATCH v1] Add memory/disk usage for WindowAgg nodes in EXPLAIN.

---
 src/backend/commands/explain.c       | 37 ++++++++++++++++++++++++++++
 src/backend/executor/nodeWindowAgg.c | 19 ++++++++++++++
 src/include/nodes/execnodes.h        |  2 ++
 3 files changed, 58 insertions(+)

diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index 1e80fd8b68..177687ea80 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -126,6 +126,7 @@ static void show_incremental_sort_info(IncrementalSortState *incrsortstate,
 									   ExplainState *es);
 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_memoize_info(MemoizeState *mstate, List *ancestors,
 							  ExplainState *es);
 static void show_hashagg_info(AggState *aggstate, ExplainState *es);
@@ -2215,6 +2216,7 @@ ExplainNode(PlanState *planstate, List *ancestors,
 										   planstate, es);
 			show_upper_qual(((WindowAgg *) plan)->runConditionOrig,
 							"Run Condition", planstate, ancestors, es);
+			show_windowagg_info(castNode(WindowAggState, planstate), es);
 			break;
 		case T_Group:
 			show_group_keys(castNode(GroupState, planstate), ancestors, es);
@@ -3362,6 +3364,41 @@ show_material_info(MaterialState *mstate, ExplainState *es)
 	}
 }
 
+/*
+ * Show information on WindowAgg node, storage method and maximum memory/disk
+ * space used.
+ */
+static void
+show_windowagg_info(WindowAggState *winstate, ExplainState *es)
+{
+	const char *storageType;
+	int64		spaceUsedKB;
+
+	/*
+	 * Nothing to show if ANALYZE option wasn't used or if execution didn't
+	 * get as far as creating the tuplestore.
+	 */
+	if (!es->analyze || winstate->storageType == NULL)
+		return;
+
+	storageType = winstate->storageType;
+	spaceUsedKB = BYTES_TO_KILOBYTES(winstate->storageSize);
+
+	if (es->format != EXPLAIN_FORMAT_TEXT)
+	{
+		ExplainPropertyText("Storage", storageType, es);
+		ExplainPropertyInteger("Maximum Storage", "kB", spaceUsedKB, es);
+	}
+	else
+	{
+		ExplainIndentText(es);
+		appendStringInfo(es->str,
+						 "Storage: %s  Maximum Storage: " INT64_FORMAT "kB\n",
+						 storageType,
+						 spaceUsedKB);
+	}
+}
+
 /*
  * Show information on memoize hits/misses/evictions and memory usage.
  */
diff --git a/src/backend/executor/nodeWindowAgg.c b/src/backend/executor/nodeWindowAgg.c
index 3221fa1522..bcfe144511 100644
--- a/src/backend/executor/nodeWindowAgg.c
+++ b/src/backend/executor/nodeWindowAgg.c
@@ -1360,7 +1360,23 @@ release_partition(WindowAggState *winstate)
 	}
 
 	if (winstate->buffer)
+	{
+		int64	spaceUsed = tuplestore_space_used(winstate->buffer);
+
+		/*
+		 * We want to track the max mem/disk usage so that we can use the info
+		 * in EXPLAIN (ANALYZE).
+		 */
+		if (spaceUsed > winstate->storageSize)
+		{
+			if (winstate->storageType != NULL)
+				pfree(winstate->storageType);
+			winstate->storageType =
+				pstrdup(tuplestore_storage_type_name(winstate->buffer));
+			winstate->storageSize = spaceUsed;
+		}
 		tuplestore_end(winstate->buffer);
+	}
 	winstate->buffer = NULL;
 	winstate->partition_spooled = false;
 }
@@ -2671,6 +2687,9 @@ ExecInitWindowAgg(WindowAgg *node, EState *estate, int eflags)
 	winstate->partition_spooled = false;
 	winstate->more_partitions = false;
 
+	winstate->storageType = NULL;
+	winstate->storageSize = 0;
+
 	return winstate;
 }
 
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index b62c96f206..7a3dfa2bc3 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -2567,6 +2567,8 @@ typedef struct WindowAggState
 	ExprState  *partEqfunction; /* equality funcs for partition columns */
 	ExprState  *ordEqfunction;	/* equality funcs for ordering columns */
 	Tuplestorestate *buffer;	/* stores rows of current partition */
+	int64		storageSize;	/* max storage size in buffer */
+	char		*storageType;	/* the storage type above */
 	int			current_ptr;	/* read pointer # for current row */
 	int			framehead_ptr;	/* read pointer # for frame head, if used */
 	int			frametail_ptr;	/* read pointer # for frame tail, if used */
-- 
2.25.1