Compare commits

..

4 Commits

Author SHA1 Message Date
Andrew Kane
f90d52d562 Added test for explicit zeros [skip ci] 2023-11-06 14:29:52 -08:00
Andrew Kane
22a1b06924 Added todo [skip ci] 2023-11-05 21:26:28 -08:00
Andrew Kane
b17b9c1ca2 Updated version [skip ci] 2023-11-05 18:21:09 -08:00
Andrew Kane
492ae1225c Added support for sparse vectors 2023-11-05 18:12:19 -08:00
20 changed files with 1214 additions and 600 deletions

View File

@@ -1,6 +1,6 @@
## 0.5.2 (unreleased) ## 0.6.0 (unreleased)
- Added support for on-disk parallel index builds for HNSW - Added support for sparse vectors
## 0.5.1 (2023-10-10) ## 0.5.1 (2023-10-10)

View File

@@ -3,8 +3,8 @@ EXTVERSION = 0.5.1
MODULE_big = vector MODULE_big = vector
DATA = $(wildcard sql/*--*.sql) 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 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/svector.o src/vector.o
HEADERS = src/vector.h HEADERS = src/svector.h src/vector.h
TESTS = $(wildcard test/sql/*.sql) TESTS = $(wildcard test/sql/*.sql)
REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS)) REGRESS = $(patsubst test/sql/%.sql,%,$(TESTS))

View File

@@ -1,8 +1,8 @@
EXTENSION = vector EXTENSION = vector
EXTVERSION = 0.5.1 EXTVERSION = 0.5.1
OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\svector.obj src\vector.obj
HEADERS = src\vector.h HEADERS = src\svector.h src\vector.h
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION) REGRESS_OPTS = --inputdir=test --load-extension=$(EXTENSION)

View File

@@ -369,6 +369,26 @@ To speed up queries with an IVFFlat index, increase the number of inverted lists
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 1000); CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 1000);
``` ```
## Sparse Vectors
Create a sparse vector column with 10 dimensions
```sql
CREATE TABLE items (id bigserial PRIMARY KEY, embedding svector(10));
```
Insert vectors
```sql
INSERT INTO items (embedding) VALUES ('(0,1),(1,2),(2,3)|10|'), ('(0,4),(1,5),(4,6)|10|');
```
Get the nearest neighbors by L2 distance
```sql
SELECT * FROM items ORDER BY embedding <-> '(0,3),(1,1),(2,2)|10|' LIMIT 5;
```
## Languages ## Languages
Use pgvector from any language with a Postgres client. You can even generate and store vectors in one language and query them in another. Use pgvector from any language with a Postgres client. You can even generate and store vectors in one language and query them in another.
@@ -678,16 +698,6 @@ SELECT extversion FROM pg_extension WHERE extname = 'vector';
## Upgrade Notes ## Upgrade Notes
### 0.6.0
If upgrading with Postgres < 13, remove this line from `sql/vector--0.5.1--0.6.0.sql`:
```sql
ALTER TYPE vector SET (STORAGE = external);
```
Then run `make install` and `ALTER EXTENSION vector UPDATE;`.
### 0.4.0 ### 0.4.0
If upgrading with Postgres < 13, remove this line from `sql/vector--0.3.2--0.4.0.sql`: If upgrading with Postgres < 13, remove this line from `sql/vector--0.3.2--0.4.0.sql`:

View File

@@ -1,5 +1,79 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION -- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.6.0'" to load this file. \quit \echo Use "ALTER EXTENSION vector UPDATE TO '0.6.0'" to load this file. \quit
-- remove this single line for Postgres < 13 CREATE TYPE svector;
ALTER TYPE vector SET (STORAGE = external);
CREATE FUNCTION svector_in(cstring, oid, integer) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_out(svector) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_typmod_in(cstring[]) RETURNS integer
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_recv(internal, oid, integer) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_send(svector) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE svector (
INPUT = svector_in,
OUTPUT = svector_out,
TYPMOD_IN = svector_typmod_in,
RECEIVE = svector_recv,
SEND = svector_send,
STORAGE = external
);
CREATE FUNCTION l2_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION jaccard_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_jaccard_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_l2_squared_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_negative_inner_product(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector(svector, integer, boolean) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_to_svector(vector, integer, boolean) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_to_vector(svector, integer, boolean) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE CAST (svector AS svector)
WITH FUNCTION svector(svector, integer, boolean) AS IMPLICIT;
CREATE CAST (svector AS vector)
WITH FUNCTION svector_to_vector(svector, integer, boolean) AS IMPLICIT;
CREATE CAST (vector AS svector)
WITH FUNCTION vector_to_svector(vector, integer, boolean) AS IMPLICIT;
CREATE OPERATOR <-> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = svector_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);

View File

@@ -26,7 +26,7 @@ CREATE TYPE vector (
TYPMOD_IN = vector_typmod_in, TYPMOD_IN = vector_typmod_in,
RECEIVE = vector_recv, RECEIVE = vector_recv,
SEND = vector_send, SEND = vector_send,
STORAGE = external STORAGE = extended
); );
-- functions -- functions
@@ -290,3 +290,92 @@ CREATE OPERATOR CLASS vector_cosine_ops
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector), FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector); FUNCTION 2 vector_norm(vector);
--- svector type
CREATE TYPE svector;
CREATE FUNCTION svector_in(cstring, oid, integer) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_out(svector) RETURNS cstring
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_typmod_in(cstring[]) RETURNS integer
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_recv(internal, oid, integer) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_send(svector) RETURNS bytea
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE TYPE svector (
INPUT = svector_in,
OUTPUT = svector_out,
TYPMOD_IN = svector_typmod_in,
RECEIVE = svector_recv,
SEND = svector_send,
STORAGE = external
);
-- svector functions
CREATE FUNCTION l2_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_l2_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION inner_product(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_inner_product' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION cosine_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_cosine_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION jaccard_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME', 'svector_jaccard_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- svector private functions
CREATE FUNCTION svector_l2_squared_distance(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_negative_inner_product(svector, svector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- svector cast functions
CREATE FUNCTION svector(svector, integer, boolean) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_to_svector(vector, integer, boolean) RETURNS svector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION svector_to_vector(svector, integer, boolean) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- svector casts
CREATE CAST (svector AS svector)
WITH FUNCTION svector(svector, integer, boolean) AS IMPLICIT;
CREATE CAST (svector AS vector)
WITH FUNCTION svector_to_vector(svector, integer, boolean) AS IMPLICIT;
CREATE CAST (vector AS svector)
WITH FUNCTION vector_to_svector(vector, integer, boolean) AS IMPLICIT;
-- svector operators
CREATE OPERATOR <-> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = l2_distance,
COMMUTATOR = '<->'
);
CREATE OPERATOR <#> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = svector_negative_inner_product,
COMMUTATOR = '<#>'
);
CREATE OPERATOR <=> (
LEFTARG = svector, RIGHTARG = svector, PROCEDURE = cosine_distance,
COMMUTATOR = '<=>'
);

View File

@@ -14,7 +14,6 @@
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
bool hnsw_enable_parallel_build;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
/* /*
@@ -40,11 +39,6 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search", 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 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* Behind a variable for now since can be slower than building in memory */
DefineCustomBoolVariable("hnsw.enable_parallel_build", "Enables or disables building indexes in parallel",
NULL, &hnsw_enable_parallel_build,
false, PGC_USERSET, 0, NULL, NULL, NULL);
} }
/* /*

View File

@@ -4,7 +4,6 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h" #include "access/generic_xlog.h"
#include "access/parallel.h"
#include "access/reloptions.h" #include "access/reloptions.h"
#include "nodes/execnodes.h" #include "nodes/execnodes.h"
#include "port.h" /* for random() */ #include "port.h" /* for random() */
@@ -15,10 +14,6 @@
#error "Requires PostgreSQL 11+" #error "Requires PostgreSQL 11+"
#endif #endif
#if PG_VERSION_NUM < 120000
#include "access/relscan.h"
#endif
#define HNSW_MAX_DIM 2000 #define HNSW_MAX_DIM 2000
/* Support functions */ /* Support functions */
@@ -64,7 +59,7 @@
#define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData)) #define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData))
#define HNSW_ELEMENT_TUPLE_SIZE(size) MAXALIGN(offsetof(HnswElementTupleData, data) + (size)) #define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData)) #define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
@@ -95,7 +90,6 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -109,7 +103,7 @@ typedef struct HnswElementData
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Datum value; Vector *vec;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
@@ -142,49 +136,6 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */ int efConstruction; /* size of dynamic candidate list */
} HnswOptions; } HnswOptions;
typedef struct HnswSpool
{
Relation heap;
Relation index;
} HnswSpool;
typedef struct HnswShared
{
/* Immutable state */
Oid heaprelid;
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
/* Worker progress */
ConditionVariable workersdonecv;
/* Mutex for mutable state */
slock_t mutex;
/* Mutable state */
int nparticipantsdone;
double reltuples;
double indtuples;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} HnswShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromHnswShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared)))
#endif
typedef struct HnswLeader
{
ParallelContext *pcxt;
int nparticipanttuplesorts;
HnswShared *hnswshared;
Snapshot snapshot;
} HnswLeader;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -212,16 +163,12 @@ typedef struct HnswBuildState
HnswElement entryPoint; HnswElement entryPoint;
double ml; double ml;
int maxLevel; int maxLevel;
long memoryLeft; double maxInMemoryElements;
bool flushed; bool flushed;
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Parallel builds */
HnswLeader *hnswleader;
HnswShared *hnswshared;
} HnswBuildState; } HnswBuildState;
typedef struct HnswMetaPageData typedef struct HnswMetaPageData
@@ -257,7 +204,7 @@ typedef struct HnswElementTupleData
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused2; uint16 unused2;
Vector data; Vector vec;
} HnswElementTupleData; } HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple; typedef HnswElementTupleData * HnswElementTuple;
@@ -342,7 +289,6 @@ void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation i
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element); 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 HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -2,16 +2,12 @@
#include <math.h> #include <math.h>
#include "access/parallel.h"
#include "access/xact.h"
#include "catalog/index.h" #include "catalog/index.h"
#include "hnsw.h" #include "hnsw.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "nodes/pg_list.h" #include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -39,23 +35,6 @@
#define UpdateProgress(index, val) ((void)val) #define UpdateProgress(index, val) ((void)val)
#endif #endif
#if PG_VERSION_NUM >= 140000
#include "utils/backend_status.h"
#include "utils/wait_event.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/table.h"
#include "optimizer/optimizer.h"
#else
#include "access/heapam.h"
#include "optimizer/planner.h"
#include "pgstat.h"
#endif
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000002)
/* /*
* Create the metapage * Create the metapage
*/ */
@@ -126,7 +105,8 @@ CreateElementPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
Size etupAllocSize; int dimensions = buildstate->dimensions;
Size etupSize;
Size maxSize; Size maxSize;
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
@@ -137,11 +117,11 @@ CreateElementPages(HnswBuildState * buildstate)
ListCell *lc; ListCell *lc;
/* Calculate sizes */ /* Calculate sizes */
etupAllocSize = BLCKSZ;
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */ /* Allocate once */
etup = palloc0(etupAllocSize); etup = palloc0(etupSize);
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
/* Prepare first page */ /* Prepare first page */
@@ -153,24 +133,15 @@ CreateElementPages(HnswBuildState * buildstate)
foreach(lc, buildstate->elements) foreach(lc, buildstate->elements)
{ {
HnswElement element = lfirst(lc); HnswElement element = lfirst(lc);
Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
/* Zero memory for each element */ HnswSetElementTuple(etup, element);
MemSet(etup, 0, etupAllocSize);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(element->value)));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
/* Initial size check */
if (etupSize > etupAllocSize)
elog(ERROR, "index tuple too large");
HnswSetElementTuple(etup, element);
/* Keep element and neighbors on the same page if possible */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize)) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
HnswBuildAppendPage(index, &buf, &page, &state, forkNum); HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
@@ -293,14 +264,13 @@ FlushPages(HnswBuildState * buildstate)
* Insert tuple * Insert tuple
*/ */
static bool static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx) InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint; HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
MemoryContext oldCtx;
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -313,9 +283,7 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
} }
/* Copy value to element so accessible outside of memory context */ /* Copy value to element so accessible outside of memory context */
oldCtx = MemoryContextSwitchTo(outerCtx); memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
@@ -345,21 +313,6 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
return *dup == NULL; return *dup == NULL;
} }
/*
* Get the memory used by an element
*/
static long
HnswElementMemory(HnswElement e, int m)
{
long elementSize = sizeof(HnswElementData);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
return elementSize;
}
/* /*
* Callback for table_index_build_scan * Callback for table_index_build_scan
*/ */
@@ -381,7 +334,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0]) if (isnull[0])
return; return;
if (buildstate->memoryLeft <= 0) if (buildstate->indtuples >= buildstate->maxInMemoryElements)
{ {
if (!buildstate->flushed) if (!buildstate->flushed)
{ {
@@ -396,18 +349,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap)) if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap))
{
if (buildstate->hnswshared)
{
HnswShared *hnswshared = buildstate->hnswshared;
SpinLockAcquire(&hnswshared->mutex);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++hnswshared->indtuples);
SpinLockRelease(&hnswshared->mutex);
}
else
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples); UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
}
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -418,12 +360,13 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Allocate necessary memory outside of memory context */ /* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel); element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
/* Use memory context since detoast can allocate */ /* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx); inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -431,21 +374,31 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
/* Add outside memory context */ /* Add outside memory context */
if (dup != NULL) if (dup != NULL)
{
HnswAddHeapTid(dup, tid); HnswAddHeapTid(dup, tid);
buildstate->memoryLeft -= sizeof(ItemPointerData);
}
/* Add to buildstate or free */ /* Add to buildstate or free */
if (inserted) if (inserted)
{
buildstate->elements = lappend(buildstate->elements, element); buildstate->elements = lappend(buildstate->elements, element);
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
}
else else
HnswFreeElement(element); HnswFreeElement(element);
} }
/*
* Get the max number of elements that fit into maintenance_work_mem
*/
static double
HnswGetMaxInMemoryElements(int m, double ml, int dimensions)
{
Size elementSize = sizeof(HnswElementData);
double avgLevel = -log(0.5) * ml;
elementSize += sizeof(HnswNeighborArray) * (avgLevel + 1);
elementSize += sizeof(HnswCandidate) * (m * (avgLevel + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VECTOR_SIZE(dimensions);
return (maintenance_work_mem * 1024L) / elementSize;
}
/* /*
* Initialize the build state * Initialize the build state
*/ */
@@ -483,7 +436,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->entryPoint = NULL; buildstate->entryPoint = NULL;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->memoryLeft = maintenance_work_mem * 1024L; buildstate->maxInMemoryElements = HnswGetMaxInMemoryElements(buildstate->m, buildstate->ml, buildstate->dimensions);
buildstate->flushed = false; buildstate->flushed = false;
/* Reuse for each tuple */ /* Reuse for each tuple */
@@ -492,9 +445,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context", "Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
buildstate->hnswleader = NULL;
buildstate->hnswshared = NULL;
} }
/* /*
@@ -507,362 +457,14 @@ FreeBuildState(HnswBuildState * buildstate)
MemoryContextDelete(buildstate->tmpCtx); MemoryContextDelete(buildstate->tmpCtx);
} }
/*
* Within leader, wait for end of heap scan
*/
static double
ParallelHeapScan(HnswBuildState * buildstate)
{
HnswShared *hnswshared = buildstate->hnswleader->hnswshared;
int nparticipanttuplesorts;
double reltuples;
nparticipanttuplesorts = buildstate->hnswleader->nparticipanttuplesorts;
for (;;)
{
SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->indtuples = hnswshared->indtuples;
reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex);
break;
}
SpinLockRelease(&hnswshared->mutex);
ConditionVariableSleep(&hnswshared->workersdonecv,
WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN);
}
ConditionVariableCancelSleep();
return reltuples;
}
/*
* Perform a worker's portion of a parallel insert
*/
static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
{
HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples;
IndexInfo *indexInfo;
/* Join parallel scan */
indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
/* TODO Support in-memory builds */
buildstate.memoryLeft = 0;
buildstate.flushed = true;
buildstate.hnswshared = hnswshared;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
ParallelTableScanFromHnswShared(hnswshared));
reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo,
true, progress, BuildCallback,
(void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Record statistics */
SpinLockAcquire(&hnswshared->mutex);
hnswshared->nparticipantsdone++;
hnswshared->reltuples += reltuples;
SpinLockRelease(&hnswshared->mutex);
/* Log statistics */
if (progress)
ereport(DEBUG1, (errmsg("leader processed " INT64_FORMAT " tuples", (int64) reltuples)));
else
ereport(DEBUG1, (errmsg("worker processed " INT64_FORMAT " tuples", (int64) reltuples)));
/* Notify leader */
ConditionVariableSignal(&hnswshared->workersdonecv);
FreeBuildState(&buildstate);
}
/*
* Perform work within a launched parallel process
*/
void
HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{
char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
LOCKMODE indexLockmode;
/* Set debug_query_string for individual workers first */
sharedquery = shm_toc_lookup(toc, PARALLEL_KEY_QUERY_TEXT, true);
debug_query_string = sharedquery;
/* Report the query string from leader */
pgstat_report_activity(STATE_RUNNING, debug_query_string);
/* Look up shared state */
hnswshared = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_SHARED, false);
/* Open relations using lock modes known to be obtained by index.c */
if (!hnswshared->isconcurrent)
{
heapLockmode = ShareLock;
indexLockmode = AccessExclusiveLock;
}
else
{
heapLockmode = ShareUpdateExclusiveLock;
indexLockmode = RowExclusiveLock;
}
/* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(hnswshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(hnswshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(hnswshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */
hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
/* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false);
/* Close relations within worker */
index_close(indexRel, indexLockmode);
#if PG_VERSION_NUM >= 120000
table_close(heapRel, heapLockmode);
#else
heap_close(heapRel, heapLockmode);
#endif
}
/*
* End parallel build
*/
static void
HnswEndParallel(HnswLeader * hnswleader)
{
/* Shutdown worker processes */
WaitForParallelWorkersToFinish(hnswleader->pcxt);
/* Free last reference to MVCC snapshot, if one was used */
if (IsMVCCSnapshot(hnswleader->snapshot))
UnregisterSnapshot(hnswleader->snapshot);
DestroyParallelContext(hnswleader->pcxt);
ExitParallelMode();
}
/*
* Return size of shared memory required for parallel index build
*/
static Size
ParallelEstimateShared(Relation heap, Snapshot snapshot)
{
#if PG_VERSION_NUM >= 120000
return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(HnswShared);
}
return add_size(offsetof(HnswShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
}
/*
* Within leader, participate as a parallel worker
*/
static void
HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
{
HnswLeader *hnswleader = buildstate->hnswleader;
HnswSpool *leaderworker;
/* Allocate memory and initialize private spool */
leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
}
/*
* Begin parallel build
*/
static void
HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
{
ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot;
Size esthnswshared;
HnswShared *hnswshared;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true;
int querylen;
#ifdef DISABLE_LEADER_PARTICIPATION
leaderparticipates = false;
#endif
/* Enter parallel mode and create context */
EnterParallelMode();
Assert(request > 0);
#if PG_VERSION_NUM >= 120000
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request;
/* Get snapshot for table scan */
if (!isconcurrent)
snapshot = SnapshotAny;
else
snapshot = RegisterSnapshot(GetTransactionSnapshot());
/* Estimate size of workspaces */
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
shm_toc_estimate_keys(&pcxt->estimator, 1);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string)
{
querylen = strlen(debug_query_string);
shm_toc_estimate_chunk(&pcxt->estimator, querylen + 1);
shm_toc_estimate_keys(&pcxt->estimator, 1);
}
else
querylen = 0; /* keep compiler quiet */
/* Everyone's had a chance to ask for space, so now create the DSM */
InitializeParallelDSM(pcxt);
/* If no DSM segment was available, back out (do serial build) */
if (pcxt->seg == NULL)
{
if (IsMVCCSnapshot(snapshot))
UnregisterSnapshot(snapshot);
DestroyParallelContext(pcxt);
ExitParallelMode();
return;
}
/* Store shared build state, for which we reserved space */
hnswshared = (HnswShared *) shm_toc_allocate(pcxt->toc, esthnswshared);
/* Initialize immutable state */
hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
hnswshared->indexrelid = RelationGetRelid(buildstate->index);
hnswshared->isconcurrent = isconcurrent;
hnswshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&hnswshared->workersdonecv);
SpinLockInit(&hnswshared->mutex);
/* Initialize mutable state */
hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0;
hnswshared->indtuples = 0;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared),
snapshot);
#else
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
/* Store query string for workers */
if (debug_query_string)
{
char *sharedquery;
sharedquery = (char *) shm_toc_allocate(pcxt->toc, querylen + 1);
memcpy(sharedquery, debug_query_string, querylen + 1);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_QUERY_TEXT, sharedquery);
}
/* Launch workers, saving status for leader/caller */
LaunchParallelWorkers(pcxt);
hnswleader->pcxt = pcxt;
hnswleader->nparticipanttuplesorts = pcxt->nworkers_launched;
if (leaderparticipates)
hnswleader->nparticipanttuplesorts++;
hnswleader->hnswshared = hnswshared;
hnswleader->snapshot = snapshot;
/* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0)
{
HnswEndParallel(hnswleader);
return;
}
/* Log participants */
ereport(DEBUG1, (errmsg("using %d parallel workers", pcxt->nworkers_launched)));
/* Save leader state now that it's clear build will be parallel */
buildstate->hnswleader = hnswleader;
/* Join heap scan ourselves */
if (leaderparticipates)
HnswLeaderParticipateAsWorker(buildstate);
/* Wait for all launched workers */
WaitForParallelWorkersToAttach(pcxt);
}
/* /*
* Build graph * Build graph
*/ */
static void static void
BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum) BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
{ {
int parallel_workers = 0;
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD); UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */
if (hnsw_enable_parallel_build)
parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to sort */
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL); true, true, BuildCallback, (void *) buildstate, NULL);
@@ -872,11 +474,6 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
#endif #endif
} }
/* End parallel build */
if (buildstate->hnswleader)
HnswEndParallel(buildstate->hnswleader);
}
/* /*
* Build the index * Build the index
*/ */

View File

@@ -123,6 +123,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Size minCombinedSize; Size minCombinedSize;
HnswElementTuple etup; HnswElementTuple etup;
BlockNumber currentPage = insertPage; BlockNumber currentPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
@@ -131,7 +132,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(e->value))); etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -404,9 +405,8 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
ItemId itemid; Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
HnswElementTuple etup; HnswElementTuple etup;
Size etupSize;
int i; int i;
/* Read page */ /* Read page */
@@ -416,9 +416,7 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
page = GenericXLogRegisterBuffer(state, buf, 0); page = GenericXLogRegisterBuffer(state, buf, 0);
/* Find space */ /* Find space */
itemid = PageGetItemId(page, dup->offno); etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, dup->offno));
etup = (HnswElementTuple) PageGetItem(page, itemid);
etupSize = ItemIdGetLength(itemid);
for (i = 0; i < HNSW_HEAPTIDS; i++) for (i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (!ItemPointerIsValid(&etup->heaptids[i])) if (!ItemPointerIsValid(&etup->heaptids[i]))
@@ -517,7 +515,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
element->value = value; element->vec = DatumGetVector(value);
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)

View File

@@ -4,7 +4,6 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "vector.h" #include "vector.h"
/* /*
@@ -188,8 +187,7 @@ HnswFreeElement(HnswElement element)
{ {
HnswFreeNeighbors(element); HnswFreeNeighbors(element);
list_free_deep(element->heaptids); list_free_deep(element->heaptids);
if (DatumGetPointer(element->value)) pfree(element->vec);
pfree(DatumGetPointer(element->value));
pfree(element); pfree(element);
} }
@@ -216,7 +214,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; element->neighbors = NULL;
element->value = PointerGetDatum(NULL); element->vec = NULL;
return element; return element;
} }
@@ -326,7 +324,7 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value))); memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
} }
/* /*
@@ -448,7 +446,10 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
} }
if (loadVec) if (loadVec)
element->value = datumCopy(PointerGetDatum(&etup->data), false, -1); {
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
} }
/* /*
@@ -475,7 +476,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -486,7 +487,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
static float static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{ {
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
} }
/* /*
@@ -749,7 +750,7 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
} }
} }
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value)); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
} }
/* /*
@@ -876,7 +877,7 @@ HnswFindDuplicate(HnswElement e)
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(e->value, neighbor->element->value, false, -1)) if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
break; break;
/* Check for space */ /* Check for space */
@@ -929,13 +930,13 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = hc->element->value; Datum q = PointerGetDatum(hc->element->vec);
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
if (DatumGetPointer(hc3->element->value) == NULL) if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
@@ -1016,7 +1017,7 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = element->value; Datum q = PointerGetDatum(element->vec);
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */ /* No neighbors if no entry point */

View File

@@ -62,8 +62,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
/* Iterate over nodes */ /* Iterate over nodes */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
ItemId itemid = PageGetItemId(page, offno); HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
int idx = 0; int idx = 0;
bool itemUpdated = false; bool itemUpdated = false;
@@ -94,7 +93,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (itemUpdated) if (itemUpdated)
{ {
Size etupSize = ItemIdGetLength(itemid); Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
/* Mark rest as invalid */ /* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++) for (int i = idx; i < HNSW_HEAPTIDS; i++)
@@ -478,8 +477,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Update element and neighbors together */ /* Update element and neighbors together */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
ItemId itemid = PageGetItemId(page, offno); HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
@@ -507,7 +505,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
continue; continue;
/* Calculate sizes */ /* Calculate sizes */
etupSize = ItemIdGetLength(itemid); etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
/* Get neighbor page */ /* Get neighbor page */
@@ -530,7 +528,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data)); MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)

705
src/svector.c Normal file
View File

@@ -0,0 +1,705 @@
#include "postgres.h"
#include <math.h>
#include "fmgr.h"
#include "libpq/pqformat.h"
#include "svector.h"
#include "utils/array.h"
#include "vector.h"
#if PG_VERSION_NUM >= 120000
#include "common/shortest_dec.h"
#include "utils/float.h"
#else
#include <float.h>
#include "utils/builtins.h"
#endif
/*
* Ensure same dimensions
*/
static inline void
CheckDims(SVector * a, SVector * b)
{
if (a->dim != b->dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("different svector dimensions %d and %d", a->dim, b->dim)));
}
/*
* Ensure expected dimensions
*/
static inline void
CheckExpectedDim(int32 typmod, int dim)
{
if (typmod != -1 && typmod != dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("expected %d dimensions, not %d", typmod, dim)));
}
/*
* Ensure valid dimensions
*/
static inline void
CheckDim(int dim)
{
if (dim < 1)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("svector must have at least 1 dimension")));
if (dim > SVECTOR_MAX_DIM)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("svector cannot have more than %d dimensions", SVECTOR_MAX_DIM)));
}
/*
* Ensure valid nnz
*/
static inline void
CheckNnz(int nnz, int dim)
{
if (nnz < 0)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("svector must have at least one element")));
if (nnz > dim)
ereport(ERROR,
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("svector cannot have more elements than dimensions")));
}
/*
* Ensure valid index
*/
static inline void
CheckIndex(int32 *indices, int i, int dim)
{
int32 index = indices[i];
if (index < 0)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("index must not be negative")));
if (index >= dim)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("index must be less than dimensions")));
if (i > 0)
{
if (index < indices[i - 1])
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("indexes must be in ascending order")));
if (index == indices[i - 1])
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("indexes must not contain duplicates")));
}
}
/*
* Ensure finite element
*/
static inline void
CheckElement(float value)
{
if (isnan(value))
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("NaN not allowed in svector")));
if (isinf(value))
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("infinite value not allowed in svector")));
}
/*
* Allocate and initialize a new sparse vector
*/
SVector *
InitSVector(int dim, int nnz)
{
SVector *result;
int size;
size = SVECTOR_SIZE(nnz);
result = (SVector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
result->nnz = nnz;
return result;
}
/*
* Convert textual representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_in);
Datum
svector_in(PG_FUNCTION_ARGS)
{
char *str = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2);
int dim;
char *pt;
SVector *result;
float *rvalues;
char *lit = pstrdup(str);
int n;
int32 *indices;
float *values;
int index;
float value;
int maxNnz;
int nnz = 0;
/* TODO Improve code and checks after deciding on format */
maxNnz = 1;
pt = str;
while (*pt != '\0')
{
if (*pt == ',')
maxNnz++;
pt++;
}
maxNnz /= 2;
indices = palloc(maxNnz * sizeof(int32));
values = palloc(maxNnz * sizeof(float));
while (sscanf(str, "(%d,%f)%n", &index, &value, &n) == 2)
{
/* TODO Better error */
if (nnz == maxNnz)
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("ran out of buffer: \"%s\"", lit)));
/* TODO Decide whether to store zero values */
indices[nnz] = index;
values[nnz] = value;
nnz++;
str += n;
if (*str == ',')
str++;
else if (*str == '|')
break;
else
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("malformed svector literal: \"%s\"", lit)));
}
if (sscanf(str, "|%d|%n", &dim, &n) != 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("malformed svector literal: \"%s\"", lit)));
str += n;
if (*str != '\0')
ereport(ERROR,
(errcode(ERRCODE_INVALID_TEXT_REPRESENTATION),
errmsg("malformed svector literal: \"%s\"", lit),
errdetail("Junk after closing pipe.")));
pfree(lit);
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitSVector(dim, nnz);
rvalues = SVECTOR_VALUES(result);
for (int i = 0; i < nnz; i++)
{
result->indices[i] = indices[i];
rvalues[i] = values[i];
CheckIndex(result->indices, i, dim);
CheckElement(rvalues[i]);
}
PG_RETURN_POINTER(result);
}
/*
* Convert internal representation to textual representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_out);
Datum
svector_out(PG_FUNCTION_ARGS)
{
SVector *svector = PG_GETARG_SVECTOR_P(0);
float *values = SVECTOR_VALUES(svector);
char *buf;
char *ptr;
int n;
/* TODO Improve code after deciding on format */
#if PG_VERSION_NUM < 120000
int ndig = FLT_DIG + extra_float_digits;
if (ndig < 1)
ndig = 1;
#define FLOAT_SHORTEST_DECIMAL_LEN (ndig + 10)
#endif
/* TODO Move */
#define APPEND_CHAR(ptr, ch) (*(ptr)++ = (ch))
/* TODO Improve */
buf = (char *) palloc((FLOAT_SHORTEST_DECIMAL_LEN + 20) * svector->nnz + 20);
ptr = buf;
for (int i = 0; i < svector->nnz; i++)
{
if (i > 0)
APPEND_CHAR(ptr, ',');
n = sprintf(ptr, "(%d,", svector->indices[i]);
ptr += n;
#if PG_VERSION_NUM >= 120000
n = float_to_shortest_decimal_bufn(values[i], ptr);
#else
n = sprintf(ptr, "%.*g", ndig, values[i]);
#endif
ptr += n;
APPEND_CHAR(ptr, ')');
}
n = sprintf(ptr, "|%d|", svector->dim);
ptr += n;
APPEND_CHAR(ptr, '\0');
PG_FREE_IF_COPY(svector, 0);
PG_RETURN_CSTRING(buf);
}
/*
* Convert type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_typmod_in);
Datum
svector_typmod_in(PG_FUNCTION_ARGS)
{
ArrayType *ta = PG_GETARG_ARRAYTYPE_P(0);
int32 *tl;
int n;
tl = ArrayGetIntegerTypmods(ta, &n);
if (n != 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("invalid type modifier")));
if (*tl < 1)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("dimensions for type svector must be at least 1")));
if (*tl > SVECTOR_MAX_DIM)
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("dimensions for type svector cannot exceed %d", SVECTOR_MAX_DIM)));
PG_RETURN_INT32(*tl);
}
/*
* Convert external binary representation to internal representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_recv);
Datum
svector_recv(PG_FUNCTION_ARGS)
{
StringInfo buf = (StringInfo) PG_GETARG_POINTER(0);
int32 typmod = PG_GETARG_INT32(2);
SVector *result;
int32 dim;
int32 nnz;
int32 unused;
float *values;
dim = pq_getmsgint(buf, sizeof(int32));
nnz = pq_getmsgint(buf, sizeof(int32));
unused = pq_getmsgint(buf, sizeof(int32));
CheckDim(dim);
CheckNnz(nnz, dim);
CheckExpectedDim(typmod, dim);
if (unused != 0)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("expected unused to be 0, not %d", unused)));
result = InitSVector(dim, nnz);
values = SVECTOR_VALUES(result);
for (int i = 0; i < nnz; i++)
{
result->indices[i] = pq_getmsgint(buf, sizeof(int32));
CheckIndex(result->indices, i, dim);
}
for (int i = 0; i < nnz; i++)
{
values[i] = pq_getmsgfloat4(buf);
CheckElement(values[i]);
}
PG_RETURN_POINTER(result);
}
/*
* Convert internal representation to the external binary representation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_send);
Datum
svector_send(PG_FUNCTION_ARGS)
{
SVector *svec = PG_GETARG_SVECTOR_P(0);
float *values = SVECTOR_VALUES(svec);
StringInfoData buf;
pq_begintypsend(&buf);
pq_sendint(&buf, svec->dim, sizeof(int32));
pq_sendint(&buf, svec->nnz, sizeof(int32));
pq_sendint(&buf, svec->unused, sizeof(int32));
for (int i = 0; i < svec->nnz; i++)
pq_sendint(&buf, svec->indices[i], sizeof(int32));
for (int i = 0; i < svec->nnz; i++)
pq_sendfloat4(&buf, values[i]);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
}
/*
* Convert sparse vector to sparse vector
* This is needed to check the type modifier
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector);
Datum
svector(PG_FUNCTION_ARGS)
{
SVector *svec = PG_GETARG_SVECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1);
CheckExpectedDim(typmod, svec->dim);
PG_RETURN_POINTER(svec);
}
/*
* Convert dense vector to sparse vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_to_svector);
Datum
vector_to_svector(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1);
SVector *result;
int dim = vec->dim;
int nnz = 0;
float *values;
int j = 0;
CheckDim(dim);
CheckExpectedDim(typmod, dim);
for (int i = 0; i < dim; i++)
{
if (vec->x[i] != 0)
nnz++;
}
result = InitSVector(dim, nnz);
values = SVECTOR_VALUES(result);
for (int i = 0; i < dim; i++)
{
if (vec->x[i] != 0)
{
/* Safety check */
if (j == nnz)
elog(ERROR, "safety check failed");
result->indices[j] = i;
values[j] = vec->x[i];
j++;
}
}
PG_RETURN_POINTER(result);
}
/*
* Get the L2 squared distance between sparse vectors
*/
static double
l2_distance_squared_internal(SVector * a, SVector * b)
{
float *ax = SVECTOR_VALUES(a);
float *bx = SVECTOR_VALUES(b);
double distance = 0.0;
int bpos = 0;
for (int i = 0; i < a->nnz; i++)
{
int ai = a->indices[i];
int bi = -1;
for (int j = bpos; j < b->nnz; j++)
{
bi = b->indices[j];
if (ai == bi)
{
double diff = ax[i] - bx[j];
distance += diff * diff;
}
else if (ai > bi)
distance += bx[j] * bx[j];
/* Update start for next iteration */
if (ai >= bi)
bpos = j + 1;
/* Found or passed it */
if (bi >= ai)
break;
}
if (ai != bi)
distance += ax[i] * ax[i];
}
for (int j = bpos; j < b->nnz; j++)
distance += bx[j] * bx[j];
return distance;
}
/*
* Get the L2 distance between sparse vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_l2_distance);
Datum
svector_l2_distance(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(sqrt(l2_distance_squared_internal(a, b)));
}
/*
* Get the L2 squared distance between sparse vectors
* This saves a sqrt calculation
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_l2_squared_distance);
Datum
svector_l2_squared_distance(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(l2_distance_squared_internal(a, b));
}
/*
* Get the inner product of two sparse vectors
*/
static double
inner_product_internal(SVector * a, SVector * b)
{
float *ax = SVECTOR_VALUES(a);
float *bx = SVECTOR_VALUES(b);
double distance = 0.0;
int bpos = 0;
for (int i = 0; i < a->nnz; i++)
{
int ai = a->indices[i];
for (int j = bpos; j < b->nnz; j++)
{
int bi = b->indices[j];
/* Only update when the same index */
if (ai == bi)
distance += ax[i] * bx[j];
/* Update start for next iteration */
if (ai >= bi)
bpos = j + 1;
/* Found or passed it */
if (bi >= ai)
break;
}
}
return distance;
}
/*
* Get the inner product of two sparse vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_inner_product);
Datum
svector_inner_product(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(inner_product_internal(a, b));
}
/*
* Get the negative inner product of two sparse vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_negative_inner_product);
Datum
svector_negative_inner_product(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
CheckDims(a, b);
PG_RETURN_FLOAT8(-inner_product_internal(a, b));
}
/*
* Get the cosine distance between two sparse vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_cosine_distance);
Datum
svector_cosine_distance(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
float *ax = SVECTOR_VALUES(a);
float *bx = SVECTOR_VALUES(b);
float norma = 0.0;
float normb = 0.0;
double similarity;
CheckDims(a, b);
similarity = inner_product_internal(a, b);
/* Auto-vectorized */
for (int i = 0; i < a->nnz; i++)
norma += ax[i] * ax[i];
/* Auto-vectorized */
for (int i = 0; i < b->nnz; i++)
normb += bx[i] * bx[i];
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
similarity /= sqrt((double) norma * (double) normb);
#ifdef _MSC_VER
/* /fp:fast may not propagate NaN */
if (isnan(similarity))
PG_RETURN_FLOAT8(NAN);
#endif
/* Keep in range */
if (similarity > 1)
similarity = 1.0;
else if (similarity < -1)
similarity = -1.0;
PG_RETURN_FLOAT8(1.0 - similarity);
}
/*
* Get the weighted Jaccard distance between two sparse vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_jaccard_distance);
Datum
svector_jaccard_distance(PG_FUNCTION_ARGS)
{
SVector *a = PG_GETARG_SVECTOR_P(0);
SVector *b = PG_GETARG_SVECTOR_P(1);
float *ax = SVECTOR_VALUES(a);
float *bx = SVECTOR_VALUES(b);
double num = 0.0;
double denom = 0.0;
int bpos = 0;
CheckDims(a, b);
/*
* Weighted Jaccard distance is not defined for vectors with negative
* values. Could check and return NaN if minimal impact on performance.
*/
for (int i = 0; i < a->nnz; i++)
{
int ai = a->indices[i];
int bi = -1;
for (int j = bpos; j < b->nnz; j++)
{
bi = b->indices[j];
if (ai == bi)
{
num += ax[i] < bx[j] ? ax[i] : bx[j];
denom += ax[i] > bx[j] ? ax[i] : bx[j];
}
else if (ai > bi)
denom += bx[j];
/* Update start for next iteration */
if (ai >= bi)
bpos = j + 1;
/* Found or passed it */
if (bi >= ai)
break;
}
if (ai != bi)
denom += ax[i];
}
for (int j = bpos; j < b->nnz; j++)
denom += bx[j];
if (denom > 0)
PG_RETURN_FLOAT8(1.0 - (num / denom));
else
PG_RETURN_FLOAT8(NAN);
}

23
src/svector.h Normal file
View File

@@ -0,0 +1,23 @@
#ifndef SVECTOR_H
#define SVECTOR_H
#define SVECTOR_MAX_DIM 100000
#define SVECTOR_SIZE(_nnz) (offsetof(SVector, indices) + (_nnz) * sizeof(int32) + (_nnz * sizeof(float)))
#define SVECTOR_VALUES(x) ((float *) (((char *) (x)) + offsetof(SVector, indices) + (x)->nnz * sizeof(int32)))
#define DatumGetSVector(x) ((SVector *) PG_DETOAST_DATUM(x))
#define PG_GETARG_SVECTOR_P(x) DatumGetSVector(PG_GETARG_DATUM(x))
#define PG_RETURN_SVECTOR_P(x) PG_RETURN_POINTER(x)
typedef struct SVector
{
int32 vl_len_; /* varlena header (do not touch directly!) */
int32 dim; /* number of dimensions */
int32 nnz;
int32 unused;
int32 indices[FLEXIBLE_ARRAY_MEMBER];
} SVector;
SVector *InitSVector(int dim, int nnz);
#endif

View File

@@ -9,6 +9,7 @@
#include "lib/stringinfo.h" #include "lib/stringinfo.h"
#include "libpq/pqformat.h" #include "libpq/pqformat.h"
#include "port.h" /* for strtof() */ #include "port.h" /* for strtof() */
#include "svector.h"
#include "utils/array.h" #include "utils/array.h"
#include "utils/builtins.h" #include "utils/builtins.h"
#include "utils/lsyscache.h" #include "utils/lsyscache.h"
@@ -1151,3 +1152,26 @@ vector_avg(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result); PG_RETURN_POINTER(result);
} }
/*
* Convert sparse vector to dense vector
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(svector_to_vector);
Datum
svector_to_vector(PG_FUNCTION_ARGS)
{
SVector *svec = PG_GETARG_SVECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1);
Vector *result;
int dim = svec->dim;
float *values = SVECTOR_VALUES(svec);
CheckDim(dim);
CheckExpectedDim(typmod, dim);
result = InitVector(dim);
for (int i = 0; i < svec->nnz; i++)
result->x[svec->indices[i]] = values[i];
PG_RETURN_POINTER(result);
}

View File

@@ -54,85 +54,85 @@ SELECT vector_norm('[3e37,4e37]')::real;
5e+37 5e+37
(1 row) (1 row)
SELECT l2_distance('[0,0]', '[3,4]'); SELECT l2_distance('[0,0]'::vector, '[3,4]');
l2_distance l2_distance
------------- -------------
5 5
(1 row) (1 row)
SELECT l2_distance('[0,0]', '[0,1]'); SELECT l2_distance('[0,0]'::vector, '[0,1]');
l2_distance l2_distance
------------- -------------
1 1
(1 row) (1 row)
SELECT l2_distance('[1,2]', '[3]'); SELECT l2_distance('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT l2_distance('[3e38]', '[-3e38]'); SELECT l2_distance('[3e38]'::vector, '[-3e38]');
l2_distance l2_distance
------------- -------------
Infinity Infinity
(1 row) (1 row)
SELECT inner_product('[1,2]', '[3,4]'); SELECT inner_product('[1,2]'::vector, '[3,4]');
inner_product inner_product
--------------- ---------------
11 11
(1 row) (1 row)
SELECT inner_product('[1,2]', '[3]'); SELECT inner_product('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT inner_product('[3e38]', '[3e38]'); SELECT inner_product('[3e38]'::vector, '[3e38]');
inner_product inner_product
--------------- ---------------
Infinity Infinity
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[2,4]'); SELECT cosine_distance('[1,2]'::vector, '[2,4]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[0,0]'); SELECT cosine_distance('[1,2]'::vector, '[0,0]');
cosine_distance cosine_distance
----------------- -----------------
NaN NaN
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[1,1]'); SELECT cosine_distance('[1,1]'::vector, '[1,1]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,0]', '[0,2]'); SELECT cosine_distance('[1,0]'::vector, '[0,2]');
cosine_distance cosine_distance
----------------- -----------------
1 1
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[-1,-1]'); SELECT cosine_distance('[1,1]'::vector, '[-1,-1]');
cosine_distance cosine_distance
----------------- -----------------
2 2
(1 row) (1 row)
SELECT cosine_distance('[1,2]', '[3]'); SELECT cosine_distance('[1,2]'::vector, '[3]');
ERROR: different vector dimensions 2 and 1 ERROR: different vector dimensions 2 and 1
SELECT cosine_distance('[1,1]', '[1.1,1.1]'); SELECT cosine_distance('[1,1]'::vector, '[1.1,1.1]');
cosine_distance cosine_distance
----------------- -----------------
0 0
(1 row) (1 row)
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]'); SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]');
cosine_distance cosine_distance
----------------- -----------------
2 2
(1 row) (1 row)
SELECT cosine_distance('[3e38]', '[3e38]'); SELECT cosine_distance('[3e38]'::vector, '[3e38]');
cosine_distance cosine_distance
----------------- -----------------
NaN NaN

140
test/expected/svector.out Normal file
View File

@@ -0,0 +1,140 @@
SELECT '(0,1.5),(2,3.5)|5|'::svector;
svector
--------------------
(0,1.5),(2,3.5)|5|
(1 row)
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector;
vector
-----------------
[1.5,0,3.5,0,0]
(1 row)
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(5);
vector
-----------------
[1.5,0,3.5,0,0]
(1 row)
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(4);
ERROR: expected 4 dimensions, not 5
SELECT '[0,1.5,0,3.5,0]'::vector::svector;
svector
--------------------
(1,1.5),(3,3.5)|5|
(1 row)
SELECT '(0,0),(1,1),(2,0)|3|'::svector;
svector
----------------------
(0,0),(1,1),(2,0)|3|
(1 row)
SELECT '|5|'::svector;
svector
---------
|5|
(1 row)
SELECT '|-1|'::svector;
ERROR: svector must have at least 1 dimension
LINE 1: SELECT '|-1|'::svector;
^
SELECT '|100001|'::svector;
ERROR: svector cannot have more than 100000 dimensions
LINE 1: SELECT '|100001|'::svector;
^
SELECT '|16001|'::svector::vector;
ERROR: vector cannot have more than 16000 dimensions
SELECT '(-1,1)|1|'::svector;
ERROR: index must not be negative
LINE 1: SELECT '(-1,1)|1|'::svector;
^
SELECT '(1,1)|1|'::svector;
ERROR: index must be less than dimensions
LINE 1: SELECT '(1,1)|1|'::svector;
^
SELECT '|1|'::svector(2);
ERROR: expected 2 dimensions, not 1
SELECT l2_distance('|2|'::svector, '(0,3),(1,4)|2|');
l2_distance
-------------
5
(1 row)
SELECT l2_distance('|2|'::svector, '(1,1)|2|');
l2_distance
-------------
1
(1 row)
SELECT '|2|'::svector <-> '(0,3),(1,4)|2|';
?column?
----------
5
(1 row)
SELECT inner_product('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
inner_product
---------------
10
(1 row)
SELECT svector_negative_inner_product('(0,1),(1,2)|2|', '(0,2),(1,4)|2|');
svector_negative_inner_product
--------------------------------
-10
(1 row)
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '|2|');
cosine_distance
-----------------
NaN
(1 row)
SELECT cosine_distance('(0,1),(1,1)|2|'::svector, '(0,-1),(1,-1)|2|');
cosine_distance
-----------------
2
(1 row)
SELECT cosine_distance('(0,1)|2|'::svector, '(1,2)|2|');
cosine_distance
-----------------
1
(1 row)
SELECT cosine_distance('|1|'::svector, '|1|');
cosine_distance
-----------------
NaN
(1 row)
SELECT cosine_distance('(0,1)|2|'::svector, '(0,1)|3|');
ERROR: different svector dimensions 2 and 3
SELECT jaccard_distance('(0,1)|2|', '(0,1)|2|');
jaccard_distance
------------------
0
(1 row)
SELECT jaccard_distance('(0,1)|2|', '(1,1)|2|');
jaccard_distance
------------------
1
(1 row)
SELECT jaccard_distance('|1|', '|1|');
jaccard_distance
------------------
NaN
(1 row)
SELECT jaccard_distance('(0,1)|2|', '(0,1)|3|');
ERROR: different svector dimensions 2 and 3

View File

@@ -13,24 +13,24 @@ SELECT vector_norm('[3,4]');
SELECT vector_norm('[0,1]'); SELECT vector_norm('[0,1]');
SELECT vector_norm('[3e37,4e37]')::real; SELECT vector_norm('[3e37,4e37]')::real;
SELECT l2_distance('[0,0]', '[3,4]'); SELECT l2_distance('[0,0]'::vector, '[3,4]');
SELECT l2_distance('[0,0]', '[0,1]'); SELECT l2_distance('[0,0]'::vector, '[0,1]');
SELECT l2_distance('[1,2]', '[3]'); SELECT l2_distance('[1,2]'::vector, '[3]');
SELECT l2_distance('[3e38]', '[-3e38]'); SELECT l2_distance('[3e38]'::vector, '[-3e38]');
SELECT inner_product('[1,2]', '[3,4]'); SELECT inner_product('[1,2]'::vector, '[3,4]');
SELECT inner_product('[1,2]', '[3]'); SELECT inner_product('[1,2]'::vector, '[3]');
SELECT inner_product('[3e38]', '[3e38]'); SELECT inner_product('[3e38]'::vector, '[3e38]');
SELECT cosine_distance('[1,2]', '[2,4]'); SELECT cosine_distance('[1,2]'::vector, '[2,4]');
SELECT cosine_distance('[1,2]', '[0,0]'); SELECT cosine_distance('[1,2]'::vector, '[0,0]');
SELECT cosine_distance('[1,1]', '[1,1]'); SELECT cosine_distance('[1,1]'::vector, '[1,1]');
SELECT cosine_distance('[1,0]', '[0,2]'); SELECT cosine_distance('[1,0]'::vector, '[0,2]');
SELECT cosine_distance('[1,1]', '[-1,-1]'); SELECT cosine_distance('[1,1]'::vector, '[-1,-1]');
SELECT cosine_distance('[1,2]', '[3]'); SELECT cosine_distance('[1,2]'::vector, '[3]');
SELECT cosine_distance('[1,1]', '[1.1,1.1]'); SELECT cosine_distance('[1,1]'::vector, '[1.1,1.1]');
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]'); SELECT cosine_distance('[1,1]'::vector, '[-1.1,-1.1]');
SELECT cosine_distance('[3e38]', '[3e38]'); SELECT cosine_distance('[3e38]'::vector, '[3e38]');
SELECT l1_distance('[0,0]', '[3,4]'); SELECT l1_distance('[0,0]', '[3,4]');
SELECT l1_distance('[0,0]', '[0,1]'); SELECT l1_distance('[0,0]', '[0,1]');

36
test/sql/svector.sql Normal file
View File

@@ -0,0 +1,36 @@
SELECT '(0,1.5),(2,3.5)|5|'::svector;
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector;
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(5);
SELECT '(0,1.5),(2,3.5)|5|'::svector::vector(4);
SELECT '[0,1.5,0,3.5,0]'::vector::svector;
SELECT '(0,0),(1,1),(2,0)|3|'::svector;
SELECT '|5|'::svector;
SELECT '|-1|'::svector;
SELECT '|100001|'::svector;
SELECT '|16001|'::svector::vector;
SELECT '(-1,1)|1|'::svector;
SELECT '(1,1)|1|'::svector;
SELECT '|1|'::svector(2);
SELECT l2_distance('|2|'::svector, '(0,3),(1,4)|2|');
SELECT l2_distance('|2|'::svector, '(1,1)|2|');
SELECT '|2|'::svector <-> '(0,3),(1,4)|2|';
SELECT inner_product('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
SELECT svector_negative_inner_product('(0,1),(1,2)|2|', '(0,2),(1,4)|2|');
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '(0,2),(1,4)|2|');
SELECT cosine_distance('(0,1),(1,2)|2|'::svector, '|2|');
SELECT cosine_distance('(0,1),(1,1)|2|'::svector, '(0,-1),(1,-1)|2|');
SELECT cosine_distance('(0,1)|2|'::svector, '(1,2)|2|');
SELECT cosine_distance('|1|'::svector, '|1|');
SELECT cosine_distance('(0,1)|2|'::svector, '(0,1)|3|');
SELECT jaccard_distance('(0,1)|2|', '(0,1)|2|');
SELECT jaccard_distance('(0,1)|2|', '(1,1)|2|');
SELECT jaccard_distance('|1|', '|1|');
SELECT jaccard_distance('(0,1)|2|', '(0,1)|3|');

View File

@@ -83,32 +83,11 @@ for my $i (0 .. $#operators)
push(@expected, $res); push(@expected, $res);
} }
# Build index serially # Add index
$node->safe_psql("postgres", qq( $node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
SET max_parallel_maintenance_workers = 0;
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
# Test approximate results
my $min = $operator eq "<#>" ? 0.80 : 0.99; my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator); test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET client_min_messages = DEBUG;
SET min_parallel_table_scan_size = 1;
SET hnsw.enable_parallel_build = on;
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
# Test approximate results
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
} }
done_testing(); done_testing();