0001-Support-extensions-wanting-to-do-more-advanced-analy.patch

From 8acb707cc859f1570046919295817e27f0d8b4f6 Mon Sep 17 00:00:00 2001
From: Mats Kindahl <mats@timescale.com>
Date: Sat, 14 Sep 2024 14:03:35 +0200
Subject: [PATCH] Support extensions wanting to do more advanced analyze

To support extensions that want to do more advanced analyzis
implementations, this commit adds two function:

`analuze_compute_vacattrstats` to compute vacuum attribute stats that
can be used to compute the number of target rows to analyze.

`read_stream_next_block` which allow an extension to just use the
provided analyze stream to sample blocks.
---
 src/backend/commands/analyze.c        | 120 ++++++++++++++------------
 src/backend/storage/aio/read_stream.c |  14 +++
 src/include/commands/vacuum.h         |   2 +
 src/include/storage/read_stream.h     |   2 +
 4 files changed, 85 insertions(+), 53 deletions(-)

diff --git a/src/backend/commands/analyze.c b/src/backend/commands/analyze.c
index c590a2adc35..8a402ad15c6 100644
--- a/src/backend/commands/analyze.c
+++ b/src/backend/commands/analyze.c
@@ -269,6 +269,72 @@ analyze_rel(Oid relid, RangeVar *relation,
 	pgstat_progress_end_command();
 }
 
+/*
+ * Determine which columns to analyze
+ *
+ * Note that system attributes are never analyzed, so we just reject them
+ * at the lookup stage.  We also reject duplicate column mentions.  (We
+ * could alternatively ignore duplicates, but analyzing a column twice
+ * won't work; we'd end up making a conflicting update in pg_statistic.)
+ */
+int
+analyze_compute_vacattrstats(Relation onerel, List *va_cols, VacAttrStats ***vacattrstats_out)
+{
+	int			tcnt,
+				i,
+				attr_cnt;
+	VacAttrStats **vacattrstats;
+
+	if (va_cols != NIL)
+	{
+		Bitmapset  *unique_cols = NULL;
+		ListCell   *le;
+
+		vacattrstats = (VacAttrStats **) palloc(list_length(va_cols) *
+												sizeof(VacAttrStats *));
+		tcnt = 0;
+		foreach(le, va_cols)
+		{
+			char	   *col = strVal(lfirst(le));
+
+			i = attnameAttNum(onerel, col, false);
+			if (i == InvalidAttrNumber)
+				ereport(ERROR,
+						(errcode(ERRCODE_UNDEFINED_COLUMN),
+						 errmsg("column \"%s\" of relation \"%s\" does not exist",
+								col, RelationGetRelationName(onerel))));
+			if (bms_is_member(i, unique_cols))
+				ereport(ERROR,
+						(errcode(ERRCODE_DUPLICATE_COLUMN),
+						 errmsg("column \"%s\" of relation \"%s\" appears more than once",
+								col, RelationGetRelationName(onerel))));
+			unique_cols = bms_add_member(unique_cols, i);
+
+			vacattrstats[tcnt] = examine_attribute(onerel, i, NULL);
+			if (vacattrstats[tcnt] != NULL)
+				tcnt++;
+		}
+		attr_cnt = tcnt;
+	}
+	else
+	{
+		attr_cnt = onerel->rd_att->natts;
+		vacattrstats = (VacAttrStats **)
+			palloc(attr_cnt * sizeof(VacAttrStats *));
+		tcnt = 0;
+		for (i = 1; i <= attr_cnt; i++)
+		{
+			vacattrstats[tcnt] = examine_attribute(onerel, i, NULL);
+			if (vacattrstats[tcnt] != NULL)
+				tcnt++;
+		}
+		attr_cnt = tcnt;
+	}
+
+	*vacattrstats_out = vacattrstats;
+	return attr_cnt;
+}
+
 /*
  *	do_analyze_rel() -- analyze one relation, recursively or not
  *
@@ -353,59 +419,7 @@ do_analyze_rel(Relation onerel, VacuumParams *params,
 		starttime = GetCurrentTimestamp();
 	}
 
-	/*
-	 * Determine which columns to analyze
-	 *
-	 * Note that system attributes are never analyzed, so we just reject them
-	 * at the lookup stage.  We also reject duplicate column mentions.  (We
-	 * could alternatively ignore duplicates, but analyzing a column twice
-	 * won't work; we'd end up making a conflicting update in pg_statistic.)
-	 */
-	if (va_cols != NIL)
-	{
-		Bitmapset  *unique_cols = NULL;
-		ListCell   *le;
-
-		vacattrstats = (VacAttrStats **) palloc(list_length(va_cols) *
-												sizeof(VacAttrStats *));
-		tcnt = 0;
-		foreach(le, va_cols)
-		{
-			char	   *col = strVal(lfirst(le));
-
-			i = attnameAttNum(onerel, col, false);
-			if (i == InvalidAttrNumber)
-				ereport(ERROR,
-						(errcode(ERRCODE_UNDEFINED_COLUMN),
-						 errmsg("column \"%s\" of relation \"%s\" does not exist",
-								col, RelationGetRelationName(onerel))));
-			if (bms_is_member(i, unique_cols))
-				ereport(ERROR,
-						(errcode(ERRCODE_DUPLICATE_COLUMN),
-						 errmsg("column \"%s\" of relation \"%s\" appears more than once",
-								col, RelationGetRelationName(onerel))));
-			unique_cols = bms_add_member(unique_cols, i);
-
-			vacattrstats[tcnt] = examine_attribute(onerel, i, NULL);
-			if (vacattrstats[tcnt] != NULL)
-				tcnt++;
-		}
-		attr_cnt = tcnt;
-	}
-	else
-	{
-		attr_cnt = onerel->rd_att->natts;
-		vacattrstats = (VacAttrStats **)
-			palloc(attr_cnt * sizeof(VacAttrStats *));
-		tcnt = 0;
-		for (i = 1; i <= attr_cnt; i++)
-		{
-			vacattrstats[tcnt] = examine_attribute(onerel, i, NULL);
-			if (vacattrstats[tcnt] != NULL)
-				tcnt++;
-		}
-		attr_cnt = tcnt;
-	}
+	attr_cnt = analyze_compute_vacattrstats(onerel, va_cols, &vacattrstats);
 
 	/*
 	 * Open all indexes of the relation, and see if there are any analyzable
diff --git a/src/backend/storage/aio/read_stream.c b/src/backend/storage/aio/read_stream.c
index 74b9bae6313..7e25793ea7e 100644
--- a/src/backend/storage/aio/read_stream.c
+++ b/src/backend/storage/aio/read_stream.c
@@ -804,3 +804,17 @@ read_stream_end(ReadStream *stream)
 	read_stream_reset(stream);
 	pfree(stream);
 }
+
+/*
+ * Transitional support for code that would like to perform a read directly,
+ * without using the stream.  Returns, and skips, the next block number that
+ * would be read by the stream's look-ahead algorithm, or InvalidBlockNumber
+ * if the end of the stream is reached.  Also reports the strategy that would
+ * be used to read it.
+ */
+BlockNumber
+read_stream_next_block(ReadStream *stream, BufferAccessStrategy *strategy)
+{
+	*strategy = stream->ios[0].op.strategy;
+	return read_stream_get_block(stream, NULL);
+}
diff --git a/src/include/commands/vacuum.h b/src/include/commands/vacuum.h
index 759f9a87d38..f456734855d 100644
--- a/src/include/commands/vacuum.h
+++ b/src/include/commands/vacuum.h
@@ -378,6 +378,8 @@ extern void parallel_vacuum_main(dsm_segment *seg, shm_toc *toc);
 extern void analyze_rel(Oid relid, RangeVar *relation,
 						VacuumParams *params, List *va_cols, bool in_outer_xact,
 						BufferAccessStrategy bstrategy);
+extern int	analyze_compute_vacattrstats(Relation onerel, List *va_cols,
+										 VacAttrStats ***vacattrstats_out);
 extern bool std_typanalyze(VacAttrStats *stats);
 
 /* in utils/misc/sampling.c --- duplicate of declarations in utils/sampling.h */
diff --git a/src/include/storage/read_stream.h b/src/include/storage/read_stream.h
index f676d2cc20a..7b9005e87bc 100644
--- a/src/include/storage/read_stream.h
+++ b/src/include/storage/read_stream.h
@@ -57,6 +57,8 @@ extern ReadStream *read_stream_begin_relation(int flags,
 											  void *callback_private_data,
 											  size_t per_buffer_data_size);
 extern Buffer read_stream_next_buffer(ReadStream *stream, void **per_buffer_data);
+extern BlockNumber read_stream_next_block(ReadStream *stream,
+										  BufferAccessStrategy *strategy);
 extern void read_stream_reset(ReadStream *stream);
 extern void read_stream_end(ReadStream *stream);
 
-- 
2.43.0