From 2252c038424ec683311d08a0a868d83ffc10d616 Mon Sep 17 00:00:00 2001
From: Melanie Plageman <melanieplageman@gmail.com>
Date: Sat, 6 Apr 2024 15:08:11 -0400
Subject: [PATCH v19 16/21] BitmapHeapScan manage iteration in heap BM
 functions

Now that therea are dedicated heap AM BitmapHeapScan functions to begin,
restart, and end Bitmap Heap Scans, it is easy to initialize and cleanup
the UnifiedTBMIterator there instead. Doing this in heap AM code will
allow us to move the current block iterator from the TableScanDesc to
the BitmapHeapScanDesc.

Author: Melanie Plageman
Discussion: https://postgr.es/m/CAAKRu_ZwCwWFeL_H3ia26bP2e7HiKLWt0ZmGXPVwPO6uXq0vaA%40mail.gmail.com
---
 src/backend/access/heap/heapam.c          | 15 ++++++++++++++-
 src/backend/executor/nodeBitmapHeapscan.c | 10 ++++------
 src/include/access/heapam.h               |  3 ++-
 src/include/access/tableam.h              | 21 ++++++++++-----------
 4 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index bcb9e4390d4..4fa4c29f90f 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -1049,6 +1049,9 @@ heap_beginscan_bm(Relation relation, Snapshot snapshot, uint32 flags)
 	scan->heap_common.rs_base.rs_parallel = NULL;
 	scan->heap_common.rs_strategy = NULL;
 
+	scan->heap_common.rs_base.rs_bhs_iterator.serial = NULL;
+	scan->heap_common.rs_base.rs_bhs_iterator.parallel = NULL;
+
 	Assert(snapshot && IsMVCCSnapshot(snapshot));
 
 	/* we only need to set this up once */
@@ -1104,7 +1107,8 @@ heap_rescan(TableScanDesc sscan, ScanKey key, bool set_params,
 }
 
 void
-heap_rescan_bm(TableScanDesc sscan)
+heap_rescan_bm(TableScanDesc sscan, TIDBitmap *tbm,
+			   ParallelBitmapHeapState *pstate, dsa_area *dsa)
 {
 	BitmapHeapScanDesc scan = (BitmapHeapScanDesc) sscan;
 
@@ -1117,10 +1121,17 @@ heap_rescan_bm(TableScanDesc sscan)
 
 	scan->empty_tuples_pending = 0;
 
+	unified_tbm_end_iterate(&scan->heap_common.rs_base.rs_bhs_iterator);
+
 	/*
 	 * reinitialize heap scan descriptor
 	 */
 	initscan(&scan->heap_common, NULL, true);
+
+	unified_tbm_begin_iterate(&scan->heap_common.rs_base.rs_bhs_iterator, tbm, dsa,
+							  pstate ?
+							  pstate->tbmiterator :
+							  InvalidDsaPointer);
 }
 
 
@@ -1171,6 +1182,8 @@ heap_endscan_bm(TableScanDesc sscan)
 	if (BufferIsValid(scan->vmbuffer))
 		ReleaseBuffer(scan->vmbuffer);
 
+	unified_tbm_end_iterate(&scan->heap_common.rs_base.rs_bhs_iterator);
+
 	/*
 	 * decrement relation reference count and free scan descriptor storage
 	 */
diff --git a/src/backend/executor/nodeBitmapHeapscan.c b/src/backend/executor/nodeBitmapHeapscan.c
index bb9a11aa51d..012cba2c125 100644
--- a/src/backend/executor/nodeBitmapHeapscan.c
+++ b/src/backend/executor/nodeBitmapHeapscan.c
@@ -156,15 +156,13 @@ BitmapHeapNext(BitmapHeapScanState *node)
 		scan = table_beginscan_bm(node->ss.ss_currentScanDesc,
 								  node->ss.ss_currentRelation,
 								  node->ss.ps.state->es_snapshot,
-								  need_tuples);
+								  need_tuples,
+								  node->tbm,
+								  node->pstate,
+								  node->ss.ps.state->es_query_dsa);
 
 		node->ss.ss_currentScanDesc = scan;
 
-		unified_tbm_begin_iterate(&scan->rs_bhs_iterator, tbm, dsa,
-								  pstate ?
-								  pstate->tbmiterator :
-								  InvalidDsaPointer);
-
 #ifdef USE_PREFETCH
 		if (node->prefetch_maximum > 0)
 		{
diff --git a/src/include/access/heapam.h b/src/include/access/heapam.h
index 1ee5195a8d5..700205b9149 100644
--- a/src/include/access/heapam.h
+++ b/src/include/access/heapam.h
@@ -303,7 +303,8 @@ extern bool heap_getnextslot_tidrange(TableScanDesc sscan,
 									  TupleTableSlot *slot);
 extern TableScanDesc heap_beginscan_bm(Relation relation, Snapshot snapshot, uint32 flags);
 extern void heap_endscan_bm(TableScanDesc scan);
-extern void heap_rescan_bm(TableScanDesc sscan);
+extern void heap_rescan_bm(TableScanDesc sscan, TIDBitmap *tbm,
+						   ParallelBitmapHeapState *pstate, dsa_area *dsa);
 extern bool heap_fetch(Relation relation, Snapshot snapshot,
 					   HeapTuple tuple, Buffer *userbuf, bool keep_buf);
 extern bool heap_hot_search_buffer(ItemPointer tid, Relation relation,
diff --git a/src/include/access/tableam.h b/src/include/access/tableam.h
index 0763ab9847c..e1df701c9a1 100644
--- a/src/include/access/tableam.h
+++ b/src/include/access/tableam.h
@@ -22,6 +22,7 @@
 #include "access/xact.h"
 #include "commands/vacuum.h"
 #include "executor/tuptable.h"
+#include "nodes/execnodes.h"
 #include "nodes/tidbitmap.h"
 #include "utils/rel.h"
 #include "utils/snapshot.h"
@@ -366,7 +367,8 @@ typedef struct TableAmRoutine
 									Snapshot snapshot,
 									uint32 flags);
 
-	void		(*scan_rescan_bm) (TableScanDesc scan);
+	void		(*scan_rescan_bm) (TableScanDesc scan, TIDBitmap *tbm,
+								   ParallelBitmapHeapState *pstate, dsa_area *dsa);
 
 	/*
 	 * Release resources and deallocate scan.
@@ -952,7 +954,8 @@ table_beginscan_strat(Relation rel, Snapshot snapshot,
  */
 static inline TableScanDesc
 table_beginscan_bm(TableScanDesc scan, Relation rel, Snapshot snapshot,
-				   bool need_tuple)
+				   bool need_tuple, TIDBitmap *tbm,
+				   ParallelBitmapHeapState *pstate, dsa_area *dsa)
 {
 	uint32		flags = SO_TYPE_BITMAPSCAN | SO_ALLOW_PAGEMODE;
 
@@ -964,9 +967,13 @@ table_beginscan_bm(TableScanDesc scan, Relation rel, Snapshot snapshot,
 	 * scan descriptor and begin the scan.
 	 */
 	if (!scan)
+	{
 		scan = rel->rd_tableam->scan_begin_bm(rel, snapshot, flags);
+		scan->rs_bhs_iterator.serial = NULL;
+		scan->rs_bhs_iterator.parallel = NULL;
+	}
 
-	scan->rs_rd->rd_tableam->scan_rescan_bm(scan);
+	scan->rs_rd->rd_tableam->scan_rescan_bm(scan, tbm, pstate, dsa);
 
 	return scan;
 }
@@ -1015,10 +1022,6 @@ table_beginscan_tid(Relation rel, Snapshot snapshot)
 static inline void
 table_endscan(TableScanDesc scan)
 {
-	if (scan->rs_flags & SO_TYPE_BITMAPSCAN &&
-		!scan->rs_bhs_iterator.exhausted)
-		unified_tbm_end_iterate(&scan->rs_bhs_iterator);
-
 	scan->rs_rd->rd_tableam->scan_end(scan);
 }
 
@@ -1038,10 +1041,6 @@ static inline void
 table_rescan(TableScanDesc scan,
 			 struct ScanKeyData *key)
 {
-	if (scan->rs_flags & SO_TYPE_BITMAPSCAN &&
-		!scan->rs_bhs_iterator.exhausted)
-		unified_tbm_end_iterate(&scan->rs_bhs_iterator);
-
 	scan->rs_rd->rd_tableam->scan_rescan(scan, key, false, false, false, false);
 }
 
-- 
2.40.1

