liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-09 14:20:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Switched to datum for HnswElement

Browse Source
This commit is contained in:
Andrew Kane
2023-11-09 17:35:39 -08:00
parent 2a69e22ca4
commit 3cf6f62900
4 changed files with 25 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/hnsw.h
View File

@@ -103,7 +103,7 @@ typedef struct HnswElementData
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Vector *vec; Datum value;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;

13
src/hnswbuild.c
View File

@@ -8,6 +8,7 @@
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "nodes/pg_list.h" #include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -264,13 +265,14 @@ FlushPages(HnswBuildState * buildstate)
* Insert tuple * Insert tuple
*/ */
static bool 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; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint; HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
MemoryContext oldCtx;
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); 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 */ /* 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 */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); 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(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2)); elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += sizeof(ItemPointerData); elementSize += sizeof(ItemPointerData);
elementSize += VARSIZE_ANY(e->vec); elementSize += VARSIZE_ANY(e->value);
return elementSize; return elementSize;
} }
@@ -375,13 +379,12 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Allocate necessary memory outside of memory context */ /* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
/* Use memory context since detoast can allocate */ /* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup); inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);

4
src/hnswinsert.c
View File

@@ -131,7 +131,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */ /* 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); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -517,7 +517,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
element->vec = DatumGetVector(value); element->value = value;
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)

25
src/hnswutils.c
View File

@@ -187,7 +187,8 @@ HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element); HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
pfree(element->vec); if (DatumGetPointer(element->value))
pfree(DatumGetPointer(element->value));
pfree(element); pfree(element);
} }
@@ -214,7 +215,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; element->neighbors = NULL;
element->vec = NULL; element->value = PointerGetDatum(NULL);
return element; return element;
} }
@@ -324,7 +325,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); 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) if (loadVec)
{ {
element->vec = palloc(VARSIZE_ANY(&etup->vec)); Vector *vec = palloc(VARSIZE_ANY(&etup->vec));
memcpy(element->vec, &etup->vec, 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 static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) 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]; HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */ /* 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; break;
/* Check for space */ /* Check for space */
@@ -930,13 +933,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = PointerGetDatum(hc->element->vec); Datum q = hc->element->value;
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
if (hc3->element->vec == NULL) if (DatumGetPointer(hc3->element->value) == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
@@ -1017,7 +1020,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = PointerGetDatum(element->vec); Datum q = element->value;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */ /* No neighbors if no entry point */