From 336670d2304f851c184a765f23891a2af0b76c73 Mon Sep 17 00:00:00 2001
From: Yuhang Qiu <iamqyh@gmail.com>
Date: Wed, 26 Aug 2026 11:57:58 +0000
Subject: [PATCH v3 3/4] heapam: Use streaming read I/O in sample scans

Sample scans read each selected block with ReadBufferExtended(), so they
cannot issue reads ahead of page processing.

Use a read stream to read blocks chosen by the sampling method.  For
methods that examine every block, let the stream generate the same
sequential block order.

The block callback needs SampleScanState, which is unavailable in
heap_beginscan(), so create the stream on first use.  Reset and reuse it on
rescan.

Reviewed-by: Nazir Bilal Yavuz <byavuz81@gmail.com>
Discussion: https://postgr.es/m/D1F99EB2-5A03-465F-A210-98F7DF256168@gmail.com
---
 contrib/tsm_system_time/tsm_system_time.c |   2 +-
 doc/src/sgml/tsm-system-time.sgml         |  13 ++-
 src/backend/access/heap/heapam.c          |   6 +-
 src/backend/access/heap/heapam_handler.c  | 128 ++++++++++++----------
 src/include/access/heapam.h               |  11 +-
 5 files changed, 91 insertions(+), 69 deletions(-)

diff --git a/contrib/tsm_system_time/tsm_system_time.c b/contrib/tsm_system_time/tsm_system_time.c
index 4814c31bc6f..c7f4aaf75e1 100644
--- a/contrib/tsm_system_time/tsm_system_time.c
+++ b/contrib/tsm_system_time/tsm_system_time.c
@@ -4,7 +4,7 @@
  *	  support routines for SYSTEM_TIME tablesample method
  *
  * The desire here is to produce a random sample with as many rows as possible
- * in no more than the specified amount of time.  We use a block-sampling
+ * within approximately the specified amount of time.  We use a block-sampling
  * approach.  To ensure that the whole relation will be visited if necessary,
  * we start at a randomly chosen block and then advance with a stride that
  * is randomly chosen but is relatively prime to the relation's nblocks.
diff --git a/doc/src/sgml/tsm-system-time.sgml b/doc/src/sgml/tsm-system-time.sgml
index 05c6be0caca..aa688ecf5c5 100644
--- a/doc/src/sgml/tsm-system-time.sgml
+++ b/doc/src/sgml/tsm-system-time.sgml
@@ -17,10 +17,10 @@
 
  <para>
   This table sampling method accepts a single floating-point argument that
-  is the maximum number of milliseconds to spend reading the table.  This
-  gives you direct control over how long the query takes, at the price that
-  the size of the sample becomes hard to predict.  The resulting sample will
-  contain as many rows as could be read in the specified time, unless the
+  specifies approximately how many milliseconds to spend reading the table.
+  This gives you approximate control over how long the query takes, at the
+  price that the size of the sample becomes hard to predict.  The resulting
+  sample will contain as many rows as could be read in that time, unless the
   whole table has been read first.
  </para>
 
@@ -64,8 +64,9 @@ SELECT * FROM my_table TABLESAMPLE SYSTEM_TIME(1000);
 
   <para>
    This command will return as large a sample of <structname>my_table</structname> as
-   it can read in 1 second (1000 milliseconds).  Of course, if the whole
-   table can be read in under 1 second, all its rows will be returned.
+   it can read in approximately 1 second (1000 milliseconds).  Of course, if
+   the whole table can be read in under 1 second, all its rows will be
+   returned.
   </para>
  </sect2>
 
diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index 2f73c0a3045..0b0fc133e67 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -352,7 +352,11 @@ bitmapheap_stream_read_next(ReadStream *pgsr, void *private_data,
 	Assert(false);
 }
 
-/* Set up a read stream for scan types that use one. */
+/*
+ * Set up a read stream for scan types that can create one without executor
+ * state.  Sample scans create their stream on first use instead, because
+ * their callback needs SampleScanState.
+ */
 static void
 heap_scan_init_stream(HeapScanDesc scan)
 {
diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index bf87430cf01..47c8d3de667 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2135,77 +2135,69 @@ heapam_scan_bitmap_next_tuple(TableScanDesc scan,
 	return true;
 }
 
-static bool
-heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
+/*
+ * Read stream API callback for sample scans.  Returns the next block selected
+ * by the sampling method, or the next block of a sequential sample scan.
+ */
+static BlockNumber
+heapam_scan_sample_stream_read_next(ReadStream *stream,
+									void *callback_private_data,
+									void *per_buffer_data)
 {
-	HeapScanDesc hscan = (HeapScanDesc) scan;
+	SampleScanState *scanstate = callback_private_data;
+	HeapScanDesc hscan = (HeapScanDesc) scanstate->ss.ss_currentScanDesc;
 	TsmRoutine *tsm = scanstate->tsmroutine;
 	BlockNumber blockno;
 
-	/* return false immediately if relation is empty */
-	if (hscan->rs_nblocks == 0)
-		return false;
-
-	/* release previous scan buffer, if any */
-	if (BufferIsValid(hscan->rs_cbuf))
+	if (tsm->NextSampleBlock)
 	{
-		ReleaseBuffer(hscan->rs_cbuf);
-		hscan->rs_cbuf = InvalidBuffer;
+		blockno = tsm->NextSampleBlock(scanstate, hscan->rs_nblocks);
+		CHECK_FOR_INTERRUPTS();
+		return blockno;
 	}
 
-	if (tsm->NextSampleBlock)
-		blockno = tsm->NextSampleBlock(scanstate, hscan->rs_nblocks);
+	/* Scan all blocks sequentially. */
+	if (unlikely(!hscan->rs_inited))
+	{
+		blockno = hscan->rs_startblock;
+		hscan->rs_inited = true;
+	}
 	else
 	{
-		/* scanning table sequentially */
+		blockno = hscan->rs_prefetch_block + 1;
 
-		if (hscan->rs_cblock == InvalidBlockNumber)
-		{
-			Assert(!hscan->rs_inited);
-			blockno = hscan->rs_startblock;
-		}
-		else
-		{
-			Assert(hscan->rs_inited);
+		if (blockno >= hscan->rs_nblocks)
+			blockno = 0;
 
-			blockno = hscan->rs_cblock + 1;
+		/* Report the position before detecting the end, as for a seqscan. */
+		if (hscan->rs_base.rs_flags & SO_ALLOW_SYNC)
+			ss_report_location(hscan->rs_base.rs_rd, blockno);
 
-			if (blockno >= hscan->rs_nblocks)
-			{
-				/* wrap to beginning of rel, might not have started at 0 */
-				blockno = 0;
-			}
+		if (blockno == hscan->rs_startblock)
+			blockno = InvalidBlockNumber;
+	}
 
-			/*
-			 * Report our new scan position for synchronization purposes.
-			 *
-			 * Note: we do this before checking for end of scan so that the
-			 * final state of the position hint is back at the start of the
-			 * rel.  That's not strictly necessary, but otherwise when you run
-			 * the same query multiple times the starting position would shift
-			 * a little bit backwards on every invocation, which is confusing.
-			 * We don't guarantee any specific ordering in general, though.
-			 */
-			if (scan->rs_flags & SO_ALLOW_SYNC)
-				ss_report_location(scan->rs_rd, blockno);
+	hscan->rs_prefetch_block = blockno;
+	return blockno;
+}
 
-			if (blockno == hscan->rs_startblock)
-			{
-				blockno = InvalidBlockNumber;
-			}
-		}
-	}
+static bool
+heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
+{
+	HeapScanDesc hscan = (HeapScanDesc) scan;
+	TsmRoutine *tsm = scanstate->tsmroutine;
 
-	hscan->rs_cblock = blockno;
+	/* return false immediately if relation is empty */
+	if (hscan->rs_nblocks == 0)
+		return false;
 
-	if (!BlockNumberIsValid(blockno))
+	/* release previous scan buffer, if any */
+	if (BufferIsValid(hscan->rs_cbuf))
 	{
-		hscan->rs_inited = false;
-		return false;
+		ReleaseBuffer(hscan->rs_cbuf);
+		hscan->rs_cbuf = InvalidBuffer;
 	}
 
-	Assert(hscan->rs_cblock < hscan->rs_nblocks);
-
 	/*
 	 * Be sure to check for interrupts at least once per page.  Checks at
 	 * higher code levels won't be able to stop a sample scan that encounters
@@ -2213,9 +2205,35 @@ heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
 	 */
 	CHECK_FOR_INTERRUPTS();
 
-	/* Read page using selected strategy */
-	hscan->rs_cbuf = ReadBufferExtended(hscan->rs_base.rs_rd, MAIN_FORKNUM,
-										blockno, RBM_NORMAL, hscan->rs_strategy);
+	/* The callback needs SampleScanState, which heap_beginscan() lacks. */
+	if (hscan->rs_read_stream == NULL)
+	{
+		int			stream_flags = READ_STREAM_DEFAULT;
+
+		/* Extension callbacks need not be safe in AIO batch mode. */
+		if (tsm->NextSampleBlock == NULL)
+			stream_flags |= READ_STREAM_SEQUENTIAL | READ_STREAM_USE_BATCHING;
+
+		hscan->rs_read_stream =
+			read_stream_begin_relation(stream_flags,
+									   hscan->rs_strategy,
+									   hscan->rs_base.rs_rd,
+									   MAIN_FORKNUM,
+									   heapam_scan_sample_stream_read_next,
+									   scanstate,
+									   0);
+	}
+
+	hscan->rs_cbuf = read_stream_next_buffer(hscan->rs_read_stream, NULL);
+	if (!BufferIsValid(hscan->rs_cbuf))
+	{
+		hscan->rs_cblock = InvalidBlockNumber;
+		hscan->rs_inited = false;
+		return false;
+	}
+
+	hscan->rs_cblock = BufferGetBlockNumber(hscan->rs_cbuf);
+	Assert(hscan->rs_cblock < hscan->rs_nblocks);
 
 	/* in pagemode, prune the page and determine visible tuple offsets */
 	if (hscan->rs_base.rs_flags & SO_ALLOW_PAGEMODE)
diff --git a/src/include/access/heapam.h b/src/include/access/heapam.h
index 5176478c295..e2b43a6e260 100644
--- a/src/include/access/heapam.h
+++ b/src/include/access/heapam.h
@@ -80,12 +80,11 @@ typedef struct HeapScanDescData
 	ReadStream *rs_read_stream;
 
 	/*
-	 * For sequential scans and TID range scans to stream reads. The read
-	 * stream is allocated at the beginning of the scan and reset on rescan or
-	 * when the scan direction changes. The scan direction is saved each time
-	 * a new page is requested. If the scan direction changes from one page to
-	 * the next, the read stream releases all previously pinned buffers and
-	 * resets the prefetch block.
+	 * For sequential, sample, and TID range scans to stream reads. The read
+	 * stream is reset on rescan or when the scan direction changes. The scan
+	 * direction is saved each time a new page is requested. If the scan
+	 * direction changes from one page to the next, the read stream releases
+	 * all previously pinned buffers and resets the prefetch block.
 	 */
 	ScanDirection rs_dir;
 	BlockNumber rs_prefetch_block;
-- 
2.43.7

