From 84ca6625a77d693b02040bc4dd51cd59b47580fb Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sun, 22 Oct 2023 15:58:40 -0700
Subject: [PATCH] Added support for index-only scans to HNSW

---
 src/hnsw.c      | 11 ++++++++++-
 src/hnsw.h      |  2 +-
 src/hnswscan.c  | 25 ++++++++++++++++++++++---
 src/hnswutils.c |  4 ++--
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/hnsw.c b/src/hnsw.c
index 758e418..a383312 100644
--- a/src/hnsw.c
+++ b/src/hnsw.c
@@ -155,6 +155,15 @@ hnswvalidate(Oid opclassoid)
 	return true;
 }
 
+/*
+ * Checks if index-only scan is supported
+ */
+static bool
+hnswcanreturn(Relation indexRelation, int attno)
+{
+	return attno == 1 && !OidIsValid(index_getprocid(indexRelation, 1, HNSW_NORM_PROC));
+}
+
 /*
  * Define index handler
  *
@@ -196,7 +205,7 @@ hnswhandler(PG_FUNCTION_ARGS)
 	amroutine->aminsert = hnswinsert;
 	amroutine->ambulkdelete = hnswbulkdelete;
 	amroutine->amvacuumcleanup = hnswvacuumcleanup;
-	amroutine->amcanreturn = NULL;
+	amroutine->amcanreturn = hnswcanreturn;
 	amroutine->amcostestimate = hnswcostestimate;
 	amroutine->amoptions = hnswoptions;
 	amroutine->amproperty = NULL;	/* TODO AMPROP_DISTANCE_ORDERABLE */
diff --git a/src/hnsw.h b/src/hnsw.h
index eb2aa9f..ecfde27 100644
--- a/src/hnsw.h
+++ b/src/hnsw.h
@@ -269,7 +269,7 @@ Buffer		HnswNewBuffer(Relation index, ForkNumber forkNum);
 void		HnswInitPage(Buffer buf, Page page);
 void		HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
 void		HnswInit(void);
-List	   *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
+List	   *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool loadVec, HnswElement skipElement);
 HnswElement HnswGetEntryPoint(Relation index);
 void		HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
 HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
diff --git a/src/hnswscan.c b/src/hnswscan.c
index 7cf2bf0..ac97250 100644
--- a/src/hnswscan.c
+++ b/src/hnswscan.c
@@ -17,6 +17,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
 	Relation	index = scan->indexRelation;
 	FmgrInfo   *procinfo = so->procinfo;
 	Oid			collation = so->collation;
+	bool		loadVec = scan->xs_want_itup;
 	List	   *ep;
 	List	   *w;
 	int			m;
@@ -28,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
 	if (entryPoint == NULL)
 		return NIL;
 
-	ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
+	ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, loadVec));
 
 	for (int lc = entryPoint->level; lc >= 1; lc--)
 	{
-		w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
+		w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, loadVec, NULL);
 		ep = w;
 	}
 
-	return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
+	return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, loadVec, NULL);
 }
 
 /*
@@ -113,6 +114,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
 
 	scan->opaque = so;
 
+	/* OK to always set since cheap */
+	scan->xs_itupdesc = RelationGetDescr(index);
+
 	return scan;
 }
 
@@ -198,6 +202,18 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
 
 		hc->element->heaptids = list_delete_last(hc->element->heaptids);
 
+		if (scan->xs_want_itup)
+		{
+			Datum		value = PointerGetDatum(hc->element->vec);
+			bool		isnull = false;
+
+			if (scan->xs_itup)
+				pfree(scan->xs_itup);
+
+			scan->xs_itup = index_form_tuple(scan->xs_itupdesc, &value, &isnull);
+			scan->xs_itup->t_tid = *heaptid;
+		}
+
 		MemoryContextSwitchTo(oldCtx);
 
 #if PG_VERSION_NUM >= 120000
@@ -226,4 +242,7 @@ hnswendscan(IndexScanDesc scan)
 
 	pfree(so);
 	scan->opaque = NULL;
+
+	if (scan->xs_itup)
+		pfree(scan->xs_itup);
 }
diff --git a/src/hnswutils.c b/src/hnswutils.c
index e7d1705..c81425d 100644
--- a/src/hnswutils.c
+++ b/src/hnswutils.c
@@ -569,7 +569,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
  * Algorithm 2 from paper
  */
 List *
-HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
+HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool loadVec, HnswElement skipElement)
 {
 	ListCell   *lc2;
 
@@ -645,7 +645,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
 				if (index == NULL)
 					eDistance = GetCandidateDistance(e, q, procinfo, collation);
 				else
-					HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
+					HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, loadVec);
 
 				Assert(!e->element->deleted);