Compare commits

..

35 Commits

Author SHA1 Message Date
Andrew Kane
7a347bc2de Added HnswGetElementTupleSize method 2024-10-09 19:38:16 -07:00
Andrew Kane
67f9a3e61c Init collation as well [skip ci] 2024-10-09 19:30:04 -07:00
Andrew Kane
3ccfab8f92 Added support for inline filtering with HNSW 2024-10-09 19:02:40 -07:00
Andrew Kane
3126fbdb6f Use double for distance [skip ci] 2024-10-09 17:04:25 -07:00
Andrew Kane
f4b67b078f DRY HNSW distance calculations 2024-10-09 17:01:49 -07:00
Andrew Kane
77688b4309 Improve total cost for cost estimation (#686) 2024-10-08 12:42:03 -07:00
Andrew Kane
d5f4a0e435 Fixed memory context leak in HnswUpdateNeighborsOnDisk - fixes #692 2024-10-08 12:21:26 -07:00
Andrew Kane
57248ba128 Use separate memory context for updating neighbors, which improves performance around 10% for larger vectors 2024-09-30 11:15:27 -07:00
Andrew Kane
ff6da4fcea Moved logic to get update neighbor on disk to separate function 2024-09-30 10:30:01 -07:00
Andrew Kane
a8b4b6675a Moved logic to get update index to separate function 2024-09-30 10:14:52 -07:00
Andrew Kane
d148b4e61b Fixed insert logic 2024-09-30 09:59:12 -07:00
Andrew Kane
658d74e2f6 Use Size for memory [skip ci] 2024-09-29 23:48:58 -07:00
Andrew Kane
7ba593c492 Improved SelectNeighbors signature [skip ci] 2024-09-29 23:03:02 -07:00
Andrew Kane
525e3b81e1 Improved HnswUpdateConnection parameters [skip ci] 2024-09-29 19:47:25 -07:00
Andrew Kane
8eb8cdf0f3 Moved insert-specific code to hnswinsert.c 2024-09-29 19:44:11 -07:00
Andrew Kane
4c72f91206 Improved variable name [skip ci] 2024-09-29 19:26:15 -07:00
Andrew Kane
4ac86f62a1 Improved variable names [skip ci] 2024-09-29 19:22:35 -07:00
Andrew Kane
648dd8af78 Moved LoadElementsForInsert to separate function and removed unused code path 2024-09-29 19:12:38 -07:00
Andrew Kane
ee43ee9b16 Use HnswLoadNeighborTids for inserts 2024-09-29 18:52:12 -07:00
Andrew Kane
5ce367e18b Removed lc from HnswUpdateConnection [skip ci] 2024-09-29 18:18:42 -07:00
Andrew Kane
f371eb119b Removed lc from SelectNeighbors [skip ci] 2024-09-29 18:14:28 -07:00
Andrew Kane
382a25aefb Split loading neighbor TIDs into separate function [skip ci] 2024-09-29 17:20:54 -07:00
Andrew Kane
0b6214aad6 Moved HnswLoadNeighbors to hnswinsert.c [skip ci] 2024-09-29 15:49:01 -07:00
Andrew Kane
f2afd11257 Use sc for search candidates [skip ci] 2024-09-29 15:09:54 -07:00
Andrew Kane
cae3458329 Updated distance to use double 2024-09-29 15:06:50 -07:00
Andrew Kane
dc23752618 Fixed uninitialized variable [skip ci] 2024-09-28 19:18:52 -07:00
Andrew Kane
158d9340bc Added distance filters to cost tests [skip ci] 2024-09-28 14:50:23 -07:00
Andrew Kane
5ee0471ead Updated readme [skip ci] 2024-09-28 09:23:41 -07:00
Andrew Kane
54f8d9733d Updated default Postgres version in Dockerfile [skip ci] 2024-09-27 16:19:57 -07:00
Andrew Kane
cf419f448b Updated Postgres version for Docker [skip ci] 2024-09-27 16:19:06 -07:00
Andrew Kane
8a2eebd6a4 Added note about Postgres 17 on Windows - #669 [skip ci] 2024-09-27 14:05:36 -07:00
Andrew Kane
daf9c5c743 Updated package versions in readme [skip ci] 2024-09-27 13:52:11 -07:00
Andrew Kane
2bca4e406b Restored quarterly package version for FreeBSD in readme [skip ci] 2024-09-27 13:50:57 -07:00
Andrew Kane
74020a90da Updated package versions in readme [skip ci] 2024-09-27 13:49:43 -07:00
Andrew Kane
44d8d28b40 Added note about postgresql@17 formula [skip ci] 2024-09-27 13:39:54 -07:00
20 changed files with 964 additions and 974 deletions

View File

@@ -1,8 +1,9 @@
## 0.8.0 (unreleased)
- Added support for iterative index scans
- Added support for inline filtering with HNSW
- Added casts for arrays to `sparsevec`
- Improved cost estimation
- Improved performance of HNSW inserts and on-disk index builds
- Reduced memory usage for HNSW index scans
- Dropped support for Postgres 12

View File

@@ -1,4 +1,4 @@
ARG PG_MAJOR=16
ARG PG_MAJOR=17
FROM postgres:$PG_MAJOR
ARG PG_MAJOR

View File

@@ -66,7 +66,7 @@ dist:
git archive --format zip --prefix=$(EXTENSION)-$(EXTVERSION)/ --output dist/$(EXTENSION)-$(EXTVERSION).zip master
# for Docker
PG_MAJOR ?= 16
PG_MAJOR ?= 17
.PHONY: docker

102
README.md
View File

@@ -52,6 +52,8 @@ nmake /F Makefile.win
nmake /F Makefile.win install
```
Note: Postgres 17 is not supported yet due to an upstream issue
See the [installation notes](#installation-notes---windows) if you run into issues
You can also install it with [Docker](#docker) or [conda-forge](#conda-forge).
@@ -100,6 +102,8 @@ Or add a vector column to an existing table
ALTER TABLE items ADD COLUMN embedding vector(3);
```
Also supports [half-precision](#half-precision-vectors), [binary](#binary-vectors), and [sparse](#sparse-vectors) vectors
Insert vectors
```sql
@@ -145,6 +149,8 @@ Supported distance functions are:
- `<#>` - (negative) inner product
- `<=>` - cosine distance
- `<+>` - L1 distance (added in 0.7.0)
- `<~>` - Hamming distance (binary vectors, added in 0.7.0)
- `<%>` - Jaccard distance (binary vectors, added in 0.7.0)
Get the nearest neighbors to a row
@@ -433,6 +439,12 @@ Create an index on one [or more](https://www.postgresql.org/docs/current/indexes
CREATE INDEX ON items (category_id);
```
Or a composite HNSW index for approximate search (added in 0.8.0)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops, category_id);
```
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql
@@ -445,63 +457,6 @@ Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html
CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id);
```
## Streaming Queries [unreleased]
*Added in 0.8.0*
With approximate indexes, you can end up with less results than expected due to filtering conditions in the query.
Starting with 0.8.0, you can enable streaming queries. If too few results from the initial index scan match the query filters, it will resume scanning until enough results are found. This can significantly improve recall (at the cost of speed).
```tsql
SET hnsw.streaming = on;
-- or
SET ivfflat.streaming = on;
```
### Streaming Options
Since scanning a large portion of the index is expensive, there are options to control when the scan ends.
#### HNSW
Specify the max number of additional tuples visited
```sql
SET hnsw.ef_stream = 10000;
```
The scan will also end if reaches `work_mem`, at which point a notice is shown
```text
NOTICE: hnsw index scan exceeded work_mem after 50000 tuples
HINT: Increase work_mem to scan more tuples.
```
Adjust this with:
```sql
SET work_mem = '8MB';
```
#### IVFFlat
Specify the max number of probes
```sql
SET ivfflat.max_probes = 100;
```
### Streaming Order
With streaming queries, its possible for rows to be slightly out of order by distance. For strict ordering, use:
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order ORDER BY distance;
```
## Half-Precision Vectors
*Added in 0.7.0*
@@ -1040,7 +995,7 @@ l2_normalize(sparsevec) → sparsevec | Normalize with Euclidean norm | 0.7.0
If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with:
```sh
export PG_CONFIG=/Library/PostgreSQL/16/bin/pg_config
export PG_CONFIG=/Library/PostgreSQL/17/bin/pg_config
```
Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use:
@@ -1051,11 +1006,11 @@ sudo --preserve-env=PG_CONFIG make install
A few common paths on Mac are:
- EDB installer - `/Library/PostgreSQL/16/bin/pg_config`
- Homebrew (arm64) - `/opt/homebrew/opt/postgresql@16/bin/pg_config`
- Homebrew (x86-64) - `/usr/local/opt/postgresql@16/bin/pg_config`
- EDB installer - `/Library/PostgreSQL/17/bin/pg_config`
- Homebrew (arm64) - `/opt/homebrew/opt/postgresql@17/bin/pg_config`
- Homebrew (x86-64) - `/usr/local/opt/postgresql@17/bin/pg_config`
Note: Replace `16` with your Postgres server version
Note: Replace `17` with your Postgres server version
### Missing Header
@@ -1064,10 +1019,10 @@ If compilation fails with `fatal error: postgres.h: No such file or directory`,
For Ubuntu and Debian, use:
```sh
sudo apt install postgresql-server-dev-16
sudo apt install postgresql-server-dev-17
```
Note: Replace `16` with your Postgres server version
Note: Replace `17` with your Postgres server version
### Missing SDK
@@ -1100,17 +1055,17 @@ If installation fails with `Access is denied`, re-run the installation instructi
Get the [Docker image](https://hub.docker.com/r/pgvector/pgvector) with:
```sh
docker pull pgvector/pgvector:pg16
docker pull pgvector/pgvector:pg17
```
This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (replace `16` with your Postgres server version, and run it the same way).
This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (replace `17` with your Postgres server version, and run it the same way).
You can also build the image manually:
```sh
git clone --branch v0.7.4 https://github.com/pgvector/pgvector.git
cd pgvector
docker build --pull --build-arg PG_MAJOR=16 -t myuser/pgvector .
docker build --pull --build-arg PG_MAJOR=17 -t myuser/pgvector .
```
### Homebrew
@@ -1121,7 +1076,7 @@ With Homebrew Postgres, you can use:
brew install pgvector
```
Note: This only adds it to the `postgresql@14` formula
Note: This only adds it to the `postgresql@17` and `postgresql@14` formulas
### PGXN
@@ -1136,22 +1091,22 @@ pgxn install vector
Debian and Ubuntu packages are available from the [PostgreSQL APT Repository](https://wiki.postgresql.org/wiki/Apt). Follow the [setup instructions](https://wiki.postgresql.org/wiki/Apt#Quickstart) and run:
```sh
sudo apt install postgresql-16-pgvector
sudo apt install postgresql-17-pgvector
```
Note: Replace `16` with your Postgres server version
Note: Replace `17` with your Postgres server version
### Yum
RPM packages are available from the [PostgreSQL Yum Repository](https://yum.postgresql.org/). Follow the [setup instructions](https://www.postgresql.org/download/linux/redhat/) for your distribution and run:
```sh
sudo yum install pgvector_16
sudo yum install pgvector_17
# or
sudo dnf install pgvector_16
sudo dnf install pgvector_17
```
Note: Replace `16` with your Postgres server version
Note: Replace `17` with your Postgres server version
### pkg
@@ -1240,6 +1195,7 @@ Thanks to:
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
- [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
- [HQANN: Efficient and Robust Similarity Search for Hybrid Queries with Structured and Unstructured Constraints](https://arxiv.org/pdf/2207.07940.pdf)
## History

View File

@@ -24,3 +24,11 @@ CREATE CAST (double precision[] AS sparsevec)
CREATE CAST (numeric[] AS sparsevec)
WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT;
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -916,3 +916,13 @@ CREATE OPERATOR CLASS sparsevec_l1_ops
OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(sparsevec, sparsevec),
FUNCTION 3 hnsw_sparsevec_support(internal);
-- hnsw attributes
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -19,8 +19,6 @@
#endif
int hnsw_ef_search;
int hnsw_ef_stream;
bool hnsw_streaming;
int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind;
@@ -71,17 +69,6 @@ HnswInit(void)
"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);
/* TODO Figure out name */
DefineCustomBoolVariable("hnsw.streaming", "Use streaming mode",
NULL, &hnsw_streaming,
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Figure out name */
/* TODO Use same value as ivfflat.max_probes for "all" */
DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search",
"-1 means all", &hnsw_ef_stream,
HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw");
}
@@ -102,33 +89,6 @@ hnswbuildphasename(int64 phasenum)
}
}
/*
* Estimate ef needed for iterative scans
*/
static int
EstimateEf(PlannerInfo *root, IndexPath *path)
{
double selectivity = 1;
ListCell *lc;
/* Cannot estimate without limit */
/* limit_tuples includes offset */
if (root->limit_tuples < 0)
return 0;
/* Get the selectivity of non-index conditions */
foreach(lc, path->indexinfo->indrestrictinfo)
{
RestrictInfo *rinfo = lfirst(lc);
/* Skip DEFAULT_INEQ_SEL since it may be a distance filter */
if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
selectivity *= rinfo->norm_selec;
}
return root->limit_tuples / Max(selectivity, 0.00001);
}
/*
* Estimate the cost of an index scan
*/
@@ -140,11 +100,8 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
{
GenericCosts costs;
int m;
int ef;
int entryLevel;
int layer0TuplesMax;
double layer0Selectivity;
double scalingFactor = 0.55;
double ratio;
double startupPages;
double spc_seq_page_cost;
Relation index;
@@ -161,12 +118,12 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
MemSet(&costs, 0, sizeof(costs));
genericcostestimate(root, path, loop_count, &costs);
index = index_open(path->indexinfo->indexoid, NoLock);
HnswGetMetaPageInfo(index, &m, NULL);
index_close(index, NoLock);
ef = hnsw_streaming ? Max(hnsw_ef_search, EstimateEf(root, path)) : hnsw_ef_search;
/*
* HNSW cost estimation follows a formula that accounts for the total
* number of tuples indexed combined with the parameters that most
@@ -194,30 +151,38 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
* at L0, accounting for previously visited tuples, multiplied by the
* "scalingFactor" (currently hardcoded).
*/
entryLevel = (int) (log(path->indexinfo->tuples + 1) * HnswGetMl(m));
layer0TuplesMax = HnswGetLayerM(m, 0) * ef;
layer0Selectivity = (scalingFactor * log(path->indexinfo->tuples + 1)) /
(log(m) * (1 + log(ef)));
if (path->indexinfo->tuples > 0)
{
double scalingFactor = 0.55;
int entryLevel = (int) (log(path->indexinfo->tuples) * HnswGetMl(m));
int layer0TuplesMax = HnswGetLayerM(m, 0) * hnsw_ef_search;
double layer0Selectivity = scalingFactor * log(path->indexinfo->tuples) / (log(m) * (1 + log(hnsw_ef_search)));
costs.numIndexTuples = (entryLevel * m) +
(layer0TuplesMax * layer0Selectivity);
ratio = (entryLevel * m + layer0TuplesMax * layer0Selectivity) / path->indexinfo->tuples;
genericcostestimate(root, path, loop_count, &costs);
if (ratio > 1)
ratio = 1;
}
else
ratio = 1;
get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
/* Startup cost is cost before returning the first row */
costs.indexStartupCost = costs.indexTotalCost * ratio;
/* Adjust cost if needed since TOAST not included in seq scan cost */
if (costs.numIndexPages > path->indexinfo->rel->pages)
startupPages = costs.numIndexPages * ratio;
if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
{
/* Change all page cost from random to sequential */
costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
costs.indexStartupCost -= startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
/* Remove cost of extra pages */
costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
}
/* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost;
*indexStartupCost = costs.indexStartupCost;
*indexTotalCost = costs.indexTotalCost;
*indexSelectivity = costs.indexSelectivity;
*indexCorrelation = costs.indexCorrelation;
@@ -262,13 +227,13 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 3;
amroutine->amsupport = 4;
amroutine->amoptsprocnum = 0;
amroutine->amcanorder = false;
amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false;
amroutine->amcanmulticol = false;
amroutine->amcanmulticol = true;
amroutine->amoptionalkey = true;
amroutine->amsearcharray = false;
amroutine->amsearchnulls = false;
@@ -320,3 +285,17 @@ hnswhandler(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(amroutine);
}
/*
* Get the distance between two int4 attributes
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_int4_attribute_distance);
Datum
hnsw_int4_attribute_distance(PG_FUNCTION_ARGS)
{
int32 a = PG_GETARG_INT32(0);
int32 b = PG_GETARG_INT32(1);
double distance = ((double) a) - ((double) b);
PG_RETURN_FLOAT8(distance);
}

View File

@@ -12,10 +12,6 @@
#include "utils/sampling.h"
#include "vector.h"
#ifdef HNSW_BENCH
#include "portability/instr_time.h"
#endif
#define HNSW_MAX_DIM 2000
#define HNSW_MAX_NNZ 1000
@@ -23,6 +19,7 @@
#define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2
#define HNSW_TYPE_INFO_PROC 3
#define HNSW_ATTRIBUTE_DISTANCE_PROC 4
#define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953
@@ -46,10 +43,6 @@
#define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000
#define HNSW_DEFAULT_STREAMING false
#define HNSW_DEFAULT_EF_STREAM -1
#define HNSW_MIN_EF_STREAM -1
#define HNSW_MAX_EF_STREAM INT_MAX
/* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1
@@ -76,21 +69,6 @@
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page))
#ifdef HNSW_BENCH
#define HnswBench(name, code) \
do { \
instr_time start; \
instr_time duration; \
INSTR_TIME_SET_CURRENT(start); \
(code); \
INSTR_TIME_SET_CURRENT(duration); \
INSTR_TIME_SUBTRACT(duration, start); \
elog(INFO, "%s: %.3f ms", name, INSTR_TIME_GET_MILLISEC(duration)); \
} while (0)
#else
#define HnswBench(name, code) (code)
#endif
#if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
@@ -127,10 +105,10 @@
#define HnswPtrPointer(hp) (hp).ptr
#define HnswPtrOffset(hp) relptr_offset((hp).relptr)
#define HnswUseIndexTuple(index) (IndexRelationGetNumberOfAttributes(index) > 1)
/* Variables */
extern int hnsw_ef_search;
extern int hnsw_ef_stream;
extern bool hnsw_streaming;
extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData;
@@ -146,6 +124,7 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr);
HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);
struct HnswElementData
{
@@ -154,7 +133,6 @@ struct HnswElementData
uint8 heaptidsLength;
uint8 level;
uint8 deleted;
uint8 version;
uint32 hash;
HnswNeighborsPtr neighbors;
BlockNumber blkno;
@@ -162,6 +140,7 @@ struct HnswElementData
OffsetNumber neighborOffno;
BlockNumber neighborPage;
DatumPtr value;
IndexTuplePtr itup;
LWLock lock;
};
@@ -186,12 +165,10 @@ typedef struct HnswSearchCandidate
pairingheap_node c_node;
pairingheap_node w_node;
HnswElementPtr element;
float distance;
double distance;
bool matches;
} HnswSearchCandidate;
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
/* HNSW index options */
typedef struct HnswOptions
{
@@ -214,8 +191,8 @@ typedef struct HnswGraph
/* Allocations state */
LWLock allocatorLock;
long memoryUsed;
long memoryTotal;
Size memoryUsed;
Size memoryTotal;
/* Flushed state */
LWLock flushLock;
@@ -285,15 +262,17 @@ typedef struct HnswBuildState
double reltuples;
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo;
Oid collation;
Oid *collation;
/* Variables */
HnswGraph graphData;
HnswGraph *graph;
double ml;
int maxLevel;
bool useIndexTuple;
TupleDesc tupdesc;
/* Memory */
MemoryContext graphCtx;
@@ -335,10 +314,10 @@ typedef struct HnswElementTupleData
uint8 type;
uint8 level;
uint8 deleted;
uint8 version;
uint8 unused;
ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid;
uint16 unused;
uint16 unused2;
Vector data;
} HnswElementTupleData;
@@ -347,37 +326,24 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData
{
uint8 type;
uint8 version;
uint8 unused;
uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData;
typedef HnswNeighborTupleData * HnswNeighborTuple;
typedef union
{
struct pointerhash_hash *pointers;
struct offsethash_hash *offsets;
struct tidhash_hash *tids;
} visited_hash;
typedef struct HnswScanOpaqueData
{
const HnswTypeInfo *typeInfo;
bool first;
List *w;
visited_hash v;
pairingheap *discarded;
Datum q;
int m;
int64 tuples;
double previousDistance;
MemoryContext tmpCtx;
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo;
Oid collation;
Oid *collation;
} HnswScanOpaqueData;
typedef HnswScanOpaqueData * HnswScanOpaque;
@@ -395,8 +361,8 @@ typedef struct HnswVacuumState
int efConstruction;
/* Support functions */
FmgrInfo *procinfo;
Oid collation;
FmgrInfo *procinfo[2];
Oid *collation;
/* Variables */
struct tidhash_hash *deleted;
@@ -417,28 +383,32 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
List *HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, int m, bool inserting, HnswElement skipElement, bool inMemory);
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfo, Oid *collation, int m, int efConstruction, bool existing, bool inMemory);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation rel, FmgrInfo **procinfo, Oid *collation, bool loadVec, bool inMemory);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
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, float *maxDistance);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index);
void HnswLoadElement(HnswElement element, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo **procinfo, Oid *collation);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
void HnswInitLockTranche(void);
const HnswTypeInfo *HnswGetTypeInfo(Relation index);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
void HnswInitProcinfo(FmgrInfo **procinfo, Oid **collation, Relation index);
Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);
/* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -148,6 +148,7 @@ CreateGraphPages(HnswBuildState * buildstate)
Page page;
HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
/* Calculate sizes */
maxSize = HNSW_MAX_SIZE;
@@ -167,7 +168,6 @@ CreateGraphPages(HnswBuildState * buildstate)
Size etupSize;
Size ntupSize;
Size combinedSize;
Pointer valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */
iter = element->next;
@@ -176,7 +176,7 @@ CreateGraphPages(HnswBuildState * buildstate)
MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE);
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
etupSize = HnswGetElementTupleSize(base, element, useIndexTuple);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -186,7 +186,7 @@ CreateGraphPages(HnswBuildState * buildstate)
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("index tuple too large")));
HnswSetElementTuple(base, etup, element);
HnswSetElementTuple(base, etup, element, useIndexTuple);
/* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -327,20 +327,29 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
* Find duplicate element
*/
static bool
FindDuplicateInMemory(char *base, HnswElement element)
FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
if (useIndexTuple)
{
/* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
/* Check for space */
if (AddDuplicateInMemory(element, neighborElement))
@@ -366,7 +375,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors
*/
static void
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m)
{
for (int lc = e->level; lc >= 0; lc--)
{
@@ -388,7 +397,7 @@ UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswEleme
Assert(neighborElement);
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, procinfo, collation);
LWLockRelease(&neighborElement->lock);
}
}
@@ -398,20 +407,20 @@ UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswEleme
* Update graph in memory
*/
static void
UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
UpdateGraphInMemory(FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
{
HnswGraph *graph = buildstate->graph;
char *base = buildstate->hnswarea;
/* Look for duplicate */
if (FindDuplicateInMemory(base, element))
if (FindDuplicateInMemory(base, element, buildstate->useIndexTuple, buildstate->tupdesc))
return;
/* Add element */
AddElementInMemory(base, graph, element);
/* Update neighbors */
UpdateNeighborsInMemory(base, procinfo, collation, element, m);
UpdateNeighborsInMemory(base, buildstate->index, procinfo, collation, element, m);
/* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -424,8 +433,9 @@ UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int
static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
Relation index = buildstate->index;
FmgrInfo **procinfo = buildstate->procinfo;
Oid *collation = buildstate->collation;
HnswGraph *graph = buildstate->graph;
HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock;
@@ -458,7 +468,7 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, true);
/* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
@@ -481,6 +491,11 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
TupleDesc tupdesc = buildstate->tupdesc;
IndexTuple itup;
Size itupSize;
IndexTuple itupPtr;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -492,10 +507,10 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation, value))
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation[0], value))
return false;
value = HnswNormValue(typeInfo, buildstate->collation, value);
value = HnswNormValue(typeInfo, buildstate->collation[0], value);
}
/* Get datum size */
@@ -546,7 +561,17 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
/* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
valuePtr = HnswAlloc(allocator, valueSize);
if (useIndexTuple)
{
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
itupSize = IndexTupleSize(itup);
itupPtr = HnswAlloc(allocator, itupSize);
}
else
valuePtr = HnswAlloc(allocator, valueSize);
/*
* We have now allocated the space needed for the element, so we don't
@@ -556,8 +581,19 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(&graph->allocatorLock);
/* Copy the datum */
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
if (useIndexTuple)
{
bool unused;
memcpy(itupPtr, itup, itupSize);
HnswPtrStore(base, element->itup, itupPtr);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupPtr, 1, tupdesc, &unused)));
}
else
{
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
}
/* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -607,7 +643,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
* Initialize the graph
*/
static void
InitGraph(HnswGraph * graph, char *base, long memoryTotal)
InitGraph(HnswGraph * graph, char *base, Size memoryTotal)
{
/* Initialize the lock tranche if needed */
HnswInitLockTranche();
@@ -684,6 +720,19 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("type not supported for hnsw index")));
/* TODO See if needed */
if (IndexRelationGetNumberOfKeyAttributes(index) > 2)
elog(ERROR, "index cannot have more than two columns");
if (!OidIsValid(index_getprocid(index, 1, HNSW_DISTANCE_PROC)))
elog(ERROR, "first column must be a vector");
for (int i = 1; i < IndexRelationGetNumberOfKeyAttributes(index); i++)
{
if (!OidIsValid(index_getprocid(index, i + 1, HNSW_ATTRIBUTE_DISTANCE_PROC)))
elog(ERROR, "column %d cannot be a vector", i + 1);
}
/* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0)
ereport(ERROR,
@@ -704,14 +753,15 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->indtuples = 0;
/* Get support functions */
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
HnswInitProcinfo(buildstate->procinfo, &buildstate->collation, index);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0];
InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
InitGraph(&buildstate->graphData, NULL, (Size) maintenance_work_mem * 1024L);
buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->useIndexTuple = HnswUseIndexTuple(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",

View File

@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
* Check for a free offset
*/
static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion)
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
{
OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,7 +98,6 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
{
*freeOffno = offno;
*freeNeighborOffno = neighborOffno;
*tupleVersion = etup->version;
return true;
}
else if (*nbuf != buf)
@@ -154,11 +153,11 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion;
char *base = NULL;
bool useIndexTuple = HnswUseIndexTuple(index);
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
etupSize = HnswGetElementTupleSize(base, e, useIndexTuple);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE;
@@ -166,7 +165,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
/* Prepare element tuple */
etup = palloc0(etupSize);
HnswSetElementTuple(base, etup, e);
HnswSetElementTuple(base, etup, e, useIndexTuple);
/* Prepare neighbor tuple */
ntup = palloc0(ntupSize);
@@ -204,7 +203,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
}
/* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
{
if (nbuf != buf)
{
@@ -214,10 +213,6 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
/* Set tuple version */
etup->version = tupleVersion;
ntup->version = tupleVersion;
break;
}
@@ -340,6 +335,107 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
*updatedInsertPage = newInsertPage;
}
/*
* Load neighbors
*/
static HnswNeighborArray *
HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
{
char *base = NULL;
HnswNeighborArray *neighbors = HnswInitNeighborArray(lm, NULL);
ItemPointerData indextids[HNSW_MAX_M * 2];
if (!HnswLoadNeighborTids(element, indextids, index, m, lm, lc))
return neighbors;
for (int i = 0; i < lm; i++)
{
ItemPointer indextid = &indextids[i];
HnswElement e;
HnswCandidate *hc;
if (!ItemPointerIsValid(indextid))
break;
e = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
hc = &neighbors->items[neighbors->length++];
HnswPtrStore(base, hc->element, e);
}
return neighbors;
}
/*
* Load elements for insert
*/
static void
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, int *idx, Relation index, FmgrInfo **procinfo, Oid *collation)
{
char *base = NULL;
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element);
double distance;
bool matches;
HnswLoadElement(element, &distance, &matches, &q, qtup, NULL, index, procinfo, collation, true, NULL);
hc->distance = distance;
/* Prune element if being deleted */
if (element->heaptidsLength == 0)
{
*idx = i;
break;
}
}
}
/*
* Get update index
*/
static int
GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, MemoryContext updateCtx)
{
char *base = NULL;
int idx = -1;
HnswNeighborArray *neighbors;
MemoryContext oldCtx = MemoryContextSwitchTo(updateCtx);
/*
* Get latest neighbors since they may have changed. Do not lock yet since
* selecting neighbors can take time. Could use optimistic locking to
* retry if another update occurs before getting exclusive lock.
*/
neighbors = HnswLoadNeighbors(element, index, m, lm, lc);
/*
* 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.
*/
if (neighbors->length < lm)
idx = -2;
else
{
Datum q = HnswGetValue(base, element);
IndexTuple qtup = HnswPtrAccess(base, element->itup);;
LoadElementsForInsert(neighbors, q, qtup, &idx, index, procinfo, collation);
if (idx == -1)
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation);
}
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(updateCtx);
return idx;
}
/*
* Check if connection already exists
*/
@@ -360,14 +456,94 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
return false;
}
/*
* Update neighbor
*/
static void
UpdateNeighborOnDisk(HnswElement element, HnswElement newElement, int idx, int m, int lm, int lc, Relation index, bool checkExisting, bool building)
{
Buffer buf;
Page page;
GenericXLogState *state;
HnswNeighborTuple ntup;
int startIdx;
OffsetNumber offno = element->neighborOffno;
/* Register page */
buf = ReadBuffer(index, element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
if (building)
{
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Get tuple */
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, offno));
/* Calculate index for update */
startIdx = (element->level - lc) * m;
/* Check for existing connection */
if (checkExisting && ConnectionExists(newElement, 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 on the buffer */
ItemPointerSet(indextid, newElement->blkno, newElement->offno);
/* Commit */
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
}
else if (!building)
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
/*
* Update neighbors
*/
void
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m, bool checkExisting, bool building)
{
char *base = NULL;
/* Use separate memory context to improve performance for larger vectors */
MemoryContext updateCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw insert update context",
#if PG_VERSION_NUM >= 150000
128 * 1024, 128 * 1024,
#endif
128 * 1024);
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
@@ -376,96 +552,20 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, Hns
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
Buffer buf;
Page page;
GenericXLogState *state;
HnswNeighborTuple ntup;
int idx = -1;
int startIdx;
HnswElement neighborElement = HnswPtrAccess(base, hc->element);
OffsetNumber offno = neighborElement->neighborOffno;
int idx;
/*
* Get latest neighbors since they may have changed. Do not lock
* yet since selecting neighbors can take time. Could use
* optimistic locking to retry if another update occurs before
* getting exclusive lock.
*/
HnswLoadNeighbors(neighborElement, index, m);
/*
* 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.
*/
/* Select neighbors */
HnswUpdateConnection(NULL, e, hc, lm, lc, &idx, index, procinfo, collation);
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation, updateCtx);
/* New element was not selected as a neighbor */
if (idx == -1)
continue;
/* Register page */
buf = ReadBuffer(index, neighborElement->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
if (building)
{
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Get tuple */
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, offno));
/* Calculate index for update */
startIdx = (neighborElement->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 on the buffer */
ItemPointerSet(indextid, e->blkno, e->offno);
/* Commit */
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
}
else if (!building)
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
UpdateNeighborOnDisk(neighborElement, e, idx, m, lm, lc, index, checkExisting, building);
}
}
MemoryContextDelete(updateCtx);
}
/*
@@ -533,16 +633,26 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
TupleDesc tupdesc = RelationGetDescr(index);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
if (HnswUseIndexTuple(index))
{
/* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true;
@@ -555,7 +665,7 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
* Update graph on disk
*/
static void
UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{
BlockNumber newInsertPage = InvalidBlockNumber;
@@ -588,11 +698,13 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
HnswElement element;
int m;
int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
FmgrInfo *procinfo[2];
Oid *collation;
LOCKMODE lockmode = ShareLock;
char *base = NULL;
HnswInitProcinfo(procinfo, &collation, index);
/*
* 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
@@ -605,7 +717,23 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
/* Create an element */
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->value, DatumGetPointer(value));
if (HnswUseIndexTuple(index))
{
/* TODO no toast */
TupleDesc tupdesc = RelationGetDescr(index);
IndexTuple itup;
bool unused;
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
}
else
HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -622,7 +750,7 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false);
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, false);
/* Update graph on disk */
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
@@ -642,7 +770,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
Datum value;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
FmgrInfo *normprocinfo;
Oid collation = index->rd_indcollation[0];
Oid *collation = index->rd_indcollation;
/* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -655,10 +783,10 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswCheckNorm(normprocinfo, collation, value))
if (!HnswCheckNorm(normprocinfo, collation[0], value))
return;
value = HnswNormValue(typeInfo, collation, value);
value = HnswNormValue(typeInfo, collation[0], value);
}
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);

View File

@@ -1,7 +1,5 @@
#include "postgres.h"
#include <float.h>
#include "access/relscan.h"
#include "hnsw.h"
#include "pgstat.h"
@@ -17,65 +15,31 @@ GetScanItems(IndexScanDesc scan, Datum q)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
FmgrInfo **procinfo = so->procinfo;
Oid *collation = so->collation;
List *ep;
List *w;
int m;
HnswElement entryPoint;
char *base = NULL;
bool inMemory = false;
ScanKeyData *keyData = scan->keyData;
/* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint);
so->q = q;
so->m = m;
if (entryPoint == NULL)
return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false));
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, NULL, keyData, index, procinfo, collation, false, inMemory));
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true, NULL);
w = HnswSearchLayer(base, q, NULL, keyData, ep, 1, lc, index, procinfo, collation, m, false, NULL, inMemory);
ep = w;
}
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true, &so->tuples);
}
/*
* Resume scan at ground level with discarded candidates
*/
static List *
ResumeScanItems(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep = NIL;
char *base = NULL;
int batch_size = hnsw_ef_search;
if (pairingheap_is_empty(so->discarded))
return NIL;
/* Get next batch of candidates */
for (int i = 0; i < batch_size; i++)
{
HnswSearchCandidate *hc;
if (pairingheap_is_empty(so->discarded))
break;
hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
ep = lappend(ep, hc);
}
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
return HnswSearchLayer(base, q, NULL, keyData, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, inMemory);
}
/*
@@ -99,7 +63,7 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */
if (so->normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->collation, value);
value = HnswNormValue(so->typeInfo, so->collation[0], value);
}
return value;
@@ -119,16 +83,13 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index);
so->first = true;
so->v.tids = NULL;
so->discarded = NULL;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
HnswInitProcinfo(so->procinfo, &so->collation, index);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->collation = index->rd_indcollation[0];
scan->opaque = so;
@@ -143,15 +104,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
if (so->v.tids != NULL)
tidhash_reset(so->v.tids);
if (so->discarded != NULL)
pairingheap_reset(so->discarded);
so->first = true;
so->tuples = 0;
so->previousDistance = -INFINITY;
MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0)
@@ -201,7 +154,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = GetScanItems(scan, value));
so->w = GetScanItems(scan, value);
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
@@ -213,97 +166,22 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
#endif
}
for (;;)
while (list_length(so->w) > 0)
{
char *base = NULL;
HnswSearchCandidate *hc;
HnswElement element;
HnswSearchCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid;
if (list_length(so->w) == 0)
{
if (!hnsw_streaming)
break;
/* Empty index */
if (so->discarded == NULL)
break;
/* Reached max number of additional tuples */
if (hnsw_ef_stream != -1 && so->tuples >= hnsw_ef_search + hnsw_ef_stream)
{
if (pairingheap_is_empty(so->discarded))
break;
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
/* Prevent scans from consuming too much memory */
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
{
if (pairingheap_is_empty(so->discarded))
{
ereport(NOTICE,
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples.")));
break;
}
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
else
{
/*
* Locking ensures when neighbors are read, the elements they
* reference will not be deleted (and replaced) during the
* iteration.
*
* Elements loaded into memory on previous iterations may have
* been deleted (and replaced), so when reading neighbors, the
* element version must be checked.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = ResumeScanItems(scan));
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif
}
if (list_length(so->w) == 0)
break;
}
hc = llast(so->w);
element = HnswPtrAccess(base, hc->element);
/* Move to next element if no valid heap TIDs */
if (element->heaptidsLength == 0)
if (!hc->matches || element->heaptidsLength == 0)
{
so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */
if (hnsw_streaming)
{
pfree(element);
pfree(hc);
}
continue;
}
heaptid = &element->heaptids[--element->heaptidsLength];
if (hc->distance < so->previousDistance)
continue;
so->previousDistance = hc->distance;
MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid;

File diff suppressed because it is too large Load Diff

View File

@@ -189,8 +189,8 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
GenericXLogState *state;
int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction;
FmgrInfo *procinfo = vacuumstate->procinfo;
Oid collation = vacuumstate->collation;
FmgrInfo **procinfo = vacuumstate->procinfo;
Oid *collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
@@ -205,7 +205,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true, false);
/* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
HnswLoadElement(highestPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
/* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine.
*/
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
HnswLoadElement(entryPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
if (NeedsUpdated(vacuumstate, entryPoint))
{
@@ -370,7 +370,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Create an element */
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true);
HnswLoadElementFromTuple(element, etup, false, true, index);
elements = lappend(elements, element);
}
@@ -440,6 +440,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
bool useIndexTuple = HnswUseIndexTuple(index);
/*
* Wait for index scans to complete. Scans before this point may contain
@@ -521,20 +522,19 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */
etup->deleted = 1;
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
if (useIndexTuple)
{
IndexTuple itup = (IndexTuple) &etup->data;
MemSet(itup, 0, IndexTupleSize(itup));
}
else
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
/* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++)
ItemPointerSetInvalid(&ntup->indextids[i]);
/* Increment version */
/* This is used to avoid incorrect reads for iterative scans */
/* Reserve some bits for future use */
etup->version++;
if (etup->version > 15)
etup->version = 1;
ntup->version = etup->version;
/*
* We modified the tuples in place, no need to call
* PageIndexTupleOverwrite
@@ -581,8 +581,7 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->callback_state = callback_state;
vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
vacuumstate->collation = index->rd_indcollation[0];
HnswInitProcinfo(vacuumstate->procinfo, &vacuumstate->collation, index);
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context",

View File

@@ -69,6 +69,8 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
GenericCosts costs;
int lists;
double ratio;
double sequentialRatio = 0.5;
double startupPages;
double spc_seq_page_cost;
Relation index;
@@ -85,6 +87,8 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
MemSet(&costs, 0, sizeof(costs));
genericcostestimate(root, path, loop_count, &costs);
index = index_open(path->indexinfo->indexoid, NoLock);
IvfflatGetMetaPageInfo(index, &lists, NULL);
index_close(index, NoLock);
@@ -94,34 +98,26 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
if (ratio > 1.0)
ratio = 1.0;
/*
* This gives us the subset of tuples to visit. This value is passed into
* the generic cost estimator to determine the number of pages to visit
* during the index scan.
*/
costs.numIndexTuples = path->indexinfo->tuples * ratio;
genericcostestimate(root, path, loop_count, &costs);
get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
/* Change some page cost from random to sequential */
costs.indexTotalCost -= sequentialRatio * costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
/* Startup cost is cost before returning the first row */
costs.indexStartupCost = costs.indexTotalCost * ratio;
/* Adjust cost if needed since TOAST not included in seq scan cost */
if (costs.numIndexPages > path->indexinfo->rel->pages && ratio < 0.5)
startupPages = costs.numIndexPages * ratio;
if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
{
/* Change all page cost from random to sequential */
costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
/* Change rest of page cost from random to sequential */
costs.indexStartupCost -= (1 - sequentialRatio) * startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
/* Remove cost of extra pages */
costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
}
else
{
/* Change some page cost from random to sequential */
costs.indexTotalCost -= 0.5 * costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
}
/* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost;
*indexStartupCost = costs.indexStartupCost;
*indexTotalCost = costs.indexTotalCost;
*indexSelectivity = costs.indexSelectivity;
*indexCorrelation = costs.indexCorrelation;

View File

@@ -41,8 +41,7 @@ my $c = int(rand() * $nc);
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
@@ -60,8 +59,7 @@ like($explain, qr/Index Scan using idx/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed like
$explain = $node->safe_psql("postgres", qq(

View File

@@ -17,12 +17,11 @@ $node->safe_psql("postgres", "CREATE EXTENSION vector;");
for my $dim (@dims)
{
my $array_sql = join(",", ('random()') x $dim);
my $n = $dim == 384 ? 2000 : 1000;
# Create table and index
$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, $n) i;"
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 2000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
@@ -40,6 +39,21 @@ for my $dim (@dims)
));
like($explain, qr/Index Scan using idx/);
# 3x the rows are needed for distance filters
# since the planner uses DEFAULT_INEQ_SEL for the selectivity (should be 1)
# Recreate index for performance
$node->safe_psql("postgres", "DROP INDEX idx;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(2001, 6000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$node->safe_psql("postgres", "DROP TABLE tst;");
}

View File

@@ -39,6 +39,11 @@ for my $dim (@dims)
));
like($explain, qr/Index Scan using idx/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$node->safe_psql("postgres", "DROP TABLE tst;");
}

View File

@@ -0,0 +1,109 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @cs = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my $nc = 50;
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $cs[0] ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
is(scalar(@actual_ids), $limit);
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('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($dim), c int4);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 20000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
push(@cs, int(rand() * $nc));
}
# Get exact results
@expected = ();
for my $i (0 .. $#queries)
{
my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <-> '$queries[$i]' LIMIT $limit;");
push(@expected, $res);
}
# Add index
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c);");
# Test recall
test_recall(0.99, '<->');
# Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
$node->safe_psql("postgres", "VACUUM tst;");
# Test columns
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_l2_ops);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c, c);");
like($stderr, qr/index cannot have more than two columns/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, v vector_l2_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();

View File

@@ -1,66 +0,0 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
foreach ((30000, 50000, 70000))
{
my $ef_stream = $_;
my $expected = $ef_stream / 10000;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET hnsw.ef_stream = $ef_stream;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '2MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
like($stderr, qr/hnsw index scan exceeded work_mem after \d+ tuples/);
done_testing();

View File

@@ -1,131 +0,0 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my @cs = (100, 1000);
sub test_recall
{
my ($c, $ef_search, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.streaming = on;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('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($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
if ($c == 100)
{
test_recall($c, 40, 0.99, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 40, 0.99, $operator);
}
else
{
test_recall($c, 40, 0.99, $operator);
}
}
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();