Added comment [skip ci]

CHANGELOG.md

@@ -1,7 +1,6 @@
 ## 0.5.1 (unreleased)
 
 - Improved performance of index scans for IVFFlat after updates and deletes
-- Fixed locking for index scans for HNSW
 
 ## 0.5.0 (2023-08-28)
 

Makefile

@@ -8,7 +8,7 @@ HEADERS = src/vector.h
 
 TESTS = $(wildcard test/sql/*.sql)
 REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))
-REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
+REGRESS_OPTS = --inputdir=test --load-extension=vector
 
 OPTFLAGS = -march=native
 

Makefile.win

@@ -5,7 +5,7 @@ OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hn
 HEADERS = src\vector.h
 
 REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
-REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)
+REGRESS_OPTS = --inputdir=test --load-extension=vector
 
 # For /arch flags
 # https://learn.microsoft.com/en-us/cpp/build/reference/arch-minimum-cpu-architecture

src/hnsw.h

@@ -218,6 +218,7 @@ typedef HnswNeighborTupleData * HnswNeighborTuple;
 typedef struct HnswScanOpaqueData
 {
 	bool		first;
+	Buffer		buf;
 	List	   *w;
 	MemoryContext tmpCtx;
 

src/hnswscan.c

@@ -101,6 +101,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 	scan = RelationGetIndexScan(index, nkeys, norderbys);
 
 	so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
+	so->buf = InvalidBuffer;
 	so->first = true;
 	so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
 									   "Hnsw scan temporary context",
@@ -113,12 +114,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
-	/*
-	 * Get a shared lock. This allows vacuum to ensure no in-flight scans
-	 * before marking tuples as deleted.
-	 */
-	LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
-
 	return scan;
 }
 
@@ -169,8 +164,17 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		/* Get scan value */
 		value = GetScanValue(scan);
 
+		/*
+		 * Get a shared lock. This allows vacuum to ensure no in-flight scans
+		 * before marking tuples as deleted.
+		 */
+		LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
+
 		so->w = GetScanItems(scan, value);
 
+		/* Release shared lock */
+		UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
+
 		so->first = false;
 	}
 
@@ -178,6 +182,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 	{
 		HnswCandidate *hc = llast(so->w);
 		ItemPointer heaptid;
+		BlockNumber indexblkno;
 
 		/* Move to next element if no valid heap TIDs */
 		if (list_length(hc->element->heaptids) == 0)
@@ -187,6 +192,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		}
 
 		heaptid = llast(hc->element->heaptids);
+		indexblkno = hc->element->blkno;
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
@@ -198,14 +204,17 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 		scan->xs_ctup.t_self = *heaptid;
 #endif
 
+		/* Unpin buffer */
+		if (BufferIsValid(so->buf))
+			ReleaseBuffer(so->buf);
+
 		/*
-		 * Typically, an index scan must maintain a pin on the index page
-		 * holding the item last returned by amgettuple. However, this is not
-		 * needed with the current vacuum strategy, which ensures scans do not
-		 * visit tuples in danger of being marked as deleted.
+		 * An index scan must maintain a pin on the index page holding the
+		 * item last returned by amgettuple
 		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
+		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -223,8 +232,9 @@ hnswendscan(IndexScanDesc scan)
 {
 	HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
 
-	/* Release shared lock */
-	UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
+	/* Release pin */
+	if (BufferIsValid(so->buf))
+		ReleaseBuffer(so->buf);
 
 	MemoryContextDelete(so->tmpCtx);
 

src/hnswvacuum.c

@@ -330,10 +330,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
 	BufferAccessStrategy bas = vacuumstate->bas;
 	BlockNumber blkno = HNSW_HEAD_BLKNO;
 
-	/*
-	 * Wait for inserts to complete. Inserts before this point may have
-	 * neighbors about to be deleted. Inserts after this point will not.
-	 */
+	/* Wait for inserts to complete */
 	LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
 	UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
 
@@ -446,11 +443,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 	Relation	index = vacuumstate->index;
 	BufferAccessStrategy bas = vacuumstate->bas;
 
-	/*
-	 * Wait for index scans to complete. Scans before this point may contain
-	 * tuples about to be deleted. Scans after this point will not, since the
-	 * graph has been repaired.
-	 */
+	/* Wait for selects to complete */
 	LockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
 	UnlockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
 

src/ivfflat.c

@@ -182,6 +182,15 @@ ivfflatvalidate(Oid opclassoid)
 	return true;
 }
 
+/*
+ * Checks if index-only scan is supported
+ */
+static bool
+ivfflatcanreturn(Relation index, int attno)
+{
+	return attno == 1 && IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC) == NULL;
+}
+
 /*
  * Define index handler
  *
@@ -223,7 +232,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
 	amroutine->aminsert = ivfflatinsert;
 	amroutine->ambulkdelete = ivfflatbulkdelete;
 	amroutine->amvacuumcleanup = ivfflatvacuumcleanup;
-	amroutine->amcanreturn = NULL;	/* tuple not included in heapsort */
+	amroutine->amcanreturn = ivfflatcanreturn;
 	amroutine->amcostestimate = ivfflatcostestimate;
 	amroutine->amoptions = ivfflatoptions;
 	amroutine->amproperty = NULL;	/* TODO AMPROP_DISTANCE_ORDERABLE */

src/ivfflat.h

@@ -40,9 +40,6 @@
 #define IVFFLAT_METAPAGE_BLKNO	0
 #define IVFFLAT_HEAD_BLKNO		1	/* first list page */
 
-/* Must correspond to page numbers since page lock is used */
-#define IVFFLAT_SCAN_LOCK		0
-
 /* IVFFlat parameters */
 #define IVFFLAT_DEFAULT_LISTS	100
 #define IVFFLAT_MIN_LISTS		1
@@ -249,9 +246,9 @@ typedef struct IvfflatScanOpaqueData
 	int			probes;
 	int			dimensions;
 	bool		first;
-	bool		hasLock;
 	Buffer		buf;
 	ItemPointerData heaptid;
+	IndexTuple	itup;
 
 	/* Sorting */
 	Tuplesortstate *sortstate;

src/ivfscan.c

@@ -9,7 +9,6 @@
 #include "miscadmin.h"
 #include "pgstat.h"
 #include "storage/bufmgr.h"
-#include "storage/lmgr.h"
 
 /*
  * Compare list distances
@@ -142,6 +141,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				IndexTuple	itup;
 				Datum		datum;
 				bool		isnull;
+				ItemPointerData indextid;
 				ItemId		itemid = PageGetItemId(page, offno);
 
 				/* Skip dead tuples */
@@ -150,6 +150,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 
 				itup = (IndexTuple) PageGetItem(page, itemid);
 				datum = index_getattr(itup, 1, tupdesc, &isnull);
+				ItemPointerSet(&indextid, searchPage, offno);
 
 				/*
 				 * Add virtual tuple
@@ -162,7 +163,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
 				slot->tts_isnull[0] = false;
 				slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
 				slot->tts_isnull[1] = false;
-				slot->tts_values[2] = Int32GetDatum((int) searchPage);
+				slot->tts_values[2] = PointerGetDatum(&indextid);
 				slot->tts_isnull[2] = false;
 				ExecStoreVirtualTuple(slot);
 
@@ -237,6 +238,34 @@ MarkPriorTupleDead(IndexScanDesc scan)
 	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
 }
 
+/*
+ * Set tuple for index-only scan
+ */
+static void
+SetIndexTuple(IndexScanDesc scan, ItemPointer indextid)
+{
+	IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
+	Buffer		buf = so->buf;
+	Page		page;
+	OffsetNumber offno = ItemPointerGetOffsetNumber(indextid);
+	IndexTuple	itup;
+	Size		itupSize;
+
+	LockBuffer(buf, BUFFER_LOCK_SHARE);
+	page = BufferGetPage(buf);
+	itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
+	itupSize = IndexTupleSize(itup);
+
+	if (so->itup == NULL)
+		so->itup = palloc(BLCKSZ);
+
+	memcpy(so->itup, itup, itupSize);
+
+	scan->xs_itup = so->itup;
+
+	LockBuffer(buf, BUFFER_LOCK_UNLOCK);
+}
+
 /*
  * Prepare for an index scan
  */
@@ -264,8 +293,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 	so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
 	so->buf = InvalidBuffer;
 	so->first = true;
-	so->hasLock = false;
 	ItemPointerSetInvalid(&so->heaptid);
+	so->itup = NULL;
 	so->probes = probes;
 	so->dimensions = dimensions;
 
@@ -282,7 +311,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 #endif
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
 	TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
-	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
+	TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indextid", TIDOID, -1, 0);
 
 	/* Prep sort */
 	so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -297,6 +326,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
+	scan->xs_itupdesc = RelationGetDescr(index);
+
 	return scan;
 }
 
@@ -349,13 +380,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		if (scan->orderByData == NULL)
 			elog(ERROR, "cannot scan ivfflat index without order");
 
-		/* Get a shared lock for non-MVCC snapshots */
-		if (!so->hasLock && !IsMVCCSnapshot(scan->xs_snapshot))
-		{
-			so->hasLock = true;
-			LockPage(scan->indexRelation, IVFFLAT_SCAN_LOCK, ShareLock);
-		}
-
 		if (scan->orderByData->sk_flags & SK_ISNULL)
 			value = PointerGetDatum(InitVector(so->dimensions));
 		else
@@ -389,7 +413,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 	if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
 	{
 		ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
-		BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
+		ItemPointer indextid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 3, &so->isnull));
 
 #if PG_VERSION_NUM >= 120000
 		scan->xs_heaptid = *heaptid;
@@ -410,7 +434,11 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
 		 *
 		 * https://www.postgresql.org/docs/current/index-locking.html
 		 */
-		so->buf = ReadBuffer(scan->indexRelation, indexblkno);
+		so->buf = ReadBuffer(scan->indexRelation, ItemPointerGetBlockNumber(indextid));
+
+		/* Set tuple for index-only scan */
+		if (scan->xs_want_itup)
+			SetIndexTuple(scan, indextid);
 
 		scan->xs_recheckorderby = false;
 		return true;
@@ -431,9 +459,8 @@ ivfflatendscan(IndexScanDesc scan)
 	if (BufferIsValid(so->buf))
 		ReleaseBuffer(so->buf);
 
-	/* Release lock */
-	if (so->hasLock)
-		UnlockPage(scan->indexRelation, IVFFLAT_SCAN_LOCK, ShareLock);
+	if (so->itup != NULL)
+		pfree(so->itup);
 
 	pairingheap_free(so->listQueue);
 	tuplesort_end(so->sortstate);

src/ivfvacuum.c

@@ -3,7 +3,6 @@
 #include "commands/vacuum.h"
 #include "ivfflat.h"
 #include "storage/bufmgr.h"
-#include "storage/lmgr.h"
 
 /*
  * Bulk delete tuples from the index
@@ -66,10 +65,14 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 
 				vacuum_delay_point();
 
-				/* Ensure no in-flight index scans for non-MVCC snapshots */
-				LockPage(index, IVFFLAT_SCAN_LOCK, ExclusiveLock);
-
 				buf = ReadBufferExtended(index, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
+
+				/*
+				 * ambulkdelete cannot delete entries from pages that are
+				 * pinned by other backends
+				 *
+				 * https://www.postgresql.org/docs/current/index-locking.html
+				 */
 				LockBufferForCleanup(buf);
 
 				state = GenericXLogStart(index);
@@ -111,8 +114,6 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
 					GenericXLogAbort(state);
 
 				UnlockReleaseBuffer(buf);
-
-				UnlockPage(index, IVFFLAT_SCAN_LOCK, ExclusiveLock);
 			}
 
 			/*