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

Added support for index-only scans to HNSW

Browse Source
This commit is contained in:
Andrew Kane
2023-10-22 15:58:40 -07:00
parent 3f49b95f01
commit 84ca6625a7
4 changed files with 35 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/hnsw.c
View File

@@ -155,6 +155,15 @@ hnswvalidate(Oid opclassoid)
return true; return true;
} }
/*
* Checks if index-only scan is supported
*/
static bool
hnswcanreturn(Relation indexRelation, int attno)
{
return attno == 1 && !OidIsValid(index_getprocid(indexRelation, 1, HNSW_NORM_PROC));
}
/* /*
* Define index handler * Define index handler
* *
@@ -196,7 +205,7 @@ hnswhandler(PG_FUNCTION_ARGS)
amroutine->aminsert = hnswinsert; amroutine->aminsert = hnswinsert;
amroutine->ambulkdelete = hnswbulkdelete; amroutine->ambulkdelete = hnswbulkdelete;
amroutine->amvacuumcleanup = hnswvacuumcleanup; amroutine->amvacuumcleanup = hnswvacuumcleanup;
amroutine->amcanreturn = NULL; amroutine->amcanreturn = hnswcanreturn;
amroutine->amcostestimate = hnswcostestimate; amroutine->amcostestimate = hnswcostestimate;
amroutine->amoptions = hnswoptions; amroutine->amoptions = hnswoptions;
amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */ amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */

2
src/hnsw.h
View File

@@ -269,7 +269,7 @@ 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 *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement); List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool loadVec, HnswElement skipElement);
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);

25
src/hnswscan.c
View File

@@ -17,6 +17,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
Relation index = scan->indexRelation; Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo; FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation; Oid collation = so->collation;
bool loadVec = scan->xs_want_itup;
List *ep; List *ep;
List *w; List *w;
int m; int m;
@@ -28,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false)); ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, loadVec));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, false, NULL); w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, loadVec, NULL);
ep = w; ep = w;
} }
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL); return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, loadVec, NULL);
} }
/* /*
@@ -113,6 +114,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
scan->opaque = so; scan->opaque = so;
/* OK to always set since cheap */
scan->xs_itupdesc = RelationGetDescr(index);
return scan; return scan;
} }
@@ -198,6 +202,18 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
hc->element->heaptids = list_delete_last(hc->element->heaptids); hc->element->heaptids = list_delete_last(hc->element->heaptids);
if (scan->xs_want_itup)
{
Datum value = PointerGetDatum(hc->element->vec);
bool isnull = false;
if (scan->xs_itup)
pfree(scan->xs_itup);
scan->xs_itup = index_form_tuple(scan->xs_itupdesc, &value, &isnull);
scan->xs_itup->t_tid = *heaptid;
}
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
@@ -226,4 +242,7 @@ hnswendscan(IndexScanDesc scan)
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;
if (scan->xs_itup)
pfree(scan->xs_itup);
} }

4
src/hnswutils.c
View File

@@ -569,7 +569,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement) HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool loadVec, HnswElement skipElement)
{ {
ListCell *lc2; ListCell *lc2;
@@ -645,7 +645,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
if (index == NULL) if (index == NULL)
eDistance = GetCandidateDistance(e, q, procinfo, collation); eDistance = GetCandidateDistance(e, q, procinfo, collation);
else else
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting); HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, loadVec);
Assert(!e->element->deleted); Assert(!e->element->deleted);