liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-30 09:41:15 +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;
Buffer buf;
List *w;
MemoryContext tmpCtx;
/* Support functions */
FmgrInfo *procinfo;

17
src/hnswscan.c
View File

@@ -4,6 +4,7 @@
#include "hnsw.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
#include "utils/memutils.h"
/*
* Algorithm 5 from paper
@@ -22,7 +23,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL)
return;
/* TODO Use memory context */
ep = lappend(ep, EntryCandidate(entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--)
@@ -31,7 +31,6 @@ GetScanItems(IndexScanDesc scan, Datum q)
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);
}
@@ -72,7 +71,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->buf = InvalidBuffer;
so->first = true;
so->w = NIL;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */
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;
so->first = true;
list_free(so->w);
so->w = NIL;
MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -110,6 +110,7 @@ bool
hnswgettuple(IndexScanDesc scan, ScanDirection dir)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/*
* 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);
MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *tid;
#else
@@ -187,9 +190,11 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->buf = ReadBuffer(scan->indexRelation, indexblkno);
scan->xs_recheckorderby = false;
return true;
}
MemoryContextSwitchTo(oldCtx);
return false;
}
@@ -205,7 +210,7 @@ hnswendscan(IndexScanDesc scan)
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
list_free(so->w);
MemoryContextDelete(so->tmpCtx);
pfree(so);
scan->opaque = NULL;