liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-01 02:02:10 +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-09-02 18:54:46 -07:00
parent 034d4acaea
commit 0a3669a680
4 changed files with 28 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 && HnswOptionalProcInfo(indexRelation, HNSW_NORM_PROC) == NULL;
}
/*
* 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

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

18
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;
HnswElement entryPoint = HnswGetEntryPoint(index);
@@ -24,15 +25,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, false, NULL);
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, loadVec, NULL);
ep = w;
}
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL);
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, loadVec, NULL);
}
/*
@@ -83,6 +84,8 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
scan->opaque = so;
scan->xs_itupdesc = RelationGetDescr(index);
return scan;
}
@@ -177,6 +180,15 @@ 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;
scan->xs_itup = index_form_tuple(scan->xs_itupdesc, &value, &isnull);
scan->xs_itup->t_tid = *tid;
}
MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000

4
src/hnswutils.c
View File

@@ -543,7 +543,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, bool inserting, HnswElement skipElement)
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, HnswElement skipElement)
{
ListCell *lc2;
@@ -619,7 +619,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);