Compare commits

..

1 Commits

Author SHA1 Message Date
Andrew Kane
ad71bba602 Removed index page pinning for HNSW 2023-09-07 15:48:22 -07:00
20 changed files with 295 additions and 282 deletions

View File

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

View File

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

View File

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

View File

@@ -1,5 +1,5 @@
EXTENSION = vector EXTENSION = vector
EXTVERSION = 0.5.1 EXTVERSION = 0.5.0
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 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 HEADERS = src\vector.h

View File

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

View File

@@ -1,2 +0,0 @@
-- 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,8 +57,6 @@
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */ /* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_HNSW_PHASE_LOAD 2 #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_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)) #define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
@@ -112,13 +110,11 @@ typedef struct HnswCandidate
{ {
HnswElement element; HnswElement element;
float distance; float distance;
bool closer;
} HnswCandidate; } HnswCandidate;
typedef struct HnswNeighborArray typedef struct HnswNeighborArray
{ {
int length; int length;
bool closerSet;
HnswCandidate *items; HnswCandidate *items;
} HnswNeighborArray; } HnswNeighborArray;
@@ -262,7 +258,7 @@ typedef struct HnswVacuumState
/* Methods */ /* Methods */
int HnswGetM(Relation index); int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *HnswOptionalProcInfo(Relation rel, uint16 procnum);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result); bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
void HnswCommitBuffer(Buffer buf, GenericXLogState *state); void HnswCommitBuffer(Buffer buf, GenericXLogState *state);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);

View File

@@ -81,6 +81,7 @@ HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf); HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */ /* Commit */
MarkBufferDirty(*buf);
GenericXLogFinish(*state); GenericXLogFinish(*state);
UnlockReleaseBuffer(*buf); UnlockReleaseBuffer(*buf);
@@ -117,12 +118,12 @@ CreateElementPages(HnswBuildState * buildstate)
ListCell *lc; ListCell *lc;
/* Calculate sizes */ /* Calculate sizes */
maxSize = HNSW_MAX_SIZE; maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */ /* Allocate once */
etup = palloc0(etupSize); etup = palloc0(etupSize);
ntup = palloc0(BLCKSZ); ntup = palloc0(maxSize);
/* Prepare first page */ /* Prepare first page */
buf = HnswNewBuffer(index, forkNum); buf = HnswNewBuffer(index, forkNum);
@@ -178,6 +179,7 @@ CreateElementPages(HnswBuildState * buildstate)
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -225,6 +227,7 @@ CreateNeighborPages(HnswBuildState * buildstate)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }

View File

@@ -135,7 +135,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
minCombinedSize = etupSize + HNSW_NEIGHBOR_TUPLE_SIZE(0, m) + sizeof(ItemIdData); minCombinedSize = etupSize + HNSW_NEIGHBOR_TUPLE_SIZE(0, m) + sizeof(ItemIdData);
/* Prepare element tuple */ /* Prepare element tuple */
@@ -202,6 +202,8 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
HnswInsertAppendPage(index, &newbuf, &newpage, state, page); HnswInsertAppendPage(index, &newbuf, &newpage, state, page);
/* Commit */ /* Commit */
MarkBufferDirty(newbuf);
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
/* Unlock previous buffer */ /* Unlock previous buffer */
@@ -268,6 +270,9 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
} }
/* Commit */ /* Commit */
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
if (nbuf != buf) if (nbuf != buf)
@@ -386,6 +391,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
} }
else else
@@ -439,6 +445,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);

View File

@@ -113,6 +113,12 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
scan->opaque = so; scan->opaque = so;
/*
* Get a shared lock. This allows vacuum to ensure no in-flight scans
* before marking tuples as deleted.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
return scan; return scan;
} }
@@ -160,25 +166,11 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL) if (scan->orderByData == NULL)
elog(ERROR, "cannot scan hnsw index without order"); 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 */ /* Get scan value */
value = GetScanValue(scan); value = GetScanValue(scan);
/*
* Get a shared lock. This allows vacuum to ensure no in-flight scans
* before marking tuples as deleted.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->w = GetScanItems(scan, value); so->w = GetScanItems(scan, value);
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->first = false; so->first = false;
} }
@@ -206,6 +198,15 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid; scan->xs_ctup.t_self = *heaptid;
#endif #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; scan->xs_recheckorderby = false;
return true; return true;
} }
@@ -222,6 +223,9 @@ hnswendscan(IndexScanDesc scan)
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
MemoryContextDelete(so->tmpCtx); MemoryContextDelete(so->tmpCtx);
pfree(so); pfree(so);

View File

@@ -38,12 +38,12 @@ HnswGetEfConstruction(Relation index)
* Get proc * Get proc
*/ */
FmgrInfo * FmgrInfo *
HnswOptionalProcInfo(Relation index, uint16 procnum) HnswOptionalProcInfo(Relation rel, uint16 procnum)
{ {
if (!OidIsValid(index_getprocid(index, 1, procnum))) if (!OidIsValid(index_getprocid(rel, 1, procnum)))
return NULL; return NULL;
return index_getprocinfo(index, 1, procnum); return index_getprocinfo(rel, 1, procnum);
} }
/* /*
@@ -117,6 +117,7 @@ HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState *
void void
HnswCommitBuffer(Buffer buf, GenericXLogState *state) HnswCommitBuffer(Buffer buf, GenericXLogState *state)
{ {
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -139,21 +140,9 @@ HnswInitNeighbors(HnswElement element, int m)
a = &element->neighbors[lc]; a = &element->neighbors[lc];
a->length = 0; a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm); a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false;
} }
} }
/*
* Free neighbors
*/
static void
HnswFreeNeighbors(HnswElement element)
{
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
}
/* /*
* Allocate an element * Allocate an element
*/ */
@@ -185,8 +174,10 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
void void
HnswFreeElement(HnswElement element) HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
pfree(element->vec); pfree(element->vec);
pfree(element); pfree(element);
} }
@@ -693,34 +684,6 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
return w; 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 * Calculate the distance between elements
*/ */
@@ -777,77 +740,33 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates) SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswCandidate * *pruned)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
pairingheap *wd; pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet;
List *added = NIL;
bool removedAny = false;
if (list_length(w) <= m) if (list_length(w) <= m)
return w; return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL); 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) while (list_length(w) > 0 && list_length(r) < m)
{ {
/* Assumes w is already ordered desc */ /* Assumes w is already ordered desc */
HnswCandidate *e = llast(w); HnswCandidate *e = llast(w);
bool closer;
w = list_delete_last(w); w = list_delete_last(w);
/* Use previous state of r and wd to skip work when possible */ closer = CheckElementCloser(e, r, lc, procinfo, collation);
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 (!e->closer) if (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); r = lappend(r, e);
else else
pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node)); 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 */ /* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m) while (!pairingheap_is_empty(wd) && list_length(r) < m)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner); r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
@@ -901,6 +820,28 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2)); 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 * Update connections
*/ */
@@ -954,12 +895,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
{ {
List *c = NIL; List *c = NIL;
/* Add candidates */ /* Add and sort candidates */
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]); c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2); c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true); SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1058,12 +1000,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
else else
lw = w; 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); AddConnections(element, neighbors, lm, lc);

View File

@@ -128,7 +128,10 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
blkno = HnswPageGetOpaque(page)->nextblkno; blkno = HnswPageGetOpaque(page)->nextblkno;
if (updated) if (updated)
{
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
}
else else
GenericXLogAbort(state); GenericXLogAbort(state);
@@ -226,6 +229,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -543,6 +547,9 @@ MarkDeleted(HnswVacuumState * vacuumstate)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state); GenericXLogFinish(state);
if (nbuf != buf) if (nbuf != buf)
UnlockReleaseBuffer(nbuf); UnlockReleaseBuffer(nbuf);

View File

@@ -11,7 +11,6 @@
#include "miscadmin.h" #include "miscadmin.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h" #include "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -66,18 +65,11 @@
static void static void
AddSample(Datum *values, IvfflatBuildState * buildstate) AddSample(Datum *values, IvfflatBuildState * buildstate)
{ {
MemoryContext oldCtx; VectorArray samples = buildstate->samples;
Datum value; int targsamples = samples->maxlen;
int targsamples = buildstate->targsamples;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Detoast once for all calls */ /* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Restore memory context */
MemoryContextSwitchTo(oldCtx);
/* /*
* Normalize with KMEANS_NORM_PROC since spherical distance function * Normalize with KMEANS_NORM_PROC since spherical distance function
@@ -89,23 +81,18 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
return; return;
} }
/* Copy datum */ if (samples->length < targsamples)
value = datumCopy(value, false, -1); {
VectorArraySet(samples, samples->length, DatumGetVector(value));
/* Reset memory context */ samples->length++;
MemoryContextReset(buildstate->tmpCtx); }
if (list_length(buildstate->samples) < targsamples)
buildstate->samples = lappend(buildstate->samples, DatumGetVector(value));
else else
{ {
if (buildstate->rowstoskip < 0) if (buildstate->rowstoskip < 0)
buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, list_length(buildstate->samples), targsamples); buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, samples->length, targsamples);
if (buildstate->rowstoskip <= 0) if (buildstate->rowstoskip <= 0)
{ {
ListCell *lc;
#if PG_VERSION_NUM >= 150000 #if PG_VERSION_NUM >= 150000
int k = (int) (targsamples * sampler_random_fract(&buildstate->rstate.randstate)); int k = (int) (targsamples * sampler_random_fract(&buildstate->rstate.randstate));
#else #else
@@ -113,8 +100,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
#endif #endif
Assert(k >= 0 && k < targsamples); Assert(k >= 0 && k < targsamples);
lc = list_nth_cell(buildstate->samples, k); VectorArraySet(samples, k, DatumGetVector(value));
lfirst(lc) = DatumGetVector(value);
} }
buildstate->rowstoskip -= 1; buildstate->rowstoskip -= 1;
@@ -129,13 +115,21 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state) bool *isnull, bool tupleIsAlive, void *state)
{ {
IvfflatBuildState *buildstate = (IvfflatBuildState *) state; IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
MemoryContext oldCtx;
/* Skip nulls */ /* Skip nulls */
if (isnull[0]) if (isnull[0])
return; return;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add sample */ /* Add sample */
AddSample(values, buildstate); AddSample(values, state);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
} }
/* /*
@@ -144,7 +138,7 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
static void static void
SampleRows(IvfflatBuildState * buildstate) SampleRows(IvfflatBuildState * buildstate)
{ {
int targsamples = buildstate->targsamples; int targsamples = buildstate->samples->maxlen;
BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap); BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
buildstate->rowstoskip = -1; buildstate->rowstoskip = -1;
@@ -455,13 +449,12 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */ /* Sample rows */
/* TODO Ensure within maintenance_work_mem */ /* TODO Ensure within maintenance_work_mem */
buildstate->samples = NIL; buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
buildstate->targsamples = numSamples;
if (buildstate->heap != NULL) if (buildstate->heap != NULL)
{ {
SampleRows(buildstate); SampleRows(buildstate);
if (list_length(buildstate->samples) < buildstate->lists) if (buildstate->samples->length < buildstate->lists)
{ {
ereport(NOTICE, ereport(NOTICE,
(errmsg("ivfflat index created with little data"), (errmsg("ivfflat index created with little data"),
@@ -474,7 +467,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers)); IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers));
/* Free samples before we allocate more memory */ /* Free samples before we allocate more memory */
list_free_deep(buildstate->samples); VectorArrayFree(buildstate->samples);
} }
/* /*
@@ -513,30 +506,29 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size listSize; OffsetNumber offno;
Size itemsz;
IvfflatList list; IvfflatList list;
listSize = MAXALIGN(IVFFLAT_LIST_SIZE(dimensions)); itemsz = MAXALIGN(IVFFLAT_LIST_SIZE(dimensions));
list = palloc(listSize); list = palloc(itemsz);
buf = IvfflatNewBuffer(index, forkNum); buf = IvfflatNewBuffer(index, forkNum);
IvfflatInitRegisterPage(index, &buf, &page, &state); IvfflatInitRegisterPage(index, &buf, &page, &state);
for (int i = 0; i < lists; i++) for (int i = 0; i < lists; i++)
{ {
OffsetNumber offno;
/* Load list */ /* Load list */
list->startPage = InvalidBlockNumber; list->startPage = InvalidBlockNumber;
list->insertPage = InvalidBlockNumber; list->insertPage = InvalidBlockNumber;
memcpy(&list->center, VectorArrayGet(centers, i), VECTOR_SIZE(dimensions)); memcpy(&list->center, VectorArrayGet(centers, i), VECTOR_SIZE(dimensions));
/* Ensure free space */ /* Ensure free space */
if (PageGetFreeSpace(page) < listSize) if (PageGetFreeSpace(page) < itemsz)
IvfflatAppendPage(index, &buf, &page, &state, forkNum); IvfflatAppendPage(index, &buf, &page, &state, forkNum);
/* Add the item */ /* Add the item */
offno = PageAddItem(page, (Item) list, listSize, InvalidOffsetNumber, false, false); offno = PageAddItem(page, (Item) list, itemsz, InvalidOffsetNumber, false, false);
if (offno == InvalidOffsetNumber) if (offno == InvalidOffsetNumber)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));

View File

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

View File

@@ -11,37 +11,36 @@
* Find the list that minimizes the distance function * Find the list that minimizes the distance function
*/ */
static void static void
FindInsertPage(Relation index, Datum *values, BlockNumber *insertPage, ListInfo * listInfo) FindInsertPage(Relation rel, Datum *values, BlockNumber *insertPage, ListInfo * listInfo)
{ {
Buffer cbuf;
Page cpage;
IvfflatList list;
double distance;
double minDistance = DBL_MAX; double minDistance = DBL_MAX;
BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO; BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO;
FmgrInfo *procinfo; FmgrInfo *procinfo;
Oid collation; Oid collation;
OffsetNumber offno;
OffsetNumber maxoffno;
/* Avoid compiler warning */ /* Avoid compiler warning */
listInfo->blkno = nextblkno; listInfo->blkno = nextblkno;
listInfo->offno = FirstOffsetNumber; listInfo->offno = FirstOffsetNumber;
procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC); procinfo = index_getprocinfo(rel, 1, IVFFLAT_DISTANCE_PROC);
collation = index->rd_indcollation[0]; collation = rel->rd_indcollation[0];
/* Search all list pages */ /* Search all list pages */
while (BlockNumberIsValid(nextblkno)) while (BlockNumberIsValid(nextblkno))
{ {
Buffer cbuf; cbuf = ReadBuffer(rel, nextblkno);
Page cpage;
OffsetNumber maxoffno;
cbuf = ReadBuffer(index, nextblkno);
LockBuffer(cbuf, BUFFER_LOCK_SHARE); LockBuffer(cbuf, BUFFER_LOCK_SHARE);
cpage = BufferGetPage(cbuf); cpage = BufferGetPage(cbuf);
maxoffno = PageGetMaxOffsetNumber(cpage); maxoffno = PageGetMaxOffsetNumber(cpage);
for (OffsetNumber offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
IvfflatList list;
double distance;
list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, offno)); list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, offno));
distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, values[0], PointerGetDatum(&list->center))); distance = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, values[0], PointerGetDatum(&list->center)));
@@ -64,7 +63,7 @@ FindInsertPage(Relation index, Datum *values, BlockNumber *insertPage, ListInfo
* Insert a tuple into the index * Insert a tuple into the index
*/ */
static void static void
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel) InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel)
{ {
IndexTuple itup; IndexTuple itup;
Datum value; Datum value;
@@ -81,33 +80,33 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */ /* Normalize if needed */
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC); normprocinfo = IvfflatOptionalProcInfo(rel, IVFFLAT_NORM_PROC);
if (normprocinfo != NULL) if (normprocinfo != NULL)
{ {
if (!IvfflatNormValue(normprocinfo, index->rd_indcollation[0], &value, NULL)) if (!IvfflatNormValue(normprocinfo, rel->rd_indcollation[0], &value, NULL))
return; return;
} }
/* Find the insert page - sets the page and list info */ /* Find the insert page - sets the page and list info */
FindInsertPage(index, values, &insertPage, &listInfo); FindInsertPage(rel, values, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;
/* Form tuple */ /* Form tuple */
itup = index_form_tuple(RelationGetDescr(index), &value, isnull); itup = index_form_tuple(RelationGetDescr(rel), &value, isnull);
itup->t_tid = *heap_tid; itup->t_tid = *heap_tid;
/* Get tuple size */ /* Get tuple size */
itemsz = MAXALIGN(IndexTupleSize(itup)); itemsz = MAXALIGN(IndexTupleSize(itup));
Assert(itemsz <= BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(IvfflatPageOpaqueData)) - sizeof(ItemIdData)); Assert(itemsz <= BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(IvfflatPageOpaqueData)));
/* Find a page to insert the item */ /* Find a page to insert the item */
for (;;) for (;;)
{ {
buf = ReadBuffer(index, insertPage); buf = ReadBuffer(rel, insertPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); state = GenericXLogStart(rel);
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
if (PageGetFreeSpace(page) >= itemsz) if (PageGetFreeSpace(page) >= itemsz)
@@ -127,9 +126,9 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
Page newpage; Page newpage;
/* Add a new page */ /* Add a new page */
LockRelationForExtension(index, ExclusiveLock); LockRelationForExtension(rel, ExclusiveLock);
newbuf = IvfflatNewBuffer(index, MAIN_FORKNUM); newbuf = IvfflatNewBuffer(rel, MAIN_FORKNUM);
UnlockRelationForExtension(index, ExclusiveLock); UnlockRelationForExtension(rel, ExclusiveLock);
/* Init new page */ /* Init new page */
newpage = GenericXLogRegisterBuffer(state, newbuf, GENERIC_XLOG_FULL_IMAGE); newpage = GenericXLogRegisterBuffer(state, newbuf, GENERIC_XLOG_FULL_IMAGE);
@@ -142,13 +141,15 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
IvfflatPageGetOpaque(page)->nextblkno = insertPage; IvfflatPageGetOpaque(page)->nextblkno = insertPage;
/* Commit */ /* Commit */
MarkBufferDirty(newbuf);
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
/* Unlock previous buffer */ /* Unlock previous buffer */
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Prepare new buffer */ /* Prepare new buffer */
state = GenericXLogStart(index); state = GenericXLogStart(rel);
buf = newbuf; buf = newbuf;
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
break; break;
@@ -157,13 +158,13 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
/* Add to next offset */ /* Add to next offset */
if (PageAddItem(page, (Item) itup, itemsz, InvalidOffsetNumber, false, false) == InvalidOffsetNumber) if (PageAddItem(page, (Item) itup, itemsz, InvalidOffsetNumber, false, false) == InvalidOffsetNumber)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(rel));
IvfflatCommitBuffer(buf, state); IvfflatCommitBuffer(buf, state);
/* Update the insert page */ /* Update the insert page */
if (insertPage != originalInsertPage) if (insertPage != originalInsertPage)
IvfflatUpdateList(index, listInfo, insertPage, originalInsertPage, InvalidBlockNumber, MAIN_FORKNUM); IvfflatUpdateList(rel, listInfo, insertPage, originalInsertPage, InvalidBlockNumber, MAIN_FORKNUM);
} }
/* /*

View File

@@ -12,38 +12,38 @@
* https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf * https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf
*/ */
static void static void
InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBound) InitCenters(Relation index, VectorArray samples, VectorArray centers, float *lowerBound)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
Oid collation; Oid collation;
int64 j; int64 j;
float *weight = palloc(list_length(samples) * sizeof(float)); double distance;
double sum;
double choice;
Vector *vec;
float *weight = palloc(samples->length * sizeof(float));
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = list_length(samples); int numSamples = samples->length;
procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC); procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
collation = index->rd_indcollation[0]; collation = index->rd_indcollation[0];
/* Choose an initial center uniformly at random */ /* Choose an initial center uniformly at random */
VectorArraySet(centers, 0, list_nth(samples, RandomInt() % list_length(samples))); VectorArraySet(centers, 0, VectorArrayGet(samples, RandomInt() % samples->length));
centers->length++; centers->length++;
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
weight[j] = FLT_MAX; weight[j] = DBL_MAX;
for (int i = 0; i < numCenters; i++) for (int i = 0; i < numCenters; i++)
{ {
double sum;
double choice;
CHECK_FOR_INTERRUPTS(); CHECK_FOR_INTERRUPTS();
sum = 0.0; sum = 0.0;
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
{ {
Vector *vec = list_nth(samples, j); vec = VectorArrayGet(samples, j);
double distance;
/* Only need to compute distance for new center */ /* Only need to compute distance for new center */
/* TODO Use triangle inequality to reduce distance calculations */ /* TODO Use triangle inequality to reduce distance calculations */
@@ -74,7 +74,7 @@ InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBoun
break; break;
} }
VectorArraySet(centers, i + 1, list_nth(samples, j)); VectorArraySet(centers, i + 1, VectorArrayGet(samples, j));
centers->length++; centers->length++;
} }
@@ -106,41 +106,26 @@ CompareVectors(const void *a, const void *b)
return vector_cmp_internal((Vector *) a, (Vector *) 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 * Quick approach if we have little data
*/ */
static void static void
QuickCenters(Relation index, List *samples, VectorArray centers) QuickCenters(Relation index, VectorArray samples, VectorArray centers)
{ {
Vector *vec;
int dimensions = centers->dim; int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
/* Copy existing vectors while avoiding duplicates */ /* Copy existing vectors while avoiding duplicates */
if (list_length(samples) > 0) if (samples->length > 0)
{ {
list_sort(samples, CompareListVectors); qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
for (int i = 0; i < list_length(samples); i++) for (int i = 0; i < samples->length; i++)
{ {
Vector *vec = list_nth(samples, i); vec = VectorArrayGet(samples, i);
if (i == 0 || CompareVectors(vec, list_nth(samples, i - 1)) != 0) if (i == 0 || CompareVectors(vec, VectorArrayGet(samples, i - 1)) != 0)
{ {
VectorArraySet(centers, centers->length, vec); VectorArraySet(centers, centers->length, vec);
centers->length++; centers->length++;
@@ -151,7 +136,7 @@ QuickCenters(Relation index, List *samples, VectorArray centers)
/* Fill remaining with random data */ /* Fill remaining with random data */
while (centers->length < centers->maxlen) while (centers->length < centers->maxlen)
{ {
Vector *vec = VectorArrayGet(centers, centers->length); vec = VectorArrayGet(centers, centers->length);
SET_VARSIZE(vec, VECTOR_SIZE(dimensions)); SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
vec->dim = dimensions; vec->dim = dimensions;
@@ -176,18 +161,19 @@ QuickCenters(Relation index, List *samples, VectorArray centers)
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
*/ */
static void static void
ElkanKmeans(Relation index, List *samples, VectorArray centers) ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid collation; Oid collation;
Vector *vec; Vector *vec;
Vector *newCenter; Vector *newCenter;
int iteration;
int64 j; int64 j;
int64 k; int64 k;
int dimensions = centers->dim; int dimensions = centers->dim;
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = list_length(samples); int numSamples = samples->length;
VectorArray newCenters; VectorArray newCenters;
int *centerCounts; int *centerCounts;
int *closestCenters; int *closestCenters;
@@ -196,9 +182,17 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
float *s; float *s;
float *halfcdist; float *halfcdist;
float *newcdist; float *newcdist;
int changes;
double minDistance;
int closestCenter;
double distance;
bool rj;
bool rjreset;
double dxcx;
double dxc;
/* Calculate allocation sizes */ /* Calculate allocation sizes */
Size samplesSize = 0; Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim); Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions); Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
Size centerCountsSize = sizeof(int) * numCenters; Size centerCountsSize = sizeof(int) * numCenters;
@@ -253,14 +247,14 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
/* Assign each x to its closest initial center c(x) = argmin d(x,c) */ /* Assign each x to its closest initial center c(x) = argmin d(x,c) */
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
{ {
float minDistance = FLT_MAX; minDistance = DBL_MAX;
int closestCenter = 0; closestCenter = 0;
/* Find closest center */ /* Find closest center */
for (k = 0; k < numCenters; k++) for (k = 0; k < numCenters; k++)
{ {
/* TODO Use Lemma 1 in k-means++ initialization */ /* TODO Use Lemma 1 in k-means++ initialization */
float distance = lowerBound[j * numCenters + k]; distance = lowerBound[j * numCenters + k];
if (distance < minDistance) if (distance < minDistance)
{ {
@@ -274,14 +268,13 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
} }
/* Give 500 iterations to converge */ /* Give 500 iterations to converge */
for (int iteration = 0; iteration < 500; iteration++) for (iteration = 0; iteration < 500; iteration++)
{ {
int changes = 0;
bool rjreset;
/* Can take a while, so ensure we can interrupt */ /* Can take a while, so ensure we can interrupt */
CHECK_FOR_INTERRUPTS(); CHECK_FOR_INTERRUPTS();
changes = 0;
/* Step 1: For all centers, compute distance */ /* Step 1: For all centers, compute distance */
for (j = 0; j < numCenters; j++) for (j = 0; j < numCenters; j++)
{ {
@@ -289,8 +282,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
for (k = j + 1; k < numCenters; k++) for (k = j + 1; k < numCenters; k++)
{ {
float distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k)))); distance = 0.5 * DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
halfcdist[j * numCenters + k] = distance; halfcdist[j * numCenters + k] = distance;
halfcdist[k * numCenters + j] = distance; halfcdist[k * numCenters + j] = distance;
} }
@@ -299,12 +291,10 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
/* For all centers c, compute s(c) */ /* For all centers c, compute s(c) */
for (j = 0; j < numCenters; j++) for (j = 0; j < numCenters; j++)
{ {
float minDistance = FLT_MAX; minDistance = DBL_MAX;
for (k = 0; k < numCenters; k++) for (k = 0; k < numCenters; k++)
{ {
float distance;
if (j == k) if (j == k)
continue; continue;
@@ -320,8 +310,6 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
{ {
bool rj;
/* Step 2: Identify all points x such that u(x) <= s(c(x)) */ /* Step 2: Identify all points x such that u(x) <= s(c(x)) */
if (upperBound[j] <= s[closestCenters[j]]) if (upperBound[j] <= s[closestCenters[j]])
continue; continue;
@@ -330,8 +318,6 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
for (k = 0; k < numCenters; k++) for (k = 0; k < numCenters; k++)
{ {
float dxcx;
/* Step 3: For all remaining points x and centers c */ /* Step 3: For all remaining points x and centers c */
if (k == closestCenters[j]) if (k == closestCenters[j])
continue; continue;
@@ -342,7 +328,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k]) if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
continue; continue;
vec = list_nth(samples, j); vec = VectorArrayGet(samples, j);
/* Step 3a */ /* Step 3a */
if (rj) if (rj)
@@ -361,7 +347,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
/* Step 3b */ /* Step 3b */
if (dxcx > lowerBound[j * numCenters + k] || dxcx > halfcdist[closestCenters[j] * numCenters + k]) if (dxcx > lowerBound[j * numCenters + k] || dxcx > halfcdist[closestCenters[j] * numCenters + k])
{ {
float dxc = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k)))); dxc = DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(vec), PointerGetDatum(VectorArrayGet(centers, k))));
/* d(x,c) calculated */ /* d(x,c) calculated */
lowerBound[j * numCenters + k] = dxc; lowerBound[j * numCenters + k] = dxc;
@@ -375,6 +361,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
changes++; changes++;
} }
} }
} }
} }
@@ -391,9 +378,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
{ {
int closestCenter; vec = VectorArrayGet(samples, j);
vec = list_nth(samples, j);
closestCenter = closestCenters[j]; closestCenter = closestCenters[j];
/* Increment sum and count of closest center */ /* Increment sum and count of closest center */
@@ -441,7 +426,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
{ {
for (k = 0; k < numCenters; k++) for (k = 0; k < numCenters; k++)
{ {
float distance = lowerBound[j * numCenters + k] - newcdist[k]; distance = lowerBound[j * numCenters + k] - newcdist[k];
if (distance < 0) if (distance < 0)
distance = 0; distance = 0;
@@ -457,7 +442,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
/* Step 7 */ /* Step 7 */
for (j = 0; j < numCenters; j++) for (j = 0; j < numCenters; j++)
VectorArraySet(centers, j, VectorArrayGet(newCenters, j)); memcpy(VectorArrayGet(centers, j), VectorArrayGet(newCenters, j), VECTOR_SIZE(dimensions));
if (changes == 0 && iteration != 0) if (changes == 0 && iteration != 0)
break; break;
@@ -480,6 +465,9 @@ static void
CheckCenters(Relation index, VectorArray centers) CheckCenters(Relation index, VectorArray centers)
{ {
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid collation;
Vector *vec;
double norm;
if (centers->length != centers->maxlen) if (centers->length != centers->maxlen)
elog(ERROR, "Not enough centers. Please report a bug."); elog(ERROR, "Not enough centers. Please report a bug.");
@@ -487,7 +475,7 @@ CheckCenters(Relation index, VectorArray centers)
/* Ensure no NaN or infinite values */ /* Ensure no NaN or infinite values */
for (int i = 0; i < centers->length; i++) for (int i = 0; i < centers->length; i++)
{ {
Vector *vec = VectorArrayGet(centers, i); vec = VectorArrayGet(centers, i);
for (int j = 0; j < vec->dim; j++) for (int j = 0; j < vec->dim; j++)
{ {
@@ -513,12 +501,11 @@ CheckCenters(Relation index, VectorArray centers)
normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC); normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
if (normprocinfo != NULL) if (normprocinfo != NULL)
{ {
Oid collation = index->rd_indcollation[0]; collation = index->rd_indcollation[0];
for (int i = 0; i < centers->length; i++) for (int i = 0; i < centers->length; i++)
{ {
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, PointerGetDatum(VectorArrayGet(centers, i)))); norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, PointerGetDatum(VectorArrayGet(centers, i))));
if (norm == 0) if (norm == 0)
elog(ERROR, "Zero norm detected. Please report a bug."); elog(ERROR, "Zero norm detected. Please report a bug.");
} }
@@ -530,9 +517,9 @@ CheckCenters(Relation index, VectorArray centers)
* We use spherical k-means for inner product and cosine * We use spherical k-means for inner product and cosine
*/ */
void void
IvfflatKmeans(Relation index, List *samples, VectorArray centers) IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
{ {
if (list_length(samples) <= centers->maxlen) if (samples->length <= centers->maxlen)
QuickCenters(index, samples, centers); QuickCenters(index, samples, centers);
else else
ElkanKmeans(index, samples, centers); ElkanKmeans(index, samples, centers);

View File

@@ -143,6 +143,10 @@ GetScanItems(IndexScanDesc scan, Datum value)
bool isnull; bool isnull;
ItemId itemid = PageGetItemId(page, offno); ItemId itemid = PageGetItemId(page, offno);
/* Skip dead tuples */
if (scan->ignore_killed_tuples && ItemIdIsDead(itemid))
continue;
itup = (IndexTuple) PageGetItem(page, itemid); itup = (IndexTuple) PageGetItem(page, itemid);
datum = index_getattr(itup, 1, tupdesc, &isnull); datum = index_getattr(itup, 1, tupdesc, &isnull);
@@ -157,6 +161,8 @@ GetScanItems(IndexScanDesc scan, Datum value)
slot->tts_isnull[0] = false; slot->tts_isnull[0] = false;
slot->tts_values[1] = PointerGetDatum(&itup->t_tid); slot->tts_values[1] = PointerGetDatum(&itup->t_tid);
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = Int32GetDatum((int) searchPage);
slot->tts_isnull[2] = false;
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
tuplesort_puttupleslot(so->sortstate, slot); tuplesort_puttupleslot(so->sortstate, slot);
@@ -181,6 +187,55 @@ GetScanItems(IndexScanDesc scan, Datum value)
tuplesort_performsort(so->sortstate); 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 * Prepare for an index scan
*/ */
@@ -206,7 +261,9 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
probes = lists; probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList)); so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
so->buf = InvalidBuffer;
so->first = true; so->first = true;
ItemPointerSetInvalid(&so->heaptid);
so->probes = probes; so->probes = probes;
so->dimensions = dimensions; so->dimensions = dimensions;
@@ -217,12 +274,13 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Create tuple description for sorting */ /* Create tuple description for sorting */
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
so->tupdesc = CreateTemplateTupleDesc(2); so->tupdesc = CreateTemplateTupleDesc(3);
#else #else
so->tupdesc = CreateTemplateTupleDesc(2, false); so->tupdesc = CreateTemplateTupleDesc(3, false);
#endif #endif
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 3, "indexblkno", INT4OID, -1, 0);
/* Prep sort */ /* Prep sort */
so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false); so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
@@ -254,6 +312,7 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
#endif #endif
so->first = true; so->first = true;
ItemPointerSetInvalid(&so->heaptid);
pairingheap_reset(so->listQueue); pairingheap_reset(so->listQueue);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -288,11 +347,6 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL) if (scan->orderByData == NULL)
elog(ERROR, "cannot scan ivfflat index without order"); 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) if (scan->orderByData->sk_flags & SK_ISNULL)
value = PointerGetDatum(InitVector(so->dimensions)); value = PointerGetDatum(InitVector(so->dimensions));
else else
@@ -316,10 +370,17 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (value != scan->orderByData->sk_argument) if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value)); 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)) if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
{ {
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull)); 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 #if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
@@ -327,6 +388,21 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
scan->xs_ctup.t_self = *heaptid; scan->xs_ctup.t_self = *heaptid;
#endif #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; scan->xs_recheckorderby = false;
return true; return true;
} }
@@ -342,6 +418,10 @@ ivfflatendscan(IndexScanDesc scan)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
/* Release pin */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);

View File

@@ -57,12 +57,12 @@ IvfflatGetLists(Relation index)
* Get proc * Get proc
*/ */
FmgrInfo * FmgrInfo *
IvfflatOptionalProcInfo(Relation index, uint16 procnum) IvfflatOptionalProcInfo(Relation rel, uint16 procnum)
{ {
if (!OidIsValid(index_getprocid(index, 1, procnum))) if (!OidIsValid(index_getprocid(rel, 1, procnum)))
return NULL; return NULL;
return index_getprocinfo(index, 1, procnum); return index_getprocinfo(rel, 1, procnum);
} }
/* /*
@@ -136,6 +136,7 @@ IvfflatInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogStat
void void
IvfflatCommitBuffer(Buffer buf, GenericXLogState *state) IvfflatCommitBuffer(Buffer buf, GenericXLogState *state)
{ {
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -159,6 +160,8 @@ IvfflatAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **st
IvfflatInitPage(newbuf, newpage); IvfflatInitPage(newbuf, newpage);
/* Commit */ /* Commit */
MarkBufferDirty(*buf);
MarkBufferDirty(newbuf);
GenericXLogFinish(*state); GenericXLogFinish(*state);
/* Unlock */ /* Unlock */

View File

@@ -107,6 +107,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
{ {
/* Delete tuples */ /* Delete tuples */
PageIndexMultiDelete(page, deletable, ndeletable); PageIndexMultiDelete(page, deletable, ndeletable);
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
} }
else else

View File

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