From 011e5438aba40086086473a1b43e8146c029f42f Mon Sep 17 00:00:00 2001
From: Tomas Vondra <tomas@vondra.me>
Date: Thu, 1 May 2025 23:17:20 +0200
Subject: [PATCH v20250501 5/7] fix: mark/restore

call AM mark/restore even with batching, so that the AM can set some
fields (e.g. moreRight/moreLeft in nbtree)
---
 src/backend/access/index/indexam.c | 14 ++++-----
 src/backend/access/nbtree/nbtree.c | 47 +++++++++++++++++++++++-------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/src/backend/access/index/indexam.c b/src/backend/access/index/indexam.c
index f0fda6d761c..20d4d65487f 100644
--- a/src/backend/access/index/indexam.c
+++ b/src/backend/access/index/indexam.c
@@ -516,11 +516,9 @@ index_markpos(IndexScanDesc scan)
 	 * Without batching, just use the ammarkpos() callback. With batching
 	 * everything is handled at this layer, without calling the AM.
 	 */
-	if (scan->xs_batches == NULL)
-	{
-		scan->indexRelation->rd_indam->ammarkpos(scan);
-	}
-	else
+	scan->indexRelation->rd_indam->ammarkpos(scan);
+
+	if (scan->xs_batches != NULL)
 	{
 		IndexScanBatches *batches = scan->xs_batches;
 		IndexScanBatchPos *pos = &batches->markPos;
@@ -588,9 +586,9 @@ index_restrpos(IndexScanDesc scan)
 	 * Without batching, just use the amrestrpos() callback. With batching
 	 * everything is handled at this layer, without calling the AM.
 	 */
-	if (scan->xs_batches == NULL)
-		scan->indexRelation->rd_indam->amrestrpos(scan);
-	else
+	scan->indexRelation->rd_indam->amrestrpos(scan);
+
+	if (scan->xs_batches != NULL)
 	{
 		IndexScanBatches *batches = scan->xs_batches;
 		IndexScanBatchPos *pos = &batches->markPos;
diff --git a/src/backend/access/nbtree/nbtree.c b/src/backend/access/nbtree/nbtree.c
index 6ea1dfcc52c..0143df993aa 100644
--- a/src/backend/access/nbtree/nbtree.c
+++ b/src/backend/access/nbtree/nbtree.c
@@ -309,10 +309,13 @@ btgetbatch(IndexScanDesc scan, ScanDirection dir)
 		IndexScanBatch batch = INDEX_SCAN_BATCH(scan, scan->xs_batches->nextBatch-1);
 		pos = (BTBatchScanPos) batch->opaque;
 
-		if (ScanDirectionIsForward(scan->xs_batches->direction))
-			pos->moreRight = true;
-		else
-			pos->moreLeft = true;
+		if (so->needPrimScan)
+		{
+			if (ScanDirectionIsForward(scan->xs_batches->direction))
+				pos->moreRight = true;
+			else
+				pos->moreLeft = true;
+		}
 	}
 
 	/* Each loop iteration performs another primitive index scan */
@@ -327,11 +330,7 @@ btgetbatch(IndexScanDesc scan, ScanDirection dir)
 			res = _bt_first_batch(scan, dir);
 		else
 		{
-			if (so->numArrayKeys)
-			{
-				_bt_start_array_keys(scan, so->currPos.dir);
-				so->needPrimScan = false;
-			}
+			so->needPrimScan = false;
 
 			/*
 			 * Now continue the scan.
@@ -656,7 +655,20 @@ btmarkpos(IndexScanDesc scan)
 	BTScanOpaque so = (BTScanOpaque) scan->opaque;
 
 	/* with batching, mark/restore is handled in indexam */
-	Assert(scan->xs_batches == NULL);
+	if (scan->xs_batches != NULL)
+	{
+		IndexScanBatch	batch = INDEX_SCAN_BATCH(scan, scan->xs_batches->firstBatch);
+		BTBatchScanPos pos = NULL;
+		pos = (BTBatchScanPos) batch->opaque;
+		if (so->needPrimScan)
+		{
+			if (ScanDirectionIsForward(scan->xs_batches->direction))
+				pos->moreRight = true;
+			else
+				pos->moreLeft = true;
+		}
+		return;
+	}
 
 	/* There may be an old mark with a pin (but no lock). */
 	BTScanPosUnpinIfPinned(so->markPos);
@@ -688,7 +700,20 @@ btrestrpos(IndexScanDesc scan)
 	BTScanOpaque so = (BTScanOpaque) scan->opaque;
 
 	/* with batching, mark/restore is handled in indexam */
-	Assert(scan->xs_batches == NULL);
+	if (scan->xs_batches != NULL)
+	{
+		if (scan->xs_batches->markPos.batch != scan->xs_batches->firstBatch)
+		{
+			/* Reset the scan's array keys (see _bt_steppage for why) */
+			if (so->numArrayKeys)
+			{
+				_bt_start_array_keys(scan, so->currPos.dir);
+				so->needPrimScan = false;
+			}
+		}
+
+		return;
+	}
 
 	if (so->markItemIndex >= 0)
 	{
-- 
2.49.0