liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-12 23: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-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; 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 * 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

@@ -266,7 +266,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, 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 HnswGetEntryPoint(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);

18
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;
HnswElement entryPoint = HnswGetEntryPoint(index); HnswElement entryPoint = HnswGetEntryPoint(index);
@@ -24,15 +25,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, false, NULL); w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, loadVec, NULL);
ep = w; 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->opaque = so;
scan->xs_itupdesc = RelationGetDescr(index);
return scan; return scan;
} }
@@ -177,6 +180,15 @@ 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;
scan->xs_itup = index_form_tuple(scan->xs_itupdesc, &value, &isnull);
scan->xs_itup->t_tid = *tid;
}
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000 #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 * Algorithm 2 from paper
*/ */
List * 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; ListCell *lc2;
@@ -619,7 +619,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);