liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-14 16:46:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added support for index-only scans to IVFFlat

Browse Source
This commit is contained in:
Andrew Kane
2023-09-03 16:56:34 -07:00
parent 3431acef94
commit 97000a7903
2 changed files with 31 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/ivfflat.c
View File

@@ -182,6 +182,15 @@ ivfflatvalidate(Oid opclassoid)
return true; return true;
} }
/*
* Checks if index-only scan is supported
*/
static bool
ivfflatcanreturn(Relation index, int attno)
{
return attno == 1 && IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC) == NULL;
}
/* /*
* Define index handler * Define index handler
* *
@@ -223,7 +232,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
amroutine->aminsert = ivfflatinsert; amroutine->aminsert = ivfflatinsert;
amroutine->ambulkdelete = ivfflatbulkdelete; amroutine->ambulkdelete = ivfflatbulkdelete;
amroutine->amvacuumcleanup = ivfflatvacuumcleanup; amroutine->amvacuumcleanup = ivfflatvacuumcleanup;
amroutine->amcanreturn = NULL; /* tuple not included in heapsort */ amroutine->amcanreturn = ivfflatcanreturn;
amroutine->amcostestimate = ivfflatcostestimate; amroutine->amcostestimate = ivfflatcostestimate;
amroutine->amoptions = ivfflatoptions; amroutine->amoptions = ivfflatoptions;
amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */ amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */

25
src/ivfscan.c
View File

@@ -141,6 +141,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
IndexTuple itup; IndexTuple itup;
Datum datum; Datum datum;
bool isnull; bool isnull;
ItemPointerData indextid;
ItemId itemid = PageGetItemId(page, offno); ItemId itemid = PageGetItemId(page, offno);
/* Skip dead tuples */ /* Skip dead tuples */
@@ -149,6 +150,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
itup = (IndexTuple) PageGetItem(page, itemid); itup = (IndexTuple) PageGetItem(page, itemid);
datum = index_getattr(itup, 1, tupdesc, &isnull); datum = index_getattr(itup, 1, tupdesc, &isnull);
ItemPointerSet(&indextid, searchPage, offno);
/* /*
* Add virtual tuple * Add virtual tuple
@@ -161,7 +163,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
slot->tts_isnull[0] = false; slot->tts_isnull[0] = false;
slot->tts_values[1] = PointerGetDatum(&itup->t_tid); slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = Int32GetDatum((int) searchPage); slot->tts_values[2] = PointerGetDatum(&indextid);
slot->tts_isnull[2] = false; slot->tts_isnull[2] = false;
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
@@ -280,7 +282,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
#endif #endif
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indextid", TIDOID, -1, 0);
/* Prep sort */ /* Prep sort */
so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false); so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -295,6 +297,8 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
scan->opaque = so; scan->opaque = so;
scan->xs_itupdesc = RelationGetDescr(index);
return scan; return scan;
} }
@@ -380,7 +384,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL)) if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
{ {
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull)); ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull)); ItemPointer indextid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 3, &so->isnull));
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
@@ -401,7 +405,20 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
* *
* https://www.postgresql.org/docs/current/index-locking.html * https://www.postgresql.org/docs/current/index-locking.html
*/ */
so->buf = ReadBuffer(scan->indexRelation, indexblkno); so->buf = ReadBuffer(scan->indexRelation, ItemPointerGetBlockNumber(indextid));
if (scan->xs_want_itup)
{
Page page;
LockBuffer(so->buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(so->buf);
/* TODO Copy tuple to IvfflatScanOpaque */
scan->xs_itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, ItemPointerGetOffsetNumber(indextid)));
LockBuffer(so->buf, BUFFER_LOCK_UNLOCK);
}
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;