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

Merge branch 'master' into ivfflat-streaming

Browse Source
This commit is contained in:
Andrew Kane
2024-09-21 19:26:12 -07:00
parent 88f56dc234 97cf990e0f
commit c950c5ffaa
2 changed files with 22 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/ivfflat.h
View File

@@ -257,8 +257,9 @@ typedef struct IvfflatScanOpaqueData
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc tupdesc;
TupleTableSlot *slot; TupleTableSlot *vslot;
bool isnull; TupleTableSlot *mslot;
BufferAccessStrategy bas;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;

33
src/ivfscan.c
View File

@@ -128,16 +128,9 @@ GetScanItems(IndexScanDesc scan, Datum value)
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; double tuples = 0;
TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual); TupleTableSlot *slot = so->vslot;
int batchProbes = 0; int batchProbes = 0;
/*
* Reuse same set of shared buffers for scan
*
* See postgres/src/backend/storage/buffer/README for description
*/
BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
tuplesort_reset(so->sortstate); tuplesort_reset(so->sortstate);
/* Search closest probes lists */ /* Search closest probes lists */
@@ -152,7 +145,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
Page page; Page page;
OffsetNumber maxoffno; OffsetNumber maxoffno;
buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas); buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, so->bas);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page); maxoffno = PageGetMaxOffsetNumber(page);
@@ -191,8 +184,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
FreeAccessStrategy(bas);
if (tuples < 100 && !ivfflat_streaming) if (tuples < 100 && !ivfflat_streaming)
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("index scan found few tuples"), (errmsg("index scan found few tuples"),
@@ -303,7 +294,16 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Prep sort */ /* Prep sort */
so->sortstate = InitScanSortState(so->tupdesc); so->sortstate = InitScanSortState(so->tupdesc);
so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple); /* Need separate slots for puttuple and gettuple */
so->vslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
so->mslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
/*
* Reuse same set of shared buffers for scan
*
* See postgres/src/backend/storage/buffer/README for description
*/
so->bas = GetAccessStrategy(BAS_BULKREAD);
/* Order by closest list for streaming */ /* Order by closest list for streaming */
if (ivfflat_streaming) if (ivfflat_streaming)
@@ -342,6 +342,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
ItemPointer heaptid; ItemPointer heaptid;
bool isnull;
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -374,7 +375,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
/* TODO clean up if we allocated a new value */ /* TODO clean up if we allocated a new value */
} }
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
if (pairingheap_is_empty(so->listQueue)) if (pairingheap_is_empty(so->listQueue))
return false; return false;
@@ -382,7 +383,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanItems", GetScanItems(scan, so->value)); IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
} }
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull)); heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
scan->xs_recheck = false; scan->xs_recheck = false;
@@ -400,6 +401,10 @@ ivfflatendscan(IndexScanDesc scan)
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
/* TODO Free vslot and mslot without freeing TupleDesc */
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;