Added support for in-memory parallel index builds for HNSW

2026-07-10 14:46:56 +08:00 · 2024-01-15 15:07:31 -08:00
parent 0ce497a1b1
commit 9d3e4e74df
10 changed files with 771 additions and 310 deletions
--- a/src/hnswbuild.c
+++ b/src/hnswbuild.c
@@ -54,7 +54,8 @@
 #endif

 #define PARALLEL_KEY_HNSW_SHARED		UINT64CONST(0xA000000000000001)
-#define PARALLEL_KEY_QUERY_TEXT			UINT64CONST(0xA000000000000002)
+#define PARALLEL_KEY_HNSW_AREA			UINT64CONST(0xA000000000000002)
+#define PARALLEL_KEY_QUERY_TEXT			UINT64CONST(0xA000000000000003)

 #if PG_VERSION_NUM < 130000
 #define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2)
@@ -135,9 +136,11 @@ CreateElementPages(HnswBuildState * buildstate)
 	HnswElementTuple etup;
 	HnswNeighborTuple ntup;
 	BlockNumber insertPage;
+	HnswElement entryPoint;
 	Buffer		buf;
 	Page		page;
-	slist_iter	iter;
+	HnswElementPtr iter = buildstate->graph->head;
+	char	   *base = buildstate->hnswarea;

 	/* Calculate sizes */
 	etupAllocSize = BLCKSZ;
@@ -152,18 +155,22 @@ CreateElementPages(HnswBuildState * buildstate)
 	page = BufferGetPage(buf);
 	HnswInitPage(buf, page);

-	slist_foreach(iter, &buildstate->graph->elements)
+	while (!HnswPtrIsNull(base, iter))
 	{
-		HnswElement element = slist_container(HnswElementData, next, iter.cur);
+		HnswElement element = HnswPtrAccess(base, iter);
 		Size		etupSize;
 		Size		ntupSize;
 		Size		combinedSize;
+		void	   *valuePtr = HnswPtrAccess(base, element->value);
+
+		/* Update iterator */
+		iter = element->next;

 		/* Zero memory for each element */
 		MemSet(etup, 0, etupAllocSize);

 		/* Calculate sizes */
-		etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(element->value)));
+		etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
 		ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
 		combinedSize = etupSize + ntupSize + sizeof(ItemIdData);

@@ -171,7 +178,7 @@ CreateElementPages(HnswBuildState * buildstate)
 		if (etupSize > etupAllocSize)
 			elog(ERROR, "index tuple too large");

-		HnswSetElementTuple(etup, element);
+		HnswSetElementTuple(base, etup, element);

 		/* Keep element and neighbors on the same page if possible */
 		if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -212,7 +219,8 @@ CreateElementPages(HnswBuildState * buildstate)
 	MarkBufferDirty(buf);
 	UnlockReleaseBuffer(buf);

-	HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->graph->entryPoint, insertPage, forkNum, true);
+	entryPoint = HnswPtrAccess(base, buildstate->graph->entryPoint);
+	HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, entryPoint, insertPage, forkNum, true);

 	pfree(etup);
 	pfree(ntup);
@@ -227,19 +235,23 @@ CreateNeighborPages(HnswBuildState * buildstate)
 	Relation	index = buildstate->index;
 	ForkNumber	forkNum = buildstate->forkNum;
 	int			m = buildstate->m;
-	slist_iter	iter;
+	HnswElementPtr iter = buildstate->graph->head;
+	char	   *base = buildstate->hnswarea;
 	HnswNeighborTuple ntup;

 	/* Allocate once */
 	ntup = palloc0(BLCKSZ);

-	slist_foreach(iter, &buildstate->graph->elements)
+	while (!HnswPtrIsNull(base, iter))
 	{
-		HnswElement e = slist_container(HnswElementData, next, iter.cur);
+		HnswElement e = HnswPtrAccess(base, iter);
 		Buffer		buf;
 		Page		page;
 		Size		ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);

+		/* Update iterator */
+		iter = e->next;
+
 		/* Can take a while, so ensure we can interrupt */
 		/* Needs to be called when no buffer locks are held */
 		CHECK_FOR_INTERRUPTS();
@@ -248,7 +260,7 @@ CreateNeighborPages(HnswBuildState * buildstate)
 		LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
 		page = BufferGetPage(buf);

-		HnswSetNeighborTuple(ntup, e, m);
+		HnswSetNeighborTuple(base, ntup, e, m);

 		if (!PageIndexTupleOverwrite(page, e->neighborOffno, (Item) ntup, ntupSize))
 			elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
@@ -261,24 +273,6 @@ CreateNeighborPages(HnswBuildState * buildstate)
 	pfree(ntup);
 }

-#ifdef HNSW_MEMORY
-/*
- * Show memory usage
- */
-static void
-ShowMemoryUsage(HnswBuildState * buildstate)
-{
-#if PG_VERSION_NUM >= 130000
-	elog(INFO, "graph memory: %zu MB, total memory: %zu MB",
-		 MemoryContextMemAllocated(buildstate->graphCtx, false) / (1024 * 1024),
-		 MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
-#else
-	MemoryContextStats(CurrentMemoryContext);
-	elog(INFO, "estimated memory: %zu MB", buildstate->memoryUsed / (1024 * 1024));
-#endif
-}
-#endif
-
 /*
 * Flush pages
 */
@@ -286,7 +280,7 @@ static void
 FlushPages(HnswBuildState * buildstate)
 {
 #ifdef HNSW_MEMORY
-	ShowMemoryUsage(buildstate);
+	elog(INFO, "memory: %zu MB", buildstate->graph->memoryUsed / (1024 * 1024));
 #endif

 	CreateMetaPage(buildstate);
@@ -297,66 +291,177 @@ FlushPages(HnswBuildState * buildstate)
 	MemoryContextReset(buildstate->graphCtx);
 }

-#if PG_VERSION_NUM < 130000
 /*
- * Get the memory used by an element
+ * Initialize a readers-writer lock
 */
-static long
-HnswElementMemory(HnswElement e, int m)
+static void
+HnswRWLockInitialize(HnswRWLock * lock)
 {
-	long		elementSize = sizeof(HnswElementData);
-
-	elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
-	elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
-	elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
-	/* Each allocation has a chunk header */
-	elementSize += (e->level + 4) * GENERATIONCHUNK_RAWSIZE;
-	/* Add an extra 5% for alignment and other overhead */
-	return elementSize * 1.05;
+	lock->readers = 0;
+	SpinLockInit(&lock->readersMutex);
+	SpinLockInit(&lock->globalMutex);
+}
+
+/*
+ * Acquire a readers-writer lock
+ */
+static void
+HnswRWLockAcquire(HnswRWLock * lock, HnswLWLockMode lockmode)
+{
+	if (lockmode == RW_EXCLUSIVE)
+		SpinLockAcquire(&lock->globalMutex);
+	else
+	{
+		SpinLockAcquire(&lock->readersMutex);
+		if (++lock->readers == 1)
+			SpinLockAcquire(&lock->globalMutex);
+		SpinLockRelease(&lock->readersMutex);
+	}
+}
+
+/*
+ * Release a readers-writer lock
+ */
+static void
+HnswRWLockRelease(HnswRWLock * lock, HnswLWLockMode lockmode)
+{
+	if (lockmode == RW_EXCLUSIVE)
+		SpinLockRelease(&lock->globalMutex);
+	else
+	{
+		SpinLockAcquire(&lock->readersMutex);
+		if (--lock->readers == 0)
+			SpinLockRelease(&lock->globalMutex);
+		SpinLockRelease(&lock->readersMutex);
+	}
+}
+
+/*
+ * Add a heap TID to an existing element
+ */
+static bool
+HnswAddDuplicateInMemory(HnswElement element, HnswElement dup)
+{
+	SpinLockAcquire(&dup->lock);
+
+	if (dup->heaptidsLength == HNSW_HEAPTIDS)
+	{
+		SpinLockRelease(&dup->lock);
+		return false;
+	}
+
+	HnswAddHeapTid(dup, &element->heaptids[0]);
+
+	SpinLockRelease(&dup->lock);
+
+	return true;
 }
-#endif

 /*
 * Find duplicate element
 */
 static bool
-HnswFindDuplicateInMemory(HnswElement element)
+HnswFindDuplicateInMemory(char *base, HnswElement element)
 {
-	HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0);
+	HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);

 	for (int i = 0; i < neighbors->length; i++)
 	{
 		HnswCandidate *neighbor = &neighbors->items[i];
+		HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
+		Datum		value = HnswGetValue(base, element);
+		Datum		neighborValue = HnswGetValue(base, neighborElement);

 		/* Exit early since ordered by distance */
-		if (!datumIsEqual(element->value, neighbor->element->value, false, -1))
+		if (!datumIsEqual(value, neighborValue, false, -1))
 			return false;

 		/* Check for space */
-		if (neighbor->element->heaptidsLength < HNSW_HEAPTIDS)
-		{
-			HnswAddHeapTid(neighbor->element, &element->heaptids[0]);
+		if (HnswAddDuplicateInMemory(element, neighborElement))
 			return true;
-		}
 	}

 	return false;
 }

 /*
- * Insert tuple into in-memory graph
+ * Add to element and neighbor pages
+ */
+static void
+WriteNewElementPagesInMemory(char *base, HnswGraph * graph, HnswElement element)
+{
+	SpinLockAcquire(&graph->lock);
+	element->next = graph->head;
+	HnswPtrStore(base, graph->head, element);
+	SpinLockRelease(&graph->lock);
+}
+
+/*
+ * Update neighbors
+ */
+static void
+HnswUpdateNeighborPagesInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
+{
+	for (int lc = e->level; lc >= 0; lc--)
+	{
+		int			lm = HnswGetLayerM(m, lc);
+		HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
+
+		for (int i = 0; i < neighbors->length; i++)
+		{
+			HnswCandidate *hc = &neighbors->items[i];
+			HnswElement neighborElement = HnswPtrAccess(base, hc->element);
+
+			/* Use element for lock instead of hc since hc can be replaced */
+			SpinLockAcquire(&neighborElement->lock);
+			HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
+			SpinLockRelease(&neighborElement->lock);
+		}
+	}
+}
+
+/*
+ * Write changes in memory
+ */
+static void
+WriteElementInMemory(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate, HnswGraph * graph, bool updateEntryPoint)
+{
+	char	   *base = buildstate->hnswarea;
+
+	/* Try to add to existing page */
+	if (HnswFindDuplicateInMemory(base, element))
+		return;
+
+	/* Write element and neighbor tuples */
+	WriteNewElementPagesInMemory(base, graph, element);
+
+	/* Update neighbors */
+	HnswUpdateNeighborPagesInMemory(base, procinfo, collation, element, m);
+
+	/* Update entry point if needed (already have lock) */
+	if (updateEntryPoint)
+		HnswPtrStore(base, graph->entryPoint, element);
+}
+
+/*
+ * Insert tuple
 */
 static bool
-InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuildState * buildstate)
+InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
 {
 	FmgrInfo   *procinfo = buildstate->procinfo;
 	Oid			collation = buildstate->collation;
 	HnswGraph  *graph = buildstate->graph;
-	HnswElement entryPoint = graph->entryPoint;
+	HnswElement entryPoint;
 	int			efConstruction = buildstate->efConstruction;
 	int			m = buildstate->m;
-	MemoryContext oldCtx;
 	HnswElement element;
+	HnswAllocator *allocator = &buildstate->allocator;
+	Size		valueSize;
+	Pointer		valuePtr;
+	bool		updateEntryPoint;
+	HnswRWLock *flushLock = &graph->flushLock;
+	char	   *base = buildstate->hnswarea;

 	/* Detoast once for all calls */
 	Datum		value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -368,69 +473,84 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
 			return false;
 	}

-	/* Allocate element in graph memory context */
-	oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
-	element = HnswInitElement(heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel);
-	element->value = datumCopy(value, false, -1);
-	MemoryContextSwitchTo(oldCtx);
+	/* Get datum size */
+	valueSize = VARSIZE_ANY(DatumGetPointer(value));

-	/* Update memory usage */
-#if PG_VERSION_NUM >= 130000
-	graph->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
-#else
-	graph->memoryUsed += HnswElementMemory(element, buildstate->m);
-#endif
+	/* Ensure graph not flushed when inserting */
+	HnswRWLockAcquire(flushLock, RW_SHARED);
+
+	if (graph->flushed)
+	{
+		HnswRWLockRelease(flushLock, RW_SHARED);
+
+		return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
+	}
+
+	/* Get lock for allocator */
+	SpinLockAcquire(&graph->allocatorLock);
+
+	/* Flush pages if needed */
+	if (graph->memoryUsed >= graph->memoryTotal)
+	{
+		SpinLockRelease(&graph->allocatorLock);
+
+		HnswRWLockRelease(flushLock, RW_SHARED);
+		HnswRWLockAcquire(flushLock, RW_EXCLUSIVE);
+
+		if (!graph->flushed)
+		{
+			ereport(NOTICE,
+					(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) graph->indtuples),
+					 errdetail("Building will take significantly more time."),
+					 errhint("Increase maintenance_work_mem to speed up builds.")));
+
+			FlushPages(buildstate);
+		}
+
+		HnswRWLockRelease(flushLock, RW_EXCLUSIVE);
+
+		return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
+	}
+
+	/* Create an element */
+	element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
+	valuePtr = HnswAlloc(allocator, valueSize);
+
+	/* Release allocator lock */
+	SpinLockRelease(&graph->allocatorLock);
+
+	/* Copy datum */
+	memcpy(valuePtr, DatumGetPointer(value), valueSize);
+	HnswPtrStore(base, element->value, valuePtr);
+
+	/* Create element lock */
+	SpinLockInit(&element->lock);
+
+	/* Get entry point */
+	SpinLockAcquire(&graph->entryLock);
+	entryPoint = HnswPtrAccess(base, graph->entryPoint);
+	updateEntryPoint = entryPoint == NULL || element->level > entryPoint->level;
+
+	/* Release lock if not updating entry point */
+	if (!updateEntryPoint)
+		SpinLockRelease(&graph->entryLock);

 	/* Insert element in graph */
-	HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
+	HnswInsertElement(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);

-	/* Look for duplicate */
-	if (HnswFindDuplicateInMemory(element))
-	{
-		/* No need to free element since memory unlikely to be reallocated */
-		return true;
-	}
+	/* Write to memory */
+	WriteElementInMemory(index, procinfo, collation, element, m, efConstruction, entryPoint, buildstate, graph, updateEntryPoint);

-	/* Add element */
-	slist_push_head(&graph->elements, &element->next);
+	/* Release lock if needed */
+	if (updateEntryPoint)
+		SpinLockRelease(&graph->entryLock);

-	/* Update neighbors */
-	for (int lc = element->level; lc >= 0; lc--)
-	{
-		int			lm = HnswGetLayerM(m, lc);
-		HnswNeighborArray *neighbors = HnswGetNeighbors(element, lc);
-
-		for (int i = 0; i < neighbors->length; i++)
-			HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
-	}
-
-	/* Update entry point if needed */
-	if (entryPoint == NULL || element->level > entryPoint->level)
-		graph->entryPoint = element;
+	/* Release flush lock */
+	HnswRWLockRelease(flushLock, RW_SHARED);

 	return true;
 }

-/*
- * Acquire a lock if needed
- */
-static inline void
-HnswLockAcquire(HnswShared * hnswshared)
-{
-	if (hnswshared)
-		SpinLockAcquire(&hnswshared->mutex);
-}
-
-/*
- * Release a lock if needed
- */
-static inline void
-HnswLockRelease(HnswShared * hnswshared)
-{
-	if (hnswshared)
-		SpinLockRelease(&hnswshared->mutex);
-}
-
 /*
 * Callback for table_index_build_scan
 */
@@ -440,9 +560,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 {
 	HnswBuildState *buildstate = (HnswBuildState *) state;
 	HnswGraph  *graph = buildstate->graph;
-	HnswShared *hnswshared = buildstate->hnswshared;
 	MemoryContext oldCtx;
-	bool		inserted;

 #if PG_VERSION_NUM < 130000
 	ItemPointer tid = &hup->t_self;
@@ -452,31 +570,16 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 	if (isnull[0])
 		return;

-	/* Flush pages if needed */
-	if (!graph->flushed && graph->memoryUsed >= graph->memoryTotal)
-	{
-		ereport(NOTICE,
-				(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) graph->indtuples),
-				 errdetail("Building will take significantly more time."),
-				 errhint("Increase maintenance_work_mem to speed up builds.")));
-
-		FlushPages(buildstate);
-	}
-
+	/* Use memory context */
 	oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);

 	/* Insert tuple */
-	if (graph->flushed)
-		inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true);
-	else
-		inserted = InsertTupleInMemory(index, values, tid, buildstate);
-
-	/* Update progress */
-	if (inserted)
+	if (InsertTuple(index, values, isnull, tid, buildstate))
 	{
-		HnswLockAcquire(hnswshared);
+		/* Update progress */
+		SpinLockAcquire(&graph->lock);
 		UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples);
-		HnswLockRelease(hnswshared);
+		SpinLockRelease(&graph->lock);
 	}

 	/* Reset memory context */
@@ -488,14 +591,59 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
 * Initialize the graph
 */
 static void
-InitGraph(HnswGraph * graph)
+InitGraph(HnswGraph * graph, char *base, long memoryTotal)
 {
-	slist_init(&graph->elements);
-	graph->entryPoint = NULL;
+	HnswPtrSetNull(base, graph->head);
+	HnswPtrSetNull(base, graph->entryPoint);
 	graph->memoryUsed = 0;
-	graph->memoryTotal = maintenance_work_mem * 1024L;
+	graph->memoryTotal = memoryTotal;
 	graph->flushed = false;
 	graph->indtuples = 0;
+	SpinLockInit(&graph->lock);
+	SpinLockInit(&graph->entryLock);
+	SpinLockInit(&graph->allocatorLock);
+	HnswRWLockInitialize(&graph->flushLock);
+}
+
+/*
+ * Initialize an allocator
+ */
+static void
+InitAllocator(HnswAllocator * allocator, void *(*alloc) (Size size, void *state), void *state)
+{
+	allocator->alloc = alloc;
+	allocator->state = state;
+}
+
+/*
+ * Memory context allocator
+ */
+static void *
+HnswMemoryContextAlloc(Size size, void *state)
+{
+	HnswBuildState *buildstate = (HnswBuildState *) state;
+	void	   *chunk = MemoryContextAlloc(buildstate->graphCtx, size);
+
+#if PG_VERSION_NUM >= 130000
+	buildstate->graphData.memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
+#else
+	buildstate->graphData.memoryUsed += MAXALIGN(size);
+#endif
+
+	return chunk;
+}
+
+/*
+ * Shared memory allocator
+ */
+static void *
+HnswSharedMemoryAlloc(Size size, void *state)
+{
+	HnswBuildState *buildstate = (HnswBuildState *) state;
+	void	   *chunk = buildstate->hnswarea + buildstate->graph->memoryUsed;
+
+	buildstate->graph->memoryUsed += MAXALIGN(size);
+	return chunk;
 }

 /*
@@ -531,7 +679,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 	buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
 	buildstate->collation = index->rd_indcollation[0];

-	InitGraph(&buildstate->graphData);
+	InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
 	buildstate->graph = &buildstate->graphData;
 	buildstate->ml = HnswGetMl(buildstate->m);
 	buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
@@ -549,8 +697,11 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
 											   "Hnsw build temporary context",
 											   ALLOCSET_DEFAULT_SIZES);

+	InitAllocator(&buildstate->allocator, &HnswMemoryContextAlloc, buildstate);
+
 	buildstate->hnswleader = NULL;
 	buildstate->hnswshared = NULL;
+	buildstate->hnswarea = NULL;
 }

 /*
@@ -581,6 +732,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
 		if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
 		{
 			buildstate->graph = &hnswshared->graphData;
+			buildstate->hnswarea = buildstate->hnswleader->hnswarea;
 			reltuples = hnswshared->reltuples;
 			SpinLockRelease(&hnswshared->mutex);
 			break;
@@ -600,7 +752,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
 * Perform a worker's portion of a parallel insert
 */
 static void
-HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
+HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, char *hnswarea, bool progress)
 {
 	HnswBuildState buildstate;
 #if PG_VERSION_NUM >= 120000
@@ -616,7 +768,8 @@ HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool p
 	indexInfo->ii_Concurrent = hnswshared->isconcurrent;
 	InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
 	buildstate.graph = &hnswshared->graphData;
-	buildstate.hnswshared = hnswshared;
+	buildstate.hnswarea = hnswarea;
+	InitAllocator(&buildstate.allocator, &HnswSharedMemoryAlloc, &buildstate);
 #if PG_VERSION_NUM >= 120000
 	scan = table_beginscan_parallel(hnswspool->heap,
 									ParallelTableScanFromHnswShared(hnswshared));
@@ -657,6 +810,7 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
 	char	   *sharedquery;
 	HnswSpool  *hnswspool;
 	HnswShared *hnswshared;
+	char	   *hnswarea;
 	Relation	heapRel;
 	Relation	indexRel;
 	LOCKMODE	heapLockmode;
@@ -697,8 +851,10 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
 	hnswspool->heap = heapRel;
 	hnswspool->index = indexRel;

+	hnswarea = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false);
+
 	/* Perform inserts */
-	HnswParallelScanAndInsert(hnswspool, hnswshared, false);
+	HnswParallelScanAndInsert(hnswspool, hnswshared, hnswarea, false);

 	/* Close relations within worker */
 	index_close(indexRel, indexLockmode);
@@ -761,7 +917,7 @@ HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
 	leaderworker->index = buildstate->index;

 	/* Perform work common to all participants */
-	HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
+	HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, hnswleader->hnswarea, true);
 }

 /*
@@ -774,7 +930,9 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
 	int			scantuplesortstates;
 	Snapshot	snapshot;
 	Size		esthnswshared;
+	Size		esthnswarea;
 	HnswShared *hnswshared;
+	char	   *hnswarea;
 	HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
 	bool		leaderparticipates = true;
 	int			querylen;
@@ -803,7 +961,9 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
 	/* Estimate size of workspaces */
 	esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
 	shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
-	shm_toc_estimate_keys(&pcxt->estimator, 1);
+	esthnswarea = maintenance_work_mem * 1024L;
+	shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea);
+	shm_toc_estimate_keys(&pcxt->estimator, 2);

 	/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
 	if (debug_query_string)
@@ -840,10 +1000,6 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
 	/* Initialize mutable state */
 	hnswshared->nparticipantsdone = 0;
 	hnswshared->reltuples = 0;
-	InitGraph(&hnswshared->graphData);
-	/* TODO Support in-memory builds */
-	hnswshared->graphData.memoryTotal = 0;
-	hnswshared->graphData.flushed = true;
 #if PG_VERSION_NUM >= 120000
 	table_parallelscan_initialize(buildstate->heap,
 								  ParallelTableScanFromHnswShared(hnswshared),
@@ -852,7 +1008,12 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
 	heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
 #endif

+	hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea);
+	/* Report less than allocated so never fails */
+	InitGraph(&hnswshared->graphData, hnswarea, esthnswarea - 1024 * 1024);
+
 	shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
+	shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_AREA, hnswarea);

 	/* Store query string for workers */
 	if (debug_query_string)
@@ -872,6 +1033,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
 		hnswleader->nparticipanttuplesorts++;
 	hnswleader->hnswshared = hnswshared;
 	hnswleader->snapshot = snapshot;
+	hnswleader->hnswarea = hnswarea;

 	/* If no workers were successfully launched, back out (do serial build) */
 	if (pcxt->nworkers_launched == 0)
@@ -926,16 +1088,12 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
 	UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);

 	/* Calculate parallel workers */
-	if (buildstate->heap != NULL && hnsw_enable_parallel_build)
+	if (buildstate->heap != NULL)
 		parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);

 	/* Attempt to launch parallel worker scan when required */
 	if (parallel_workers > 0)
-	{
-		/* TODO Support in-memory builds */
-		FlushPages(buildstate);
 		HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
-	}

 	/* Add tuples to graph */
 	if (buildstate->heap != NULL)