Compare commits

..

18 Commits

Author SHA1 Message Date
Andrew Kane
64223989cd Use List for samples 2023-10-16 15:32:51 -07:00
Andrew Kane
e630efd195 Version bump to 0.5.1 [skip ci] 2023-10-10 17:40:57 -07:00
Andrew Kane
b5b912906b Added check for MVCC-compliant snapshot and removed marking tuples as dead for IVFFlat index scans - closes #260 2023-10-10 17:28:48 -07:00
Andrew Kane
4b5db94307 Disable closer caching for new elements for now 2023-10-06 14:27:09 -07:00
Andrew Kane
65e70326b8 Updated comment [skip ci] 2023-10-06 14:07:35 -07:00
Andrew Kane
71641ed84e Updated comment [skip ci] 2023-10-06 13:58:07 -07:00
Andrew Kane
f3dba25036 Added comment [skip ci] 2023-10-06 13:56:25 -07:00
Andrew Kane
5588ba6410 Improved variable name [skip ci] 2023-10-06 13:46:19 -07:00
Andrew Kane
ec9fac5456 Improved closerSet logic 2023-10-06 13:39:55 -07:00
Andrew Kane
8085d3e538 Moved sorting logic into SelectNeighbors 2023-10-06 12:56:15 -07:00
Andrew Kane
cae162ffc6 Ensure order is deterministic for SelectNeighbors closer caching 2023-10-06 12:26:53 -07:00
Andrew Kane
62482e3760 Use e for consistency 2023-10-05 16:15:13 -07:00
Heikki Linnakangas
c81302b835 Improve HNSW index build performance more (#295)
This takes the approach from commit a713e2acaa further. Once we have
remove a candidate from the "closer" set, we still don't need to
recalculate everything that follows. Any candidates that were in the
closer set before still only need to be compared with any new
candidates that we have added.
2023-10-05 16:04:50 -07:00
Andrew Kane
a713e2acaa Improved performance of HNSW index builds - closes #292
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-10-05 13:21:26 -07:00
Andrew Kane
6e1312ddbe DRY max size 2023-10-04 21:43:34 -07:00
Andrew Kane
4ef5bca275 Use BLCKSZ for consistency 2023-10-04 21:37:55 -07:00
Xiaoran Wang
1ecf6ada76 Include ItemIdData when computing the maxSize for the data in a page (#274)
As the data is aligned, for hnsw, the combined size won't be in the range
(8156 (maxSize exlucding `ItemIdData`), 8160]. So even if the
ItemIdData is not included in the maxSize, it works well now, but I
think it's better to make it correct.
2023-10-04 21:26:46 -07:00
Andrew Kane
564a3d45fc Added check for MVCC-compliant snapshot for HNSW index scans - closes #281 2023-10-04 20:14:50 -07:00
17 changed files with 190 additions and 187 deletions

View File

@@ -1,6 +1,7 @@
## 0.5.1 (unreleased)
## 0.5.1 (2023-10-10)
- Improved performance of index scans for IVFFlat after updates and deletes
- Improved performance of HNSW index builds
- Added check for MVCC-compliant snapshot for index scans
## 0.5.0 (2023-08-28)

View File

@@ -2,7 +2,7 @@
"name": "vector",
"abstract": "Open-source vector similarity search for Postgres",
"description": "Supports L2 distance, inner product, and cosine distance",
"version": "0.5.0",
"version": "0.5.1",
"maintainer": [
"Andrew Kane <andrew@ankane.org>"
],
@@ -20,7 +20,7 @@
"vector": {
"file": "sql/vector.sql",
"docfile": "README.md",
"version": "0.5.0",
"version": "0.5.1",
"abstract": "Open-source vector similarity search for Postgres"
}
},

View File

@@ -1,5 +1,5 @@
EXTENSION = vector
EXTVERSION = 0.5.0
EXTVERSION = 0.5.1
MODULE_big = vector
DATA = $(wildcard sql/*--*.sql)

View File

@@ -1,5 +1,5 @@
EXTENSION = vector
EXTVERSION = 0.5.0
EXTVERSION = 0.5.1
OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
HEADERS = src\vector.h

View File

@@ -18,7 +18,7 @@ Compile and install the extension (supports Postgres 11+)
```sh
cd /tmp
git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector
make
make install # may need sudo
@@ -509,7 +509,7 @@ Then use `nmake` to build:
```cmd
set "PGROOT=C:\Program Files\PostgreSQL\15"
git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector
nmake /F Makefile.win
nmake /F Makefile.win install
@@ -530,7 +530,7 @@ This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (r
You can also build the image manually:
```sh
git clone --branch v0.5.0 https://github.com/pgvector/pgvector.git
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector
docker build --build-arg PG_MAJOR=15 -t myuser/pgvector .
```

View File

@@ -0,0 +1,2 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.5.1'" to load this file. \quit

View File

@@ -57,6 +57,8 @@
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_HNSW_PHASE_LOAD 2
#define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData))
#define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
@@ -110,11 +112,13 @@ typedef struct HnswCandidate
{
HnswElement element;
float distance;
bool closer;
} HnswCandidate;
typedef struct HnswNeighborArray
{
int length;
bool closerSet;
HnswCandidate *items;
} HnswNeighborArray;

View File

@@ -117,12 +117,12 @@ CreateElementPages(HnswBuildState * buildstate)
ListCell *lc;
/* Calculate sizes */
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
maxSize = HNSW_MAX_SIZE;
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */
etup = palloc0(etupSize);
ntup = palloc0(maxSize);
ntup = palloc0(BLCKSZ);
/* Prepare first page */
buf = HnswNewBuffer(index, forkNum);

View File

@@ -135,7 +135,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
maxSize = HNSW_MAX_SIZE;
minCombinedSize = etupSize + HNSW_NEIGHBOR_TUPLE_SIZE(0, m) + sizeof(ItemIdData);
/* Prepare element tuple */

View File

@@ -160,6 +160,11 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL)
elog(ERROR, "cannot scan hnsw index without order");
/* Requires MVCC-compliant snapshot as not able to maintain a pin */
/* https://www.postgresql.org/docs/current/index-locking.html */
if (!IsMVCCSnapshot(scan->xs_snapshot))
elog(ERROR, "non-MVCC snapshots are not supported with hnsw");
/* Get scan value */
value = GetScanValue(scan);
@@ -201,15 +206,6 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid;
#endif
/*
* Typically, an index scan must maintain a pin on the index page
* holding the item last returned by amgettuple. However, this is not
* needed with the current vacuum strategy, which ensures scans do not
* visit tuples in danger of being marked as deleted.
*
* https://www.postgresql.org/docs/current/index-locking.html
*/
scan->xs_recheckorderby = false;
return true;
}

View File

@@ -139,6 +139,7 @@ HnswInitNeighbors(HnswElement element, int m)
a = &element->neighbors[lc];
a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false;
}
}
@@ -692,6 +693,34 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
return w;
}
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
if (hca->element < hcb->element)
return 1;
if (hca->element > hcb->element)
return -1;
return 0;
}
/*
* Calculate the distance between elements
*/
@@ -748,33 +777,77 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper
*/
static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswCandidate * *pruned)
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{
List *r = NIL;
List *w = list_copy(c);
pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet;
List *added = NIL;
bool removedAny = false;
if (list_length(w) <= m)
return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL);
/* Ensure order of candidates is deterministic for closer caching */
if (sortCandidates)
list_sort(w, CompareCandidateDistances);
while (list_length(w) > 0 && list_length(r) < m)
{
/* Assumes w is already ordered desc */
HnswCandidate *e = llast(w);
bool closer;
w = list_delete_last(w);
closer = CheckElementCloser(e, r, lc, procinfo, collation);
/* Use previous state of r and wd to skip work when possible */
if (mustCalculate)
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
else if (list_length(added) > 0)
{
/*
* If the current candidate was closer, we only need to compare it
* with the other candidates that we have added.
*/
if (e->closer)
{
e->closer = CheckElementCloser(e, added, lc, procinfo, collation);
if (closer)
if (!e->closer)
removedAny = true;
}
else
{
/*
* If we have removed any candidates from closer, a candidate
* that was not closer earlier might now be.
*/
if (removedAny)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
if (e->closer)
added = lappend(added, e);
}
}
}
else if (e == newCandidate)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
if (e->closer)
added = lappend(added, e);
}
if (e->closer)
r = lappend(r, e);
else
pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node));
}
/* Cached value can only be used in future if sorted deterministically */
e2->neighbors[lc].closerSet = sortCandidates;
/* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
@@ -828,28 +901,6 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
}
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
return 0;
}
/*
* Update connections
*/
@@ -903,13 +954,12 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
{
List *c = NIL;
/* Add and sort candidates */
/* Add candidates */
for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true);
/* Should not happen */
if (pruned == NULL)
@@ -1008,7 +1058,12 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
else
lw = w;
neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, NULL);
/*
* Candidates are sorted, but not deterministically. Could set
* sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance.
*/
neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
AddConnections(element, neighbors, lm, lc);

View File

@@ -11,6 +11,7 @@
#include "miscadmin.h"
#include "storage/bufmgr.h"
#include "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000
@@ -65,11 +66,18 @@
static void
AddSample(Datum *values, IvfflatBuildState * buildstate)
{
VectorArray samples = buildstate->samples;
int targsamples = samples->maxlen;
MemoryContext oldCtx;
Datum value;
int targsamples = buildstate->targsamples;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Restore memory context */
MemoryContextSwitchTo(oldCtx);
/*
* Normalize with KMEANS_NORM_PROC since spherical distance function
@@ -81,18 +89,23 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
return;
}
if (samples->length < targsamples)
{
VectorArraySet(samples, samples->length, DatumGetVector(value));
samples->length++;
}
/* Copy datum */
value = datumCopy(value, false, -1);
/* Reset memory context */
MemoryContextReset(buildstate->tmpCtx);
if (list_length(buildstate->samples) < targsamples)
buildstate->samples = lappend(buildstate->samples, DatumGetVector(value));
else
{
if (buildstate->rowstoskip < 0)
buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, samples->length, targsamples);
buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, list_length(buildstate->samples), targsamples);
if (buildstate->rowstoskip <= 0)
{
ListCell *lc;
#if PG_VERSION_NUM >= 150000
int k = (int) (targsamples * sampler_random_fract(&buildstate->rstate.randstate));
#else
@@ -100,7 +113,8 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
#endif
Assert(k >= 0 && k < targsamples);
VectorArraySet(samples, k, DatumGetVector(value));
lc = list_nth_cell(buildstate->samples, k);
lfirst(lc) = DatumGetVector(value);
}
buildstate->rowstoskip -= 1;
@@ -115,21 +129,13 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state)
{
IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
MemoryContext oldCtx;
/* Skip nulls */
if (isnull[0])
return;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add sample */
AddSample(values, state);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
AddSample(values, buildstate);
}
/*
@@ -138,7 +144,7 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
static void
SampleRows(IvfflatBuildState * buildstate)
{
int targsamples = buildstate->samples->maxlen;
int targsamples = buildstate->targsamples;
BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
buildstate->rowstoskip = -1;
@@ -449,12 +455,13 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */
/* TODO Ensure within maintenance_work_mem */
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
buildstate->samples = NIL;
buildstate->targsamples = numSamples;
if (buildstate->heap != NULL)
{
SampleRows(buildstate);
if (buildstate->samples->length < buildstate->lists)
if (list_length(buildstate->samples) < buildstate->lists)
{
ereport(NOTICE,
(errmsg("ivfflat index created with little data"),
@@ -467,7 +474,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers));
/* Free samples before we allocate more memory */
VectorArrayFree(buildstate->samples);
list_free_deep(buildstate->samples);
}
/*

View File

@@ -80,6 +80,10 @@
#define RandomInt() random()
#endif
#if PG_VERSION_NUM < 130000
#define list_sort(list, cmp) list_qsort(list, cmp)
#endif
/* Variables */
extern int ivfflat_probes;
@@ -178,7 +182,8 @@ typedef struct IvfflatBuildState
Oid collation;
/* Variables */
VectorArray samples;
List *samples;
int targsamples;
VectorArray centers;
ListInfo *listInfo;
Vector *normvec;
@@ -246,8 +251,6 @@ typedef struct IvfflatScanOpaqueData
int probes;
int dimensions;
bool first;
Buffer buf;
ItemPointerData heaptid;
/* Sorting */
Tuplesortstate *sortstate;
@@ -276,7 +279,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
VectorArray VectorArrayInit(int maxlen, int dimensions);
void VectorArrayFree(VectorArray arr);
void PrintVectorArray(char *msg, VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
void IvfflatKmeans(Relation index, List *samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
int IvfflatGetLists(Relation index);

View File

@@ -99,7 +99,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
/* Get tuple size */
itemsz = MAXALIGN(IndexTupleSize(itup));
Assert(itemsz <= BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(IvfflatPageOpaqueData)));
Assert(itemsz <= BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(IvfflatPageOpaqueData)) - sizeof(ItemIdData));
/* Find a page to insert the item */
for (;;)

View File

@@ -12,20 +12,20 @@
* https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf
*/
static void
InitCenters(Relation index, VectorArray samples, VectorArray centers, float *lowerBound)
InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBound)
{
FmgrInfo *procinfo;
Oid collation;
int64 j;
float *weight = palloc(samples->length * sizeof(float));
float *weight = palloc(list_length(samples) * sizeof(float));
int numCenters = centers->maxlen;
int numSamples = samples->length;
int numSamples = list_length(samples);
procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
collation = index->rd_indcollation[0];
/* Choose an initial center uniformly at random */
VectorArraySet(centers, 0, VectorArrayGet(samples, RandomInt() % samples->length));
VectorArraySet(centers, 0, list_nth(samples, RandomInt() % list_length(samples)));
centers->length++;
for (j = 0; j < numSamples; j++)
@@ -42,7 +42,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
for (j = 0; j < numSamples; j++)
{
Vector *vec = VectorArrayGet(samples, j);
Vector *vec = list_nth(samples, j);
double distance;
/* Only need to compute distance for new center */
@@ -74,7 +74,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
break;
}
VectorArraySet(centers, i + 1, VectorArrayGet(samples, j));
VectorArraySet(centers, i + 1, list_nth(samples, j));
centers->length++;
}
@@ -106,25 +106,41 @@ CompareVectors(const void *a, const void *b)
return vector_cmp_internal((Vector *) a, (Vector *) b);
}
/*
* Compare list vectors
*/
static int
#if PG_VERSION_NUM >= 130000
CompareListVectors(const ListCell *a, const ListCell *b)
#else
CompareListVectors(const void *a, const void *b)
#endif
{
Vector *va = lfirst((ListCell *) a);
Vector *vb = lfirst((ListCell *) b);
return CompareVectors(va, vb);
}
/*
* Quick approach if we have little data
*/
static void
QuickCenters(Relation index, VectorArray samples, VectorArray centers)
QuickCenters(Relation index, List *samples, VectorArray centers)
{
int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0];
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
/* Copy existing vectors while avoiding duplicates */
if (samples->length > 0)
if (list_length(samples) > 0)
{
qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
for (int i = 0; i < samples->length; i++)
list_sort(samples, CompareListVectors);
for (int i = 0; i < list_length(samples); i++)
{
Vector *vec = VectorArrayGet(samples, i);
Vector *vec = list_nth(samples, i);
if (i == 0 || CompareVectors(vec, VectorArrayGet(samples, i - 1)) != 0)
if (i == 0 || CompareVectors(vec, list_nth(samples, i - 1)) != 0)
{
VectorArraySet(centers, centers->length, vec);
centers->length++;
@@ -160,7 +176,7 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
*/
static void
ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
ElkanKmeans(Relation index, List *samples, VectorArray centers)
{
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
@@ -171,7 +187,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
int64 k;
int dimensions = centers->dim;
int numCenters = centers->maxlen;
int numSamples = samples->length;
int numSamples = list_length(samples);
VectorArray newCenters;
int *centerCounts;
int *closestCenters;
@@ -182,7 +198,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
float *newcdist;
/* Calculate allocation sizes */
Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
Size samplesSize = 0;
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
Size centerCountsSize = sizeof(int) * numCenters;
@@ -326,7 +342,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
continue;
vec = VectorArrayGet(samples, j);
vec = list_nth(samples, j);
/* Step 3a */
if (rj)
@@ -377,7 +393,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
{
int closestCenter;
vec = VectorArrayGet(samples, j);
vec = list_nth(samples, j);
closestCenter = closestCenters[j];
/* Increment sum and count of closest center */
@@ -514,9 +530,9 @@ CheckCenters(Relation index, VectorArray centers)
* We use spherical k-means for inner product and cosine
*/
void
IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
IvfflatKmeans(Relation index, List *samples, VectorArray centers)
{
if (samples->length <= centers->maxlen)
if (list_length(samples) <= centers->maxlen)
QuickCenters(index, samples, centers);
else
ElkanKmeans(index, samples, centers);

View File

@@ -143,10 +143,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
bool isnull;
ItemId itemid = PageGetItemId(page, offno);
/* Skip dead tuples */
if (scan->ignore_killed_tuples && ItemIdIsDead(itemid))
continue;
itup = (IndexTuple) PageGetItem(page, itemid);
datum = index_getattr(itup, 1, tupdesc, &isnull);
@@ -161,8 +157,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
slot->tts_isnull[0] = false;
slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
slot->tts_isnull[1] = false;
slot->tts_values[2] = Int32GetDatum((int) searchPage);
slot->tts_isnull[2] = false;
ExecStoreVirtualTuple(slot);
tuplesort_puttupleslot(so->sortstate, slot);
@@ -187,55 +181,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
tuplesort_performsort(so->sortstate);
}
/*
* Mark prior tuple as dead
*/
static void
MarkPriorTupleDead(IndexScanDesc scan)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
Buffer buf = so->buf;
Page page;
OffsetNumber maxoffno;
/* Safety check */
if (!BufferIsValid(so->buf) || !ItemPointerIsValid(&so->heaptid))
return;
/* Only a shared locked is needed for ItemIdMarkDead */
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
ItemId itemid = PageGetItemId(page, offno);
IndexTuple itup = (IndexTuple) PageGetItem(page, itemid);
/*
* Find tuple. Since buffer has been pinned, tuple cannot have been
* vacuumed (and heap TID reused).
*/
if (ItemPointerEquals(&itup->t_tid, &so->heaptid))
{
/*
* Make sure tuple has not already been marked dead to avoid extra
* WAL if wal_log_hints or data checksums enabled
*/
if (!ItemIdIsDead(itemid))
{
ItemIdMarkDead(itemid);
MarkBufferDirtyHint(buf, true);
}
break;
}
}
/* Unlock buffer */
LockBuffer(buf, BUFFER_LOCK_UNLOCK);
}
/*
* Prepare for an index scan
*/
@@ -261,9 +206,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
so->buf = InvalidBuffer;
so->first = true;
ItemPointerSetInvalid(&so->heaptid);
so->probes = probes;
so->dimensions = dimensions;
@@ -274,13 +217,12 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Create tuple description for sorting */
#if PG_VERSION_NUM >= 120000
so->tupdesc = CreateTemplateTupleDesc(3);
so->tupdesc = CreateTemplateTupleDesc(2);
#else
so->tupdesc = CreateTemplateTupleDesc(3, false);
so->tupdesc = CreateTemplateTupleDesc(2, false);
#endif
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
/* Prep sort */
so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -312,7 +254,6 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
#endif
so->first = true;
ItemPointerSetInvalid(&so->heaptid);
pairingheap_reset(so->listQueue);
if (keys && scan->numberOfKeys > 0)
@@ -347,6 +288,11 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL)
elog(ERROR, "cannot scan ivfflat index without order");
/* Requires MVCC-compliant snapshot as not able to pin during sorting */
/* https://www.postgresql.org/docs/current/index-locking.html */
if (!IsMVCCSnapshot(scan->xs_snapshot))
elog(ERROR, "non-MVCC snapshots are not supported with ivfflat");
if (scan->orderByData->sk_flags & SK_ISNULL)
value = PointerGetDatum(InitVector(so->dimensions));
else
@@ -370,17 +316,10 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value));
}
else
{
/* Mark prior tuple as dead */
if (scan->kill_prior_tuple)
MarkPriorTupleDead(scan);
}
if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
{
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
BlockNumber indexblkno = DatumGetInt32(slot_getattr(so->slot, 3, &so->isnull));
#if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid;
@@ -388,21 +327,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid;
#endif
/* Keep track of info needed to mark tuple as dead */
so->heaptid = *heaptid;
/* Unpin buffer */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
/*
* An index scan must maintain a pin on the index page holding the
* item last returned by amgettuple
*
* https://www.postgresql.org/docs/current/index-locking.html
*/
so->buf = ReadBuffer(scan->indexRelation, indexblkno);
scan->xs_recheckorderby = false;
return true;
}
@@ -418,10 +342,6 @@ ivfflatendscan(IndexScanDesc scan)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
/* Release pin */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
pairingheap_free(so->listQueue);
tuplesort_end(so->sortstate);

View File

@@ -1,5 +1,4 @@
comment = 'vector data type and ivfflat and hnsw access methods'
default_version = '0.5.0'
default_version = '0.5.1'
module_pathname = '$libdir/vector'
relocatable = true
trusted = true