Compare commits

..

20 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
Andrew Kane
8d7abb6590 Revert "Fixed locking for index scans for HNSW - #256"
This reverts commit d032726976.
2023-09-26 23:00:14 -07:00
jeff-davis
b247b688a8 No need to MarkBufferDirty(); GenericXLogFinish() does that. (#265) 2023-09-15 13:14:10 -07:00
20 changed files with 239 additions and 270 deletions

View File

@@ -1,7 +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
- Fixed locking for index scans for HNSW - Added check for MVCC-compliant snapshot for index scans
## 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.0", "version": "0.5.1",
"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.0", "version": "0.5.1",
"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.0 EXTVERSION = 0.5.1
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.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 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.0 https://github.com/pgvector/pgvector.git git clone --branch v0.5.1 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.0 https://github.com/pgvector/pgvector.git git clone --branch v0.5.1 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.0 https://github.com/pgvector/pgvector.git git clone --branch v0.5.1 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

@@ -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,7 +57,9 @@
/* 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_ELEMENT_TUPLE_SIZE(_datum) MAXALIGN(offsetof(HnswElementTupleData, value) + VARSIZE_ANY(_datum)) #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)) #define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
@@ -96,13 +98,12 @@ typedef struct HnswElementData
List *heaptids; List *heaptids;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
bool loaded;
HnswNeighborArray *neighbors; HnswNeighborArray *neighbors;
BlockNumber blkno; BlockNumber blkno;
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Datum value; Vector *vec;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
@@ -111,11 +112,13 @@ 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;
@@ -201,7 +204,7 @@ typedef struct HnswElementTupleData
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused2; uint16 unused2;
char value[FLEXIBLE_ARRAY_MEMBER]; Vector vec;
} HnswElementTupleData; } HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple; typedef HnswElementTupleData * HnswElementTuple;

View File

@@ -8,7 +8,6 @@
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "nodes/pg_list.h" #include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -82,7 +81,6 @@ 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);
@@ -107,6 +105,8 @@ CreateElementPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
int dimensions = buildstate->dimensions;
Size etupSize;
Size maxSize; Size maxSize;
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
@@ -117,11 +117,12 @@ CreateElementPages(HnswBuildState * buildstate)
ListCell *lc; ListCell *lc;
/* Calculate sizes */ /* Calculate sizes */
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)); maxSize = HNSW_MAX_SIZE;
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */ /* Allocate once */
etup = palloc0(maxSize); etup = palloc0(etupSize);
ntup = palloc0(maxSize); ntup = palloc0(BLCKSZ);
/* Prepare first page */ /* Prepare first page */
buf = HnswNewBuffer(index, forkNum); buf = HnswNewBuffer(index, forkNum);
@@ -132,14 +133,12 @@ CreateElementPages(HnswBuildState * buildstate)
foreach(lc, buildstate->elements) foreach(lc, buildstate->elements)
{ {
HnswElement element = lfirst(lc); HnswElement element = lfirst(lc);
Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
HnswSetElementTuple(etup, element); HnswSetElementTuple(etup, element);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(element->value);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -179,7 +178,6 @@ CreateElementPages(HnswBuildState * buildstate)
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
/* Commit */ /* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state); GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -227,7 +225,6 @@ 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);
} }
@@ -276,15 +273,18 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
int m = buildstate->m; int m = buildstate->m;
/* Detoast once for all calls */ /* Detoast once for all calls */
element->value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */ /* Normalize if needed */
if (buildstate->normprocinfo != NULL) if (buildstate->normprocinfo != NULL)
{ {
if (!HnswNormValue(buildstate->normprocinfo, collation, &element->value, buildstate->normvec)) if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
return false; return false;
} }
/* Copy value to element so accessible outside of memory context */
memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
@@ -360,6 +360,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Allocate necessary memory outside of memory context */ /* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
/* Use memory context since detoast can allocate */ /* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
@@ -367,8 +368,9 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup); inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Switch memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
/* Add outside memory context */ /* Add outside memory context */
if (dup != NULL) if (dup != NULL)
@@ -376,16 +378,9 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Add to buildstate or free */ /* Add to buildstate or free */
if (inserted) if (inserted)
{
element->value = datumCopy(element->value, false, -1);
element->loaded = true;
buildstate->elements = lappend(buildstate->elements, element); buildstate->elements = lappend(buildstate->elements, element);
}
else else
HnswFreeElement(element); HnswFreeElement(element);
/* Reset memory context */
MemoryContextReset(buildstate->tmpCtx);
} }
/* /*

View File

@@ -123,6 +123,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Size minCombinedSize; Size minCombinedSize;
HnswElementTuple etup; HnswElementTuple etup;
BlockNumber currentPage = insertPage; BlockNumber currentPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
@@ -131,10 +132,10 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(e->value); 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 = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)); maxSize = HNSW_MAX_SIZE;
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 */
@@ -201,8 +202,6 @@ 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 */
@@ -269,9 +268,6 @@ 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)
@@ -390,7 +386,6 @@ 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
@@ -410,7 +405,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->value); Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
HnswElementTuple etup; HnswElementTuple etup;
int i; int i;
@@ -444,7 +439,6 @@ 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);
@@ -521,7 +515,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
element->value = value; element->vec = DatumGetVector(value);
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)

View File

@@ -113,12 +113,6 @@ 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;
} }
@@ -166,11 +160,25 @@ 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;
} }
@@ -198,15 +206,6 @@ 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;
} }
@@ -223,9 +222,6 @@ 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

@@ -4,7 +4,6 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "vector.h" #include "vector.h"
/* /*
@@ -118,7 +117,6 @@ 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);
} }
@@ -141,6 +139,7 @@ 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;
} }
} }
@@ -188,8 +187,7 @@ HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element); HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
if (element->loaded) pfree(element->vec);
pfree(DatumGetPointer(element->value));
pfree(element); pfree(element);
} }
@@ -216,7 +214,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; element->neighbors = NULL;
element->loaded = false; element->vec = NULL;
return element; return element;
} }
@@ -326,7 +324,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->value, DatumGetPointer(element->value), VARSIZE_ANY(element->value)); memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
} }
/* /*
@@ -449,10 +447,8 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
if (loadVec) if (loadVec)
{ {
Datum value = PointerGetDatum(&etup->value); element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
element->value = datumCopy(value, false, -1);
element->loaded = true;
} }
} }
@@ -480,7 +476,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->value))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -491,7 +487,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
static float static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{ {
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
} }
/* /*
@@ -697,6 +693,34 @@ 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
*/ */
@@ -726,7 +750,7 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
} }
} }
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
} }
/* /*
@@ -753,33 +777,77 @@ 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, 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 *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);
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); 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);
@@ -809,7 +877,7 @@ HnswFindDuplicate(HnswElement e)
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(e->value, neighbor->element->value, false, -1)) if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
break; break;
/* Check for space */ /* Check for space */
@@ -833,28 +901,6 @@ 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
*/ */
@@ -884,13 +930,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = hc->element->value; Datum q = PointerGetDatum(hc->element->vec);
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
if (!hc3->element->loaded) if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
@@ -908,13 +954,12 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
{ {
List *c = NIL; List *c = NIL;
/* Add and sort candidates */ /* Add 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, &pruned); SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -972,7 +1017,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = element->value; Datum q = PointerGetDatum(element->vec);
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */ /* No neighbors if no entry point */
@@ -1013,7 +1058,12 @@ 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

@@ -93,7 +93,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (itemUpdated) if (itemUpdated)
{ {
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(PointerGetDatum(&etup->value)); Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
/* Mark rest as invalid */ /* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++) for (int i = idx; i < HNSW_HEAPTIDS; i++)
@@ -128,10 +128,7 @@ 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);
@@ -229,7 +226,6 @@ 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);
@@ -485,7 +481,6 @@ MarkDeleted(HnswVacuumState * vacuumstate)
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
Datum value;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
BlockNumber neighborPage; BlockNumber neighborPage;
@@ -509,11 +504,8 @@ MarkDeleted(HnswVacuumState * vacuumstate)
if (ItemPointerIsValid(&etup->heaptids[0])) if (ItemPointerIsValid(&etup->heaptids[0]))
continue; continue;
/* Get datum */
value = PointerGetDatum(&etup->value);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(value); etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
/* Get neighbor page */ /* Get neighbor page */
@@ -536,7 +528,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
MemSet(&etup->value, 0, VARSIZE_ANY(value)); MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)
@@ -551,9 +543,6 @@ 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,6 +11,7 @@
#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
@@ -65,11 +66,18 @@
static void static void
AddSample(Datum *values, IvfflatBuildState * buildstate) AddSample(Datum *values, IvfflatBuildState * buildstate)
{ {
VectorArray samples = buildstate->samples; MemoryContext oldCtx;
int targsamples = samples->maxlen; Datum value;
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 */
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 * Normalize with KMEANS_NORM_PROC since spherical distance function
@@ -81,18 +89,23 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
return; return;
} }
if (samples->length < targsamples) /* Copy datum */
{ value = datumCopy(value, false, -1);
VectorArraySet(samples, samples->length, DatumGetVector(value));
samples->length++; /* Reset memory context */
} 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, samples->length, targsamples); buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, list_length(buildstate->samples), 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
@@ -100,7 +113,8 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
#endif #endif
Assert(k >= 0 && k < targsamples); Assert(k >= 0 && k < targsamples);
VectorArraySet(samples, k, DatumGetVector(value)); lc = list_nth_cell(buildstate->samples, k);
lfirst(lc) = DatumGetVector(value);
} }
buildstate->rowstoskip -= 1; buildstate->rowstoskip -= 1;
@@ -115,21 +129,13 @@ 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, state); AddSample(values, buildstate);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
} }
/* /*
@@ -138,7 +144,7 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
static void static void
SampleRows(IvfflatBuildState * buildstate) SampleRows(IvfflatBuildState * buildstate)
{ {
int targsamples = buildstate->samples->maxlen; int targsamples = buildstate->targsamples;
BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap); BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
buildstate->rowstoskip = -1; buildstate->rowstoskip = -1;
@@ -449,12 +455,13 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */ /* Sample rows */
/* TODO Ensure within maintenance_work_mem */ /* TODO Ensure within maintenance_work_mem */
buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions); buildstate->samples = NIL;
buildstate->targsamples = numSamples;
if (buildstate->heap != NULL) if (buildstate->heap != NULL)
{ {
SampleRows(buildstate); SampleRows(buildstate);
if (buildstate->samples->length < buildstate->lists) if (list_length(buildstate->samples) < buildstate->lists)
{ {
ereport(NOTICE, ereport(NOTICE,
(errmsg("ivfflat index created with little data"), (errmsg("ivfflat index created with little data"),
@@ -467,7 +474,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 */
VectorArrayFree(buildstate->samples); list_free_deep(buildstate->samples);
} }
/* /*
@@ -506,9 +513,11 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size listSize;
IvfflatList list; IvfflatList list;
list = palloc0(BLCKSZ); listSize = MAXALIGN(IVFFLAT_LIST_SIZE(dimensions));
list = palloc(listSize);
buf = IvfflatNewBuffer(index, forkNum); buf = IvfflatNewBuffer(index, forkNum);
IvfflatInitRegisterPage(index, &buf, &page, &state); IvfflatInitRegisterPage(index, &buf, &page, &state);
@@ -516,13 +525,11 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
for (int i = 0; i < lists; i++) for (int i = 0; i < lists; i++)
{ {
OffsetNumber offno; OffsetNumber offno;
Datum center = PointerGetDatum(VectorArrayGet(centers, i));
Size listSize = MAXALIGN(IVFFLAT_LIST_SIZE(center));
/* Load list */ /* Load list */
list->startPage = InvalidBlockNumber; list->startPage = InvalidBlockNumber;
list->insertPage = InvalidBlockNumber; list->insertPage = InvalidBlockNumber;
memcpy(&list->center, DatumGetPointer(center), VARSIZE_ANY(center)); memcpy(&list->center, VectorArrayGet(centers, i), VECTOR_SIZE(dimensions));
/* Ensure free space */ /* Ensure free space */
if (PageGetFreeSpace(page) < listSize) if (PageGetFreeSpace(page) < listSize)

View File

@@ -52,7 +52,7 @@
#define PROGRESS_IVFFLAT_PHASE_ASSIGN 3 #define PROGRESS_IVFFLAT_PHASE_ASSIGN 3
#define PROGRESS_IVFFLAT_PHASE_LOAD 4 #define PROGRESS_IVFFLAT_PHASE_LOAD 4
#define IVFFLAT_LIST_SIZE(_datum) (offsetof(IvfflatListData, center) + VARSIZE_ANY(_datum)) #define IVFFLAT_LIST_SIZE(_dim) (offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
#define IvfflatPageGetOpaque(page) ((IvfflatPageOpaque) PageGetSpecialPointer(page)) #define IvfflatPageGetOpaque(page) ((IvfflatPageOpaque) PageGetSpecialPointer(page))
#define IvfflatPageGetMeta(page) ((IvfflatMetaPageData *) PageGetContents(page)) #define IvfflatPageGetMeta(page) ((IvfflatMetaPageData *) PageGetContents(page))
@@ -80,6 +80,10 @@
#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;
@@ -178,7 +182,8 @@ typedef struct IvfflatBuildState
Oid collation; Oid collation;
/* Variables */ /* Variables */
VectorArray samples; List *samples;
int targsamples;
VectorArray centers; VectorArray centers;
ListInfo *listInfo; ListInfo *listInfo;
Vector *normvec; Vector *normvec;
@@ -229,7 +234,7 @@ typedef struct IvfflatListData
{ {
BlockNumber startPage; BlockNumber startPage;
BlockNumber insertPage; BlockNumber insertPage;
char center[FLEXIBLE_ARRAY_MEMBER]; Vector center;
} IvfflatListData; } IvfflatListData;
typedef IvfflatListData * IvfflatList; typedef IvfflatListData * IvfflatList;
@@ -246,8 +251,6 @@ 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;
@@ -276,7 +279,7 @@ 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, VectorArray samples, VectorArray centers); void IvfflatKmeans(Relation index, List *samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *IvfflatOptionalProcInfo(Relation index, 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);

View File

@@ -99,7 +99,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
/* Get tuple size */ /* Get tuple size */
itemsz = MAXALIGN(IndexTupleSize(itup)); 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 */ /* Find a page to insert the item */
for (;;) for (;;)
@@ -142,8 +142,6 @@ 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 */

View File

@@ -12,20 +12,20 @@
* 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, VectorArray samples, VectorArray centers, float *lowerBound) InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBound)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
Oid collation; Oid collation;
int64 j; int64 j;
float *weight = palloc(samples->length * sizeof(float)); float *weight = palloc(list_length(samples) * sizeof(float));
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = samples->length; int numSamples = list_length(samples);
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, VectorArrayGet(samples, RandomInt() % samples->length)); VectorArraySet(centers, 0, list_nth(samples, RandomInt() % list_length(samples)));
centers->length++; centers->length++;
for (j = 0; j < numSamples; j++) 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++) for (j = 0; j < numSamples; j++)
{ {
Vector *vec = VectorArrayGet(samples, j); Vector *vec = list_nth(samples, j);
double distance; double distance;
/* Only need to compute distance for new center */ /* Only need to compute distance for new center */
@@ -74,7 +74,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
break; break;
} }
VectorArraySet(centers, i + 1, VectorArrayGet(samples, j)); VectorArraySet(centers, i + 1, list_nth(samples, j));
centers->length++; centers->length++;
} }
@@ -106,25 +106,41 @@ 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, VectorArray samples, VectorArray centers) QuickCenters(Relation index, List *samples, VectorArray centers)
{ {
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 (samples->length > 0) if (list_length(samples) > 0)
{ {
qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors); list_sort(samples, CompareListVectors);
for (int i = 0; i < samples->length; i++) 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); VectorArraySet(centers, centers->length, vec);
centers->length++; centers->length++;
@@ -160,7 +176,7 @@ QuickCenters(Relation index, VectorArray 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, VectorArray samples, VectorArray centers) ElkanKmeans(Relation index, List *samples, VectorArray centers)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -171,7 +187,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
int64 k; int64 k;
int dimensions = centers->dim; int dimensions = centers->dim;
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = samples->length; int numSamples = list_length(samples);
VectorArray newCenters; VectorArray newCenters;
int *centerCounts; int *centerCounts;
int *closestCenters; int *closestCenters;
@@ -182,7 +198,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
float *newcdist; float *newcdist;
/* Calculate allocation sizes */ /* 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 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;
@@ -326,7 +342,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k]) if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
continue; continue;
vec = VectorArrayGet(samples, j); vec = list_nth(samples, j);
/* Step 3a */ /* Step 3a */
if (rj) if (rj)
@@ -377,7 +393,7 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
{ {
int closestCenter; 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 */
@@ -514,9 +530,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, 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); QuickCenters(index, samples, centers);
else else
ElkanKmeans(index, samples, centers); ElkanKmeans(index, samples, centers);

View File

@@ -143,10 +143,6 @@ 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);
@@ -161,8 +157,6 @@ 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);
@@ -187,55 +181,6 @@ 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
*/ */
@@ -261,9 +206,7 @@ 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;
@@ -274,13 +217,12 @@ 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(3); so->tupdesc = CreateTemplateTupleDesc(2);
#else #else
so->tupdesc = CreateTemplateTupleDesc(3, false); so->tupdesc = CreateTemplateTupleDesc(2, 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);
@@ -312,7 +254,6 @@ 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)
@@ -347,6 +288,11 @@ 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
@@ -370,17 +316,10 @@ 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;
@@ -388,21 +327,6 @@ 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;
} }
@@ -418,10 +342,6 @@ 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

@@ -136,7 +136,6 @@ 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);
} }
@@ -160,8 +159,6 @@ 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,7 +107,6 @@ 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.0' default_version = '0.5.1'
module_pathname = '$libdir/vector' module_pathname = '$libdir/vector'
relocatable = true relocatable = true