liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-30 09:41:15 +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;
}
/*
* 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
*
@@ -196,7 +205,7 @@ hnswhandler(PG_FUNCTION_ARGS)
amroutine->aminsert = hnswinsert;
amroutine->ambulkdelete = hnswbulkdelete;
amroutine->amvacuumcleanup = hnswvacuumcleanup;
amroutine->amcanreturn = NULL;
amroutine->amcanreturn = hnswcanreturn;
amroutine->amcostestimate = hnswcostestimate;
amroutine->amoptions = hnswoptions;
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 HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
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);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
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;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
bool loadVec = scan->xs_want_itup;
List *ep;
List *w;
int m;
@@ -28,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL)
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--)
{
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;
}
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;
/* OK to always set since cheap */
scan->xs_itupdesc = RelationGetDescr(index);
return scan;
}
@@ -198,6 +202,18 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
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);
#if PG_VERSION_NUM >= 120000
@@ -226,4 +242,7 @@ hnswendscan(IndexScanDesc scan)
pfree(so);
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
*/
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;
@@ -645,7 +645,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
if (index == NULL)
eDistance = GetCandidateDistance(e, q, procinfo, collation);
else
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, loadVec);
Assert(!e->element->deleted);