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

Use memory context for index scan [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2023-08-04 19:59:09 -07:00
parent 4f4f444396
commit 3a6cd544ff
2 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/hnsw.h
View File

@@ -220,6 +220,7 @@ typedef struct HnswScanOpaqueData
bool first; bool first;
Buffer buf; Buffer buf;
List *w; List *w;
MemoryContext tmpCtx;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;

17
src/hnswscan.c
View File

@@ -4,6 +4,7 @@
#include "hnsw.h" #include "hnsw.h"
#include "pgstat.h" #include "pgstat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/memutils.h"
/* /*
* Algorithm 5 from paper * Algorithm 5 from paper
@@ -22,7 +23,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL) if (entryPoint == NULL)
return; return;
/* TODO Use memory context */
ep = lappend(ep, EntryCandidate(entryPoint, q, index, procinfo, collation, false)); ep = lappend(ep, EntryCandidate(entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
@@ -31,7 +31,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
ep = w; ep = w;
} }
/* TODO Return all visited elements at level 0, not just ef search */
so->w = SearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL, NULL); so->w = SearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL, NULL);
} }
@@ -72,7 +71,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->buf = InvalidBuffer; so->buf = InvalidBuffer;
so->first = true; so->first = true;
so->w = NIL; so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
@@ -93,8 +94,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
so->first = true; so->first = true;
list_free(so->w); MemoryContextReset(so->tmpCtx);
so->w = NIL;
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));
@@ -110,6 +110,7 @@ bool
hnswgettuple(IndexScanDesc scan, ScanDirection dir) hnswgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -169,6 +170,8 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
hc->element->heaptids = list_delete_last(hc->element->heaptids); hc->element->heaptids = list_delete_last(hc->element->heaptids);
MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *tid; scan->xs_heaptid = *tid;
#else #else
@@ -187,9 +190,11 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->buf = ReadBuffer(scan->indexRelation, indexblkno); so->buf = ReadBuffer(scan->indexRelation, indexblkno);
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
} }
MemoryContextSwitchTo(oldCtx);
return false; return false;
} }
@@ -205,7 +210,7 @@ hnswendscan(IndexScanDesc scan)
if (BufferIsValid(so->buf)) if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf); ReleaseBuffer(so->buf);
list_free(so->w); MemoryContextDelete(so->tmpCtx);
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;