mirror of
https://github.com/pgvector/pgvector.git
synced 2026-06-30 17:51:18 +08:00
Switched to datum for HnswElement
This commit is contained in:
@@ -103,7 +103,7 @@ typedef struct HnswElementData
|
||||
OffsetNumber offno;
|
||||
OffsetNumber neighborOffno;
|
||||
BlockNumber neighborPage;
|
||||
Vector *vec;
|
||||
Datum value;
|
||||
} HnswElementData;
|
||||
|
||||
typedef HnswElementData * HnswElement;
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 */
|
||||
|
||||
Reference in New Issue
Block a user