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

Use memory context for building index

Browse Source
This commit is contained in:
Andrew Kane
2023-01-10 00:53:33 -08:00
parent 4a2f20f5fc
commit ec800ef903
2 changed files with 42 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
src/ivfbuild.c
View File

@@ -6,6 +6,7 @@
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
#include "utils/backend_progress.h" #include "utils/backend_progress.h"
@@ -118,27 +119,20 @@ SampleRows(IvfflatBuildState * buildstate)
} }
/* /*
* Callback for table_index_build_scan * Add tuple to sort
*/ */
static void static void
BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values, AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate)
bool *isnull, bool tupleIsAlive, void *state)
{ {
IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
double distance; double distance;
double minDistance = DBL_MAX; double minDistance = DBL_MAX;
int closestCenter = -1; int closestCenter = -1;
VectorArray centers = buildstate->centers; VectorArray centers = buildstate->centers;
TupleTableSlot *slot = buildstate->slot; TupleTableSlot *slot = buildstate->slot;
Datum value = values[0];
int i; int i;
#if PG_VERSION_NUM < 130000 /* Detoast once for all calls */
ItemPointer tid = &hup->t_self; Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
#endif
if (isnull[0])
return;
/* Normalize if needed */ /* Normalize if needed */
if (buildstate->normprocinfo != NULL) if (buildstate->normprocinfo != NULL)
@@ -186,6 +180,34 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
buildstate->indtuples++; buildstate->indtuples++;
} }
/*
* Callback for table_index_build_scan
*/
static void
BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state)
{
IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
MemoryContext oldCtx;
#if PG_VERSION_NUM < 130000
ItemPointer tid = &hup->t_self;
#endif
if (isnull[0])
return;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add tuple to sort */
AddTupleToSort(index, tid, values, buildstate);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
}
/* /*
* Get index tuple from sort state * Get index tuple from sort state
*/ */
@@ -331,6 +353,10 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
/* Reuse for each tuple */ /* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions); buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat build temporary context",
ALLOCSET_DEFAULT_SIZES);
#ifdef IVFFLAT_KMEANS_DEBUG #ifdef IVFFLAT_KMEANS_DEBUG
buildstate->inertia = 0; buildstate->inertia = 0;
buildstate->listSums = palloc0(sizeof(double) * buildstate->lists); buildstate->listSums = palloc0(sizeof(double) * buildstate->lists);
@@ -352,6 +378,8 @@ FreeBuildState(IvfflatBuildState * buildstate)
pfree(buildstate->listSums); pfree(buildstate->listSums);
pfree(buildstate->listCounts); pfree(buildstate->listCounts);
#endif #endif
MemoryContextDelete(buildstate->tmpCtx);
} }
/* /*

3
src/ivfflat.h
View File

@@ -145,6 +145,9 @@ typedef struct IvfflatBuildState
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc tupdesc;
TupleTableSlot *slot; TupleTableSlot *slot;
/* Memory */
MemoryContext tmpCtx;
} IvfflatBuildState; } IvfflatBuildState;
typedef struct IvfflatMetaPageData typedef struct IvfflatMetaPageData