From 68350a7a21de67e97664fb669ec11ed0d7009f36 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Wed, 9 Oct 2024 15:25:03 -0400
Subject: [PATCH v25 09/11] Move BitmapTableScan per-scan setup into a helper

Add BitmapTableScanSetup() which contains all of the code that must be
done on every scan of the table in a bitmap table scan.
This includes scanning the index, building the bitmap, and setting up
the scan descriptors.

Much of BitmapHeapNext() code was this setup, so it made sense to put it
in a helper function.
---
 src/backend/executor/nodeBitmapHeapscan.c | 227 +++++++++++-----------
 1 file changed, 116 insertions(+), 111 deletions(-)

diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index fcdbbb821ab..f522ba4959a 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -54,148 +54,152 @@ static inline void BitmapDoneInitializingSharedState(ParallelBitmapHeapState *ps
 static bool BitmapShouldInitializeSharedState(ParallelBitmapHeapState *pstate);
 
 
-/* ----------------------------------------------------------------
- *		BitmapHeapNext
+/*
+ * Do the underlying index scan, build the bitmap, set up the parallel state
+ * needed for parallel workers to iterate through the bitmap, and set up the
+ * underlying table scan descriptor.
  *
- *		Retrieve next tuple from the BitmapHeapScan node's currentRelation
- * ----------------------------------------------------------------
+ * While prefetching is the responsibility of the table AM, the prefetch
+ * iterator is set up here along with the main iterator. Two iterators are
+ * used -- one for the pages actually being scanned now and the other to
+ * run ahead and inform the storage layer what blocks to prefetch.
  */
-static TupleTableSlot *
-BitmapHeapNext(BitmapHeapScanState *node)
+static void
+BitmapTableScanSetup(BitmapHeapScanState *node)
 {
-	ExprContext *econtext;
-	TableScanDesc scan;
-	TIDBitmap  *tbm;
-	TupleTableSlot *slot;
+	TBMIterator tbmiterator = {0};
+	TBMIterator prefetch_iterator = {0};
+	int			prefetch_max;
 	ParallelBitmapHeapState *pstate = node->pstate;
 	dsa_area   *dsa = node->ss.ps.state->es_query_dsa;
+	Relation	rel = node->ss.ss_currentRelation;
 
 	/*
-	 * extract necessary information from index scan node
+	 * Maximum number of prefetches for the tablespace if configured,
+	 * otherwise the current value of the effective_io_concurrency GUC.
 	 */
-	econtext = node->ss.ps.ps_ExprContext;
-	slot = node->ss.ss_ScanTupleSlot;
-	scan = node->ss.ss_currentScanDesc;
-	tbm = node->tbm;
+	prefetch_max =
+		get_tablespace_io_concurrency(rel->rd_rel->reltablespace);
 
-	/*
-	 * If we haven't yet performed the underlying index scan, do it, and begin
-	 * the iteration over the bitmap.
-	 *
-	 * While prefetching is the responsibility of the table AM, the prefetch
-	 * iterator is set up here along with the main iterator. Two iterators are
-	 * used -- one for the pages actually being scanned now and the other to
-	 * run ahead and inform the storage layer what blocks to prefetch.
-	 */
-	if (!node->initialized)
+	if (!pstate)
 	{
-		TBMIterator tbmiterator;
-		int			prefetch_max;
-		Relation	rel = node->ss.ss_currentRelation;
-#ifdef USE_PREFETCH
-		TBMIterator prefetch_iterator = {0};
-#endif							/* USE_PREFETCH */
+		node->tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
 
+		if (!node->tbm || !IsA(node->tbm, TIDBitmap))
+			elog(ERROR, "unrecognized result from subplan");
+	}
+	else if (BitmapShouldInitializeSharedState(pstate))
+	{
 		/*
-		 * Maximum number of prefetches for the tablespace if configured,
-		 * otherwise the current value of the effective_io_concurrency GUC.
+		 * The leader will immediately come out of the function, but others
+		 * will be blocked until leader populates the TBM and wakes them up.
 		 */
-		prefetch_max =
-			get_tablespace_io_concurrency(rel->rd_rel->reltablespace);
-
-		if (!pstate)
-		{
-			tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
+		node->tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
+		if (!node->tbm || !IsA(node->tbm, TIDBitmap))
+			elog(ERROR, "unrecognized result from subplan");
 
-			if (!tbm || !IsA(tbm, TIDBitmap))
-				elog(ERROR, "unrecognized result from subplan");
+		/*
+		 * Prepare to iterate over the TBM. This will return the dsa_pointer
+		 * of the iterator state which will be used by multiple processes to
+		 * iterate jointly.
+		 */
+		pstate->tbmiterator = tbm_prepare_shared_iterate(node->tbm);
 
-			node->tbm = tbm;
-		}
-		else if (BitmapShouldInitializeSharedState(pstate))
-		{
-			/*
-			 * The leader will immediately come out of the function, but
-			 * others will be blocked until leader populates the TBM and wakes
-			 * them up.
-			 */
-			tbm = (TIDBitmap *) MultiExecProcNode(outerPlanState(node));
-			if (!tbm || !IsA(tbm, TIDBitmap))
-				elog(ERROR, "unrecognized result from subplan");
+		/*
+		 * For prefetching, we use *two* iterators, one for the pages we are
+		 * actually scanning and another that runs ahead of the first for
+		 * prefetching.  node->prefetch_pages tracks exactly how many pages
+		 * ahead the prefetch iterator is.  Also, node->prefetch_target tracks
+		 * the desired prefetch distance, which starts small and increases up
+		 * to the node->prefetch_maximum.  This is to avoid doing a lot of
+		 * prefetching in a scan that stops after a few tuples because of a
+		 * LIMIT.
+		 */
+#ifdef USE_PREFETCH
+		if (prefetch_max > 0)
+			pstate->prefetch_iterator =
+				tbm_prepare_shared_iterate(node->tbm);
+#endif
 
-			node->tbm = tbm;
+		/* We have initialized the shared state so wake up others. */
+		BitmapDoneInitializingSharedState(pstate);
+	}
 
-			/*
-			 * Prepare to iterate over the TBM. This will return the
-			 * dsa_pointer of the iterator state which will be used by
-			 * multiple processes to iterate jointly.
-			 */
-			pstate->tbmiterator = tbm_prepare_shared_iterate(tbm);
+	tbmiterator = tbm_begin_iterate(node->tbm, dsa,
+									pstate ?
+									pstate->tbmiterator :
+									InvalidDsaPointer);
 
 #ifdef USE_PREFETCH
-			if (prefetch_max > 0)
-			{
-				pstate->prefetch_iterator =
-					tbm_prepare_shared_iterate(tbm);
-			}
+	if (prefetch_max > 0)
+		prefetch_iterator = tbm_begin_iterate(node->tbm, dsa,
+											  pstate ?
+											  pstate->prefetch_iterator :
+											  InvalidDsaPointer);
 #endif							/* USE_PREFETCH */
 
-			/* We have initialized the shared state so wake up others. */
-			BitmapDoneInitializingSharedState(pstate);
-		}
-
-		tbmiterator = tbm_begin_iterate(tbm, dsa,
-										pstate ?
-										pstate->tbmiterator :
-										InvalidDsaPointer);
-#ifdef USE_PREFETCH
-		if (prefetch_max > 0)
-			prefetch_iterator = tbm_begin_iterate(tbm, dsa,
-												  pstate ?
-												  pstate->prefetch_iterator :
-												  InvalidDsaPointer);
-#endif							/* USE_PREFETCH */
+	/*
+	 * If this is the first scan of the underlying table, create the table
+	 * scan descriptor and begin the scan.
+	 */
+	if (!node->ss.ss_currentScanDesc)
+	{
+		bool		need_tuples = false;
 
 		/*
-		 * If this is the first scan of the underlying table, create the table
-		 * scan descriptor and begin the scan.
+		 * We can potentially skip fetching heap pages if we do not need any
+		 * columns of the table, either for checking non-indexable quals or
+		 * for returning data.  This test is a bit simplistic, as it checks
+		 * the stronger condition that there's no qual or return tlist at all.
+		 * But in most cases it's probably not worth working harder than that.
 		 */
-		if (!scan)
-		{
-			bool		need_tuples = false;
+		need_tuples = (node->ss.ps.plan->qual != NIL ||
+					   node->ss.ps.plan->targetlist != NIL);
+
+		node->ss.ss_currentScanDesc =
+			table_beginscan_bm(node->ss.ss_currentRelation,
+							   node->ss.ps.state->es_snapshot,
+							   pstate,
+							   0,
+							   NULL,
+							   need_tuples,
+							   prefetch_max);
+	}
 
-			/*
-			 * We can potentially skip fetching heap pages if we do not need
-			 * any columns of the table, either for checking non-indexable
-			 * quals or for returning data.  This test is a bit simplistic, as
-			 * it checks the stronger condition that there's no qual or return
-			 * tlist at all. But in most cases it's probably not worth working
-			 * harder than that.
-			 */
-			need_tuples = (node->ss.ps.plan->qual != NIL ||
-						   node->ss.ps.plan->targetlist != NIL);
-
-			scan = table_beginscan_bm(node->ss.ss_currentRelation,
-									  node->ss.ps.state->es_snapshot,
-									  pstate,
-									  0,
-									  NULL,
-									  need_tuples,
-									  prefetch_max);
-
-			node->ss.ss_currentScanDesc = scan;
-		}
+	node->ss.ss_currentScanDesc->st.bitmap.rs_iterator = tbmiterator;
+	node->ss.ss_currentScanDesc->st.bitmap.rs_prefetch_iterator =
+		prefetch_iterator;
+	node->initialized = true;
+}
 
-		scan->st.bitmap.rs_iterator = tbmiterator;
-		scan->st.bitmap.rs_prefetch_iterator = prefetch_iterator;
-		node->initialized = true;
+/* ----------------------------------------------------------------
+ *		BitmapHeapNext
+ *
+ *		Retrieve next tuple from the BitmapHeapScan node's currentRelation
+ * ----------------------------------------------------------------
+ */
+static TupleTableSlot *
+BitmapHeapNext(BitmapHeapScanState *node)
+{
+	ExprContext *econtext = node->ss.ps.ps_ExprContext;
+	TupleTableSlot *slot = node->ss.ss_ScanTupleSlot;
 
+	/*
+	 * If we haven't yet performed the underlying index scan, do it, and begin
+	 * the iteration over the bitmap.
+	 */
+	if (!node->initialized)
+	{
+		/* BitmapTableScanSetup sets node->ss.ss_currentScanDesc */
+		BitmapTableScanSetup(node);
 		goto new_page;
 	}
 
 	for (;;)
 	{
-		while (table_scan_bitmap_next_tuple(scan, slot))
+		while (table_scan_bitmap_next_tuple(
+											node->ss.ss_currentScanDesc,
+											slot))
 		{
 			/*
 			 * Continuing in previously obtained page.
@@ -229,7 +233,8 @@ new_page:
 		 * Returns false if the bitmap is exhausted and there are no further
 		 * blocks we need to scan.
 		 */
-		if (!table_scan_bitmap_next_block(scan, &node->recheck,
+		if (!table_scan_bitmap_next_block(node->ss.ss_currentScanDesc,
+										  &node->recheck,
 										  &node->stats.lossy_pages,
 										  &node->stats.exact_pages))
 			break;
-- 
2.45.2

