diff --git a/CHANGELOG.md b/CHANGELOG.md index 757f998..9ea71ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +## 0.8.1 (unreleased) + +- Improved performance of HNSW index scans for Postgres 17 + ## 0.8.0 (2024-10-30) - Added support for iterative index scans diff --git a/src/hnsw.h b/src/hnsw.h index 5102bfb..b224f4d 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -362,6 +362,13 @@ typedef union ItemPointerData indextid; } HnswUnvisited; +typedef struct HnswReadStreamData +{ + HnswUnvisited *unvisited; + int unvisitedLength; + int index; +} HnswReadStreamData; + typedef struct HnswScanOpaqueData { const HnswTypeInfo *typeInfo; diff --git a/src/hnswutils.c b/src/hnswutils.c index c52d2c7..d62cf77 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -15,6 +15,10 @@ #include "utils/memdebug.h" #include "utils/rel.h" +#if PG_VERSION_NUM >= 170000 +#include "storage/read_stream.h" +#endif + #if PG_VERSION_NUM < 170000 static inline uint64 murmurhash64(uint64 data) @@ -525,14 +529,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); @@ -553,7 +555,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); } @@ -567,7 +569,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); } /* @@ -807,6 +811,26 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u } } +#if PG_VERSION_NUM >= 170000 +/* + * 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->index == streamData->unvisitedLength) + return InvalidBlockNumber; + + uv = &streamData->unvisited[streamData->index++]; + *offno = ItemPointerGetOffsetNumber(&uv->indextid); + return ItemPointerGetBlockNumber(&uv->indextid); +} +#endif + /* * Algorithm 2 from paper */ @@ -825,6 +849,10 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); int unvisitedLength; bool inMemory = index == NULL; +#if PG_VERSION_NUM >= 170000 + HnswReadStreamData streamData; + ReadStream *stream = inMemory ? NULL : read_stream_begin_relation(READ_STREAM_DEFAULT, NULL, index, MAIN_FORKNUM, HnswReadStreamNextBlock, &streamData, sizeof(OffsetNumber)); +#endif if (v == NULL) { @@ -888,8 +916,17 @@ 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 >= 170000 + read_stream_reset(stream); + streamData.unvisited = unvisited; + streamData.unvisitedLength = unvisitedLength; + streamData.index = 0; +#endif + } + /* OK to count elements instead of tuples */ if (tuples != NULL) (*tuples) += unvisitedLength; @@ -910,13 +947,24 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in } else { + Buffer buf; + OffsetNumber offno; + +#if PG_VERSION_NUM >= 170000 + OffsetNumber *offnoPtr; + + buf = read_stream_next_buffer(stream, (void **) &offnoPtr); + offno = *offnoPtr; +#else ItemPointer indextid = &unvisited[i].indextid; - BlockNumber blkno = ItemPointerGetBlockNumber(indextid); - OffsetNumber offno = ItemPointerGetOffsetNumber(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; @@ -972,6 +1020,11 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in w = lappend(w, sc); } +#if PG_VERSION_NUM >= 170000 + if (stream) + read_stream_end(stream); +#endif + return w; }