From 4ff36af67eae561b495b97df07d2e2c4713b12bb Mon Sep 17 00:00:00 2001 From: Andrew Kane Date: Thu, 11 Apr 2024 13:47:36 -0700 Subject: [PATCH] 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 40661b3..34af580 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 3ca6fe9..ce364b4 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 fde328b..97acfa0 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 2a40ed6..402d46b 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; }