From 3cf6f6290019ea60ea166f3b4dd4841900b4ccea Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 9 Nov 2023 17:35:39 -0800 Subject: [PATCH] Switched to datum for HnswElement --- src/hnsw.h | 2 +- src/hnswbuild.c | 13 ++++++++----- src/hnswinsert.c | 4 ++-- src/hnswutils.c | 25 ++++++++++++++----------- 4 files changed, 25 insertions(+), 19 deletions(-) diff --git a/src/hnsw.h b/src/hnsw.h index ff7a8c7..e6b45a4 100644 --- a/src/hnsw.h +++ b/src/hnsw.h @@ -103,7 +103,7 @@ typedef struct HnswElementData OffsetNumber offno; OffsetNumber neighborOffno; BlockNumber neighborPage; - Vector *vec; + Datum value; } HnswElementData; typedef HnswElementData * HnswElement; diff --git a/src/hnswbuild.c b/src/hnswbuild.c index 9621cc2..7035840 100644 --- a/src/hnswbuild.c +++ b/src/hnswbuild.c @@ -8,6 +8,7 @@ #include "lib/pairingheap.h" #include "nodes/pg_list.h" #include "storage/bufmgr.h" +#include "utils/datum.h" #include "utils/memutils.h" #if PG_VERSION_NUM >= 140000 @@ -264,13 +265,14 @@ FlushPages(HnswBuildState * buildstate) * Insert tuple */ static bool -InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup) +InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx) { FmgrInfo *procinfo = buildstate->procinfo; Oid collation = buildstate->collation; HnswElement entryPoint = buildstate->entryPoint; int efConstruction = buildstate->efConstruction; int m = buildstate->m; + MemoryContext oldCtx; /* Detoast once for all calls */ Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); @@ -283,7 +285,9 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * } /* Copy value to element so accessible outside of memory context */ - memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions)); + oldCtx = MemoryContextSwitchTo(outerCtx); + element->value = datumCopy(value, false, -1); + MemoryContextSwitchTo(oldCtx); /* Insert element in graph */ HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); @@ -324,7 +328,7 @@ HnswElementMemory(HnswElement e, int m) elementSize += sizeof(HnswNeighborArray) * (e->level + 1); elementSize += sizeof(HnswCandidate) * (m * (e->level + 2)); elementSize += sizeof(ItemPointerData); - elementSize += VARSIZE_ANY(e->vec); + elementSize += VARSIZE_ANY(e->value); return elementSize; } @@ -375,13 +379,12 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values, /* Allocate necessary memory outside of memory context */ element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); - element->vec = palloc(VECTOR_SIZE(buildstate->dimensions)); /* Use memory context since detoast can allocate */ oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); /* Insert tuple */ - inserted = InsertTuple(index, values, element, buildstate, &dup); + inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx); /* Reset memory context */ MemoryContextSwitchTo(oldCtx); diff --git a/src/hnswinsert.c b/src/hnswinsert.c index 705b1f1..c9414af 100644 --- a/src/hnswinsert.c +++ b/src/hnswinsert.c @@ -131,7 +131,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag BlockNumber newInsertPage = InvalidBlockNumber; /* Calculate sizes */ - etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(e->vec)); + etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(e->value)); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); combinedSize = etupSize + ntupSize + sizeof(ItemIdData); maxSize = HNSW_MAX_SIZE; @@ -517,7 +517,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti /* Create an element */ element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); - element->vec = DatumGetVector(value); + element->value = value; /* Prevent concurrent inserts when likely updating entry point */ if (entryPoint == NULL || element->level > entryPoint->level) diff --git a/src/hnswutils.c b/src/hnswutils.c index 0c82636..ea52b7f 100644 --- a/src/hnswutils.c +++ b/src/hnswutils.c @@ -187,7 +187,8 @@ HnswFreeElement(HnswElement element) { HnswFreeNeighbors(element); list_free_deep(element->heaptids); - pfree(element->vec); + if (DatumGetPointer(element->value)) + pfree(DatumGetPointer(element->value)); pfree(element); } @@ -214,7 +215,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno) element->blkno = blkno; element->offno = offno; element->neighbors = NULL; - element->vec = NULL; + element->value = PointerGetDatum(NULL); return element; } @@ -324,7 +325,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element) else ItemPointerSetInvalid(&etup->heaptids[i]); } - memcpy(&etup->vec, element->vec, VARSIZE_ANY(element->vec)); + memcpy(&etup->vec, DatumGetPointer(element->value), VARSIZE_ANY(element->value)); } /* @@ -447,8 +448,10 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe if (loadVec) { - element->vec = palloc(VARSIZE_ANY(&etup->vec)); - memcpy(element->vec, &etup->vec, VARSIZE_ANY(&etup->vec)); + Vector *vec = palloc(VARSIZE_ANY(&etup->vec)); + + memcpy(vec, &etup->vec, VARSIZE_ANY(&etup->vec)); + element->value = PointerGetDatum(vec); } } @@ -487,7 +490,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, static float GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) { - return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec))); + return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value)); } /* @@ -750,7 +753,7 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co } } - return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec))); + return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value)); } /* @@ -877,7 +880,7 @@ HnswFindDuplicate(HnswElement e) HnswCandidate *neighbor = &neighbors->items[i]; /* Exit early since ordered by distance */ - if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0) + if (vector_cmp_internal(DatumGetVector(e->value), DatumGetVector(neighbor->element->value)) != 0) break; /* Check for space */ @@ -930,13 +933,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int /* Load elements on insert */ if (index != NULL) { - Datum q = PointerGetDatum(hc->element->vec); + Datum q = hc->element->value; for (int i = 0; i < currentNeighbors->length; i++) { HnswCandidate *hc3 = ¤tNeighbors->items[i]; - if (hc3->element->vec == NULL) + if (DatumGetPointer(hc3->element->value) == NULL) HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); else hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); @@ -1017,7 +1020,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F List *w; int level = element->level; int entryLevel; - Datum q = PointerGetDatum(element->vec); + Datum q = element->value; HnswElement skipElement = existing ? element : NULL; /* No neighbors if no entry point */