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

Improved approach [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-09-21 19:55:55 -07:00
parent c950c5ffaa
commit ff6267917e
2 changed files with 15 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/ivfflat.h
View File

@@ -269,6 +269,8 @@ typedef struct IvfflatScanOpaqueData
/* Lists */
pairingheap *listQueue;
BlockNumber *startPages;
int currentIndex;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
} IvfflatScanOpaqueData;

36
src/ivfscan.c
View File

@@ -33,21 +33,6 @@ CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
return 0;
}
/*
* Compare list distances for streaming
*/
static int
CompareListsStreaming(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
return -1;
if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
return 1;
return 0;
}
/*
* Get lists and sort by distance
*/
@@ -117,6 +102,11 @@ GetScanLists(IndexScanDesc scan, Datum value)
UnlockReleaseBuffer(cbuf);
}
for (int i = listCount - 1; i >= 0; i--)
so->startPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
Assert(pairingheap_is_empty(so->listQueue));
}
/*
@@ -134,9 +124,9 @@ GetScanItems(IndexScanDesc scan, Datum value)
tuplesort_reset(so->sortstate);
/* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue) && (++batchProbes) <= so->probes)
while (so->currentIndex < so->maxProbes && (++batchProbes) <= so->probes)
{
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
BlockNumber searchPage = so->startPages[so->currentIndex++];
/* Search all entry pages for list */
while (BlockNumberIsValid(searchPage))
@@ -305,11 +295,9 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
*/
so->bas = GetAccessStrategy(BAS_BULKREAD);
/* Order by closest list for streaming */
if (ivfflat_streaming)
so->listQueue = pairingheap_allocate(CompareListsStreaming, scan);
else
so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listQueue = pairingheap_allocate(CompareLists, scan);
so->startPages = palloc(maxProbes * sizeof(BlockNumber));
so->currentIndex = 0;
scan->opaque = so;
@@ -326,6 +314,7 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
so->first = true;
pairingheap_reset(so->listQueue);
so->currentIndex = 0;
if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -377,7 +366,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{
if (pairingheap_is_empty(so->listQueue))
if (so->currentIndex == so->maxProbes)
return false;
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
@@ -400,6 +389,7 @@ ivfflatendscan(IndexScanDesc scan)
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue);
pfree(so->startPages);
tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);