From cc6c927708ab6e111d2abbb467a6f8dd798d0992 Mon Sep 17 00:00:00 2001
From: Yuhang Qiu <iamqyh@gmail.com>
Date: Wed, 26 Aug 2026 11:57:58 +0000
Subject: [PATCH 2/3] 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.
---
 src/backend/access/heap/heapam.c         |  10 +-
 src/backend/access/heap/heapam_handler.c | 127 +++++++++++++----------
 src/include/access/heapam.h              |  11 +-
 3 files changed, 88 insertions(+), 60 deletions(-)

diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index 72d6541734c..a43d9fb87b1 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -412,10 +412,17 @@ initscan(HeapScanDesc scan, ScanKey key, bool keep_startblock)
 	else
 	{
 		if (scan->rs_strategy != NULL)
+		{
+			if (scan->rs_read_stream != NULL)
+				read_stream_set_strategy(scan->rs_read_stream, NULL);
 			FreeAccessStrategy(scan->rs_strategy);
+		}
 		scan->rs_strategy = NULL;
 	}
 
+	if (scan->rs_read_stream != NULL && scan->rs_strategy != NULL)
+		read_stream_set_strategy(scan->rs_read_stream, scan->rs_strategy);
+
 	if (scan->rs_base.rs_parallel != NULL)
 	{
 		/* For parallel scan, believe whatever ParallelTableScanDesc says. */
@@ -1204,6 +1211,7 @@ heap_beginscan(Relation relation, Snapshot snapshot,
 	scan->rs_base.rs_parallel = parallel_scan;
 	scan->rs_base.rs_instrument = NULL;
 	scan->rs_strategy = NULL;	/* set in initscan */
+	scan->rs_read_stream = NULL;
 	scan->rs_cbuf = InvalidBuffer;
 
 	/*
@@ -1270,8 +1278,6 @@ heap_beginscan(Relation relation, Snapshot snapshot,
 
 	initscan(scan, key, false);
 
-	scan->rs_read_stream = NULL;
-
 	/*
 	 * Set up a read stream for sequential scans and TID range scans. This
 	 * should be done after initscan() because initscan() allocates the
diff --git a/src/backend/access/heap/heapam_handler.c b/src/backend/access/heap/heapam_handler.c
index bf87430cf01..e804bb0db11 100644
--- a/src/backend/access/heap/heapam_handler.c
+++ b/src/backend/access/heap/heapam_handler.c
@@ -2135,11 +2135,58 @@ heapam_scan_bitmap_next_tuple(TableScanDesc scan,
 	return true;
 }
 
+/*
+ * 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)
+{
+	SampleScanState *scanstate = callback_private_data;
+	HeapScanDesc hscan = (HeapScanDesc) scanstate->ss.ss_currentScanDesc;
+	TsmRoutine *tsm = scanstate->tsmroutine;
+	BlockNumber blockno;
+
+	if (tsm->NextSampleBlock)
+	{
+		blockno = tsm->NextSampleBlock(scanstate, hscan->rs_nblocks);
+		CHECK_FOR_INTERRUPTS();
+		return blockno;
+	}
+
+	/* Scan all blocks sequentially. */
+	if (unlikely(!hscan->rs_inited))
+	{
+		blockno = hscan->rs_startblock;
+		hscan->rs_inited = true;
+	}
+	else
+	{
+		blockno = hscan->rs_prefetch_block + 1;
+
+		if (blockno >= hscan->rs_nblocks)
+			blockno = 0;
+
+		/* 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_startblock)
+			blockno = InvalidBlockNumber;
+	}
+
+	hscan->rs_prefetch_block = blockno;
+	return blockno;
+}
+
 static bool
 heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
 {
 	HeapScanDesc hscan = (HeapScanDesc) scan;
 	TsmRoutine *tsm = scanstate->tsmroutine;
+	int			stream_flags = READ_STREAM_DEFAULT;
 	BlockNumber blockno;
 
 	/* return false immediately if relation is empty */
@@ -2153,70 +2200,46 @@ heapam_scan_sample_next_block(TableScanDesc scan, SampleScanState *scanstate)
 		hscan->rs_cbuf = InvalidBuffer;
 	}
 
-	if (tsm->NextSampleBlock)
-		blockno = tsm->NextSampleBlock(scanstate, hscan->rs_nblocks);
-	else
-	{
-		/* scanning table sequentially */
-
-		if (hscan->rs_cblock == InvalidBlockNumber)
-		{
-			Assert(!hscan->rs_inited);
-			blockno = hscan->rs_startblock;
-		}
-		else
-		{
-			Assert(hscan->rs_inited);
-
-			blockno = hscan->rs_cblock + 1;
-
-			if (blockno >= hscan->rs_nblocks)
-			{
-				/* wrap to beginning of rel, might not have started at 0 */
-				blockno = 0;
-			}
-
-			/*
-			 * 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);
+	/*
+	 * 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
+	 * many pages' worth of consecutive dead tuples.
+	 */
+	CHECK_FOR_INTERRUPTS();
 
-			if (blockno == hscan->rs_startblock)
-			{
-				blockno = InvalidBlockNumber;
-			}
-		}
+	/* The callback needs SampleScanState, which heap_beginscan() lacks. */
+	if (hscan->rs_read_stream == NULL)
+	{
+		/* 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_cblock = blockno;
+	hscan->rs_cbuf = read_stream_next_buffer(hscan->rs_read_stream, NULL);
+	if (BufferIsValid(hscan->rs_cbuf))
+		blockno = BufferGetBlockNumber(hscan->rs_cbuf);
+	else
+		blockno = InvalidBlockNumber;
 
 	if (!BlockNumberIsValid(blockno))
 	{
+		hscan->rs_cblock = InvalidBlockNumber;
 		hscan->rs_inited = false;
 		return false;
 	}
 
+	hscan->rs_cblock = blockno;
 	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
-	 * many pages' worth of consecutive dead tuples.
-	 */
-	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);
-
 	/* in pagemode, prune the page and determine visible tuple offsets */
 	if (hscan->rs_base.rs_flags & SO_ALLOW_PAGEMODE)
 		heap_prepare_pagescan(scan);
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

