pgvector.git / summary / log / commit / refs
commit 4ff36af67eae561b495b97df07d2e2c4713b12bb
Author: Andrew Kane <andrew@ankane.org>
Date: Thu Apr 11 20:47:36 2024 +0000
Added itemsize to VectorArray [skip ci]
src/ivfbuild.c | 10 +++++-----
src/ivfflat.h | 9 +++++----
src/ivfkmeans.c | 12 ++++++------
src/ivfutils.c | 5 +++--
4 files changed, 19 insertions(+), 17 deletions(-)
diff --git a/src/ivfbuild.c b/src/ivfbuild.c
index 40661b39..34af5803 100644
--- a/src/ivfbuild.c
+++ b/src/ivfbuild.c
@@ -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);
diff --git a/src/ivfflat.h b/src/ivfflat.h
index 3ca6fe92..ce364b43 100644
--- a/src/ivfflat.h
+++ b/src/ivfflat.h
@@ -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);
diff --git a/src/ivfkmeans.c b/src/ivfkmeans.c
index fde328b9..97acfa0d 100644
--- a/src/ivfkmeans.c
+++ b/src/ivfkmeans.c
@@ -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)
diff --git a/src/ivfutils.c b/src/ivfutils.c
index 2a40ed6e..402d46b0 100644
--- a/src/ivfutils.c
+++ b/src/ivfutils.c
@@ -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;
}
[parent: 5834b58c5ae0]