From c5b2f3ac8bd404017ab6dcd18b01b3396a386882 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sun, 13 Aug 2023 16:50:36 -0700
Subject: [PATCH] Improved HNSW vacuum performance

---
 src/hnsw.h       |  3 ++-
 src/hnswinsert.c | 33 ++++++++++++++++++++++++++++++---
 src/hnswvacuum.c | 10 +++++-----
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/src/hnsw.h b/src/hnsw.h
index c9aabf6..b6e680d 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -268,11 +268,12 @@ void		HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
 void		HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
 void		HnswInitNeighbors(HnswElement element, int m);
 bool		HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
-void		UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m);
+void		UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, HTAB *deleted);
 void		HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
 void		HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
 void		HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
 void		HnswLoadNeighbors(HnswElement element, Relation index);
+bool		HnswDeletedContains(HTAB *deleted, ItemPointer indextid);
 
 /* Index access methods */
 IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
diff --git a/src/hnswinsert.c b/src/hnswinsert.c
index b397941..fea75aa 100644
--- a/src/hnswinsert.c
+++ b/src/hnswinsert.c
@@ -283,11 +283,33 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
 		*updatedInsertPage = newInsertPage;
 }
 
+/*
+ * Get index of deleted neighbor
+ */
+static int
+GetDeletedIndex(HnswElement element, HTAB *deleted, int lc)
+{
+	HnswNeighborArray *neighbors = &element->neighbors[lc];
+
+	for (int i = 0; i < neighbors->length; i++)
+	{
+		HnswCandidate *hc = &neighbors->items[i];
+		ItemPointerData ipData;
+
+		ItemPointerSet(&ipData, hc->element->blkno, hc->element->offno);
+
+		if (HnswDeletedContains(deleted, &ipData))
+			return i;
+	}
+
+	return -1;
+}
+
 /*
  * Update neighbors
  */
 void
-UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
+UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, HTAB *deleted)
 {
 	for (int lc = e->level; lc >= 0; lc--)
 	{
@@ -311,8 +333,13 @@ UpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswEleme
 			/* Do not lock yet since selecting neighbors can take time */
 			HnswLoadNeighbors(hc->element, index);
 
+			/* Find deleted index without loading elements if possible */
+			if (deleted != NULL)
+				idx = GetDeletedIndex(hc->element, deleted, lc);
+
 			/* Select neighbors */
-			HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
+			if (idx == -1)
+				HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
 
 			/* New element was not selected as a neighbor */
 			if (idx == -1)
@@ -484,7 +511,7 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
 		HnswUpdateMetaPage(index, false, NULL, newInsertPage, MAIN_FORKNUM);
 
 	/* Update neighbors */
-	UpdateNeighborPages(index, procinfo, collation, element, m);
+	UpdateNeighborPages(index, procinfo, collation, element, m, NULL);
 
 	/* Update metapage if needed */
 	if (element->level > entryPoint->level)
diff --git a/src/hnswvacuum.c b/src/hnswvacuum.c
index e2f412d..a476b1b 100644
--- a/src/hnswvacuum.c
+++ b/src/hnswvacuum.c
@@ -10,8 +10,8 @@
 /*
  * Check if deleted list contains an index tid
  */
-static bool
-DeletedContains(HTAB *deleted, ItemPointer indextid)
+bool
+HnswDeletedContains(HTAB *deleted, ItemPointer indextid)
 {
 	bool		found;
 
@@ -167,7 +167,7 @@ NeedsUpdated(HnswVacuumState * vacuumstate, HnswElement element)
 			continue;
 
 		/* Check if in deleted list */
-		if (DeletedContains(vacuumstate->deleted, indextid))
+		if (HnswDeletedContains(vacuumstate->deleted, indextid))
 		{
 			needsUpdated = true;
 			break;
@@ -255,7 +255,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element)
 	UnlockReleaseBuffer(buf);
 
 	/* Update neighbors */
-	UpdateNeighborPages(index, procinfo, collation, element, m);
+	UpdateNeighborPages(index, procinfo, collation, element, m, vacuumstate->deleted);
 }
 
 /*
@@ -293,7 +293,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
 
 		ItemPointerSet(&epData, entryPoint->blkno, entryPoint->offno);
 
-		if (DeletedContains(vacuumstate->deleted, &epData))
+		if (HnswDeletedContains(vacuumstate->deleted, &epData))
 			HnswUpdateMetaPage(index, true, highestPoint, InvalidBlockNumber, MAIN_FORKNUM);
 		else
 		{