From ed65970ccd7953aebc86111333f55f03260885e1 Mon Sep 17 00:00:00 2001
From: Nikhil Kumar Veldanda <nikhilkv@amazon.com>
Date: Mon, 14 Apr 2025 21:51:43 +0000
Subject: [PATCH v11 5/7] generate and cleanup dictionaries using vacuum

---
 src/backend/catalog/pg_zstd_dictionaries.c | 56 ++++++++++++++++++++++
 src/backend/commands/vacuum.c              |  4 ++
 src/include/catalog/pg_zstd_dictionaries.h |  1 +
 3 files changed, 61 insertions(+)

diff --git a/src/backend/catalog/pg_zstd_dictionaries.c b/src/backend/catalog/pg_zstd_dictionaries.c
index 5ae8ed71e48..08a6883ecd4 100644
--- a/src/backend/catalog/pg_zstd_dictionaries.c
+++ b/src/backend/catalog/pg_zstd_dictionaries.c
@@ -568,6 +568,62 @@ cleanup_unused_zstd_dictionaries_internal(void)
 	return dropped_count;
 }
 
+/*
+ * generate_or_cleanup_zstd_dictionaries_for_relation
+ *
+ * Opens the relation identified by relid, iterates over its attributes,
+ * and for each valid (non-dropped, user-defined) attribute, calls
+ * build_zstd_dictionary_internal.
+ */
+void
+generate_or_cleanup_zstd_dictionaries_for_relation(Oid relid)
+{
+	Relation	rel;
+	TupleDesc	tupdesc;
+
+	/* Start a new transaction */
+	StartTransactionCommand();
+
+	/* Push an active snapshot toast want snapshot */
+	PushActiveSnapshot(GetTransactionSnapshot());
+
+	/* Open the relation using table_open (or relation_open) */
+	rel = table_open(relid, AccessShareLock);
+	tupdesc = RelationGetDescr(rel);
+
+	/* Iterate over all attributes of the relation */
+	for (int i = 0; i < tupdesc->natts; i++)
+	{
+		Form_pg_attribute attr = TupleDescAttr(tupdesc, i);
+
+		/* Skip dropped attributes and system columns (attnum <= 0) */
+		if (attr->attisdropped || attr->attnum <= 0)
+			continue;
+
+		/* Call your dictionary-building function for this attribute */
+		build_zstd_dictionary_internal(relid, attr->attnum);
+
+		/*
+		 * If build_zstd_dictionary_internal performs modifications that
+		 * subsequent iterations must see, use CommandCounterIncrement to
+		 * update the visibility of those changes.
+		 */
+		CommandCounterIncrement();
+	}
+
+	/* Close the relation and release the lock */
+	table_close(rel, NoLock);
+
+	/* Cleanup unused zstd dictionaries */
+	cleanup_unused_zstd_dictionaries_internal();
+
+	/* Pop the snapshot to clean up */
+	PopActiveSnapshot();
+
+	/* Commit the transaction */
+	CommitTransactionCommand();
+}
+
 /*
  * get_zstd_dict - Fetches the ZSTD dictionary from the catalog
  *
diff --git a/src/backend/commands/vacuum.c b/src/backend/commands/vacuum.c
index db5da3ce826..87a5708788e 100644
--- a/src/backend/commands/vacuum.c
+++ b/src/backend/commands/vacuum.c
@@ -37,6 +37,7 @@
 #include "catalog/namespace.h"
 #include "catalog/pg_database.h"
 #include "catalog/pg_inherits.h"
+#include "catalog/pg_zstd_dictionaries.h"
 #include "commands/cluster.h"
 #include "commands/defrem.h"
 #include "commands/progress.h"
@@ -2312,6 +2313,9 @@ vacuum_rel(Oid relid, RangeVar *relation, VacuumParams *params,
 		vacuum_rel(toast_relid, NULL, &toast_vacuum_params, bstrategy);
 	}
 
+	/* Generate or cleanup unused ZSTD dictionaries. */
+	generate_or_cleanup_zstd_dictionaries_for_relation(relid);
+
 	/*
 	 * Now release the session-level lock on the main table.
 	 */
diff --git a/src/include/catalog/pg_zstd_dictionaries.h b/src/include/catalog/pg_zstd_dictionaries.h
index cf847ee2801..a1e3948933c 100644
--- a/src/include/catalog/pg_zstd_dictionaries.h
+++ b/src/include/catalog/pg_zstd_dictionaries.h
@@ -44,5 +44,6 @@ DECLARE_UNIQUE_INDEX_PKEY(pg_zstd_dictionaries_dictid_index, 9949, ZstdDictidInd
 MAKE_SYSCACHE(ZSTDDICTIDOID, pg_zstd_dictionaries_dictid_index, 128);
 
 extern bytea *get_zstd_dict(Oid dictid);
+extern void generate_or_cleanup_zstd_dictionaries_for_relation(Oid relid);
 
 #endif							/* PG_ZSTD_DICTIONARIES_H */
-- 
2.47.1