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

DRY code for forming index value

Browse Source
This commit is contained in:
Andrew Kane
2024-10-09 20:50:17 -07:00
parent 3126fbdb6f
commit 57c05c59a2
4 changed files with 43 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/hnsw.h
View File

@@ -388,10 +388,11 @@ void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, in
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator); HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building); bool HnswInsertTupleOnDisk(Relation index, Datum value, ItemPointer heaptid, bool building);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building); void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, double *maxDistance); void HnswLoadElement(HnswElement element, double *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, double *maxDistance);
bool HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, FmgrInfo *normprocinfo, Oid collation);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation); void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc); bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);

24
src/hnswbuild.c
View File

@@ -473,7 +473,6 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
static bool static bool
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate) InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
{ {
const HnswTypeInfo *typeInfo = buildstate->typeInfo;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswElement element; HnswElement element;
HnswAllocator *allocator = &buildstate->allocator; HnswAllocator *allocator = &buildstate->allocator;
@@ -481,22 +480,11 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
Pointer valuePtr; Pointer valuePtr;
LWLock *flushLock = &graph->flushLock; LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
Datum value;
/* Detoast once for all calls */ /* Form index value */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, buildstate->normprocinfo, buildstate->collation))
return false;
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation, value))
return false;
value = HnswNormValue(typeInfo, buildstate->collation, value);
}
/* Get datum size */ /* Get datum size */
valueSize = VARSIZE_ANY(DatumGetPointer(value)); valueSize = VARSIZE_ANY(DatumGetPointer(value));
@@ -509,7 +497,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
{ {
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true); return HnswInsertTupleOnDisk(index, value, heaptid, true);
} }
/* /*
@@ -541,7 +529,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true); return HnswInsertTupleOnDisk(index, value, heaptid, true);
} }
/* Ok, we can proceed to allocate the element */ /* Ok, we can proceed to allocate the element */

29
src/hnswinsert.c
View File

@@ -679,7 +679,7 @@ UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building) HnswInsertTupleOnDisk(Relation index, Datum value, ItemPointer heaptid, bool building)
{ {
HnswElement entryPoint; HnswElement entryPoint;
HnswElement element; HnswElement element;
@@ -701,7 +701,7 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
/* Create an element */ /* Create an element */
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL); element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->value, DatumGetPointer(value)); HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
@@ -734,31 +734,18 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
* Insert a tuple into the index * Insert a tuple into the index
*/ */
static void static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid) HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
{ {
Datum value; Datum value;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index); const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
/* Detoast once for all calls */ /* Form index value */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); if (!HnswFormIndexValue(&value, values, isnull, typeInfo, normprocinfo, collation))
return;
/* Check value */ HnswInsertTupleOnDisk(index, value, heaptid, false);
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswCheckNorm(normprocinfo, collation, value))
return;
value = HnswNormValue(typeInfo, collation, value);
}
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);
} }
/* /*

27
src/hnswutils.c
View File

@@ -394,6 +394,33 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
/*
* Form index value
*/
bool
HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, FmgrInfo *normprocinfo, Oid collation)
{
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (normprocinfo != NULL)
{
if (!HnswCheckNorm(normprocinfo, collation, value))
return false;
value = HnswNormValue(typeInfo, collation, value);
}
*out = value;
return true;
}
/* /*
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */