Use separate memory context for updating neighbors, which improves performance around 10% for larger vectors

ankane · ankane · commit 57248ba128f1 · 2024-09-30T11:15:27.000-07:00
diff --git 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
@@ -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)
