liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-29 17:21:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added itemsize to VectorArray [skip ci]

Browse Source
This commit is contained in:
Andrew Kane
2024-04-11 13:47:36 -07:00
parent 5834b58c5a
commit 4ff36af67e
4 changed files with 19 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/ivfbuild.c
View File

@@ -367,7 +367,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, VECTOR_SIZE(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
@@ -422,7 +422,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */
/* TODO Ensure within maintenance_work_mem */
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions, buildstate->centers->itemsize);
if (buildstate->heap != NULL)
{
SampleRows(buildstate);
@@ -495,7 +495,7 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
/* Load list */
list->startPage = InvalidBlockNumber;
list->insertPage = InvalidBlockNumber;
memcpy(&list->center, VectorArrayGet(centers, i), VECTOR_SIZE(dimensions));
memcpy(&list->center, VectorArrayGet(centers, i), centers->itemsize);
/* Ensure free space */
if (PageGetFreeSpace(page) < listSize)
@@ -624,7 +624,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
indexInfo = BuildIndexInfo(ivfspool->index);
indexInfo->ii_Concurrent = ivfshared->isconcurrent;
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, VECTOR_SIZE(buildstate.centers->dim) * buildstate.centers->maxlen);
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = tuplesort_begin_heap(buildstate.tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, sortmem, coordinate, false);
buildstate.sortstate = ivfspool->sortstate;
@@ -813,7 +813,7 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
shm_toc_estimate_chunk(&pcxt->estimator, estivfshared);
estsort = tuplesort_estimate_shared(scantuplesortstates);
shm_toc_estimate_chunk(&pcxt->estimator, estsort);
estcenters = VECTOR_SIZE(buildstate->dimensions) * buildstate->lists;
estcenters = buildstate->centers->itemsize * buildstate->lists;
shm_toc_estimate_chunk(&pcxt->estimator, estcenters);
shm_toc_estimate_keys(&pcxt->estimator, 3);

9
src/ivfflat.h
View File

@@ -90,6 +90,7 @@ typedef struct VectorArrayData
int length;
int maxlen;
int dim;
Size itemsize;
Vector *items;
} VectorArrayData;
@@ -261,13 +262,13 @@ typedef struct IvfflatScanOpaqueData
typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
#define VECTOR_ARRAY_SIZE(_length, _dim) (sizeof(VectorArrayData) + (_length) * VECTOR_SIZE(_dim))
#define VECTOR_ARRAY_OFFSET(_arr, _offset) ((char*) (_arr)->items + (_offset) * VECTOR_SIZE((_arr)->dim))
#define VECTOR_ARRAY_SIZE(_length, _size) (sizeof(VectorArrayData) + (_length) * _size)
#define VECTOR_ARRAY_OFFSET(_arr, _offset) ((char*) (_arr)->items + (_offset) * (_arr)->itemsize)
#define VectorArrayGet(_arr, _offset) ((Vector *) VECTOR_ARRAY_OFFSET(_arr, _offset))
#define VectorArraySet(_arr, _offset, _val) memcpy(VECTOR_ARRAY_OFFSET(_arr, _offset), _val, VECTOR_SIZE((_arr)->dim))
#define VectorArraySet(_arr, _offset, _val) memcpy(VECTOR_ARRAY_OFFSET(_arr, _offset), _val, (_arr)->itemsize)
/* Methods */
VectorArray VectorArrayInit(int maxlen, int dimensions);
VectorArray VectorArrayInit(int maxlen, int dimensions, Size itemsize);
void VectorArrayFree(VectorArray arr);
void PrintVectorArray(char *msg, VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);

12
src/ivfkmeans.c
View File

@@ -120,7 +120,7 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
/* Copy existing vectors while avoiding duplicates */
if (samples->length > 0)
{
qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
qsort(samples->items, samples->length, samples->itemsize, CompareVectors);
for (int i = 0; i < samples->length; i++)
{
Vector *vec = VectorArrayGet(samples, i);
@@ -198,9 +198,9 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
MemoryContext oldCtx;
/* Calculate allocation sizes */
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->itemsize);
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->itemsize);
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, centers->itemsize);
Size centerCountsSize = sizeof(int) * numCenters;
Size closestCentersSize = sizeof(int) * numSamples;
Size lowerBoundSize = sizeof(float) * numSamples * numCenters;
@@ -245,7 +245,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
halfcdist = palloc_extended(halfcdistSize, MCXT_ALLOC_HUGE);
newcdist = palloc(newcdistSize);
newCenters = VectorArrayInit(numCenters, dimensions);
newCenters = VectorArrayInit(numCenters, dimensions, centers->itemsize);
for (int64 j = 0; j < numCenters; j++)
{
Vector *vec = VectorArrayGet(newCenters, j);
@@ -506,7 +506,7 @@ CheckCenters(Relation index, VectorArray centers)
/* Ensure no duplicate centers */
/* Fine to sort in-place */
qsort(centers->items, centers->length, VECTOR_SIZE(centers->dim), CompareVectors);
qsort(centers->items, centers->length, centers->itemsize, CompareVectors);
for (int i = 1; i < centers->length; i++)
{
if (CompareVectors(VectorArrayGet(centers, i), VectorArrayGet(centers, i - 1)) == 0)

5
src/ivfutils.c
View File

@@ -11,14 +11,15 @@
* Allocate a vector array
*/
VectorArray
VectorArrayInit(int maxlen, int dimensions)
VectorArrayInit(int maxlen, int dimensions, Size itemsize)
{
VectorArray res = palloc(sizeof(VectorArrayData));
res->length = 0;
res->maxlen = maxlen;
res->dim = dimensions;
res->items = palloc_extended(maxlen * VECTOR_SIZE(dimensions), MCXT_ALLOC_ZERO | MCXT_ALLOC_HUGE);
res->itemsize = itemsize;
res->items = palloc_extended(maxlen * itemsize, MCXT_ALLOC_ZERO | MCXT_ALLOC_HUGE);
return res;
}