From dbc07adf5b381282cb75ad92e7c0d62310bb668f Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Mon, 5 Oct 2026 19:08:39 -0700 Subject: [PATCH] Added support for async I/O --- src/hnsw.h | 26 +++++++++- src/hnswbuild.c | 2 +- src/hnswinsert.c | 2 +- src/hnswscan.c | 9 ++-- src/hnswutils.c | 127 +++++++++++++++++++++++++++++++++++++++++------ src/hnswvacuum.c | 2 +- 6 files changed, 146 insertions(+), 22 deletions(-) diff --git a/src/hnsw.h b/src/hnsw.h index 7900711..72f8298 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -18,6 +18,10 @@ #include "utils/sampling.h" #include "vector.h" +#if PG_VERSION_NUM >= 190000 +#include "storage/read_stream.h" +#endif + #if PG_VERSION_NUM < 190000 #include "storage/shmem.h" /* for add_size()/mul_size() in some versions */ #endif @@ -407,6 +411,21 @@ typedef union ItemPointerData indextid; } HnswUnvisited; +typedef struct HnswReadStreamData +{ + HnswUnvisited *unvisited; + int unvisitedLength; + int visited; +} HnswReadStreamData; + +typedef struct HnswSearchState +{ +#if PG_VERSION_NUM >= 190000 + HnswReadStreamData streamData; + ReadStream *stream; +#endif +} HnswSearchState; + typedef struct HnswScanOpaqueData { const HnswTypeInfo *typeInfo; @@ -417,6 +436,7 @@ typedef struct HnswScanOpaqueData HnswQuery q; int m; int64 tuples; + HnswSearchState searchState; double previousDistance; Size maxMemory; MemoryContext tmpCtx; @@ -463,13 +483,15 @@ bool HnswCheckNorm(HnswSupport * support, Datum value); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); void HnswInitPage(Buffer buf, Page page); void HnswInit(void); -List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples); +void HnswInitSearchState(HnswSearchState * searchState, Relation index, bool inMemory, bool maintenance); +void HnswFreeSearchState(HnswSearchState * searchState); +List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples, HnswSearchState * searchState); HnswElement HnswGetEntryPoint(Relation index); void HnswGetMetaPageInfo(Relation index, int *m, int *dimensions, HnswElement * entryPoint); void *HnswAlloc(HnswAllocator * allocator, Size size); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); -void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing); +void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool maintenance); HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); diff --git a/src/hnswbuild.c b/src/hnswbuild.c index bb2e9dc..711454e 100644 --- a/src/hnswbuild.c +++ b/src/hnswbuild.c @@ -470,7 +470,7 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element) } /* Find neighbors for element */ - HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false); + HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false, true); /* Update graph in memory */ UpdateGraphInMemory(support, element, m, entryPoint, buildstate); diff --git a/src/hnswinsert.c b/src/hnswinsert.c index c7afa7e..71a984e 100644 --- a/src/hnswinsert.c +++ b/src/hnswinsert.c @@ -735,7 +735,7 @@ HnswInsertTupleOnDisk(Relation index, const HnswTypeInfo * typeInfo, HnswSupport } /* Find neighbors for element */ - HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false); + HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, building); /* Update graph on disk */ UpdateGraphOnDisk(index, support, element, m, entryPoint, building); diff --git a/src/hnswscan.c b/src/hnswscan.c index 925a71c..9ee0588 100644 --- a/src/hnswscan.c +++ b/src/hnswscan.c @@ -53,11 +53,11 @@ GetScanItems(IndexScanDesc scan, Datum value) for (int lc = entryPoint->level; lc >= 1; lc--) { - w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL); + w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL, &so->searchState); ep = w; } - return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF ? &so->discarded : NULL, true, &so->tuples); + return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF ? &so->discarded : NULL, true, &so->tuples, &so->searchState); } /* @@ -88,7 +88,7 @@ ResumeScanItems(IndexScanDesc scan) ep = lappend(ep, sc); } - return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples); + return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples, &so->searchState); } /* @@ -147,6 +147,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys) /* Set support functions */ HnswInitSupport(&so->support, index); + HnswInitSearchState(&so->searchState, index, false, false); + /* * Use a lower max allocation size than default to allow scanning more * tuples for iterative search before exceeding work_mem @@ -343,6 +345,7 @@ hnswendscan(IndexScanDesc scan) { HnswScanOpaque so = (HnswScanOpaque) scan->opaque; + HnswFreeSearchState(&so->searchState); MemoryContextDelete(so->tmpCtx); pfree(so); diff --git a/src/hnswutils.c b/src/hnswutils.c index 8e6ac2a..cc8246c 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -22,6 +22,10 @@ #include "varatt.h" #endif +#if PG_VERSION_NUM >= 190000 +#include "storage/read_stream.h" +#endif + #if PG_VERSION_NUM < 170000 static inline uint64 murmurhash64(uint64 data) @@ -536,14 +540,12 @@ HnswGetDistance(Datum a, Datum b, HnswSupport * support) * Load an element and optionally get its distance from q */ static void -HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element) +HnswLoadElementImpl(Buffer buf, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element) { - Buffer buf; Page page; HnswElementTuple etup; /* Read vector */ - buf = ReadBuffer(index, blkno); LockBuffer(buf, BUFFER_LOCK_SHARE); page = BufferGetPage(buf); @@ -567,7 +569,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns if (distance == NULL || maxDistance == NULL || *distance < *maxDistance) { if (*element == NULL) - *element = HnswInitElementFromBlock(blkno, offno); + *element = HnswInitElementFromBlock(BufferGetBlockNumber(buf), offno); HnswLoadElementFromTuple(*element, etup, true, loadVec); } @@ -581,7 +583,9 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, Hns void HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance) { - HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element); + Buffer buf = ReadBuffer(index, element->blkno); + + HnswLoadElementImpl(buf, element->offno, distance, q, index, support, loadVec, maxDistance, &element); } /* @@ -821,11 +825,64 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u } } +#if PG_VERSION_NUM >= 190000 +/* + * Get next block number for read stream + */ +static BlockNumber +HnswReadStreamNextBlock(ReadStream *stream, void *callback_private_data, void *per_buffer_data) +{ + HnswReadStreamData *streamData = callback_private_data; + OffsetNumber *offno = per_buffer_data; + HnswUnvisited *uv; + + if (streamData->visited == streamData->unvisitedLength) + return read_stream_pause(stream); + + uv = &streamData->unvisited[streamData->visited++]; + *offno = ItemPointerGetOffsetNumber(&uv->indextid); + return ItemPointerGetBlockNumber(&uv->indextid); +} +#endif + +/* + * Initialize the search state + */ +void +HnswInitSearchState(HnswSearchState * searchState, Relation index, bool inMemory, bool maintenance) +{ +#if PG_VERSION_NUM >= 190000 + if (!inMemory) + { + int flags = READ_STREAM_DEFAULT; + + if (maintenance) + flags |= READ_STREAM_MAINTENANCE; + + searchState->stream = read_stream_begin_relation(flags, NULL, index, MAIN_FORKNUM, HnswReadStreamNextBlock, &searchState->streamData, sizeof(OffsetNumber)); + } + else + searchState->stream = NULL; +#endif +} + +/* + * Free the search state + */ +void +HnswFreeSearchState(HnswSearchState * searchState) +{ +#if PG_VERSION_NUM >= 190000 + if (searchState->stream != NULL) + read_stream_end(searchState->stream); +#endif +} + /* * Algorithm 2 from paper */ List * -HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples) +HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples, HnswSearchState * searchState) { List *w = NIL; pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); @@ -840,6 +897,11 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in int unvisitedLength; bool inMemory = index == NULL; +#if PG_VERSION_NUM >= 190000 + HnswReadStreamData *streamData = &searchState->streamData; + ReadStream *stream = searchState->stream; +#endif + if (v == NULL) { v = &vh; @@ -902,13 +964,23 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in if (inMemory) HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize); else + { HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc); +#if PG_VERSION_NUM >= 190000 + streamData->unvisited = unvisited; + streamData->unvisitedLength = unvisitedLength; + streamData->visited = 0; + + read_stream_resume(stream); +#endif + } + /* OK to count elements instead of tuples */ if (tuples != NULL) (*tuples) += unvisitedLength; - for (int i = 0; i < unvisitedLength; i++) + for (int i = 0;; i++) { HnswElement eElement; HnswSearchCandidate *e; @@ -919,18 +991,40 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in if (inMemory) { + if (i == unvisitedLength) + break; + eElement = unvisited[i].element; eDistance = GetElementDistance(base, eElement, q, support); } else { - ItemPointer indextid = &unvisited[i].indextid; - BlockNumber blkno = ItemPointerGetBlockNumber(indextid); - OffsetNumber offno = ItemPointerGetOffsetNumber(indextid); + Buffer buf; + OffsetNumber offno; + +#if PG_VERSION_NUM >= 190000 + void *offnoPtr; + + buf = read_stream_next_buffer(stream, &offnoPtr); + + if (!BufferIsValid(buf)) + break; + + offno = *((OffsetNumber *) offnoPtr); +#else + ItemPointer indextid; + + if (i == unvisitedLength) + break; + + indextid = &unvisited[i].indextid; + buf = ReadBuffer(index, ItemPointerGetBlockNumber(indextid)); + offno = ItemPointerGetOffsetNumber(indextid); +#endif /* Avoid any allocations if not adding */ eElement = NULL; - HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement); + HnswLoadElementImpl(buf, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement); if (eElement == NULL) continue; @@ -1281,7 +1375,7 @@ PrecomputeHash(char *base, HnswElement element) * Algorithm 1 from paper */ void -HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing) +HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool maintenance) { List *ep; List *w; @@ -1290,6 +1384,9 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint HnswQuery q; HnswElement skipElement = existing ? element : NULL; bool inMemory = index == NULL; + HnswSearchState searchState; + + HnswInitSearchState(&searchState, index, inMemory, maintenance); q.value = HnswGetValue(base, element); @@ -1308,7 +1405,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint /* 1st phase: greedy search to insert level */ for (int lc = entryLevel; lc >= level + 1; lc--) { - w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL); + w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL, &searchState); ep = w; } @@ -1327,7 +1424,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint List *lw = NIL; ListCell *lc2; - w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL); + w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL, &searchState); /* Convert search candidates to candidates */ foreach(lc2, w) @@ -1357,6 +1454,8 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint ep = w; } + + HnswFreeSearchState(&searchState); } /* diff --git a/src/hnswvacuum.c b/src/hnswvacuum.c index eb993ca..75e7ade 100644 --- a/src/hnswvacuum.c +++ b/src/hnswvacuum.c @@ -246,7 +246,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme element->heaptidsLength = 0; /* Find neighbors for element, skipping itself */ - HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true); + HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, true); /* Zero memory for each element */ MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);