Improved performance of HNSW inserts with sparsevec - resolves #1025 and resolves #1026

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.9.0 (unreleased)
 
 - Added `vector_dims` function for `sparsevec`
+- Improved performance of HNSW inserts with `sparsevec`
 - Improved `amvalidate` for indexes
 - Dropped support for Postgres 13 and 14
 

src/hnsw.h

@@ -281,6 +281,7 @@ typedef struct HnswTypeInfo
 	Datum		(*dimensions) (PG_FUNCTION_ARGS);
 	Datum		(*normalize) (PG_FUNCTION_ARGS);
 	void		(*checkValue) (Pointer v);
+	bool		useFsm;
 }			HnswTypeInfo;
 
 typedef struct HnswSupport
@@ -472,6 +473,7 @@ void		HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element
 void		HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
 bool		HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
 void		HnswInitLockTranche(void);
+void		HnswRecordPageWithFreeSpace(Relation index, BlockNumber blkno, Page page);
 void		HnswCheckDim(int expected, const HnswTypeInfo * typeInfo, Oid collation, Datum value);
 const		HnswTypeInfo *HnswGetTypeInfo(Relation index);
 PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);

src/hnswbuild.c

@@ -55,6 +55,7 @@
 #include "optimizer/optimizer.h"
 #include "storage/bufmgr.h"
 #include "storage/condition_variable.h"
+#include "storage/freespace.h"
 #include "tcop/tcopprot.h"
 #include "utils/backend_progress.h"
 #include "utils/backend_status.h"
@@ -1106,6 +1107,34 @@ BuildGraph(HnswBuildState * buildstate)
 		HnswEndParallel(buildstate->hnswleader);
 }
 
+/*
+ * Record free space
+ */
+static void
+RecordFreeSpace(Relation index)
+{
+	BlockNumber blkno = HNSW_HEAD_BLKNO;
+
+	while (BlockNumberIsValid(blkno))
+	{
+		Buffer		buf;
+		Page		page;
+
+		buf = ReadBuffer(index, blkno);
+		LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
+		page = BufferGetPage(buf);
+
+		/* Update free space map */
+		HnswRecordPageWithFreeSpace(index, blkno, page);
+
+		blkno = HnswPageGetOpaque(page)->nextblkno;
+		UnlockReleaseBuffer(buf);
+	}
+
+	/* Ensure free space becomes visible */
+	FreeSpaceMapVacuum(index);
+}
+
 /*
  * Build the index
  */
@@ -1124,6 +1153,9 @@ BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
 	if (RelationNeedsWAL(index) || forkNum == INIT_FORKNUM)
 		log_newpage_range(index, forkNum, 0, RelationGetNumberOfBlocksInFork(index, forkNum), true);
 
+	if (buildstate->typeInfo->useFsm && forkNum == MAIN_FORKNUM)
+		RecordFreeSpace(index);
+
 	FreeBuildState(buildstate);
 }
 

src/hnswinsert.c

@@ -2,9 +2,11 @@
 
 #include "access/genam.h"
 #include "access/generic_xlog.h"
+#include "access/htup_details.h"
 #include "hnsw.h"
 #include "nodes/execnodes.h"
 #include "storage/bufmgr.h"
+#include "storage/freespace.h"
 #include "storage/lmgr.h"
 #include "storage/lwlock.h"
 #include "utils/datum.h"
@@ -15,6 +17,8 @@
 #include "varatt.h"
 #endif
 
+#define HNSW_MAX_FSM_TRIES 3
+
 /*
  * Get the insert page
  */
@@ -42,7 +46,7 @@ GetInsertPage(Relation index)
  * Check for a free offset
  */
 static bool
-HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion)
+HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion, bool useFsm)
 {
 	OffsetNumber offno;
 	OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -65,7 +69,7 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
 			Size		pageFree;
 			Size		npageFree;
 
-			if (!BlockNumberIsValid(*newInsertPage))
+			if (!BlockNumberIsValid(*newInsertPage) && !useFsm)
 				*newInsertPage = elementPage;
 
 			if (neighborPage == elementPage)
@@ -142,7 +146,7 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
  * Add to element and neighbor pages
  */
 static void
-AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building)
+AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber *updatedInsertPage, bool building)
 {
 	Buffer		buf;
 	Page		page;
@@ -153,7 +157,8 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 	Size		maxSize;
 	Size		minCombinedSize;
 	HnswElementTuple etup;
-	BlockNumber currentPage = insertPage;
+	BlockNumber insertPage = InvalidBlockNumber;
+	BlockNumber currentPage = InvalidBlockNumber;
 	HnswNeighborTuple ntup;
 	Buffer		nbuf;
 	Page		npage;
@@ -162,6 +167,11 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 	BlockNumber newInsertPage = InvalidBlockNumber;
 	uint8		tupleVersion;
 	char	   *base = NULL;
+	bool		useFsm = !building && HnswGetTypeInfo(index)->useFsm;
+	bool		tryFsm;
+	int			fsmTries = 0;
+	bool		fsmPage;
+	BlockNumber fsmHistory[HNSW_MAX_FSM_TRIES];
 
 	/* Calculate sizes */
 	etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
@@ -178,9 +188,48 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 	ntup = palloc0(ntupSize);
 	HnswSetNeighborTuple(base, ntup, e, m);
 
+	/* Ensure within MaxFSMRequestSize */
+	tryFsm = useFsm && combinedSize <= MaxHeapTupleSize;
+
 	/* Find a page (or two if needed) to insert the tuples */
 	for (;;)
 	{
+		/* Try free space map */
+		if (tryFsm && fsmTries < HNSW_MAX_FSM_TRIES)
+		{
+			currentPage = GetPageWithFreeSpace(index, combinedSize);
+			fsmHistory[fsmTries++] = currentPage;
+
+			/*
+			 * Even if a page is returned, it may not be possible to add the
+			 * tuples (see HnswFreeOffset for details), so avoid trying the
+			 * same page again
+			 */
+			if (BlockNumberIsValid(currentPage))
+			{
+				for (int i = 0; i < fsmTries - 1; i++)
+				{
+					if (fsmHistory[i] == currentPage)
+					{
+						currentPage = InvalidBlockNumber;
+						break;
+					}
+				}
+			}
+
+			fsmPage = BlockNumberIsValid(currentPage);
+			tryFsm = BlockNumberIsValid(currentPage);
+		}
+		else
+			fsmPage = false;
+
+		/* Start at insert page if free space map yields nothing */
+		if (!BlockNumberIsValid(currentPage))
+		{
+			insertPage = GetInsertPage(index);
+			currentPage = insertPage;
+		}
+
 		buf = ReadBuffer(index, currentPage);
 		LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
 
@@ -196,7 +245,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 		}
 
 		/* Keep track of first page where element at level 0 can fit */
-		if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize)
+		if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize && !useFsm)
 			newInsertPage = currentPage;
 
 		/* First, try the fastest path */
@@ -210,7 +259,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 		}
 
 		/* Next, try space from a deleted element */
-		if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
+		if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion, useFsm))
 		{
 			if (nbuf != buf)
 			{
@@ -227,6 +276,14 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 			break;
 		}
 
+		/* Try another page */
+		if (fsmPage)
+		{
+			UnlockReleaseBuffer(buf);
+			currentPage = InvalidBlockNumber;
+			continue;
+		}
+
 		/* Finally, try space for element only if last page */
 		/* Skip if both tuples can fit on the same page */
 		if (combinedSize > maxSize && PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
@@ -291,7 +348,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 
 	/* Added tuple to new page if newInsertPage is not set */
 	/* So can set to neighbor page instead of element page */
-	if (!BlockNumberIsValid(newInsertPage))
+	if (!BlockNumberIsValid(newInsertPage) && !fsmPage)
 		newInsertPage = e->neighborPage;
 
 	if (OffsetNumberIsValid(freeOffno))
@@ -328,6 +385,15 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 			elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
 	}
 
+	/* Update free space map */
+	if (useFsm)
+	{
+		HnswRecordPageWithFreeSpace(index, e->blkno, page);
+
+		if (npage != page)
+			HnswRecordPageWithFreeSpace(index, e->neighborPage, npage);
+	}
+
 	/* Commit */
 	if (building)
 	{
@@ -342,7 +408,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
 		UnlockReleaseBuffer(nbuf);
 
 	/* Update the insert page */
-	if (BlockNumberIsValid(newInsertPage) && newInsertPage != insertPage)
+	if (BlockNumberIsValid(newInsertPage) && BlockNumberIsValid(insertPage) && newInsertPage != insertPage)
 		*updatedInsertPage = newInsertPage;
 }
 
@@ -673,7 +739,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
 		return;
 
 	/* Add element */
-	AddElementOnDisk(index, element, m, GetInsertPage(index), &newInsertPage, building);
+	AddElementOnDisk(index, element, m, &newInsertPage, building);
 
 	/* Update insert page if needed */
 	if (BlockNumberIsValid(newInsertPage))

src/hnswutils.c

@@ -12,6 +12,7 @@
 #include "port/atomics.h"
 #include "sparsevec.h"
 #include "storage/bufmgr.h"
+#include "storage/freespace.h"
 #include "utils/datum.h"
 #include "utils/fmgrprotos.h"
 #include "utils/memdebug.h"
@@ -1363,6 +1364,48 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
 	}
 }
 
+/*
+ * Record free space
+ */
+void
+HnswRecordPageWithFreeSpace(Relation index, BlockNumber blkno, Page page)
+{
+	Size		freeSpace = 0;
+	OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
+
+	for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
+	{
+		ItemId		eitemid = PageGetItemId(page, offno);
+		HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, eitemid);
+		Size		itemsize;
+		BlockNumber neighborPage;
+		OffsetNumber neighborOffno;
+
+		/* Skip neighbor tuples */
+		if (!HnswIsElementTuple(etup))
+			continue;
+
+		/* Skip live tuples */
+		if (!etup->deleted)
+			continue;
+
+		/* Get neighbor page */
+		neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
+		neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
+
+		/* Keep track of largest tuple pair that will fit */
+		itemsize = ItemIdGetLength(eitemid);
+		if (neighborPage == blkno)
+			itemsize += ItemIdGetLength(PageGetItemId(page, neighborOffno)) + sizeof(ItemIdData);
+
+		freeSpace = Max(freeSpace, itemsize);
+	}
+
+	freeSpace += PageGetFreeSpace(page);
+
+	RecordPageWithFreeSpace(index, blkno, freeSpace);
+}
+
 /*
  * Check dimensions match index
  */
@@ -1410,7 +1453,8 @@ HnswGetTypeInfo(Relation index)
 			.maxDimensions = HNSW_MAX_DIM,
 			.dimensions = vector_dims,
 			.normalize = l2_normalize,
-			.checkValue = NULL
+			.checkValue = NULL,
+			.useFsm = false
 		};
 
 		return (&typeInfo);
@@ -1427,7 +1471,8 @@ hnsw_halfvec_support(PG_FUNCTION_ARGS)
 		.maxDimensions = HNSW_MAX_DIM * 2,
 		.dimensions = halfvec_vector_dims,
 		.normalize = halfvec_l2_normalize,
-		.checkValue = NULL
+		.checkValue = NULL,
+		.useFsm = false
 	};
 
 	PG_RETURN_POINTER(&typeInfo);
@@ -1441,7 +1486,8 @@ hnsw_bit_support(PG_FUNCTION_ARGS)
 		.maxDimensions = HNSW_MAX_DIM * 32,
 		.dimensions = bitlength,
 		.normalize = NULL,
-		.checkValue = NULL
+		.checkValue = NULL,
+		.useFsm = false
 	};
 
 	PG_RETURN_POINTER(&typeInfo);
@@ -1455,7 +1501,8 @@ hnsw_sparsevec_support(PG_FUNCTION_ARGS)
 		.maxDimensions = SPARSEVEC_MAX_DIM,
 		.dimensions = sparsevec_vector_dims,
 		.normalize = sparsevec_l2_normalize,
-		.checkValue = SparsevecCheckValue
+		.checkValue = SparsevecCheckValue,
+		.useFsm = true
 	};
 
 	PG_RETURN_POINTER(&typeInfo);

src/hnswvacuum.c

@@ -6,6 +6,7 @@
 #include "hnsw.h"
 #include "nodes/pg_list.h"
 #include "storage/bufmgr.h"
+#include "storage/freespace.h"
 #include "storage/lmgr.h"
 #include "utils/memutils.h"
 #include "utils/rel.h"
@@ -598,6 +599,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 	BlockNumber insertPage = InvalidBlockNumber;
 	Relation	index = vacuumstate->index;
 	BufferAccessStrategy bas = vacuumstate->bas;
+	bool		useFsm = HnswGetTypeInfo(index)->useFsm;
 
 	/*
 	 * Wait for inserts and index scans to complete. Inserts and scans before
@@ -710,7 +712,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 				UnlockReleaseBuffer(nbuf);
 
 			/* Set to first free page */
-			if (!BlockNumberIsValid(insertPage))
+			if (!useFsm && !BlockNumberIsValid(insertPage))
 				insertPage = blkno;
 
 			/* Prepare new xlog */
@@ -718,12 +720,25 @@ MarkDeleted(HnswVacuumState * vacuumstate)
 			page = GenericXLogRegisterBuffer(state, buf, 0);
 		}
 
+		if (useFsm)
+		{
+			/* Update free space map */
+			HnswRecordPageWithFreeSpace(index, blkno, page);
+
+			/* Set to last page */
+			insertPage = blkno;
+		}
+
 		blkno = HnswPageGetOpaque(page)->nextblkno;
 
 		GenericXLogAbort(state);
 		UnlockReleaseBuffer(buf);
 	}
 
+	/* Ensure increases in free space become visible */
+	if (useFsm)
+		FreeSpaceMapVacuum(index);
+
 	/* Update insert page last, after everything has been marked as deleted */
 	HnswUpdateMetaPage(index, 0, NULL, insertPage, MAIN_FORKNUM, false);
 }

test/t/038_hnsw_sparsevec_vacuum_insert.pl

@@ -20,9 +20,9 @@ for (1 .. 3)
 	{
 		my @elements;
 		my %indices;
-		for (1 .. int(rand() * 100))
+		for (1 .. int(rand(1001)))
 		{
-			my $index = int(rand() * (100000 - 1)) + 1;
+			my $index = int(rand(100000)) + 1;
 			if (!exists($indices{$index}))
 			{
 				my $value = rand();