Compare commits

..

7 Commits

Author SHA1 Message Date
Andrew Kane
a5bb59d9f6 Use normalize_l2 for ivfflat 2023-10-16 17:56:50 -07:00
Andrew Kane
dd609f200b Merge branch 'normalize_l2' into hnsw-normalize 2023-10-16 17:09:20 -07:00
Andrew Kane
9514c152bc Fixed precision on Windows 2023-10-16 17:08:19 -07:00
Andrew Kane
b391e40765 Always normalize 2023-10-16 16:53:40 -07:00
Andrew Kane
0054a9c40a Use normalize_l2 for normalization 2023-10-16 16:42:40 -07:00
Andrew Kane
9ed7e63fb7 Prevent overflow at cost to speed 2023-08-09 12:36:29 -07:00
Andrew Kane
47e361a93d Added normalize_l2 function 2023-08-09 11:29:14 -07:00
51 changed files with 697 additions and 1108 deletions

View File

@@ -78,7 +78,6 @@ jobs:
nmake /NOLOGO /F Makefile.win uninstall
shell: cmd
i386:
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
runs-on: ubuntu-latest
container:
image: debian:11

View File

@@ -3,6 +3,7 @@
- Added HNSW index type
- Added support for parallel index builds
- Added `l1_distance` function
- Added `normalize_l2` function
- Added element-wise multiplication for vectors
- Added `sum` aggregate
- Improved performance of distance functions

View File

@@ -4,7 +4,6 @@ EXTVERSION = 0.4.4
MODULE_big = vector
DATA = $(wildcard sql/*--*.sql)
OBJS = src/hnsw.o src/hnswbuild.o src/hnswinsert.o src/hnswscan.o src/hnswutils.o src/hnswvacuum.o src/ivfbuild.o src/ivfflat.o src/ivfinsert.o src/ivfkmeans.o src/ivfscan.o src/ivfutils.o src/ivfvacuum.o src/vector.o
HEADERS = src/vector.h
TESTS = $(wildcard test/sql/*.sql)
REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))

View File

@@ -2,7 +2,6 @@ EXTENSION = vector
EXTVERSION = 0.4.4
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
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
REGRESS_OPTS = --inputdir=test --load-extension=vector
@@ -55,8 +54,6 @@ install:
copy $(SHLIB) "$(PKGLIBDIR)"
copy $(EXTENSION).control "$(SHAREDIR)\extension"
copy sql\$(EXTENSION)--*.sql "$(SHAREDIR)\extension"
mkdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
copy $(HEADERS) "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
installcheck:
"$(BINDIR)\pg_regress" --bindir="$(BINDIR)" $(REGRESS_OPTS) $(REGRESS)
@@ -64,9 +61,7 @@ installcheck:
uninstall:
del /f "$(PKGLIBDIR)\$(SHLIB)"
del /f "$(SHAREDIR)\extension\$(EXTENSION).control"
del /f "$(SHAREDIR)\extension\$(EXTENSION)--*.sql"
del /f "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)\*.h"
rmdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
del /f "$(SHAREDIR)\extension\vector--*.sql"
clean:
del /f $(SHLIB) $(EXTENSION).lib $(EXTENSION).exp

View File

@@ -2,7 +2,7 @@
Open-source vector similarity search for Postgres
Store your vectors with the rest of your data. Supports:
Store all of your application data in one place. Supports:
- exact and approximate nearest neighbor search
- L2 distance, inner product, and cosine distance
@@ -392,6 +392,7 @@ cosine_distance(vector, vector) → double precision | cosine distance
inner_product(vector, vector) → double precision | inner product
l2_distance(vector, vector) → double precision | Euclidean distance
l1_distance(vector, vector) → double precision | taxicab distance [unreleased]
normalize_l2(vector) → vector | normalize with Euclidean norm [unreleased]
vector_dims(vector) → integer | number of dimensions
vector_norm(vector) → double precision | Euclidean norm

View File

@@ -4,6 +4,9 @@
CREATE FUNCTION l1_distance(vector, vector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION normalize_l2(vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

View File

@@ -49,6 +49,9 @@ CREATE FUNCTION vector_dims(vector) RETURNS integer
CREATE FUNCTION vector_norm(vector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION normalize_l2(vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_add(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
@@ -265,15 +268,15 @@ CREATE OPERATOR CLASS vector_ip_ops
OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 3 vector_spherical_distance(vector, vector),
FUNCTION 4 vector_norm(vector);
FUNCTION 6 normalize_l2(vector);
CREATE OPERATOR CLASS vector_cosine_ops
FOR TYPE vector USING ivfflat AS
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector),
FUNCTION 3 vector_spherical_distance(vector, vector),
FUNCTION 4 vector_norm(vector);
FUNCTION 5 normalize_l2(vector),
FUNCTION 6 normalize_l2(vector);
CREATE OPERATOR CLASS vector_l2_ops
FOR TYPE vector USING hnsw AS
@@ -289,4 +292,4 @@ CREATE OPERATOR CLASS vector_cosine_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector);
FUNCTION 3 normalize_l2(vector);

View File

@@ -37,7 +37,7 @@ HnswInit(void)
);
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"Valid range is 1..1000.", &hnsw_ef_search,
"Valid range is 10..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
}
@@ -167,7 +167,7 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 2;
amroutine->amsupport = 3;
#if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0;
#endif

View File

@@ -19,40 +19,32 @@
/* Support functions */
#define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2
#define HNSW_NORMALIZE_PROC 3
#define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953
#define HNSW_PAGE_ID 0xFF90
#define HNSW_PAGE_ID 0xFF85
/* Preserved page numbers */
#define HNSW_METAPAGE_BLKNO 0
#define HNSW_HEAD_BLKNO 1 /* first element page */
/* Must correspond to page numbers since page lock is used */
#define HNSW_UPDATE_LOCK 0
#define HNSW_SCAN_LOCK 1
/* HNSW parameters */
#define HNSW_DEFAULT_M 16
#define HNSW_MIN_M 2
#define HNSW_MIN_M 4
#define HNSW_MAX_M 100
#define HNSW_DEFAULT_EF_CONSTRUCTION 40
#define HNSW_MIN_EF_CONSTRUCTION 4
#define HNSW_MIN_EF_CONSTRUCTION 10
#define HNSW_MAX_EF_CONSTRUCTION 1000
#define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1
#define HNSW_MIN_EF_SEARCH 10
#define HNSW_MAX_EF_SEARCH 1000
/* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1
#define HNSW_NEIGHBOR_TUPLE_TYPE 2
/* Make graph robust against non-HOT updates */
#define HNSW_HEAPTIDS 10
#define HNSW_UPDATE_ENTRY_GREATER 1
#define HNSW_UPDATE_ENTRY_ALWAYS 2
/* Build phases */
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_HNSW_PHASE_LOAD 2
@@ -77,13 +69,10 @@
#define HnswIsElementTuple(tup) ((tup)->type == HNSW_ELEMENT_TUPLE_TYPE)
#define HnswIsNeighborTuple(tup) ((tup)->type == HNSW_NEIGHBOR_TUPLE_TYPE)
/* 2 * M connections for ground layer */
#define HnswGetLayerM(m, layer) (layer == 0 ? (m) * 2 : (m))
/* Optimal ML from paper */
#define HnswGetLayerM(m, layer) (layer == 0 ? m * 2 : m)
#define HnswGetMl(m) (1 / log(m))
/* Ensure fits on page and in uint8 */
/* Ensure fits in uint8 */
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / m) - 2, 255)
/* Variables */
@@ -118,6 +107,13 @@ typedef struct HnswNeighborArray
HnswCandidate *items;
} HnswNeighborArray;
typedef struct HnswUpdate
{
HnswCandidate hc;
int level;
int index;
} HnswUpdate;
typedef struct HnswPairingHeapNode
{
pairingheap_node ph_node;
@@ -152,6 +148,7 @@ typedef struct HnswBuildState
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation;
/* Variables */
@@ -225,6 +222,7 @@ typedef struct HnswScanOpaqueData
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation;
} HnswScanOpaqueData;
@@ -260,31 +258,25 @@ typedef struct HnswVacuumState
int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation rel, uint16 procnum);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
void HnswNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result);
void HnswCommitBuffer(Buffer buf, GenericXLogState *state);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void HnswInit(void);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage, OffsetNumber *skipOffno);
HnswElement HnswGetEntryPoint(Relation index);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswElement HnswFindDuplicate(HnswElement e);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
HnswElement HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadvec);
void HnswUpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
void HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index);
/* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -60,7 +60,6 @@ CreateMetaPage(HnswBuildState * buildstate)
metap->efConstruction = buildstate->efConstruction;
metap->entryBlkno = InvalidBlockNumber;
metap->entryOffno = InvalidOffsetNumber;
metap->entryLevel = -1;
metap->insertPage = InvalidBlockNumber;
((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
@@ -183,7 +182,7 @@ CreateElementPages(HnswBuildState * buildstate)
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum);
HnswUpdateMetaPage(index, true, buildstate->entryPoint, insertPage, forkNum);
pfree(etup);
pfree(ntup);
@@ -279,33 +278,13 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
return false;
}
HnswNormValue(buildstate->normprocinfo, buildstate->normalizeprocinfo, collation, &value, buildstate->normvec);
/* Copy value to element so accessible outside of memory context */
memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
/* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
*dup = HnswFindDuplicate(element);
/* Update neighbors if needed */
if (*dup == NULL)
{
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
}
}
*dup = HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, NULL, false);
/* Update entry point if needed */
if (*dup == NULL && (entryPoint == NULL || element->level > entryPoint->level))
@@ -424,15 +403,13 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
if (buildstate->dimensions > HNSW_MAX_DIM)
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", HNSW_MAX_DIM);
if (buildstate->efConstruction < 2 * buildstate->m)
elog(ERROR, "ef_construction must be greater than or equal to 2 * m");
buildstate->reltuples = 0;
buildstate->indtuples = 0;
/* Get support functions */
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
buildstate->collation = index->rd_indcollation[0];
buildstate->elements = NIL;

View File

@@ -34,7 +34,7 @@ GetInsertPage(Relation index)
* Check for a free offset
*/
static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *firstFreePage)
{
OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -49,15 +49,14 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
if (etup->deleted)
{
BlockNumber elementPage = BufferGetBlockNumber(buf);
BlockNumber neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
OffsetNumber neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
ItemId itemid;
if (!BlockNumberIsValid(*newInsertPage))
*newInsertPage = elementPage;
if (!BlockNumberIsValid(*firstFreePage))
*firstFreePage = neighborPage;
if (neighborPage == elementPage)
if (neighborPage == BufferGetBlockNumber(buf))
{
*nbuf = buf;
*npage = page;
@@ -111,7 +110,7 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
* Add to element and neighbor pages
*/
static void
WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage)
WriteNewElementPages(Relation index, HnswElement e, int m)
{
Buffer buf;
Page page;
@@ -119,24 +118,21 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Size etupSize;
Size ntupSize;
Size combinedSize;
Size maxSize;
Size minCombinedSize;
HnswElementTuple etup;
BlockNumber currentPage = insertPage;
BlockNumber insertPage = GetInsertPage(index);
BlockNumber originalInsertPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup;
Buffer nbuf;
Page npage;
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
BlockNumber firstFreePage = InvalidBlockNumber;
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
minCombinedSize = etupSize + HNSW_NEIGHBOR_TUPLE_SIZE(0, m) + sizeof(ItemIdData);
/* Prepare element tuple */
etup = palloc0(etupSize);
@@ -146,22 +142,16 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
ntup = palloc0(ntupSize);
HnswSetNeighborTuple(ntup, e, m);
/* Find a page (or two if needed) to insert the tuples */
/* Find a page to insert the item */
for (;;)
{
buf = ReadBuffer(index, currentPage);
buf = ReadBuffer(index, insertPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Keep track of first page where element at level 0 can fit */
if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize)
newInsertPage = currentPage;
/* First, try the fastest path */
/* Space for both tuples on the current page */
/* This can split existing tuples in rare cases */
/* Space for both */
if (PageGetFreeSpace(page) >= combinedSize)
{
nbuf = buf;
@@ -169,8 +159,15 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
break;
}
/* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
/* Space for element but not neighbors and last page */
if (PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
{
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
break;
}
/* Space from deleted item */
if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &firstFreePage))
{
if (nbuf != buf)
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
@@ -178,17 +175,9 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
break;
}
/* Finally, try space for element only if last page */
/* Skip if both tuples can fit on the same page */
if (combinedSize > maxSize && PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
{
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
break;
}
insertPage = HnswPageGetOpaque(page)->nextblkno;
currentPage = HnswPageGetOpaque(page)->nextblkno;
if (BlockNumberIsValid(currentPage))
if (BlockNumberIsValid(insertPage))
{
/* Move to next page */
GenericXLogAbort(state);
@@ -230,10 +219,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
e->blkno = BufferGetBlockNumber(buf);
e->neighborPage = BufferGetBlockNumber(nbuf);
/* Added tuple to new page if newInsertPage is not set */
/* So can set to neighbor page instead of element page */
if (!BlockNumberIsValid(newInsertPage))
newInsertPage = e->neighborPage;
insertPage = e->neighborPage;
if (OffsetNumberIsValid(freeOffno))
{
@@ -279,126 +265,75 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
UnlockReleaseBuffer(nbuf);
/* Update the insert page */
if (BlockNumberIsValid(newInsertPage) && newInsertPage != insertPage)
*updatedInsertPage = newInsertPage;
if (insertPage != originalInsertPage && (!OffsetNumberIsValid(freeOffno) || firstFreePage == insertPage))
HnswUpdateMetaPage(index, false, NULL, insertPage, MAIN_FORKNUM);
}
/*
* Check if connection already exists
* Calculate index for update
*/
static bool
ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
static int
HnswGetIndex(HnswUpdate * update, int m)
{
for (int i = 0; i < lm; i++)
{
ItemPointer indextid = &ntup->indextids[startIdx + i];
if (!ItemPointerIsValid(indextid))
break;
if (ItemPointerGetBlockNumber(indextid) == e->blkno && ItemPointerGetOffsetNumber(indextid) == e->offno)
return true;
}
return false;
return (update->hc.element->level - update->level) * m + update->index;
}
/*
* Update neighbors
*/
void
HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting)
static void
UpdateNeighborPages(Relation index, HnswElement e, int m, List *updates)
{
for (int lc = e->level; lc >= 0; lc--)
ListCell *lc;
/* Could update multiple at once for same element */
/* but should only happen a low percent of time, so keep simple for now */
foreach(lc, updates)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc];
Buffer buf;
Page page;
GenericXLogState *state;
HnswUpdate *update = lfirst(lc);
ItemId itemid;
HnswNeighborTuple ntup;
Size ntupSize;
int idx;
OffsetNumber offno = update->hc.element->neighborOffno;
for (int i = 0; i < neighbors->length; i++)
/* Register page */
buf = ReadBuffer(index, update->hc.element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Get tuple */
itemid = PageGetItemId(page, offno);
ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
ntupSize = ItemIdGetLength(itemid);
/* Calculate index */
idx = HnswGetIndex(update, m);
/* Make robust to issues */
if (idx < ntup->count)
{
HnswCandidate *hc = &neighbors->items[i];
Buffer buf;
Page page;
GenericXLogState *state;
ItemId itemid;
HnswNeighborTuple ntup;
Size ntupSize;
int idx = -1;
int startIdx;
OffsetNumber offno = hc->element->neighborOffno;
ItemPointer indextid = &ntup->indextids[idx];
/* Get latest neighbors since they may have changed */
/* Do not lock yet since selecting neighbors can take time */
HnswLoadNeighbors(hc->element, index);
/* Update neighbor */
ItemPointerSet(indextid, e->blkno, e->offno);
/*
* Could improve performance for vacuuming by checking neighbors
* against list of elements being deleted to find index. It's
* important to exclude already deleted elements for this since
* they can be replaced at any time.
*/
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, offno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Select neighbors */
HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
/* New element was not selected as a neighbor */
if (idx == -1)
continue;
/* Register page */
buf = ReadBuffer(index, hc->element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Get tuple */
itemid = PageGetItemId(page, offno);
ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
ntupSize = ItemIdGetLength(itemid);
/* Calculate index for update */
startIdx = (hc->element->level - lc) * m;
/* Check for existing connection */
if (checkExisting && ConnectionExists(e, ntup, startIdx, lm))
idx = -1;
else if (idx == -2)
{
/* Find free offset if still exists */
/* TODO Retry updating connections if not */
for (int j = 0; j < lm; j++)
{
if (!ItemPointerIsValid(&ntup->indextids[startIdx + j]))
{
idx = startIdx + j;
break;
}
}
}
else
idx += startIdx;
/* Make robust to issues */
if (idx >= 0 && idx < ntup->count)
{
ItemPointer indextid = &ntup->indextids[idx];
/* Update neighbor */
ItemPointerSet(indextid, e->blkno, e->offno);
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, offno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
}
else
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
}
else
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
}
@@ -456,10 +391,8 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
* Write changes to disk
*/
static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint)
WriteElement(Relation index, HnswElement element, int m, List *updates, HnswElement dup, HnswElement entryPoint)
{
BlockNumber newInsertPage = InvalidBlockNumber;
/* Try to add to existing page */
if (dup != NULL)
{
@@ -467,19 +400,13 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
return;
}
/* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage);
/* Update insert page if needed */
if (BlockNumberIsValid(newInsertPage))
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false);
/* If fails, take this path */
WriteNewElementPages(index, element, m);
UpdateNeighborPages(index, element, m, updates);
/* Update metapage if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM);
HnswUpdateMetaPage(index, true, element, InvalidBlockNumber, MAIN_FORKNUM);
}
/*
@@ -490,6 +417,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
{
Datum value;
FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
HnswElement entryPoint;
HnswElement element;
int m = HnswGetM(index);
@@ -497,59 +425,29 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
double ml = HnswGetMl(m);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
List *updates = NIL;
HnswElement dup;
LOCKMODE lockmode = ShareLock;
/* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswNormValue(normprocinfo, collation, &value, NULL))
return false;
}
normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
HnswNormValue(normprocinfo, normalizeprocinfo, collation, &value, NULL);
/* Create an element */
element = HnswInitElement(heap_tid, m, ml, HnswGetMaxLevel(m));
element->vec = DatumGetVector(value);
/*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
* before repairing graph. Use a page lock so it does not interfere with
* buffer lock (or reads when vacuuming).
*/
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get entry point */
entryPoint = HnswGetEntryPoint(index);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get exclusive lock */
lockmode = ExclusiveLock;
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get latest entry point after lock is acquired */
entryPoint = HnswGetEntryPoint(index);
}
/* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
dup = HnswFindDuplicate(element);
dup = HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, &updates, false);
/* Write to disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
WriteElement(index, element, m, updates, dup, entryPoint);
return true;
}

View File

@@ -4,35 +4,34 @@
#include "hnsw.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
* Algorithm 5 from paper
*/
static List *
static void
GetScanItems(IndexScanDesc scan, Datum q)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep;
List *ep = NIL;
List *w;
HnswElement entryPoint = HnswGetEntryPoint(index);
if (entryPoint == NULL)
return NIL;
return;
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
ep = lappend(ep, HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL);
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL, NULL);
ep = w;
}
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL);
so->w = HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL, NULL);
}
/*
@@ -79,6 +78,7 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
/* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->normalizeprocinfo = HnswOptionalProcInfo(index, HNSW_NORMALIZE_PROC);
so->collation = index->rd_indcollation[0];
scan->opaque = so;
@@ -141,21 +141,10 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
HnswNormValue(so->normprocinfo, so->normalizeprocinfo, so->collation, &value, NULL);
}
/*
* 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);
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
GetScanItems(scan, value);
so->first = false;
}

View File

@@ -7,7 +7,7 @@
#include "vector.h"
/*
* Get the max number of connections in an upper layer for each element in the index
* Get the number of connection in the index
*/
int
HnswGetM(Relation index)
@@ -47,17 +47,26 @@ HnswOptionalProcInfo(Relation rel, uint16 procnum)
}
/*
* Divide by the norm
*
* Returns false if value should not be indexed
* Normalize a vector
*
* The caller needs to free the pointer stored in value
* if it's different than the original value
*/
bool
HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
void
HnswNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result)
{
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
double norm;
if (normalizeprocinfo != NULL)
{
*value = FunctionCall1Coll(normalizeprocinfo, collation, *value);
return;
}
if (procinfo == NULL)
return;
norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
@@ -70,11 +79,7 @@ HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
result->x[i] = v->x[i] / norm;
*value = PointerGetDatum(result);
return true;
}
return false;
}
/*
@@ -197,8 +202,8 @@ HnswAddHeapTid(HnswElement element, ItemPointer heaptid)
/*
* Allocate an element from block and offset numbers
*/
HnswElement
HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
static HnswElement
InitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
{
HnswElement element = palloc(sizeof(HnswElementData));
@@ -226,7 +231,7 @@ HnswGetEntryPoint(Relation index)
metap = HnswPageGetMeta(page);
if (BlockNumberIsValid(metap->entryBlkno))
entryPoint = HnswInitElementFromBlock(metap->entryBlkno, metap->entryOffno);
entryPoint = InitElementFromBlock(metap->entryBlkno, metap->entryOffno);
UnlockReleaseBuffer(buf);
@@ -234,12 +239,22 @@ HnswGetEntryPoint(Relation index)
}
/*
* Update the metapage info
* Update the metapage
*/
static void
HnswUpdateMetaPageInfo(Page page, int updateEntry, HnswElement entryPoint, BlockNumber insertPage)
void
HnswUpdateMetaPage(Relation index, bool updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum)
{
HnswMetaPage metap = HnswPageGetMeta(page);
Buffer buf;
Page page;
GenericXLogState *state;
HnswMetaPage metap;
buf = ReadBufferExtended(index, forkNum, HNSW_METAPAGE_BLKNO, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
metap = HnswPageGetMeta(page);
if (updateEntry)
{
@@ -249,7 +264,7 @@ HnswUpdateMetaPageInfo(Page page, int updateEntry, HnswElement entryPoint, Block
metap->entryOffno = InvalidOffsetNumber;
metap->entryLevel = -1;
}
else if (entryPoint->level > metap->entryLevel || updateEntry == HNSW_UPDATE_ENTRY_ALWAYS)
else
{
metap->entryBlkno = entryPoint->blkno;
metap->entryOffno = entryPoint->offno;
@@ -259,24 +274,6 @@ HnswUpdateMetaPageInfo(Page page, int updateEntry, HnswElement entryPoint, Block
if (BlockNumberIsValid(insertPage))
metap->insertPage = insertPage;
}
/*
* Update the metapage
*/
void
HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum)
{
Buffer buf;
Page page;
GenericXLogState *state;
buf = ReadBufferExtended(index, forkNum, HNSW_METAPAGE_BLKNO, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage);
HnswCommitBuffer(buf, state);
}
@@ -364,7 +361,7 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page)
if (!ItemPointerIsValid(indextid))
continue;
e = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
e = InitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
/* Calculate level based on offset */
level = element->level - i / m;
@@ -380,8 +377,8 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page)
/*
* Load neighbors
*/
void
HnswLoadNeighbors(HnswElement element, Relation index)
static void
LoadNeighbors(HnswElement element, Relation index)
{
Buffer buf;
Page page;
@@ -395,37 +392,6 @@ HnswLoadNeighbors(HnswElement element, Relation index)
UnlockReleaseBuffer(buf);
}
/*
* Load an element from a tuple
*/
void
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
{
element->level = etup->level;
element->deleted = etup->deleted;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptids = NIL;
if (loadHeaptids)
{
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
/* Can stop at first invalid */
if (!ItemPointerIsValid(&etup->heaptids[i]))
break;
HnswAddHeapTid(element, &etup->heaptids[i]);
}
}
if (loadVec)
{
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
}
/*
* Load an element and optionally get its distance from q
*/
@@ -446,7 +412,25 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
Assert(HnswIsElementTuple(etup));
/* Load element */
HnswLoadElementFromTuple(element, etup, true, loadVec);
element->heaptids = NIL;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
/* Can stop at first invalid */
if (!ItemPointerIsValid(&etup->heaptids[i]))
break;
HnswAddHeapTid(element, &etup->heaptids[i]);
}
element->level = etup->level;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->deleted = etup->deleted;
if (loadVec)
{
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
/* Calculate distance */
if (distance != NULL)
@@ -468,7 +452,7 @@ GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collat
* Create a candidate for the entry point
*/
HnswCandidate *
HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadvec)
{
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
@@ -476,7 +460,7 @@ HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *pr
if (index == NULL)
hc->distance = GetCandidateDistance(hc, q, procinfo, collation);
else
HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadVec);
HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadvec);
return hc;
}
@@ -543,7 +527,7 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement)
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, BlockNumber *skipPage, OffsetNumber *skipOffno)
{
ListCell *lc2;
@@ -579,13 +563,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
/*
* Do not count elements being deleted towards ef when vacuuming. It
* would be ideal to do this for inserts as well, but this could
* affect insert performance.
*/
if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
wlen++;
wlen++;
}
while (!pairingheap_is_empty(C))
@@ -598,7 +576,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
break;
if (c->element->neighbors == NULL)
HnswLoadNeighbors(c->element, index);
LoadNeighbors(c->element, index);
/* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc];
@@ -621,10 +599,16 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
else
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
Assert(!e->element->deleted);
/* Skip if fully deleted */
if (e->element->deleted)
continue;
/* Skip for inserts if deleting */
if (inserting && list_length(e->element->heaptids) == 0)
continue;
/* Skip self for vacuuming update */
if (skipElement != NULL && e->element->blkno == skipElement->blkno && e->element->offno == skipElement->offno)
if (skipPage != NULL && e->element->neighborPage == *skipPage && e->element->neighborOffno == *skipOffno)
continue;
/* Make robust to issues */
@@ -641,20 +625,11 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(ec)->ph_node));
wlen++;
/*
* Do not count elements being deleted towards ef when
* vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance.
*/
if (skipElement == NULL || list_length(e->element->heaptids) != 0)
{
wlen++;
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
}
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
}
}
}
@@ -733,7 +708,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswC
List *w = list_copy(c);
pairingheap *wd;
if (list_length(w) <= m)
if (list_length(w) < m)
return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -773,14 +748,14 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswC
/*
* Find duplicate element
*/
HnswElement
HnswFindDuplicate(HnswElement e)
static HnswElement
HnswFindDuplicate(HnswElement e, List *neighbors)
{
HnswNeighborArray *neighbors = &e->neighbors[0];
ListCell *lc;
for (int i = 0; i < neighbors->length; i++)
foreach(lc, neighbors)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswCandidate *neighbor = lfirst(lc);
/* Exit early since ordered by distance */
if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
@@ -829,134 +804,144 @@ CompareCandidateDistances(const void *a, const void *b)
return 0;
}
/*
* Create update
*/
static HnswUpdate *
CreateUpdate(HnswCandidate * hc, int level, int index)
{
HnswUpdate *update = palloc(sizeof(HnswUpdate));
update->hc = *hc;
update->level = level;
update->index = index;
return update;
}
/*
* Update connections
*/
void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
static void
UpdateConnections(HnswElement element, List *neighbors, int m, int lc, List **updates, Relation index, FmgrInfo *procinfo, Oid collation)
{
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
ListCell *lc2;
HnswCandidate hc2;
hc2.element = element;
hc2.distance = hc->distance;
if (currentNeighbors->length < m)
foreach(lc2, neighbors)
{
currentNeighbors->items[currentNeighbors->length++] = hc2;
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
/* Track update */
if (updateIdx != NULL)
*updateIdx = -2;
}
else
{
/* Shrink connections */
HnswCandidate *pruned = NULL;
HnswCandidate hc2;
/* Load elements on insert */
if (index != NULL)
hc2.element = element;
hc2.distance = hc->distance;
if (currentNeighbors->length < m)
{
Datum q = PointerGetDatum(hc->element->vec);
currentNeighbors->items[currentNeighbors->length++] = hc2;
/* Track updates */
if (updates != NULL)
*updates = lappend(*updates, CreateUpdate(hc, lc, currentNeighbors->length - 1));
}
else
{
/* Shrink connections */
HnswCandidate *pruned = NULL;
List *c = NIL;
/* Load elements on insert */
if (index != NULL)
{
Datum q = PointerGetDatum(hc->element->vec);
for (int i = 0; i < currentNeighbors->length; i++)
{
HnswCandidate *hc3 = &currentNeighbors->items[i];
if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0)
{
pruned = &currentNeighbors->items[i];
break;
}
}
}
if (pruned == NULL)
{
/* Add and sort candidates */
for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */
if (pruned == NULL)
continue;
}
/* Find and replace the pruned element */
for (int i = 0; i < currentNeighbors->length; i++)
{
HnswCandidate *hc3 = &currentNeighbors->items[i];
if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0)
if (currentNeighbors->items[i].element == pruned->element)
{
pruned = &currentNeighbors->items[i];
currentNeighbors->items[i] = hc2;
/* Track updates */
if (updates != NULL)
*updates = lappend(*updates, CreateUpdate(hc, lc, i));
break;
}
}
}
if (pruned == NULL)
{
List *c = NIL;
/* Add and sort candidates */
for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */
if (pruned == NULL)
return;
}
/* Find and replace the pruned element */
for (int i = 0; i < currentNeighbors->length; i++)
{
if (currentNeighbors->items[i].element == pruned->element)
{
currentNeighbors->items[i] = hc2;
/* Track update */
if (updateIdx != NULL)
*updateIdx = i;
break;
}
}
}
}
/*
* Remove elements being deleted
*/
static List *
RemoveElementsBeingDeleted(List *w)
{
ListCell *lc2;
List *w2 = NIL;
foreach(lc2, w)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
if (list_length(hc->element->heaptids) != 0)
w2 = lappend(w2, hc);
}
return w2;
}
/*
* Algorithm 1 from paper
*/
void
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
HnswElement
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, List **updates, bool vacuuming)
{
List *ep;
List *ep = NIL;
List *w;
int level = element->level;
int entryLevel;
List **newNeighbors = palloc(sizeof(List *) * (level + 1));
Datum q = PointerGetDatum(element->vec);
HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */
if (entryPoint == NULL)
return;
HnswElement dup;
BlockNumber *skipPage = vacuuming ? &element->neighborPage : NULL;
OffsetNumber *skipOffno = vacuuming ? &element->neighborOffno : NULL;
bool removeEntryPoint;
HnswCandidate *entryCandidate;
/* Get entry point and level */
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, true));
entryLevel = entryPoint->level;
if (entryPoint != NULL)
{
entryCandidate = HnswEntryCandidate(entryPoint, q, index, procinfo, collation, true);
ep = lappend(ep, entryCandidate);
entryLevel = entryPoint->level;
removeEntryPoint = vacuuming && list_length(entryPoint->heaptids) == 0;
}
else
{
entryLevel = -1;
removeEntryPoint = false;
}
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, true, skipElement);
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, true, skipPage, skipOffno);
ep = w;
}
@@ -967,19 +952,35 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
for (int lc = level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
List *neighbors;
w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipElement);
w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipPage, skipOffno);
/* Elements being deleted can help with search */
/* but should be removed before selecting neighbors */
if (index != NULL)
w = RemoveElementsBeingDeleted(w);
neighbors = SelectNeighbors(w, lm, lc, procinfo, collation, NULL);
AddConnections(element, neighbors, lm, lc);
/* Remove entry point if it's being deleted */
if (removeEntryPoint)
w = list_delete_ptr(w, entryCandidate);
newNeighbors[lc] = SelectNeighbors(w, lm, lc, procinfo, collation, NULL);
ep = w;
}
/* Look for duplicate */
if (level >= 0 && !vacuuming)
{
dup = HnswFindDuplicate(element, newNeighbors[0]);
if (dup != NULL)
return dup;
}
/* Update connections */
for (int lc = level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
AddConnections(element, newNeighbors[lc], lm, lc);
if (!vacuuming)
UpdateConnections(element, newNeighbors[lc], lm, lc, updates, index, procinfo, collation);
}
return NULL;
}

View File

@@ -5,11 +5,10 @@
#include "commands/vacuum.h"
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
* Check if deleted list contains an index TID
* Check if deleted list contains an index tid
*/
static bool
DeletedContains(HTAB *deleted, ItemPointer indextid)
@@ -33,7 +32,6 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
HnswElement entryPoint = HnswGetEntryPoint(vacuumstate->index);
IndexBulkDeleteResult *stats = vacuumstate->stats;
/* Store separately since highestPoint.level is uint8 */
int highestLevel = -1;
@@ -79,15 +77,11 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
break;
if (vacuumstate->callback(&etup->heaptids[i], vacuumstate->callback_state))
{
itemUpdated = true;
stats->tuples_removed++;
}
else
{
/* Move to front of list */
etup->heaptids[idx++] = etup->heaptids[i];
stats->num_index_tuples++;
}
}
@@ -115,7 +109,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
(void) hash_search(vacuumstate->deleted, &ip, HASH_ENTER, NULL);
}
else if (etup->level > highestLevel && !(entryPoint != NULL && blkno == entryPoint->blkno && offno == entryPoint->offno))
else if (etup->level > highestLevel && !(blkno == entryPoint->blkno && offno == entryPoint->offno))
{
/* Keep track of highest non-entry point */
highestPoint->blkno = blkno;
@@ -175,11 +169,6 @@ NeedsUpdated(HnswVacuumState * vacuumstate, HnswElement element)
}
}
/* Also update if layer 0 is not full */
/* This could indicate too many candidates being deleted during insert */
if (!needsUpdated)
needsUpdated = !ItemPointerIsValid(&ntup->indextids[ntup->count - 1]);
UnlockReleaseBuffer(buf);
return needsUpdated;
@@ -189,7 +178,7 @@ NeedsUpdated(HnswVacuumState * vacuumstate, HnswElement element)
* Repair graph for a single element
*/
static void
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint)
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element)
{
Relation index = vacuumstate->index;
Buffer buf;
@@ -199,20 +188,42 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
int efConstruction = vacuumstate->efConstruction;
FmgrInfo *procinfo = vacuumstate->procinfo;
Oid collation = vacuumstate->collation;
HnswElement entryPoint;
BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
/* Skip if element is entry point */
if (entryPoint != NULL && element->blkno == entryPoint->blkno && element->offno == entryPoint->offno)
/* Check if any neighbors point to deleted values */
if (!NeedsUpdated(vacuumstate, element))
return;
/* Refresh entry point for each element */
entryPoint = HnswGetEntryPoint(index);
/* Special case for entry point */
if (element->blkno == entryPoint->blkno && element->offno == entryPoint->offno)
{
if (BlockNumberIsValid(vacuumstate->highestPoint.blkno))
{
/* Already updated */
if (vacuumstate->highestPoint.blkno == element->blkno && vacuumstate->highestPoint.offno == element->offno)
return;
entryPoint = &vacuumstate->highestPoint;
/* Reset neighbors from previous update */
entryPoint->neighbors = NULL;
}
else
entryPoint = NULL;
}
/* Init fields */
HnswInitNeighbors(element, m);
element->heaptids = NIL;
/* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, NULL, true);
/* Update neighbor tuple */
/* Do this before getting page to minimize locking */
@@ -232,9 +243,6 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, true);
}
/*
@@ -248,35 +256,16 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
HnswElement entryPoint;
MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
if (!BlockNumberIsValid(highestPoint->blkno))
highestPoint = NULL;
/*
* Repair graph for highest non-entry point. Highest point may be outdated
* due to inserts that happen during and after RemoveHeapTids.
*/
if (highestPoint != NULL)
/* Repair graph for highest non-entry point */
/* This may not be the highest with new inserts, but should be fine */
if (BlockNumberIsValid(highestPoint->blkno))
{
/* Get a shared lock */
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
/* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint))
RepairGraphElement(vacuumstate, highestPoint, HnswGetEntryPoint(index));
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, ShareLock);
RepairGraphElement(vacuumstate, highestPoint);
}
/* Prevent concurrent inserts when possibly updating entry point */
LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Get latest entry point */
/* See if entry point needs updated */
entryPoint = HnswGetEntryPoint(index);
if (entryPoint != NULL)
{
ItemPointerData epData;
@@ -284,37 +273,15 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
ItemPointerSet(&epData, entryPoint->blkno, entryPoint->offno);
if (DeletedContains(vacuumstate->deleted, &epData))
{
/*
* Replace the entry point with the highest point. If highest
* point is outdated and empty, the entry point will be empty
* until an element is repaired.
*/
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, highestPoint, InvalidBlockNumber, MAIN_FORKNUM);
}
HnswUpdateMetaPage(index, true, highestPoint, InvalidBlockNumber, MAIN_FORKNUM);
else
{
/*
* Repair the entry point with the highest point. If highest point
* is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine.
*/
/* Highest point will be used to repair */
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
if (NeedsUpdated(vacuumstate, entryPoint))
{
/* Reset neighbors from previous update */
if (highestPoint != NULL)
highestPoint->neighbors = NULL;
RepairGraphElement(vacuumstate, entryPoint, highestPoint);
}
RepairGraphElement(vacuumstate, entryPoint);
}
}
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(vacuumstate->tmpCtx);
@@ -330,11 +297,6 @@ RepairGraph(HnswVacuumState * vacuumstate)
BufferAccessStrategy bas = vacuumstate->bas;
BlockNumber blkno = HNSW_HEAD_BLKNO;
/* Wait for inserts to complete */
LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Repair entry point first */
RepairGraphEntryPoint(vacuumstate);
while (BlockNumberIsValid(blkno))
@@ -371,8 +333,14 @@ RepairGraph(HnswVacuumState * vacuumstate)
continue;
/* Create an element */
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true);
element = palloc(sizeof(HnswElementData));
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->level = etup->level;
element->blkno = blkno;
element->offno = offno;
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
elements = lappend(elements, element);
}
@@ -383,48 +351,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Update neighbor pages */
foreach(lc2, elements)
{
HnswElement element = (HnswElement) lfirst(lc2);
HnswElement entryPoint;
LOCKMODE lockmode = ShareLock;
/* Check if any neighbors point to deleted values */
if (!NeedsUpdated(vacuumstate, element))
continue;
/* Get a shared lock */
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Refresh entry point for each element */
entryPoint = HnswGetEntryPoint(index);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get exclusive lock */
lockmode = ExclusiveLock;
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get latest entry point after lock is acquired */
entryPoint = HnswGetEntryPoint(index);
}
/* Repair connections */
RepairGraphElement(vacuumstate, element, entryPoint);
/*
* Update metapage if needed. Should only happen if entry point
* was replaced and highest point was outdated.
*/
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
}
RepairGraphElement(vacuumstate, (HnswElement) lfirst(lc2));
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
@@ -442,10 +369,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
/* Wait for selects to complete */
LockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
UnlockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
IndexBulkDeleteResult *stats = vacuumstate->stats;
while (BlockNumberIsValid(blkno))
{
@@ -489,17 +413,17 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Skip deleted tuples */
if (etup->deleted)
{
/* Set to first free page */
if (!BlockNumberIsValid(insertPage))
insertPage = blkno;
continue;
}
/* Skip live tuples */
if (ItemPointerIsValid(&etup->heaptids[0]))
{
stats->num_index_tuples++;
continue;
}
/* Update stats */
stats->tuples_removed++;
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
@@ -562,8 +486,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
UnlockReleaseBuffer(buf);
}
/* Update insert page last, after everything has been marked as deleted */
HnswUpdateMetaPage(index, 0, NULL, insertPage, MAIN_FORKNUM);
HnswUpdateMetaPage(index, false, NULL, insertPage, MAIN_FORKNUM);
}
/*
@@ -611,67 +534,6 @@ FreeVacuumState(HnswVacuumState * vacuumstate)
MemoryContextDelete(vacuumstate->tmpCtx);
}
/*
* Print graph
*/
#ifdef HNSW_DEBUG
static void
PrintGraph(HnswVacuumState * vacuumstate)
{
BlockNumber blkno = HNSW_HEAD_BLKNO;
Relation index = vacuumstate->index;
while (BlockNumberIsValid(blkno))
{
Buffer buf;
Page page;
OffsetNumber offno;
OffsetNumber maxoffno;
buf = ReadBuffer(index, blkno);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElement element;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
/* Skip deleted tuples */
if (etup->deleted)
continue;
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true);
HnswLoadNeighbors(element, index);
elog(INFO, "element (%d,%d)", element->blkno, element->offno);
for (int lc = element->level; lc >= 0; lc--)
{
HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
{
HnswElement e = neighbors->items[i].element;
elog(INFO, "%d: (%d,%d)", lc, e->blkno, e->offno);
}
}
}
blkno = HnswPageGetOpaque(page)->nextblkno;
UnlockReleaseBuffer(buf);
}
}
#endif
/*
* Bulk delete tuples from the index
*/
@@ -692,10 +554,6 @@ hnswbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
/* Pass 3: Mark as deleted */
MarkDeleted(&vacuumstate);
#ifdef HNSW_DEBUG
PrintGraph(&vacuumstate);
#endif
FreeVacuumState(&vacuumstate);
return vacuumstate.stats;

View File

@@ -75,11 +75,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
* Normalize with KMEANS_NORM_PROC since spherical distance function
* expects unit vectors
*/
if (buildstate->kmeansnormprocinfo != NULL)
{
if (!IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->collation, &value, buildstate->normvec))
return;
}
IvfflatNormValue(buildstate->kmeansnormprocinfo, buildstate->kmeansnormalizeprocinfo, buildstate->collation, &value, buildstate->normvec);
if (samples->length < targsamples)
{
@@ -176,11 +172,7 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!IvfflatNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
return;
}
IvfflatNormValue(buildstate->normprocinfo, buildstate->normalizeprocinfo, buildstate->collation, &value, buildstate->normvec);
/* Find the list that minimizes the distance */
for (int i = 0; i < centers->length; i++)
@@ -368,6 +360,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
buildstate->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
buildstate->kmeansnormprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
buildstate->normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
buildstate->kmeansnormalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORMALIZE_PROC);
buildstate->collation = index->rd_indcollation[0];
/* Require more than one dimension for spherical k-means */

View File

@@ -24,7 +24,7 @@ IvfflatInit(void)
{
ivfflat_relopt_kind = add_reloption_kind();
add_int_reloption(ivfflat_relopt_kind, "lists", "Number of inverted lists",
IVFFLAT_DEFAULT_LISTS, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS
IVFFLAT_DEFAULT_LISTS, 1, IVFFLAT_MAX_LISTS
#if PG_VERSION_NUM >= 130000
,AccessExclusiveLock
#endif
@@ -32,7 +32,7 @@ IvfflatInit(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
"Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
1, 1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
}
/*
@@ -194,7 +194,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 4;
amroutine->amsupport = 6;
#if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0;
#endif

View File

@@ -31,6 +31,8 @@
#define IVFFLAT_NORM_PROC 2
#define IVFFLAT_KMEANS_DISTANCE_PROC 3
#define IVFFLAT_KMEANS_NORM_PROC 4
#define IVFFLAT_NORMALIZE_PROC 5
#define IVFFLAT_KMEANS_NORMALIZE_PROC 6
#define IVFFLAT_VERSION 1
#define IVFFLAT_MAGIC_NUMBER 0x14FF1A7
@@ -40,11 +42,8 @@
#define IVFFLAT_METAPAGE_BLKNO 0
#define IVFFLAT_HEAD_BLKNO 1 /* first list page */
/* IVFFlat parameters */
#define IVFFLAT_DEFAULT_LISTS 100
#define IVFFLAT_MIN_LISTS 1
#define IVFFLAT_MAX_LISTS 32768
#define IVFFLAT_DEFAULT_PROBES 1
/* Build phases */
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
@@ -175,6 +174,8 @@ typedef struct IvfflatBuildState
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
FmgrInfo *kmeansnormprocinfo;
FmgrInfo *normalizeprocinfo;
FmgrInfo *kmeansnormalizeprocinfo;
Oid collation;
/* Variables */
@@ -256,6 +257,7 @@ typedef struct IvfflatScanOpaqueData
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Oid collation;
/* Lists */
@@ -276,7 +278,7 @@ void VectorArrayFree(VectorArray arr);
void PrintVectorArray(char *msg, VectorArray arr);
void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation rel, uint16 procnum);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
void IvfflatNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result);
int IvfflatGetLists(Relation index);
void IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);
void IvfflatCommitBuffer(Buffer buf, GenericXLogState *state);

View File

@@ -68,6 +68,7 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
IndexTuple itup;
Datum value;
FmgrInfo *normprocinfo;
FmgrInfo *normalizeprocinfo;
Buffer buf;
Page page;
GenericXLogState *state;
@@ -81,11 +82,8 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
/* Normalize if needed */
normprocinfo = IvfflatOptionalProcInfo(rel, IVFFLAT_NORM_PROC);
if (normprocinfo != NULL)
{
if (!IvfflatNormValue(normprocinfo, rel->rd_indcollation[0], &value, NULL))
return;
}
normalizeprocinfo = IvfflatOptionalProcInfo(rel, IVFFLAT_NORMALIZE_PROC);
IvfflatNormValue(normprocinfo, normalizeprocinfo, rel->rd_indcollation[0], &value, NULL);
/* Find the insert page - sets the page and list info */
FindInsertPage(rel, values, &insertPage, &listInfo);

View File

@@ -232,6 +232,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Set support functions */
so->procinfo = index_getprocinfo(index, 1, IVFFLAT_DISTANCE_PROC);
so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
so->normalizeprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORMALIZE_PROC);
so->collation = index->rd_indcollation[0];
/* Create tuple description for sorting */
@@ -319,8 +320,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL);
IvfflatNormValue(so->normprocinfo, so->normalizeprocinfo, so->collation, &value, NULL);
}
IvfflatBench("GetScanLists", GetScanLists(scan, value));

View File

@@ -66,17 +66,26 @@ IvfflatOptionalProcInfo(Relation rel, uint16 procnum)
}
/*
* Divide by the norm
*
* Returns false if value should not be indexed
* Normalize a vector
*
* The caller needs to free the pointer stored in value
* if it's different than the original value
*/
bool
IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
void
IvfflatNormValue(FmgrInfo *procinfo, FmgrInfo *normalizeprocinfo, Oid collation, Datum *value, Vector * result)
{
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
double norm;
if (normalizeprocinfo != NULL)
{
*value = FunctionCall1Coll(normalizeprocinfo, collation, *value);
return;
}
if (procinfo == NULL)
return;
norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
@@ -89,11 +98,7 @@ IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * resul
result->x[i] = v->x[i] / norm;
*value = PointerGetDatum(result);
return true;
}
return false;
}
/*

View File

@@ -15,10 +15,6 @@
#include "utils/numeric.h"
#include "vector.h"
#if PG_VERSION_NUM >= 160000
#include "varatt.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "common/shortest_dec.h"
#include "utils/float.h"
@@ -39,7 +35,6 @@ PG_MODULE_MAGIC;
/*
* Initialize index options and variables
*/
PGDLLEXPORT void _PG_init(void);
void
_PG_init(void)
{
@@ -105,23 +100,6 @@ CheckElement(float value)
errmsg("infinite value not allowed in vector")));
}
/*
* Allocate and initialize a new vector
*/
Vector *
InitVector(int dim)
{
Vector *result;
int size;
size = VECTOR_SIZE(dim);
result = (Vector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
/*
* Check for whitespace, since array_isspace() is static
*/
@@ -755,16 +733,62 @@ vector_dims(PG_FUNCTION_ARGS)
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_norm);
Datum
vector_norm(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
float *ax = a->x;
float norm = 0.0;
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
norm += ax[i] * ax[i];
PG_RETURN_FLOAT8(sqrt((double) norm));
}
/*
* Normalize a vector with the L2 norm
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(normalize_l2);
Datum
normalize_l2(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
float *ax = a->x;
double norm = 0.0;
#ifdef _MSC_VER
/* Fix precision on Windows */
double normf;
#else
float normf;
#endif
Vector *result;
float *rx;
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
norm += (double) ax[i] * (double) ax[i];
PG_RETURN_FLOAT8(sqrt(norm));
norm = sqrt(norm);
normf = norm;
if (normf > 0)
{
/* Auto-vectorized */
for (int i = 0, imax = a->dim; i < imax; i++)
rx[i] = ax[i] / normf;
/* Check for overflow */
for (int i = 0, imax = a->dim; i < imax; i++)
{
if (isinf(rx[i]))
float_overflow_error();
}
}
PG_RETURN_POINTER(result);
}
/*

View File

@@ -1,6 +1,12 @@
#ifndef VECTOR_H
#define VECTOR_H
#include "postgres.h"
#if PG_VERSION_NUM >= 160000
#include "varatt.h"
#endif
#define VECTOR_MAX_DIM 16000
#define VECTOR_SIZE(_dim) (offsetof(Vector, x) + sizeof(float)*(_dim))
@@ -8,6 +14,9 @@
#define PG_GETARG_VECTOR_P(x) DatumGetVector(PG_GETARG_DATUM(x))
#define PG_RETURN_VECTOR_P(x) PG_RETURN_POINTER(x)
/* Exported functions */
PGDLLEXPORT void _PG_init(void);
typedef struct Vector
{
int32 vl_len_; /* varlena header (do not touch directly!) */
@@ -16,8 +25,24 @@ typedef struct Vector
float x[FLEXIBLE_ARRAY_MEMBER];
} Vector;
Vector *InitVector(int dim);
void PrintVector(char *msg, Vector * vector);
int vector_cmp_internal(Vector * a, Vector * b);
/*
* Allocate and initialize a new vector
*/
static inline Vector *
InitVector(int dim)
{
Vector *result;
int size;
size = VECTOR_SIZE(dim);
result = (Vector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
#endif

View File

@@ -48,10 +48,34 @@ SELECT vector_norm('[0,1]');
1
(1 row)
SELECT vector_norm('[3e37,4e37]')::real;
vector_norm
-------------
5e+37
SELECT normalize_l2('[3,4]');
normalize_l2
--------------
[0.6,0.8]
(1 row)
SELECT normalize_l2('[3,0]');
normalize_l2
--------------
[1,0]
(1 row)
SELECT normalize_l2('[0,0.1]');
normalize_l2
--------------
[0,1]
(1 row)
SELECT normalize_l2('[0,0]');
normalize_l2
--------------
[0,0]
(1 row)
SELECT normalize_l2('[3e38]');
normalize_l2
--------------
[1]
(1 row)
SELECT l2_distance('[0,0]', '[3,4]');

View File

@@ -9,18 +9,19 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
[1,1,1]
[1,2,3]
[1,2,4]
(3 rows)
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
count
-------
3
4
(1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
count
-------
3
4
(1 row)
DROP TABLE t;

View File

@@ -27,10 +27,4 @@ SELECT COUNT(*) FROM t;
5
(1 row)
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)
DROP TABLE t;

View File

@@ -1,26 +1,25 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 1);
ERROR: value 1 out of bounds for option "m"
DETAIL: Valid values are between "2" and "100".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 3);
ERROR: value 3 out of bounds for option "m"
DETAIL: Valid values are between "4" and "100".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 101);
ERROR: value 101 out of bounds for option "m"
DETAIL: Valid values are between "2" and "100".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 3);
ERROR: value 3 out of bounds for option "ef_construction"
DETAIL: Valid values are between "4" and "1000".
DETAIL: Valid values are between "4" and "100".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 9);
ERROR: value 9 out of bounds for option "ef_construction"
DETAIL: Valid values are between "10" and "1000".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 1001);
ERROR: value 1001 out of bounds for option "ef_construction"
DETAIL: Valid values are between "4" and "1000".
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 16, ef_construction = 31);
ERROR: ef_construction must be greater than or equal to 2 * m
DETAIL: Valid values are between "10" and "1000".
SHOW hnsw.ef_search;
hnsw.ef_search
----------------
40
(1 row)
SET hnsw.ef_search = 0;
ERROR: 0 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SET hnsw.ef_search = 9;
ERROR: 9 is outside the valid range for parameter "hnsw.ef_search" (10 .. 1000)
SET hnsw.ef_search = 1001;
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (10 .. 1000)
DROP TABLE t;

View File

@@ -9,18 +9,19 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
[1,1,1]
[1,2,3]
[1,2,4]
(3 rows)
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
count
-------
3
4
(1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
count
-------
3
4
(1 row)
DROP TABLE t;

View File

@@ -1,7 +1,7 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
@@ -27,13 +27,4 @@ SELECT COUNT(*) FROM t;
5
(1 row)
TRUNCATE t;
NOTICE: ivfflat index created with little data
DETAIL: This will cause low recall.
HINT: Drop the index until the table has more data.
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)
DROP TABLE t;

View File

@@ -1,8 +1,9 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 0);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 0);
ERROR: value 0 out of bounds for option "lists"
DETAIL: Valid values are between "1" and "32768".
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 32769);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 32769);
ERROR: value 32769 out of bounds for option "lists"
DETAIL: Valid values are between "1" and "32768".
SHOW ivfflat.probes;

View File

@@ -1,7 +1,7 @@
SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 1);
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------

View File

@@ -11,7 +11,12 @@ SELECT vector_dims('[1,2,3]');
SELECT round(vector_norm('[1,1]')::numeric, 5);
SELECT vector_norm('[3,4]');
SELECT vector_norm('[0,1]');
SELECT vector_norm('[3e37,4e37]')::real;
SELECT normalize_l2('[3,4]');
SELECT normalize_l2('[3,0]');
SELECT normalize_l2('[0,0.1]');
SELECT normalize_l2('[0,0]');
SELECT normalize_l2('[3e38]');
SELECT l2_distance('[0,0]', '[3,4]');
SELECT l2_distance('[0,0]', '[0,1]');

View File

@@ -10,7 +10,4 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
SELECT COUNT(*) FROM t;
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -1,13 +1,14 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 1);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 3);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 101);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 3);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 9);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (ef_construction = 1001);
CREATE INDEX ON t USING hnsw (val vector_l2_ops) WITH (m = 16, ef_construction = 31);
SHOW hnsw.ef_search;
SET hnsw.ef_search = 0;
SET hnsw.ef_search = 9;
SET hnsw.ef_search = 1001;
DROP TABLE t;

View File

@@ -2,7 +2,7 @@ SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
@@ -10,7 +10,4 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
SELECT COUNT(*) FROM t;
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -1,6 +1,8 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 0);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 32769);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 0);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 32769);
SHOW ivfflat.probes;

View File

@@ -2,7 +2,7 @@ SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
CREATE INDEX ON t USING ivfflat (val) WITH (lists = 1);
SELECT * FROM t ORDER BY val <-> '[3,3,3]';

View File

@@ -26,8 +26,7 @@ sub test_index_replay
or die "Timed out while waiting for replica 1 to catch up";
my @r = ();
for (1 .. $dim)
{
for (1 .. $dim) {
push(@r, rand());
}
my $sql = join(",", @r);
@@ -53,13 +52,11 @@ my $array_sql = join(",", ('random()') x $dim);
# Initialize primary node
$node_primary = get_new_node('primary');
$node_primary->init(allows_streaming => 1);
if ($dim > 32)
{
if ($dim > 32) {
# TODO use wal_keep_segments for Postgres < 13
$node_primary->append_conf('postgresql.conf', qq(wal_keep_size = 1GB));
}
if ($dim > 1500)
{
if ($dim > 1500) {
$node_primary->append_conf('postgresql.conf', qq(maintenance_work_mem = 128MB));
}
$node_primary->start;
@@ -70,7 +67,8 @@ $node_primary->backup($backup_name);
# Create streaming replica linking to primary
$node_replica = get_new_node('replica');
$node_replica->init_from_backup($node_primary, $backup_name, has_streaming => 1);
$node_replica->init_from_backup($node_primary, $backup_name,
has_streaming => 1);
$node_replica->start;
# Create ivfflat index on primary
@@ -79,7 +77,7 @@ $node_primary->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));
$node_primary->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node_primary->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
$node_primary->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
# Test that queries give same result
test_index_replay('initial');

View File

@@ -7,8 +7,7 @@ use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim)
{
for (1 .. $dim) {
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
@@ -25,7 +24,7 @@ $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
# Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");

View File

@@ -22,8 +22,7 @@ sub test_recall
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
for my $i (0 .. $#queries) {
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
@@ -34,10 +33,8 @@ sub test_recall
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
foreach (@expected_ids) {
if (exists($actual_set{$_})) {
$correct++;
}
$total++;
@@ -60,8 +57,7 @@ $node->safe_psql("postgres",
);
# Generate queries
for (1 .. 20)
{
for (1..20) {
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
@@ -70,21 +66,26 @@ for (1 .. 20)
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
foreach (@operators) {
my $operator = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
foreach (@queries) {
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
push(@expected, $res);
}
my $opclass;
if ($operator eq "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator eq "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
# Build index serially
$node->safe_psql("postgres", qq(
SET max_parallel_maintenance_workers = 0;
@@ -92,14 +93,13 @@ for my $i (0 .. $#operators)
));
# Test approximate results
if ($operator ne "<#>")
{
if ($operator ne "<#>") {
# TODO fix test
test_recall(1, 0.71, $operator);
test_recall(1, 0.75, $operator);
test_recall(10, 0.95, $operator);
}
# Account for equal distances
test_recall(100, 0.9925, $operator);
test_recall(100, 0.995, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
@@ -113,14 +113,13 @@ for my $i (0 .. $#operators)
like($stderr, qr/using \d+ parallel workers/);
# Test approximate results
if ($operator ne "<#>")
{
if ($operator ne "<#>") {
# TODO fix test
test_recall(1, 0.71, $operator);
test_recall(1, 0.75, $operator);
test_recall(10, 0.95, $operator);
}
# Account for equal distances
test_recall(100, 0.9925, $operator);
test_recall(100, 0.995, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
}

View File

@@ -20,7 +20,7 @@ sub test_centers
{
my ($lists, $min) = @_;
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops) WITH (lists = $lists);");
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v) WITH (lists = $lists);");
is($ret, 0, $stderr);
}

View File

@@ -18,21 +18,24 @@ $node->safe_psql("postgres",
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
foreach (@operators) {
my $operator = $_;
# Add index
my $opclass;
if ($operator eq "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator eq "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
# Test 100% recall
for (1 .. 20)
{
my $id = int(rand() * 100000);
my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $id;");
for (1..20) {
my $i = int(rand() * 100000);
my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $i;");
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT v FROM tst ORDER BY v <-> '$query' LIMIT 1;

View File

@@ -16,8 +16,8 @@ $node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX lists50 ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 50);");
$node->safe_psql("postgres", "CREATE INDEX lists100 ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 100);");
$node->safe_psql("postgres", "CREATE INDEX lists50 ON tst USING ivfflat (v) WITH (lists = 50);");
$node->safe_psql("postgres", "CREATE INDEX lists100 ON tst USING ivfflat (v) WITH (lists = 100);");
# Test prefers more lists
my $res = $node->safe_psql("postgres", "EXPLAIN SELECT v FROM tst ORDER BY v <-> '[0.5,0.5,0.5]' LIMIT 10;");
@@ -26,7 +26,7 @@ unlike($res, qr/lists50/);
# Test errors with too much memory
my ($ret, $stdout, $stderr) = $node->psql("postgres",
"CREATE INDEX lists10000 ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 10000);"
"CREATE INDEX lists10000 ON tst USING ivfflat (v) WITH (lists = 10000);"
);
like($stderr, qr/memory required is/);

View File

@@ -19,7 +19,7 @@ $node->safe_psql("postgres", "CREATE TABLE tst (v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
$node->pgbench(
"--no-vacuum --client=5 --transactions=100",
@@ -28,7 +28,7 @@ $node->pgbench(
[qr{^$}],
"concurrent INSERTs",
{
"007_ivfflat_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
"007_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
}
);

View File

@@ -30,8 +30,7 @@ sub test_aggregate
# Test matches real for avg
# Cannot test sum since sum(real) varies between calls
if ($agg eq 'avg')
{
if ($agg eq 'avg') {
my $r1 = $node->safe_psql("postgres", "SELECT $agg(r1)::float4 FROM tst;");
my $r2 = $node->safe_psql("postgres", "SELECT $agg(r2)::float4 FROM tst;");
my $r3 = $node->safe_psql("postgres", "SELECT $agg(r3)::float4 FROM tst;");

View File

@@ -26,8 +26,7 @@ sub test_index_replay
or die "Timed out while waiting for replica 1 to catch up";
my @r = ();
for (1 .. $dim)
{
for (1 .. $dim) {
push(@r, rand());
}
my $sql = join(",", @r);
@@ -53,13 +52,11 @@ my $array_sql = join(",", ('random()') x $dim);
# Initialize primary node
$node_primary = get_new_node('primary');
$node_primary->init(allows_streaming => 1);
if ($dim > 32)
{
if ($dim > 32) {
# TODO use wal_keep_segments for Postgres < 13
$node_primary->append_conf('postgresql.conf', qq(wal_keep_size = 1GB));
}
if ($dim > 1500)
{
if ($dim > 1500) {
$node_primary->append_conf('postgresql.conf', qq(maintenance_work_mem = 128MB));
}
$node_primary->start;
@@ -70,7 +67,8 @@ $node_primary->backup($backup_name);
# Create streaming replica linking to primary
$node_replica = get_new_node('replica');
$node_replica->init_from_backup($node_primary, $backup_name, has_streaming => 1);
$node_replica->init_from_backup($node_primary, $backup_name,
has_streaming => 1);
$node_replica->start;
# Create hnsw index on primary

View File

@@ -7,8 +7,7 @@ use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim)
{
for (1 .. $dim) {
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
@@ -23,7 +22,7 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
@@ -34,21 +33,11 @@ my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_id
$node->safe_psql("postgres", "DELETE FROM tst;");
$node->safe_psql("postgres", "VACUUM tst;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
# Check size
# May increase some due to different levels
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
cmp_ok($new_size, "<=", $size * 1.02, "size does not increase too much");
# Delete all but one
$node->safe_psql("postgres", "DELETE FROM tst WHERE i != 123;");
$node->safe_psql("postgres", "VACUUM tst;");
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst ORDER BY v <-> '[0,0,0]' LIMIT 10;
));
is($res, 123);
cmp_ok($new_size, "<=", $size * 1.01, "size does not increase too much");
done_testing();

View File

@@ -21,8 +21,7 @@ sub test_recall
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
for my $i (0 .. $#queries) {
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
@@ -32,10 +31,8 @@ sub test_recall
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
foreach (@expected_ids) {
if (exists($actual_set{$_})) {
$correct++;
}
$total++;
@@ -58,8 +55,7 @@ $node->safe_psql("postgres",
);
# Generate queries
for (1 .. 20)
{
for (1..20) {
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
@@ -68,26 +64,33 @@ for (1 .. 20)
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
foreach (@operators) {
my $operator = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
foreach (@queries) {
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
push(@expected, $res);
}
# Add index
my $opclass;
if ($operator eq "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator eq "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator);
if ($operator eq "<#>") {
test_recall(0.80, $operator);
} else {
test_recall(0.99, $operator);
}
}
done_testing();

View File

@@ -21,8 +21,7 @@ sub test_recall
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
for my $i (0 .. $#queries) {
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
@@ -32,10 +31,8 @@ sub test_recall
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
foreach (@expected_ids) {
if (exists($actual_set{$_})) {
$correct++;
}
$total++;
@@ -52,11 +49,10 @@ $node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i serial, v vector(3));");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
# Generate queries
for (1 .. 20)
{
for (1..20) {
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
@@ -65,32 +61,28 @@ for (1 .. 20)
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
foreach (@operators) {
my $operator = $_;
# Add index
my $opclass;
if ($operator eq "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator eq "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
# Use concurrent inserts
$node->pgbench(
"--no-vacuum --client=10 --transactions=1000",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"013_hnsw_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES (ARRAY[random(), random(), random()]);"
}
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 10000) i;"
);
# Get exact results
@expected = ();
foreach (@queries)
{
foreach (@queries) {
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
@@ -98,8 +90,11 @@ for my $i (0 .. $#operators)
push(@expected, $res);
}
my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator);
if ($operator eq "<#>") {
test_recall(0.80, $operator);
} else {
test_recall(0.99, $operator);
}
$node->safe_psql("postgres", "DROP INDEX idx;");
$node->safe_psql("postgres", "TRUNCATE tst;");

View File

@@ -17,8 +17,22 @@ $node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 100) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
$node->pgbench(
"--no-vacuum --client=5 --transactions=100",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"007_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
}
);
sub idx_scan
{
# Stats do not update instantaneously
@@ -27,48 +41,18 @@ sub idx_scan
$node->safe_psql("postgres", "SELECT idx_scan FROM pg_stat_user_indexes WHERE indexrelid = 'tst_v_idx'::regclass;");
}
for my $i (1 .. 20)
{
$node->pgbench(
"--no-vacuum --client=10 --transactions=1",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"014_hnsw_inserts_$i" => "INSERT INTO tst VALUES (ARRAY[$array_sql]);"
}
);
my $expected = 100 + 5 * 100 * 10;
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
is($count, 10);
my $count = $node->safe_psql("postgres", "SELECT COUNT(*) FROM tst;");
is($count, $expected);
is(idx_scan(), 0);
$node->safe_psql("postgres", "TRUNCATE tst;");
}
$node->pgbench(
"--no-vacuum --client=20 --transactions=5",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"014_hnsw_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
}
);
my $count = $node->safe_psql("postgres", qq(
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1000;
SET hnsw.ef_search = 400;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
# Elements may lose all incoming connections with the HNSW algorithm
# Vacuuming can fix this if one of the elements neighbors is deleted
cmp_ok($count, ">=", 997);
is(idx_scan(), 21);
is($count, 400);
is(idx_scan(), 1);
done_testing();

View File

@@ -1,58 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
sub insert_vectors
{
for my $i (1 .. 20)
{
$node->safe_psql("postgres", "INSERT INTO tst VALUES ('[1,1,1]');");
}
}
sub test_duplicates
{
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1;
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
));
is($res, 10);
}
# Test duplicates with build
insert_vectors();
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
test_duplicates();
# Reset
$node->safe_psql("postgres", "TRUNCATE tst;");
# Test duplicates with inserts
insert_vectors();
test_duplicates();
# Test fallback path for inserts
$node->pgbench(
"--no-vacuum --client=5 --transactions=100",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"015_hnsw_duplicates" => "INSERT INTO tst VALUES ('[1,1,1]');"
}
);
done_testing();

View File

@@ -1,97 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
sub test_recall
{
my ($min, $ef_search, $test_name) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SELECT i FROM tst ORDER BY v <-> '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $test_name);
}
# Initialize node
$node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
$node->safe_psql("postgres", "ALTER TABLE tst SET (autovacuum_enabled = false);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 10000) i;"
);
# Add index
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops) WITH (m = 4, ef_construction = 8);");
# Delete data
$node->safe_psql("postgres", "DELETE FROM tst WHERE i > 2500;");
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v <-> '$_' LIMIT $limit;
));
push(@expected, $res);
}
test_recall(0.20, $limit, "before vacuum");
test_recall(0.95, 100, "before vacuum");
# TODO test concurrent inserts with vacuum
$node->safe_psql("postgres", "VACUUM tst;");
test_recall(0.95, $limit, "after vacuum");
done_testing();