liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-11 23:26:55 +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 */ /* Lists */
pairingheap *listQueue; pairingheap *listQueue;
BlockNumber *startPages;
int currentIndex;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;

36
src/ivfscan.c
View File

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