liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-06 12:40:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Reduced index size

Browse Source
This commit is contained in:
Andrew Kane
2023-08-01 19:17:39 -07:00
parent 95eded091f
commit d3e08fdf99
8 changed files with 430 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
src/hnsw.h
View File

@@ -38,13 +38,18 @@
#define HNSW_MIN_EF_SEARCH 10 #define HNSW_MIN_EF_SEARCH 10
#define HNSW_MAX_EF_SEARCH 1000 #define HNSW_MAX_EF_SEARCH 1000
#define HNSW_ELEMENT_TUPLE_TYPE 1
#define HNSW_NEIGHBOR_TUPLE_TYPE 2
#define HNSW_HEAPTIDS 10 #define HNSW_HEAPTIDS 10
/* Build phases */ /* Build phases */
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */ /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_HNSW_PHASE_LOAD 2 #define PROGRESS_HNSW_PHASE_LOAD 2
#define HNSW_ELEMENT_TUPLE_SIZE(_dim) (offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim)) #define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, neighbors) + ((level) + 2) * (m) * sizeof(HnswNeighborTupleItem))
#define HNSW_NEIGHBOR_COUNT(itemid) ((ItemIdGetLength(itemid) - offsetof(HnswNeighborTupleData, neighbors)) / sizeof(HnswNeighborTupleItem))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page)) #define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page))
@@ -60,6 +65,9 @@
#define list_sort(list, cmp) list_qsort(list, cmp) #define list_sort(list, cmp) list_qsort(list, cmp)
#endif #endif
#define HnswIsElementTuple(tup) ((tup)->type == HNSW_ELEMENT_TUPLE_TYPE)
#define HnswIsNeighborTuple(tup) ((tup)->type == HNSW_NEIGHBOR_TUPLE_TYPE)
#define GetLayerM(m, layer) (layer == 0 ? m * 2 : m) #define GetLayerM(m, layer) (layer == 0 ? m * 2 : m)
#define HnswGetMl(m) (1 / log(m)) #define HnswGetMl(m) (1 / log(m))
@@ -77,6 +85,7 @@ typedef struct HnswElementData
BlockNumber blkno; BlockNumber blkno;
OffsetNumber offno; OffsetNumber offno;
BlockNumber neighborPage; BlockNumber neighborPage;
OffsetNumber neighborOffno;
Vector *vec; Vector *vec;
} HnswElementData; } HnswElementData;
@@ -175,21 +184,32 @@ typedef HnswPageOpaqueData * HnswPageOpaque;
typedef struct HnswElementTupleData typedef struct HnswElementTupleData
{ {
ItemPointerData heaptids[HNSW_HEAPTIDS]; uint8 type;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint16 unused; uint8 unused;
BlockNumber neighborPage; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid;
uint16 unused2;
Vector vec; Vector vec;
} HnswElementTupleData; } HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple; typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData typedef struct HnswNeighborTupleItem
{ {
ItemPointerData indextid; ItemPointerData indextid;
uint16 unused; uint16 unused;
float distance; float distance;
} HnswNeighborTupleItem;
typedef struct HnswNeighborTupleData
{
uint8 type;
uint8 unused;
uint16 unused2;
uint32 unused3;
HnswNeighborTupleItem neighbors[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData; } HnswNeighborTupleData;
typedef HnswNeighborTupleData * HnswNeighborTuple; typedef HnswNeighborTupleData * HnswNeighborTuple;
@@ -236,14 +256,14 @@ Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state); void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void HnswInit(void); void HnswInit(void);
List *SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage); List *SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage, OffsetNumber *skipOffno);
HnswElement GetEntryPoint(Relation index); HnswElement GetEntryPoint(Relation index);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element); void HnswFreeElement(HnswElement element);
HnswElement HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming); HnswElement HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming);
HnswCandidate *EntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadvec); HnswCandidate *EntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadvec);
void UpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum); void UpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
void AddNeighborsToPage(Relation index, Page page, HnswElement e, HnswNeighborTuple neighbor, Size neighborsz, int m); void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m); void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel); bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
@@ -267,6 +287,6 @@ bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan); void hnswendscan(IndexScanDesc scan);
/* Ensure fits in uint8 */ /* Ensure fits in uint8 */
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData))) / (MAXALIGN(sizeof(HnswNeighborTupleData)) + sizeof(ItemIdData)) / m) - 2, 255) #define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, neighbors) - sizeof(ItemIdData)) / (sizeof(HnswNeighborTupleItem)) / m) - 2, 255)
#endif #endif

140
src/hnswbuild.c
View File

@@ -67,6 +67,30 @@ CreateMetaPage(HnswBuildState * buildstate)
HnswCommitBuffer(buf, state); HnswCommitBuffer(buf, state);
} }
/*
* Add a new page
*/
static void
HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, ForkNumber forkNum)
{
/* Add a new page */
Buffer newbuf = HnswNewBuffer(index, forkNum);
/* Update previous page */
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */
MarkBufferDirty(*buf);
GenericXLogFinish(*state);
UnlockReleaseBuffer(*buf);
/* Prepare new page */
*buf = newbuf;
*state = GenericXLogStart(index);
*page = GenericXLogRegisterBuffer(*state, *buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*buf, *page);
}
/* /*
* Create element pages * Create element pages
*/ */
@@ -76,23 +100,23 @@ CreateElementPages(HnswBuildState * buildstate)
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
int dimensions = buildstate->dimensions; int dimensions = buildstate->dimensions;
Size elementsz; Size etupSize;
HnswElementTuple element; Size maxSize;
int elementsPerPage; HnswElementTuple etup;
BlockNumber neighborPage; HnswNeighborTuple ntup;
BlockNumber insertPage; BlockNumber insertPage;
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
ListCell *lc; ListCell *lc;
/* Allocate once */ /* Calculate sizes */
elementsz = MAXALIGN(HNSW_ELEMENT_TUPLE_SIZE(dimensions)); maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
element = palloc0(elementsz); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Calculate starting neighbor page */ /* Allocate once */
elementsPerPage = (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData))) / (elementsz + sizeof(ItemIdData)); etup = palloc0(etupSize);
neighborPage = HNSW_HEAD_BLKNO + (int) ceil(list_length(buildstate->elements) / (double) elementsPerPage); ntup = palloc0(maxSize);
/* Prepare first page */ /* Prepare first page */
buf = HnswNewBuffer(index, forkNum); buf = HnswNewBuffer(index, forkNum);
@@ -102,45 +126,46 @@ CreateElementPages(HnswBuildState * buildstate)
foreach(lc, buildstate->elements) foreach(lc, buildstate->elements)
{ {
HnswElement e = lfirst(lc); HnswElement element = lfirst(lc);
Size ntupSize;
Size combinedSize;
/* Calculate neighbor page */ HnswSetElementTuple(etup, element);
/* Will be rechecked later */
e->neighborPage = neighborPage++;
/* Set item data */ /* Calculate sizes */
HnswSetElementTuple(element, e); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
element->neighborPage = e->neighborPage; combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
/* Ensure free space */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < elementsz) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
/* Calculate offsets */
element->blkno = BufferGetBlockNumber(buf);
element->offno = OffsetNumberNext(PageGetMaxOffsetNumber(page));
if (combinedSize <= maxSize)
{ {
/* Add a new page */ element->neighborPage = element->blkno;
Buffer newbuf = HnswNewBuffer(index, forkNum); element->neighborOffno = OffsetNumberNext(element->offno);
}
/* Update previous page */ else
HnswPageGetOpaque(page)->nextblkno = BufferGetBlockNumber(newbuf); {
element->neighborPage = element->blkno + 1;
/* Commit */ element->neighborOffno = FirstOffsetNumber;
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
/* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */
CHECK_FOR_INTERRUPTS();
/* Prepare new page */
buf = newbuf;
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(buf, page);
} }
/* Add the item */ ItemPointerSet(&etup->neighbortid, element->neighborPage, element->neighborOffno);
e->blkno = BufferGetBlockNumber(buf);
e->offno = PageAddItem(page, (Item) element, elementsz, InvalidOffsetNumber, false, false); /* Add element */
if (e->offno == InvalidOffsetNumber) if (PageAddItem(page, (Item) etup, etupSize, InvalidOffsetNumber, false, false) != element->offno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Add new page if needed */
if (PageGetFreeSpace(page) < ntupSize)
HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
/* Add placeholder for neighbors */
if (PageAddItem(page, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != element->neighborOffno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
} }
@@ -162,13 +187,12 @@ CreateNeighborPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
Size neighborsz; int m = buildstate->m;
HnswNeighborTuple neighbor;
ListCell *lc; ListCell *lc;
HnswNeighborTuple ntup;
/* Allocate once */ /* Allocate once */
neighborsz = MAXALIGN(sizeof(HnswNeighborTupleData)); ntup = palloc0(BLCKSZ);
neighbor = palloc0(neighborsz);
foreach(lc, buildstate->elements) foreach(lc, buildstate->elements)
{ {
@@ -176,23 +200,17 @@ CreateNeighborPages(HnswBuildState * buildstate)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size neighborsz = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
/* Can take a while, so ensure we can interrupt */ buf = ReadBufferExtended(index, forkNum, e->neighborPage, RBM_NORMAL, NULL);
/* Needs to be called when no buffer locks are held */ LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
CHECK_FOR_INTERRUPTS();
buf = HnswNewBuffer(index, forkNum);
/* Check block number */
if (BufferGetBlockNumber(buf) != e->neighborPage)
elog(ERROR, "expected neighbor page %d, got %d", e->neighborPage, BufferGetBlockNumber(buf));
/* Prepare page */
state = GenericXLogStart(index); state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE); page = GenericXLogRegisterBuffer(state, buf, 0);
HnswInitPage(buf, page);
AddNeighborsToPage(index, page, e, neighbor, neighborsz, buildstate->m); HnswSetNeighborTuple(ntup, e, m);
if (!PageIndexTupleOverwrite(page, e->neighborOffno, (Item) ntup, neighborsz))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf); MarkBufferDirty(buf);

235
src/hnswinsert.c
View File

@@ -34,7 +34,7 @@ GetInsertPage(Relation index)
* Check for a free offset * Check for a free offset
*/ */
static bool static bool
HnswFreeOffset(Page page, OffsetNumber *freeOffno, BlockNumber *neighborPage) HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size neighborsz, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno)
{ {
OffsetNumber offno; OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page); OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -43,17 +43,67 @@ HnswFreeOffset(Page page, OffsetNumber *freeOffno, BlockNumber *neighborPage)
{ {
HnswElementTuple item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno)); HnswElementTuple item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
if (item->deleted) /* Skip neighbor tuples */
if (!HnswIsElementTuple(item))
continue;
/* TODO Remove level check */
if (item->deleted && item->level == element->level)
{ {
*freeOffno = offno; BlockNumber neighborPage = ItemPointerGetBlockNumber(&item->neighbortid);
*neighborPage = item->neighborPage; OffsetNumber neighborOffno = ItemPointerGetOffsetNumber(&item->neighbortid);
return true; ItemId itemid;
if (neighborPage == BufferGetBlockNumber(buf))
{
*nbuf = buf;
*npage = page;
}
else
{
*nbuf = ReadBuffer(index, neighborPage);
LockBuffer(*nbuf, BUFFER_LOCK_EXCLUSIVE);
/* Skip WAL for now */
*npage = BufferGetPage(*nbuf);
}
itemid = PageGetItemId(*npage, neighborOffno);
/* Check for space on neighbor tuple page */
if (PageGetFreeSpace(*npage) + ItemIdGetLength(itemid) - sizeof(ItemIdData) >= neighborsz)
{
*freeOffno = offno;
*freeNeighborOffno = neighborOffno;
return true;
}
else if (*nbuf != buf)
UnlockReleaseBuffer(*nbuf);
} }
} }
return false; return false;
} }
/*
* Add a new page
*/
static void
HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState *state, Page page)
{
/* Add a new page */
LockRelationForExtension(index, ExclusiveLock);
*nbuf = HnswNewBuffer(index, MAIN_FORKNUM);
UnlockRelationForExtension(index, ExclusiveLock);
/* Init new page */
*npage = GenericXLogRegisterBuffer(state, *nbuf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*nbuf, *npage);
/* Update previous buffer */
HnswPageGetOpaque(page)->nextblkno = BufferGetBlockNumber(*nbuf);
}
/* /*
* Add to element and neighbor pages * Add to element and neighbor pages
*/ */
@@ -63,25 +113,32 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size esize; Size etupSize;
Size ntupSize;
Size combinedSize;
HnswElementTuple etup; HnswElementTuple etup;
BlockNumber insertPage = GetInsertPage(index); BlockNumber insertPage = GetInsertPage(index);
BlockNumber originalInsertPage = insertPage; BlockNumber originalInsertPage = insertPage;
int dimensions = e->vec->dim; int dimensions = e->vec->dim;
Size nsize = MAXALIGN(sizeof(HnswNeighborTupleData)); HnswNeighborTuple ntup;
HnswNeighborTuple ntup = palloc0(nsize);
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
OffsetNumber freeOffno = InvalidOffsetNumber; OffsetNumber freeOffno = InvalidOffsetNumber;
BlockNumber neighborPage = InvalidBlockNumber; OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
/* Get tuple size */ /* Calculate sizes */
esize = MAXALIGN(HNSW_ELEMENT_TUPLE_SIZE(dimensions)); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
/* Prepare tuple */ /* Prepare element tuple */
etup = palloc0(esize); etup = palloc0(etupSize);
HnswSetElementTuple(etup, e); HnswSetElementTuple(etup, e);
/* Prepare neighbor tuple */
ntup = palloc0(ntupSize);
HnswSetNeighborTuple(ntup, e, m);
/* Find a page to insert the item */ /* Find a page to insert the item */
for (;;) for (;;)
{ {
@@ -91,8 +148,29 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
state = GenericXLogStart(index); state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
if (HnswFreeOffset(page, &freeOffno, &neighborPage) || PageGetFreeSpace(page) >= esize) /* Space for both */
if (PageGetFreeSpace(page) >= combinedSize)
{
nbuf = buf;
npage = page;
break; break;
}
/* Space for element but not neighbors and last page */
if (PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
{
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
break;
}
/* Space from deleted item */
if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno))
{
if (nbuf != buf)
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
break;
}
insertPage = HnswPageGetOpaque(page)->nextblkno; insertPage = HnswPageGetOpaque(page)->nextblkno;
@@ -107,28 +185,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
Buffer newbuf; Buffer newbuf;
Page newpage; Page newpage;
/* HnswInsertAppendPage(index, &newbuf, &newpage, state, page);
* From ReadBufferExtended: Caller is responsible for ensuring
* that only one backend tries to extend a relation at the same
* time!
*/
LockRelationForExtension(index, ExclusiveLock);
/* Add a new page */
newbuf = HnswNewBuffer(index, MAIN_FORKNUM);
/* Unlock extend relation lock as early as possible */
UnlockRelationForExtension(index, ExclusiveLock);
/* Init new page */
newpage = GenericXLogRegisterBuffer(state, newbuf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(newbuf, newpage);
/* Update insert page */
insertPage = BufferGetBlockNumber(newbuf);
/* Update previous buffer */
HnswPageGetOpaque(page)->nextblkno = insertPage;
/* Commit */ /* Commit */
MarkBufferDirty(newbuf); MarkBufferDirty(newbuf);
@@ -142,58 +199,67 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
state = GenericXLogStart(index); state = GenericXLogStart(index);
buf = newbuf; buf = newbuf;
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
/* Create new page for neighbors if needed */
if (PageGetFreeSpace(page) < combinedSize)
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
else
{
nbuf = buf;
npage = page;
}
break; break;
} }
} }
if (OffsetNumberIsValid(freeOffno))
{
/* Reuse existing page */
nbuf = ReadBuffer(index, neighborPage);
LockBuffer(nbuf, BUFFER_LOCK_EXCLUSIVE);
}
else
{
/* Add new page */
LockRelationForExtension(index, ExclusiveLock);
nbuf = HnswNewBuffer(index, MAIN_FORKNUM);
UnlockRelationForExtension(index, ExclusiveLock);
}
npage = GenericXLogRegisterBuffer(state, nbuf, GENERIC_XLOG_FULL_IMAGE);
/* Overwrites existing page via InitPage */
HnswInitPage(nbuf, npage);
/* Update neighbors */
AddNeighborsToPage(index, npage, e, ntup, nsize, m);
e->blkno = BufferGetBlockNumber(buf); e->blkno = BufferGetBlockNumber(buf);
e->neighborPage = BufferGetBlockNumber(nbuf); e->neighborPage = BufferGetBlockNumber(nbuf);
/* Set neighbor page for element */ insertPage = e->neighborPage;
etup->neighborPage = e->neighborPage;
/* Add to next offset */
if (OffsetNumberIsValid(freeOffno)) if (OffsetNumberIsValid(freeOffno))
{ {
e->offno = freeOffno; e->offno = freeOffno;
if (!PageIndexTupleOverwrite(page, freeOffno, (Item) etup, esize)) e->neighborOffno = freeNeighborOffno;
}
else
{
e->offno = OffsetNumberNext(PageGetMaxOffsetNumber(page));
if (nbuf == buf)
e->neighborOffno = OffsetNumberNext(e->offno);
else
e->neighborOffno = FirstOffsetNumber;
}
ItemPointerSet(&etup->neighbortid, e->neighborPage, e->neighborOffno);
/* Add element and neighbors */
if (OffsetNumberIsValid(freeOffno))
{
if (!PageIndexTupleOverwrite(page, e->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
if (!PageIndexTupleOverwrite(npage, e->neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
} }
else else
{ {
e->offno = PageAddItem(page, (Item) etup, esize, InvalidOffsetNumber, false, false); if (PageAddItem(page, (Item) etup, etupSize, InvalidOffsetNumber, false, false) != e->offno)
if (e->offno == InvalidOffsetNumber) elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
if (PageAddItem(npage, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != e->neighborOffno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
} }
/* Commit */ /* Commit */
MarkBufferDirty(buf); MarkBufferDirty(buf);
MarkBufferDirty(nbuf); if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
UnlockReleaseBuffer(nbuf); if (nbuf != buf)
UnlockReleaseBuffer(nbuf);
/* Update the insert page */ /* Update the insert page */
if (insertPage != originalInsertPage) if (insertPage != originalInsertPage)
@@ -201,12 +267,12 @@ WriteNewElementPages(Relation index, HnswElement e, int m)
} }
/* /*
* Calculate offset number for update * Calculate index for update
*/ */
static OffsetNumber static int
HnswGetOffsetNumber(HnswUpdate * update, int m) HnswGetIndex(HnswUpdate * update, int m)
{ {
return FirstOffsetNumber + (update->hc.element->level - update->level) * m + update->index; return (update->hc.element->level - update->level) * m + update->index;
} }
/* /*
@@ -215,36 +281,45 @@ HnswGetOffsetNumber(HnswUpdate * update, int m)
static void static void
UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates) UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates)
{ {
Buffer buf;
Page page;
GenericXLogState *state;
ListCell *lc; ListCell *lc;
OffsetNumber offno;
Size neighborsz = MAXALIGN(sizeof(HnswNeighborTupleData));
HnswNeighborTuple neighbor = palloc0(neighborsz);
/* Could update multiple at once for same element */ /* Could update multiple at once for same element */
/* but should only happen a low percent of time, so keep simple for now */ /* but should only happen a low percent of time, so keep simple for now */
foreach(lc, updates) foreach(lc, updates)
{ {
Buffer buf;
Page page;
GenericXLogState *state;
HnswUpdate *update = lfirst(lc); HnswUpdate *update = lfirst(lc);
ItemId itemid;
Size neighborsz;
int idx;
OffsetNumber offno = update->hc.element->neighborOffno;
/* Register page */ /* Register page */
buf = ReadBuffer(index, update->hc.element->neighborPage); buf = ReadBuffer(index, update->hc.element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
offno = HnswGetOffsetNumber(update, m);
itemid = PageGetItemId(page, offno);
neighborsz = ItemIdGetLength(itemid);
idx = HnswGetIndex(update, m);
/* Make robust against issues */ /* Make robust against issues */
if (offno <= PageGetMaxOffsetNumber(page)) if (idx < HNSW_NEIGHBOR_COUNT(itemid))
{ {
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
HnswNeighborTupleItem *neighbor = &ntup->neighbors[idx];
/* Set item data */ /* Set item data */
ItemPointerSet(&neighbor->indextid, e->blkno, e->offno); ItemPointerSet(&neighbor->indextid, e->blkno, e->offno);
neighbor->distance = update->hc.distance; neighbor->distance = update->hc.distance;
/* Update connections */ /* Update connections */
if (!PageIndexTupleOverwrite(page, offno, (Item) neighbor, neighborsz)) if (!PageIndexTupleOverwrite(page, offno, (Item) ntup, neighborsz))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
@@ -267,7 +342,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size esize = MAXALIGN(HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim)); Size esize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
HnswElementTuple etup; HnswElementTuple etup;
int i; int i;

4
src/hnswscan.c
View File

@@ -27,12 +27,12 @@ GetScanItems(IndexScanDesc scan, Datum q)
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = SearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL); w = SearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL, NULL);
ep = w; ep = w;
} }
/* TODO Return all visited elements at level 0, not just ef search */ /* TODO Return all visited elements at level 0, not just ef search */
so->w = SearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL); so->w = SearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL, NULL);
} }
/* /*

136
src/hnswutils.c
View File

@@ -212,6 +212,8 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
/* Look for cached distance */ /* Look for cached distance */
if (a->neighbors != NULL) if (a->neighbors != NULL)
{ {
Assert(a->level >= lc);
for (int i = 0; i < a->neighbors[lc].length; i++) for (int i = 0; i < a->neighbors[lc].length; i++)
{ {
if (a->neighbors[lc].items[i].element == b) if (a->neighbors[lc].items[i].element == b)
@@ -221,6 +223,8 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
if (b->neighbors != NULL) if (b->neighbors != NULL)
{ {
Assert(b->level >= lc);
for (int i = 0; i < b->neighbors[lc].length; i++) for (int i = 0; i < b->neighbors[lc].length; i++)
{ {
if (b->neighbors[lc].items[i].element == a) if (b->neighbors[lc].items[i].element == a)
@@ -359,6 +363,55 @@ HnswAddHeapTid(HnswElement element, ItemPointer heaptid)
element->heaptids = lappend(element->heaptids, copy); element->heaptids = lappend(element->heaptids, copy);
} }
/*
* Load neighbors from page
*/
static void
LoadNeighborsFromPage(HnswElement element, Relation index, Page page)
{
int m = HnswGetM(index);
ItemId itemid = PageGetItemId(page, element->neighborOffno);
int neighborCount = (element->level + 2) * m;
HnswInitNeighbors(element, m);
/* If not, neighbor page represents new item */
/* Only caught if item has a different level */
/* TODO Use versioning to fix this? */
if (HNSW_NEIGHBOR_COUNT(itemid) == neighborCount)
{
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
Assert(HnswIsNeighborTuple(ntup));
for (int i = 0; i < neighborCount; i++)
{
HnswElement e;
int level;
HnswCandidate *hc;
HnswNeighborTupleItem *neighbor;
HnswNeighborArray *neighbors;
neighbor = &ntup->neighbors[i];
if (!ItemPointerIsValid(&neighbor->indextid))
continue;
e = CreateElementFromBlock(ItemPointerGetBlockNumber(&neighbor->indextid), ItemPointerGetOffsetNumber(&neighbor->indextid));
/* Calculate level based on offset */
level = element->level - i / m;
if (level < 0)
level = 0;
neighbors = &element->neighbors[level];
hc = &neighbors->items[neighbors->length++];
hc->element = e;
hc->distance = neighbor->distance;
}
}
}
/* /*
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
@@ -376,6 +429,8 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, element->offno)); item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, element->offno));
Assert(HnswIsElementTuple(item));
/* Load element */ /* Load element */
element->heaptids = NIL; element->heaptids = NIL;
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
@@ -387,7 +442,8 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
HnswAddHeapTid(element, &item->heaptids[i]); HnswAddHeapTid(element, &item->heaptids[i]);
} }
element->level = item->level; element->level = item->level;
element->neighborPage = item->neighborPage; element->neighborPage = ItemPointerGetBlockNumber(&item->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&item->neighbortid);
element->deleted = item->deleted; element->deleted = item->deleted;
if (loadvec) if (loadvec)
@@ -400,6 +456,10 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&item->vec))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&item->vec)));
/* Load neighbors if on same page */
if (element->neighborPage == element->blkno)
LoadNeighborsFromPage(element, index, page);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -512,53 +572,16 @@ HnswInitNeighbors(HnswElement element, int m)
* Load neighbors * Load neighbors
*/ */
static void static void
LoadNeighbors(HnswCandidate * c, Relation index) LoadNeighbors(HnswElement element, Relation index)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
OffsetNumber offno;
OffsetNumber maxoffno;
HnswNeighborTuple neighbor;
HnswNeighborArray *neighbors;
int m = HnswGetM(index);
buf = ReadBuffer(index, c->element->neighborPage); buf = ReadBuffer(index, element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
HnswInitNeighbors(c->element, m); LoadNeighborsFromPage(element, index, page);
/* If not, neighbor page represents new item */
/* Only caught if item has a different level */
/* TODO Use versioning to fix this? */
if (maxoffno == (c->element->level + 2) * m)
{
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElement element;
int level;
HnswCandidate *hc;
neighbor = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, offno));
if (!ItemPointerIsValid(&neighbor->indextid))
continue;
element = CreateElementFromBlock(ItemPointerGetBlockNumber(&neighbor->indextid), ItemPointerGetOffsetNumber(&neighbor->indextid));
/* Calculate level based on offset */
level = c->element->level - (offno - FirstOffsetNumber) / m;
if (level < 0)
level = 0;
neighbors = &c->element->neighbors[level];
hc = &neighbors->items[neighbors->length];
hc->element = element;
hc->distance = neighbor->distance;
neighbors->length++;
}
}
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -603,11 +626,14 @@ HnswFreeElement(HnswElement element)
} }
/* /*
* Set element tuple, except for neighbor page * Set element tuple, except for neighbor info
*/ */
void void
HnswSetElementTuple(HnswElementTuple etup, HnswElement element) HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
{ {
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level;
etup->deleted = 0;
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (i < list_length(element->heaptids)) if (i < list_length(element->heaptids))
@@ -615,8 +641,6 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
etup->level = element->level;
etup->deleted = 0;
memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim)); memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
} }
@@ -650,7 +674,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage) SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage, OffsetNumber *skipOffno)
{ {
ListCell *lc2; ListCell *lc2;
@@ -699,7 +723,7 @@ SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinf
break; break;
if (c->element->neighbors == NULL) if (c->element->neighbors == NULL)
LoadNeighbors(c, index); LoadNeighbors(c->element, index);
/* Get the neighborhood at layer lc */ /* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc]; neighborhood = &c->element->neighbors[lc];
@@ -731,7 +755,7 @@ SearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinf
continue; continue;
/* Skip self for vacuuming update */ /* Skip self for vacuuming update */
if (skipPage != NULL && e->element->neighborPage == *skipPage) if (skipPage != NULL && e->element->neighborPage == *skipPage && e->element->neighborOffno == *skipOffno)
continue; continue;
/* Stale read */ /* Stale read */
@@ -825,6 +849,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
Datum q = PointerGetDatum(element->vec); Datum q = PointerGetDatum(element->vec);
HnswElement dup; HnswElement dup;
BlockNumber *skipPage = vacuuming ? &element->neighborPage : NULL; BlockNumber *skipPage = vacuuming ? &element->neighborPage : NULL;
OffsetNumber *skipOffno = vacuuming ? &element->neighborOffno : NULL;
/* Get entry point and level */ /* Get entry point and level */
if (entryPoint != NULL) if (entryPoint != NULL)
@@ -837,7 +862,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
for (int lc = entryLevel; lc >= level + 1; lc--) for (int lc = entryLevel; lc >= level + 1; lc--)
{ {
w = SearchLayer(q, ep, 1, lc, index, procinfo, collation, true, skipPage); w = SearchLayer(q, ep, 1, lc, index, procinfo, collation, true, skipPage, skipOffno);
ep = w; ep = w;
} }
@@ -848,7 +873,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
{ {
int lm = GetLayerM(m, lc); int lm = GetLayerM(m, lc);
w = SearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipPage); w = SearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipPage, skipOffno);
newNeighbors[lc] = SelectNeighbors(w, lm, lc, procinfo, collation, NULL); newNeighbors[lc] = SelectNeighbors(w, lm, lc, procinfo, collation, NULL);
ep = w; ep = w;
} }
@@ -913,11 +938,15 @@ UpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNu
} }
/* /*
* Add neighbors to page * Set neighbor tuple
*/ */
void void
AddNeighborsToPage(Relation index, Page page, HnswElement e, HnswNeighborTuple neighbor, Size neighborsz, int m) HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
{ {
int idx = 0;
ntup->type = HNSW_NEIGHBOR_TUPLE_TYPE;
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
HnswNeighborArray *neighbors = &e->neighbors[lc]; HnswNeighborArray *neighbors = &e->neighbors[lc];
@@ -925,6 +954,8 @@ AddNeighborsToPage(Relation index, Page page, HnswElement e, HnswNeighborTuple n
for (int i = 0; i < lm; i++) for (int i = 0; i < lm; i++)
{ {
HnswNeighborTupleItem *neighbor = &ntup->neighbors[idx++];
if (i < neighbors->length) if (i < neighbors->length)
{ {
HnswCandidate *hc = &neighbors->items[i]; HnswCandidate *hc = &neighbors->items[i];
@@ -937,9 +968,6 @@ AddNeighborsToPage(Relation index, Page page, HnswElement e, HnswNeighborTuple n
ItemPointerSetInvalid(&neighbor->indextid); ItemPointerSetInvalid(&neighbor->indextid);
neighbor->distance = NAN; neighbor->distance = NAN;
} }
if (PageAddItem(page, (Item) neighbor, neighborsz, InvalidOffsetNumber, false, false) == InvalidOffsetNumber)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
} }
} }
} }

101
src/hnswvacuum.c
View File

@@ -62,6 +62,10 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
int idx = 0; int idx = 0;
bool itemUpdated = false; bool itemUpdated = false;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(item))
continue;
if (ItemPointerIsValid(&item->heaptids[0])) if (ItemPointerIsValid(&item->heaptids[0]))
{ {
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
@@ -81,7 +85,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (itemUpdated) if (itemUpdated)
{ {
Size itemsz = MAXALIGN(HNSW_ELEMENT_TUPLE_SIZE(item->vec.dim)); Size itemsz = HNSW_ELEMENT_TUPLE_SIZE(item->vec.dim);
/* Mark rest as invalid */ /* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++) for (int i = idx; i < HNSW_HEAPTIDS; i++)
@@ -137,25 +141,30 @@ NeedsUpdated(HnswVacuumState * vacuumstate, HnswElement element)
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
Buffer buf; Buffer buf;
Page page; Page page;
OffsetNumber offno; ItemId itemid;
OffsetNumber maxoffno; int neighborCount;
HnswNeighborTuple ntup;
bool needsUpdated = false; bool needsUpdated = false;
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas); buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page); itemid = PageGetItemId(page, element->neighborOffno);
ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
neighborCount = HNSW_NEIGHBOR_COUNT(itemid);
Assert(HnswIsNeighborTuple(ntup));
/* Check neighbors */ /* Check neighbors */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (int i = 0; i < neighborCount; i++)
{ {
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, offno)); HnswNeighborTupleItem *neighbor = &ntup->neighbors[i];
if (!ItemPointerIsValid(&ntup->indextid)) if (!ItemPointerIsValid(&neighbor->indextid))
continue; continue;
/* Check if in deleted list */ /* Check if in deleted list */
if (DeletedContains(vacuumstate->deleted, &ntup->indextid)) if (DeletedContains(vacuumstate->deleted, &neighbor->indextid))
{ {
needsUpdated = true; needsUpdated = true;
break; break;
@@ -184,7 +193,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element)
HnswElement entryPoint; HnswElement entryPoint;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup; HnswNeighborTuple ntup = vacuumstate->ntup;
Size nsize = vacuumstate->nsize; Size neighborsz = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
/* Check if any neighbors point to deleted values */ /* Check if any neighbors point to deleted values */
if (!NeedsUpdated(vacuumstate, element)) if (!NeedsUpdated(vacuumstate, element))
@@ -217,13 +226,13 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element)
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas); buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE); page = GenericXLogRegisterBuffer(state, buf, 0);
/* Overwrites existing page via InitPage */
HnswInitPage(buf, page);
/* Update neighbors */ /* Update neighbors */
AddNeighborsToPage(index, page, element, ntup, nsize, m); HnswSetNeighborTuple(ntup, element, m);
if (!PageIndexTupleOverwrite(page, element->neighborOffno, (Item) ntup, neighborsz))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf); MarkBufferDirty(buf);
@@ -309,13 +318,18 @@ RepairGraph(HnswVacuumState * vacuumstate)
HnswElementTuple item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno)); HnswElementTuple item = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElement element; HnswElement element;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(item))
continue;
/* Skip updating neighbors if being deleted */ /* Skip updating neighbors if being deleted */
if (!ItemPointerIsValid(&item->heaptids[0])) if (!ItemPointerIsValid(&item->heaptids[0]))
continue; continue;
/* Create an element */ /* Create an element */
element = palloc(sizeof(HnswElementData)); element = palloc(sizeof(HnswElementData));
element->neighborPage = item->neighborPage; element->neighborPage = ItemPointerGetBlockNumber(&item->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&item->neighbortid);
element->level = item->level; element->level = item->level;
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
@@ -381,30 +395,68 @@ MarkDeleted(HnswVacuumState * vacuumstate)
Size itemsz; Size itemsz;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
BlockNumber neighborPage;
OffsetNumber neighborOffno;
Size ntupsz;
HnswNeighborTuple ntup;
int neighborCount;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(item))
continue;
if (ItemPointerIsValid(&item->heaptids[0])) if (ItemPointerIsValid(&item->heaptids[0]))
continue; continue;
/* Calculate sizes */
itemsz = HNSW_ELEMENT_TUPLE_SIZE(item->vec.dim);
ntupsz = HNSW_NEIGHBOR_TUPLE_SIZE(item->level, vacuumstate->m);
neighborCount = (item->level + 2) * vacuumstate->m;
/* Get neighbor page */
neighborPage = ItemPointerGetBlockNumber(&item->neighbortid);
neighborOffno = ItemPointerGetOffsetNumber(&item->neighbortid);
if (neighborPage == blkno)
{
nbuf = buf;
npage = page;
}
else
{
nbuf = ReadBufferExtended(index, MAIN_FORKNUM, neighborPage, RBM_NORMAL, bas);
LockBuffer(nbuf, BUFFER_LOCK_EXCLUSIVE);
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
ntup = (HnswNeighborTuple) PageGetItem(npage, PageGetItemId(npage, neighborOffno));
/* Overwrite element */ /* Overwrite element */
/* TODO Increment version? */ /* TODO Increment version? */
item->deleted = 1; item->deleted = 1;
MemSet(&item->vec.x, 0, item->vec.dim * sizeof(float)); MemSet(&item->vec.x, 0, item->vec.dim * sizeof(float));
itemsz = MAXALIGN(HNSW_ELEMENT_TUPLE_SIZE(item->vec.dim)); /* Overwrite neighbors */
for (int i = 0; i < neighborCount; i++)
{
ItemPointerSetInvalid(&ntup->neighbors[i].indextid);
ntup->neighbors[i].distance = NAN;
}
if (!PageIndexTupleOverwrite(page, offno, (Item) item, itemsz)) if (!PageIndexTupleOverwrite(page, offno, (Item) item, itemsz))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Overwrite neighbors */ if (!PageIndexTupleOverwrite(npage, neighborOffno, (Item) ntup, ntupsz))
nbuf = ReadBufferExtended(index, MAIN_FORKNUM, item->neighborPage, RBM_NORMAL, bas); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
LockBuffer(nbuf, BUFFER_LOCK_EXCLUSIVE);
npage = GenericXLogRegisterBuffer(state, nbuf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(nbuf, npage);
/* Commit */ /* Commit */
MarkBufferDirty(buf); MarkBufferDirty(buf);
MarkBufferDirty(nbuf); if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(nbuf); if (nbuf != buf)
UnlockReleaseBuffer(nbuf);
/* Set to first free page */ /* Set to first free page */
if (!BlockNumberIsValid(insertPage)) if (!BlockNumberIsValid(insertPage))
@@ -445,8 +497,7 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD); vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
vacuumstate->collation = index->rd_indcollation[0]; vacuumstate->collation = index->rd_indcollation[0];
vacuumstate->nsize = MAXALIGN(sizeof(HnswNeighborTupleData)); vacuumstate->ntup = palloc0(BLCKSZ);
vacuumstate->ntup = palloc0(vacuumstate->nsize);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context", "Hnsw vacuum temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);

3
test/t/011_hnsw_vacuum.pl
View File

@@ -38,6 +38,7 @@ $node->safe_psql("postgres",
# Check size # Check size
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');"); my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
is($size, $new_size, "size does not change"); # TODO Improve
cmp_ok($new_size, "<=", $size * 1.75, "size does not increase too much");
done_testing(); done_testing();

8
test/t/014_hnsw_inserts.pl
View File

@@ -4,7 +4,8 @@ use PostgresNode;
use TestLib; use TestLib;
use Test::More; use Test::More;
my $dim = 768; # Ensures elements and neighbors on both same and different pages
my $dim = 1900;
my $array_sql = join(",", ('random()') x $dim); my $array_sql = join(",", ('random()') x $dim);
@@ -16,6 +17,9 @@ $node->start;
# Create table and index # Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;"); $node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector($dim));"); $node->safe_psql("postgres", "CREATE TABLE tst (v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 100) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);"); $node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
$node->pgbench( $node->pgbench(
@@ -37,7 +41,7 @@ sub idx_scan
$node->safe_psql("postgres", "SELECT idx_scan FROM pg_stat_user_indexes WHERE indexrelid = 'tst_v_idx'::regclass;"); $node->safe_psql("postgres", "SELECT idx_scan FROM pg_stat_user_indexes WHERE indexrelid = 'tst_v_idx'::regclass;");
} }
my $expected = 5 * 100 * 10; my $expected = 100 + 5 * 100 * 10;
my $count = $node->safe_psql("postgres", "SELECT COUNT(*) FROM tst;"); my $count = $node->safe_psql("postgres", "SELECT COUNT(*) FROM tst;");
is($count, $expected); is($count, $expected);