From 57248ba128f1f4b3773d0575de8a8c936e3eee59 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 30 Sep 2024 11:15:27 -0700
Subject: [PATCH] Use separate memory context for updating neighbors, which
 improves performance around 10% for larger vectors

---
 CHANGELOG.md     |  1 +
 src/hnswinsert.c | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e5a17b3..a7d9924 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 - Added casts for arrays to `sparsevec`
 - Improved cost estimation
+- Improved performance of HNSW inserts and on-disk index builds
 - Reduced memory usage for HNSW index scans
 - Dropped support for Postgres 12
 
diff --git a/src/hnswinsert.c b/src/hnswinsert.c
index a3949b2..2f18b12 100644
--- a/src/hnswinsert.c
+++ b/src/hnswinsert.c
@@ -394,11 +394,12 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation
  * Get update index
  */
 static int
-GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo *procinfo, Oid collation)
+GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo *procinfo, Oid collation, MemoryContext updateCtx)
 {
 	char	   *base = NULL;
 	int			idx = -1;
 	HnswNeighborArray *neighbors;
+	MemoryContext oldCtx = MemoryContextSwitchTo(updateCtx);
 
 	/*
 	 * Get latest neighbors since they may have changed. Do not lock yet since
@@ -426,6 +427,9 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
 			HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation);
 	}
 
+	MemoryContextSwitchTo(oldCtx);
+	MemoryContextReset(updateCtx);
+
 	return idx;
 }
 
@@ -529,6 +533,14 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
 {
 	char	   *base = NULL;
 
+	/* Use separate memory context to improve performance for larger vectors */
+	MemoryContext updateCtx = GenerationContextCreate(CurrentMemoryContext,
+													  "Hnsw insert update context",
+#if PG_VERSION_NUM >= 150000
+													  128 * 1024, 128 * 1024,
+#endif
+													  128 * 1024);
+
 	for (int lc = e->level; lc >= 0; lc--)
 	{
 		int			lm = HnswGetLayerM(m, lc);
@@ -540,7 +552,7 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
 			HnswElement neighborElement = HnswPtrAccess(base, hc->element);
 			int			idx;
 
-			idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation);
+			idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation, updateCtx);
 
 			/* New element was not selected as a neighbor */
 			if (idx == -1)