Use separate memory context for updating neighbors, which improves performance around 10% for larger vectors

This commit is contained in:
Andrew Kane
2024-09-30 11:15:27 -07:00
parent ff6da4fcea
commit 57248ba128
2 changed files with 15 additions and 2 deletions

View File

@@ -2,6 +2,7 @@
- Added casts for arrays to `sparsevec` - Added casts for arrays to `sparsevec`
- Improved cost estimation - Improved cost estimation
- Improved performance of HNSW inserts and on-disk index builds
- Reduced memory usage for HNSW index scans - Reduced memory usage for HNSW index scans
- Dropped support for Postgres 12 - Dropped support for Postgres 12

View File

@@ -394,11 +394,12 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, int *idx, Relation
* Get update index * Get update index
*/ */
static int static int
GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo *procinfo, Oid collation) GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo *procinfo, Oid collation, MemoryContext updateCtx)
{ {
char *base = NULL; char *base = NULL;
int idx = -1; int idx = -1;
HnswNeighborArray *neighbors; HnswNeighborArray *neighbors;
MemoryContext oldCtx = MemoryContextSwitchTo(updateCtx);
/* /*
* Get latest neighbors since they may have changed. Do not lock yet since * Get latest neighbors since they may have changed. Do not lock yet since
@@ -426,6 +427,9 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation); HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation);
} }
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(updateCtx);
return idx; return idx;
} }
@@ -529,6 +533,14 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
{ {
char *base = NULL; char *base = NULL;
/* Use separate memory context to improve performance for larger vectors */
MemoryContext updateCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw insert update context",
#if PG_VERSION_NUM >= 150000
128 * 1024, 128 * 1024,
#endif
128 * 1024);
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
@@ -540,7 +552,7 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
HnswElement neighborElement = HnswPtrAccess(base, hc->element); HnswElement neighborElement = HnswPtrAccess(base, hc->element);
int idx; int idx;
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation); idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation, updateCtx);
/* New element was not selected as a neighbor */ /* New element was not selected as a neighbor */
if (idx == -1) if (idx == -1)