mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 03:57:34 +08:00
Compare commits
20 Commits
hqann2
...
hnsw-strea
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c01e76f2fa | ||
|
|
ab57217f48 | ||
|
|
5f6e031ccc | ||
|
|
1a1221f905 | ||
|
|
40c3e402c7 | ||
|
|
058248fdcc | ||
|
|
73c5145b77 | ||
|
|
ec4a23fe49 | ||
|
|
38207f5640 | ||
|
|
4e35c6abe3 | ||
|
|
11e4d040d9 | ||
|
|
b2fa625255 | ||
|
|
a8e699c927 | ||
|
|
91541fece6 | ||
|
|
f3de487da2 | ||
|
|
721d4b7e3f | ||
|
|
28066d8fe4 | ||
|
|
495041e43b | ||
|
|
52c385c03a | ||
|
|
80cbd32dab |
2
.github/workflows/build.yml
vendored
2
.github/workflows/build.yml
vendored
@@ -49,7 +49,7 @@ jobs:
|
||||
- postgres: 16
|
||||
os: macos-14
|
||||
- postgres: 14
|
||||
os: macos-13
|
||||
os: macos-12
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ankane/setup-postgres@v1
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
## 0.8.0 (unreleased)
|
||||
|
||||
- Added support for inline filtering with HNSW
|
||||
- Added support for iterative index scans
|
||||
- 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
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG PG_MAJOR=17
|
||||
ARG PG_MAJOR=16
|
||||
FROM postgres:$PG_MAJOR
|
||||
ARG PG_MAJOR
|
||||
|
||||
|
||||
2
Makefile
2
Makefile
@@ -66,7 +66,7 @@ dist:
|
||||
git archive --format zip --prefix=$(EXTENSION)-$(EXTVERSION)/ --output dist/$(EXTENSION)-$(EXTVERSION).zip master
|
||||
|
||||
# for Docker
|
||||
PG_MAJOR ?= 17
|
||||
PG_MAJOR ?= 16
|
||||
|
||||
.PHONY: docker
|
||||
|
||||
|
||||
102
README.md
102
README.md
@@ -52,8 +52,6 @@ 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).
|
||||
@@ -102,8 +100,6 @@ 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
|
||||
@@ -149,8 +145,6 @@ 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
|
||||
|
||||
@@ -439,12 +433,6 @@ 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
|
||||
@@ -457,6 +445,63 @@ 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, it’s 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*
|
||||
@@ -995,7 +1040,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/17/bin/pg_config
|
||||
export PG_CONFIG=/Library/PostgreSQL/16/bin/pg_config
|
||||
```
|
||||
|
||||
Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use:
|
||||
@@ -1006,11 +1051,11 @@ sudo --preserve-env=PG_CONFIG make install
|
||||
|
||||
A few common paths on Mac are:
|
||||
|
||||
- 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`
|
||||
- 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`
|
||||
|
||||
Note: Replace `17` with your Postgres server version
|
||||
Note: Replace `16` with your Postgres server version
|
||||
|
||||
### Missing Header
|
||||
|
||||
@@ -1019,10 +1064,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-17
|
||||
sudo apt install postgresql-server-dev-16
|
||||
```
|
||||
|
||||
Note: Replace `17` with your Postgres server version
|
||||
Note: Replace `16` with your Postgres server version
|
||||
|
||||
### Missing SDK
|
||||
|
||||
@@ -1055,17 +1100,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:pg17
|
||||
docker pull pgvector/pgvector:pg16
|
||||
```
|
||||
|
||||
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).
|
||||
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).
|
||||
|
||||
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=17 -t myuser/pgvector .
|
||||
docker build --pull --build-arg PG_MAJOR=16 -t myuser/pgvector .
|
||||
```
|
||||
|
||||
### Homebrew
|
||||
@@ -1076,7 +1121,7 @@ With Homebrew Postgres, you can use:
|
||||
brew install pgvector
|
||||
```
|
||||
|
||||
Note: This only adds it to the `postgresql@17` and `postgresql@14` formulas
|
||||
Note: This only adds it to the `postgresql@14` formula
|
||||
|
||||
### PGXN
|
||||
|
||||
@@ -1091,22 +1136,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-17-pgvector
|
||||
sudo apt install postgresql-16-pgvector
|
||||
```
|
||||
|
||||
Note: Replace `17` with your Postgres server version
|
||||
Note: Replace `16` 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_17
|
||||
sudo yum install pgvector_16
|
||||
# or
|
||||
sudo dnf install pgvector_17
|
||||
sudo dnf install pgvector_16
|
||||
```
|
||||
|
||||
Note: Replace `17` with your Postgres server version
|
||||
Note: Replace `16` with your Postgres server version
|
||||
|
||||
### pkg
|
||||
|
||||
@@ -1195,7 +1240,6 @@ 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
|
||||
|
||||
|
||||
@@ -24,11 +24,3 @@ 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);
|
||||
|
||||
@@ -916,13 +916,3 @@ 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);
|
||||
|
||||
115
src/hnsw.c
115
src/hnsw.c
@@ -18,16 +18,9 @@
|
||||
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
|
||||
#endif
|
||||
|
||||
static const struct config_enum_entry hnsw_iterative_search_options[] = {
|
||||
{"off", HNSW_ITERATIVE_SEARCH_OFF, false},
|
||||
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
|
||||
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
|
||||
{NULL, 0, false}
|
||||
};
|
||||
|
||||
int hnsw_ef_search;
|
||||
int hnsw_iterative_search_max_tuples;
|
||||
int hnsw_iterative_search;
|
||||
int hnsw_ef_stream;
|
||||
bool hnsw_streaming;
|
||||
int hnsw_lock_tranche_id;
|
||||
static relopt_kind hnsw_relopt_kind;
|
||||
|
||||
@@ -78,14 +71,16 @@ 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);
|
||||
|
||||
DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search",
|
||||
NULL, &hnsw_iterative_search,
|
||||
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, 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 Ensure ivfflat.max_probes uses same value for "all" */
|
||||
DefineCustomIntVariable("hnsw.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search",
|
||||
"-1 means all", &hnsw_iterative_search_max_tuples,
|
||||
-1, -1, INT_MAX, 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");
|
||||
}
|
||||
@@ -107,6 +102,33 @@ 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
|
||||
*/
|
||||
@@ -118,8 +140,11 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
||||
{
|
||||
GenericCosts costs;
|
||||
int m;
|
||||
double ratio;
|
||||
double startupPages;
|
||||
int ef;
|
||||
int entryLevel;
|
||||
int layer0TuplesMax;
|
||||
double layer0Selectivity;
|
||||
double scalingFactor = 0.55;
|
||||
double spc_seq_page_cost;
|
||||
Relation index;
|
||||
|
||||
@@ -136,12 +161,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
|
||||
@@ -169,38 +194,30 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
||||
* at L0, accounting for previously visited tuples, multiplied by the
|
||||
* "scalingFactor" (currently hardcoded).
|
||||
*/
|
||||
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)));
|
||||
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)));
|
||||
|
||||
ratio = (entryLevel * m + layer0TuplesMax * layer0Selectivity) / path->indexinfo->tuples;
|
||||
costs.numIndexTuples = (entryLevel * m) +
|
||||
(layer0TuplesMax * layer0Selectivity);
|
||||
|
||||
if (ratio > 1)
|
||||
ratio = 1;
|
||||
}
|
||||
else
|
||||
ratio = 1;
|
||||
genericcostestimate(root, path, loop_count, &costs);
|
||||
|
||||
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 */
|
||||
startupPages = costs.numIndexPages * ratio;
|
||||
if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
|
||||
if (costs.numIndexPages > path->indexinfo->rel->pages)
|
||||
{
|
||||
/* Change all page cost from random to sequential */
|
||||
costs.indexStartupCost -= startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
|
||||
/* Remove cost of extra pages */
|
||||
costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
|
||||
costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
|
||||
}
|
||||
|
||||
*indexStartupCost = costs.indexStartupCost;
|
||||
/* Use total cost since most work happens before first tuple is returned */
|
||||
*indexStartupCost = costs.indexTotalCost;
|
||||
*indexTotalCost = costs.indexTotalCost;
|
||||
*indexSelectivity = costs.indexSelectivity;
|
||||
*indexCorrelation = costs.indexCorrelation;
|
||||
@@ -245,13 +262,13 @@ hnswhandler(PG_FUNCTION_ARGS)
|
||||
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
|
||||
|
||||
amroutine->amstrategies = 0;
|
||||
amroutine->amsupport = 4;
|
||||
amroutine->amsupport = 3;
|
||||
amroutine->amoptsprocnum = 0;
|
||||
amroutine->amcanorder = false;
|
||||
amroutine->amcanorderbyop = true;
|
||||
amroutine->amcanbackward = false; /* can change direction mid-scan */
|
||||
amroutine->amcanunique = false;
|
||||
amroutine->amcanmulticol = true;
|
||||
amroutine->amcanmulticol = false;
|
||||
amroutine->amoptionalkey = true;
|
||||
amroutine->amsearcharray = false;
|
||||
amroutine->amsearchnulls = false;
|
||||
@@ -303,17 +320,3 @@ 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);
|
||||
}
|
||||
|
||||
114
src/hnsw.h
114
src/hnsw.h
@@ -12,6 +12,10 @@
|
||||
#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
|
||||
|
||||
@@ -19,7 +23,6 @@
|
||||
#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
|
||||
@@ -43,6 +46,10 @@
|
||||
#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
|
||||
@@ -69,6 +76,21 @@
|
||||
#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)
|
||||
@@ -89,9 +111,6 @@
|
||||
/* Ensure fits on page and in uint8 */
|
||||
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
|
||||
|
||||
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
|
||||
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
|
||||
|
||||
#define HnswGetValue(base, element) PointerGetDatum(HnswPtrAccess(base, (element)->value))
|
||||
|
||||
#if PG_VERSION_NUM < 140005
|
||||
@@ -108,21 +127,12 @@
|
||||
#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_iterative_search;
|
||||
extern int hnsw_iterative_search_max_tuples;
|
||||
extern int hnsw_ef_stream;
|
||||
extern bool hnsw_streaming;
|
||||
extern int hnsw_lock_tranche_id;
|
||||
|
||||
typedef enum HnswIterativeSearchType
|
||||
{
|
||||
HNSW_ITERATIVE_SEARCH_OFF,
|
||||
HNSW_ITERATIVE_SEARCH_RELAXED,
|
||||
HNSW_ITERATIVE_SEARCH_STRICT
|
||||
} HnswIterativeSearchType;
|
||||
|
||||
typedef struct HnswElementData HnswElementData;
|
||||
typedef struct HnswNeighborArray HnswNeighborArray;
|
||||
|
||||
@@ -136,7 +146,6 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
|
||||
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
|
||||
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
|
||||
HnswPtrDeclare(char, DatumRelptr, DatumPtr);
|
||||
HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);
|
||||
|
||||
struct HnswElementData
|
||||
{
|
||||
@@ -153,7 +162,6 @@ struct HnswElementData
|
||||
OffsetNumber neighborOffno;
|
||||
BlockNumber neighborPage;
|
||||
DatumPtr value;
|
||||
IndexTuplePtr itup;
|
||||
LWLock lock;
|
||||
};
|
||||
|
||||
@@ -178,10 +186,12 @@ typedef struct HnswSearchCandidate
|
||||
pairingheap_node c_node;
|
||||
pairingheap_node w_node;
|
||||
HnswElementPtr element;
|
||||
double distance;
|
||||
bool matches;
|
||||
float distance;
|
||||
} 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
|
||||
{
|
||||
@@ -204,8 +214,8 @@ typedef struct HnswGraph
|
||||
|
||||
/* Allocations state */
|
||||
LWLock allocatorLock;
|
||||
Size memoryUsed;
|
||||
Size memoryTotal;
|
||||
long memoryUsed;
|
||||
long memoryTotal;
|
||||
|
||||
/* Flushed state */
|
||||
LWLock flushLock;
|
||||
@@ -256,20 +266,6 @@ typedef struct HnswTypeInfo
|
||||
void (*checkValue) (Pointer v);
|
||||
} HnswTypeInfo;
|
||||
|
||||
typedef struct HnswSupport
|
||||
{
|
||||
FmgrInfo *procinfo[2];
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid *collation;
|
||||
} HnswSupport;
|
||||
|
||||
typedef struct HnswQuery
|
||||
{
|
||||
Datum value;
|
||||
IndexTuple itup;
|
||||
ScanKeyData *keyData;
|
||||
} HnswQuery;
|
||||
|
||||
typedef struct HnswBuildState
|
||||
{
|
||||
/* Info */
|
||||
@@ -289,15 +285,15 @@ typedef struct HnswBuildState
|
||||
double reltuples;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation;
|
||||
|
||||
/* Variables */
|
||||
HnswGraph graphData;
|
||||
HnswGraph *graph;
|
||||
double ml;
|
||||
int maxLevel;
|
||||
bool useIndexTuple;
|
||||
TupleDesc tupdesc;
|
||||
|
||||
/* Memory */
|
||||
MemoryContext graphCtx;
|
||||
@@ -365,12 +361,6 @@ typedef union
|
||||
struct tidhash_hash *tids;
|
||||
} visited_hash;
|
||||
|
||||
typedef union
|
||||
{
|
||||
HnswElement element;
|
||||
ItemPointerData indextid;
|
||||
} HnswUnvisited;
|
||||
|
||||
typedef struct HnswScanOpaqueData
|
||||
{
|
||||
const HnswTypeInfo *typeInfo;
|
||||
@@ -378,14 +368,16 @@ typedef struct HnswScanOpaqueData
|
||||
List *w;
|
||||
visited_hash v;
|
||||
pairingheap *discarded;
|
||||
HnswQuery q;
|
||||
Datum q;
|
||||
int m;
|
||||
int64 tuples;
|
||||
double previousDistance;
|
||||
MemoryContext tmpCtx;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation;
|
||||
} HnswScanOpaqueData;
|
||||
|
||||
typedef HnswScanOpaqueData * HnswScanOpaque;
|
||||
@@ -403,7 +395,8 @@ typedef struct HnswVacuumState
|
||||
int efConstruction;
|
||||
|
||||
/* Support functions */
|
||||
HnswSupport support;
|
||||
FmgrInfo *procinfo;
|
||||
Oid collation;
|
||||
|
||||
/* Variables */
|
||||
struct tidhash_hash *deleted;
|
||||
@@ -419,38 +412,33 @@ typedef struct HnswVacuumState
|
||||
int HnswGetM(Relation index);
|
||||
int HnswGetEfConstruction(Relation index);
|
||||
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
|
||||
void HnswInitSupport(HnswSupport * support, Relation index);
|
||||
Datum HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value);
|
||||
bool HnswCheckNorm(HnswSupport * support, Datum value);
|
||||
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, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
|
||||
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);
|
||||
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, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory);
|
||||
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec, bool inMemory);
|
||||
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 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, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc);
|
||||
void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, 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, HnswQuery * q, Relation index, HnswSupport * support, 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, HnswSupport * support);
|
||||
bool HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc);
|
||||
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
|
||||
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 HnswInitLockTranche(void);
|
||||
const HnswTypeInfo *HnswGetTypeInfo(Relation index);
|
||||
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
|
||||
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);
|
||||
|
||||
@@ -148,7 +148,6 @@ CreateGraphPages(HnswBuildState * buildstate)
|
||||
Page page;
|
||||
HnswElementPtr iter = buildstate->graph->head;
|
||||
char *base = buildstate->hnswarea;
|
||||
bool useIndexTuple = buildstate->useIndexTuple;
|
||||
|
||||
/* Calculate sizes */
|
||||
maxSize = HNSW_MAX_SIZE;
|
||||
@@ -168,6 +167,7 @@ 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 = HnswGetElementTupleSize(base, element, useIndexTuple);
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
|
||||
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, useIndexTuple);
|
||||
HnswSetElementTuple(base, etup, element);
|
||||
|
||||
/* Keep element and neighbors on the same page if possible */
|
||||
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
|
||||
@@ -327,18 +327,19 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
|
||||
* Find duplicate element
|
||||
*/
|
||||
static bool
|
||||
FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
|
||||
FindDuplicateInMemory(char *base, HnswElement element)
|
||||
{
|
||||
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
|
||||
IndexTuple itup = HnswPtrAccess(base, element->itup);
|
||||
Datum value = HnswGetValue(base, element);
|
||||
|
||||
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 (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
|
||||
if (!datumIsEqual(value, neighborValue, false, -1))
|
||||
return false;
|
||||
|
||||
/* Check for space */
|
||||
@@ -365,7 +366,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
|
||||
* Update neighbors
|
||||
*/
|
||||
static void
|
||||
UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswElement e, int m)
|
||||
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
|
||||
{
|
||||
for (int lc = e->level; lc >= 0; lc--)
|
||||
{
|
||||
@@ -387,7 +388,7 @@ UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswE
|
||||
Assert(neighborElement);
|
||||
|
||||
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
|
||||
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, support);
|
||||
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
|
||||
LWLockRelease(&neighborElement->lock);
|
||||
}
|
||||
}
|
||||
@@ -397,20 +398,20 @@ UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswE
|
||||
* Update graph in memory
|
||||
*/
|
||||
static void
|
||||
UpdateGraphInMemory(HnswSupport * support, 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, buildstate->useIndexTuple, buildstate->tupdesc))
|
||||
if (FindDuplicateInMemory(base, element))
|
||||
return;
|
||||
|
||||
/* Add element */
|
||||
AddElementInMemory(base, graph, element);
|
||||
|
||||
/* Update neighbors */
|
||||
UpdateNeighborsInMemory(base, buildstate->index, support, element, m);
|
||||
UpdateNeighborsInMemory(base, procinfo, collation, element, m);
|
||||
|
||||
/* Update entry point if needed (already have lock) */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
@@ -423,9 +424,9 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
|
||||
static void
|
||||
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
||||
{
|
||||
Relation index = buildstate->index;
|
||||
FmgrInfo *procinfo = buildstate->procinfo;
|
||||
Oid collation = buildstate->collation;
|
||||
HnswGraph *graph = buildstate->graph;
|
||||
HnswSupport *support = &buildstate->support;
|
||||
HnswElement entryPoint;
|
||||
LWLock *entryLock = &graph->entryLock;
|
||||
LWLock *entryWaitLock = &graph->entryWaitLock;
|
||||
@@ -457,10 +458,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
||||
}
|
||||
|
||||
/* Find neighbors for element */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, true);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
|
||||
|
||||
/* Update graph in memory */
|
||||
UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
|
||||
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
|
||||
|
||||
/* Release entry lock */
|
||||
LWLockRelease(entryLock);
|
||||
@@ -472,24 +473,33 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
|
||||
static bool
|
||||
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
|
||||
{
|
||||
const HnswTypeInfo *typeInfo = buildstate->typeInfo;
|
||||
HnswGraph *graph = buildstate->graph;
|
||||
HnswElement element;
|
||||
HnswAllocator *allocator = &buildstate->allocator;
|
||||
HnswSupport *support = &buildstate->support;
|
||||
Size valueSize;
|
||||
Pointer valuePtr;
|
||||
LWLock *flushLock = &graph->flushLock;
|
||||
char *base = buildstate->hnswarea;
|
||||
TupleDesc tupdesc = buildstate->tupdesc;
|
||||
IndexTuple itup;
|
||||
Size itupSize;
|
||||
IndexTuple itupShared;
|
||||
bool unused;
|
||||
|
||||
/* Form index tuple */
|
||||
if (!HnswFormIndexTuple(&itup, values, isnull, buildstate->typeInfo, support, tupdesc))
|
||||
return false;
|
||||
/* Detoast once for all calls */
|
||||
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
|
||||
/* Get tuple size */
|
||||
itupSize = IndexTupleSize(itup);
|
||||
/* Check value */
|
||||
if (typeInfo->checkValue != NULL)
|
||||
typeInfo->checkValue(DatumGetPointer(value));
|
||||
|
||||
/* Normalize if needed */
|
||||
if (buildstate->normprocinfo != NULL)
|
||||
{
|
||||
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation, value))
|
||||
return false;
|
||||
|
||||
value = HnswNormValue(typeInfo, buildstate->collation, value);
|
||||
}
|
||||
|
||||
/* Get datum size */
|
||||
valueSize = VARSIZE_ANY(DatumGetPointer(value));
|
||||
|
||||
/* Ensure graph not flushed when inserting */
|
||||
LWLockAcquire(flushLock, LW_SHARED);
|
||||
@@ -499,7 +509,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
|
||||
{
|
||||
LWLockRelease(flushLock);
|
||||
|
||||
return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
|
||||
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -531,12 +541,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
|
||||
|
||||
LWLockRelease(flushLock);
|
||||
|
||||
return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc);
|
||||
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true);
|
||||
}
|
||||
|
||||
/* Ok, we can proceed to allocate the element */
|
||||
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
|
||||
itupShared = HnswAlloc(allocator, itupSize);
|
||||
valuePtr = HnswAlloc(allocator, valueSize);
|
||||
|
||||
/*
|
||||
* We have now allocated the space needed for the element, so we don't
|
||||
@@ -545,10 +555,9 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
|
||||
*/
|
||||
LWLockRelease(&graph->allocatorLock);
|
||||
|
||||
/* Copy the tuple */
|
||||
memcpy(itupShared, itup, itupSize);
|
||||
HnswPtrStore(base, element->itup, itupShared);
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupShared, 1, tupdesc, &unused)));
|
||||
/* Copy the datum */
|
||||
memcpy(valuePtr, DatumGetPointer(value), valueSize);
|
||||
HnswPtrStore(base, element->value, valuePtr);
|
||||
|
||||
/* Create a lock for the element */
|
||||
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
|
||||
@@ -598,7 +607,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
|
||||
* Initialize the graph
|
||||
*/
|
||||
static void
|
||||
InitGraph(HnswGraph * graph, char *base, Size memoryTotal)
|
||||
InitGraph(HnswGraph * graph, char *base, long memoryTotal)
|
||||
{
|
||||
/* Initialize the lock tranche if needed */
|
||||
HnswInitLockTranche();
|
||||
@@ -675,19 +684,6 @@ 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,
|
||||
@@ -708,14 +704,14 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
|
||||
buildstate->indtuples = 0;
|
||||
|
||||
/* Get support functions */
|
||||
HnswInitSupport(&buildstate->support, index);
|
||||
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
buildstate->collation = index->rd_indcollation[0];
|
||||
|
||||
InitGraph(&buildstate->graphData, NULL, (Size) maintenance_work_mem * 1024L);
|
||||
InitGraph(&buildstate->graphData, NULL, 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",
|
||||
|
||||
345
src/hnswinsert.c
345
src/hnswinsert.c
@@ -156,10 +156,9 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
|
||||
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||
uint8 tupleVersion;
|
||||
char *base = NULL;
|
||||
bool useIndexTuple = HnswUseIndexTuple(index);
|
||||
|
||||
/* Calculate sizes */
|
||||
etupSize = HnswGetElementTupleSize(base, e, useIndexTuple);
|
||||
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
|
||||
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
|
||||
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
|
||||
maxSize = HNSW_MAX_SIZE;
|
||||
@@ -167,7 +166,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
|
||||
|
||||
/* Prepare element tuple */
|
||||
etup = palloc0(etupSize);
|
||||
HnswSetElementTuple(base, etup, e, useIndexTuple);
|
||||
HnswSetElementTuple(base, etup, e);
|
||||
|
||||
/* Prepare neighbor tuple */
|
||||
ntup = palloc0(ntupSize);
|
||||
@@ -341,110 +340,6 @@ 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, HnswQuery * q, int *idx, Relation index, HnswSupport * support)
|
||||
{
|
||||
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, index, support, 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, HnswSupport * support, 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
|
||||
{
|
||||
HnswQuery q;
|
||||
|
||||
q.value = HnswGetValue(base, element);
|
||||
q.itup = HnswPtrAccess(base, element->itup);
|
||||
q.keyData = NULL;
|
||||
|
||||
LoadElementsForInsert(neighbors, &q, &idx, index, support);
|
||||
|
||||
if (idx == -1)
|
||||
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, support);
|
||||
}
|
||||
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
MemoryContextReset(updateCtx);
|
||||
|
||||
return idx;
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if connection already exists
|
||||
*/
|
||||
@@ -465,94 +360,14 @@ 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, HnswSupport * support, 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);
|
||||
@@ -561,20 +376,96 @@ HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e,
|
||||
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);
|
||||
int idx;
|
||||
OffsetNumber offno = neighborElement->neighborOffno;
|
||||
|
||||
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, support, 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.
|
||||
*/
|
||||
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);
|
||||
|
||||
/* New element was not selected as a neighbor */
|
||||
if (idx == -1)
|
||||
continue;
|
||||
|
||||
UpdateNeighborOnDisk(neighborElement, e, idx, m, lm, lc, index, checkExisting, building);
|
||||
/* 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);
|
||||
}
|
||||
}
|
||||
|
||||
MemoryContextDelete(updateCtx);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -637,30 +528,21 @@ AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool bu
|
||||
* Find duplicate element
|
||||
*/
|
||||
static bool
|
||||
FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDesc tupdesc)
|
||||
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);
|
||||
|
||||
for (int i = 0; i < neighbors->length; i++)
|
||||
{
|
||||
HnswCandidate *neighbor = &neighbors->items[i];
|
||||
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
|
||||
Datum neighborValue = HnswGetValue(base, neighborElement);
|
||||
|
||||
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;
|
||||
}
|
||||
/* Exit early since ordered by distance */
|
||||
if (!datumIsEqual(value, neighborValue, false, -1))
|
||||
return false;
|
||||
|
||||
if (AddDuplicateOnDisk(index, element, neighborElement, building))
|
||||
return true;
|
||||
@@ -673,12 +555,12 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDes
|
||||
* Update graph on disk
|
||||
*/
|
||||
static void
|
||||
UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building, TupleDesc tupdesc)
|
||||
UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
|
||||
{
|
||||
BlockNumber newInsertPage = InvalidBlockNumber;
|
||||
|
||||
/* Look for duplicate */
|
||||
if (FindDuplicateOnDisk(index, element, building, tupdesc))
|
||||
if (FindDuplicateOnDisk(index, element, building))
|
||||
return;
|
||||
|
||||
/* Add element */
|
||||
@@ -689,7 +571,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
|
||||
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
|
||||
|
||||
/* Update neighbors */
|
||||
HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
|
||||
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building);
|
||||
|
||||
/* Update entry point if needed */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
@@ -700,15 +582,16 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
|
||||
* Insert a tuple into the index
|
||||
*/
|
||||
bool
|
||||
HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc)
|
||||
HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building)
|
||||
{
|
||||
HnswElement entryPoint;
|
||||
HnswElement element;
|
||||
int m;
|
||||
int efConstruction = HnswGetEfConstruction(index);
|
||||
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
LOCKMODE lockmode = ShareLock;
|
||||
char *base = NULL;
|
||||
bool unused;
|
||||
|
||||
/*
|
||||
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
|
||||
@@ -721,9 +604,8 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
|
||||
HnswGetMetaPageInfo(index, &m, &entryPoint);
|
||||
|
||||
/* Create an element */
|
||||
element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
|
||||
HnswPtrStore(base, element->itup, itup);
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
|
||||
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(value));
|
||||
|
||||
/* Prevent concurrent inserts when likely updating entry point */
|
||||
if (entryPoint == NULL || element->level > entryPoint->level)
|
||||
@@ -740,10 +622,10 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
|
||||
}
|
||||
|
||||
/* Find neighbors for element */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, false);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false);
|
||||
|
||||
/* Update graph on disk */
|
||||
UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building, tupdesc);
|
||||
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
|
||||
|
||||
/* Release lock */
|
||||
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
|
||||
@@ -755,20 +637,31 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
|
||||
* Insert a tuple into the index
|
||||
*/
|
||||
static void
|
||||
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
|
||||
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid)
|
||||
{
|
||||
IndexTuple itup;
|
||||
Datum value;
|
||||
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
HnswSupport support;
|
||||
FmgrInfo *normprocinfo;
|
||||
Oid collation = index->rd_indcollation[0];
|
||||
|
||||
HnswInitSupport(&support, index);
|
||||
/* Detoast once for all calls */
|
||||
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
|
||||
/* Form index tuple */
|
||||
if (!HnswFormIndexTuple(&itup, values, isnull, typeInfo, &support, tupdesc))
|
||||
return;
|
||||
/* Check value */
|
||||
if (typeInfo->checkValue != NULL)
|
||||
typeInfo->checkValue(DatumGetPointer(value));
|
||||
|
||||
HnswInsertTupleOnDisk(index, &support, itup, heaptid, false, tupdesc);
|
||||
/* Normalize if needed */
|
||||
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
if (normprocinfo != NULL)
|
||||
{
|
||||
if (!HnswCheckNorm(normprocinfo, collation, value))
|
||||
return;
|
||||
|
||||
value = HnswNormValue(typeInfo, collation, value);
|
||||
}
|
||||
|
||||
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);
|
||||
}
|
||||
|
||||
/*
|
||||
|
||||
@@ -1,51 +1,48 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include <float.h>
|
||||
|
||||
#include "access/relscan.h"
|
||||
#include "hnsw.h"
|
||||
#include "pgstat.h"
|
||||
#include "storage/bufmgr.h"
|
||||
#include "storage/lmgr.h"
|
||||
#include "utils/float.h"
|
||||
#include "utils/memutils.h"
|
||||
|
||||
/*
|
||||
* Algorithm 5 from paper
|
||||
*/
|
||||
static List *
|
||||
GetScanItems(IndexScanDesc scan, Datum value)
|
||||
GetScanItems(IndexScanDesc scan, Datum q)
|
||||
{
|
||||
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
|
||||
Relation index = scan->indexRelation;
|
||||
HnswSupport *support = &so->support;
|
||||
FmgrInfo *procinfo = so->procinfo;
|
||||
Oid collation = so->collation;
|
||||
List *ep;
|
||||
List *w;
|
||||
int m;
|
||||
HnswElement entryPoint;
|
||||
char *base = NULL;
|
||||
bool inMemory = false;
|
||||
HnswQuery *q = &so->q;
|
||||
|
||||
q->value = value;
|
||||
q->itup = NULL;
|
||||
q->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, support, false, inMemory));
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false));
|
||||
|
||||
for (int lc = entryPoint->level; lc >= 1; lc--)
|
||||
{
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, inMemory, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true, NULL);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, inMemory, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
|
||||
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);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -56,6 +53,8 @@ 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;
|
||||
@@ -66,17 +65,17 @@ ResumeScanItems(IndexScanDesc scan)
|
||||
/* Get next batch of candidates */
|
||||
for (int i = 0; i < batch_size; i++)
|
||||
{
|
||||
HnswSearchCandidate *sc;
|
||||
HnswSearchCandidate *hc;
|
||||
|
||||
if (pairingheap_is_empty(so->discarded))
|
||||
break;
|
||||
|
||||
sc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
|
||||
hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
|
||||
|
||||
ep = lappend(ep, sc);
|
||||
ep = lappend(ep, hc);
|
||||
}
|
||||
|
||||
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, false, &so->v, &so->discarded, false, &so->tuples);
|
||||
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -99,8 +98,8 @@ GetScanValue(IndexScanDesc scan)
|
||||
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
|
||||
|
||||
/* Normalize if needed */
|
||||
if (so->support.normprocinfo != NULL)
|
||||
value = HnswNormValue(so->typeInfo, so->support.collation[0], value);
|
||||
if (so->normprocinfo != NULL)
|
||||
value = HnswNormValue(so->typeInfo, so->collation, value);
|
||||
}
|
||||
|
||||
return value;
|
||||
@@ -127,7 +126,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
|
||||
/* Set support functions */
|
||||
HnswInitSupport(&so->support, index);
|
||||
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
so->collation = index->rd_indcollation[0];
|
||||
|
||||
scan->opaque = so;
|
||||
|
||||
@@ -150,7 +151,7 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
|
||||
|
||||
so->first = true;
|
||||
so->tuples = 0;
|
||||
so->previousDistance = -get_float8_infinity();
|
||||
so->previousDistance = -INFINITY;
|
||||
MemoryContextReset(so->tmpCtx);
|
||||
|
||||
if (keys && scan->numberOfKeys > 0)
|
||||
@@ -200,7 +201,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
*/
|
||||
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||
|
||||
so->w = GetScanItems(scan, value);
|
||||
HnswBench("scan iteration", so->w = GetScanItems(scan, value));
|
||||
|
||||
/* Release shared lock */
|
||||
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||
@@ -215,13 +216,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
for (;;)
|
||||
{
|
||||
char *base = NULL;
|
||||
HnswSearchCandidate *sc;
|
||||
HnswSearchCandidate *hc;
|
||||
HnswElement element;
|
||||
ItemPointer heaptid;
|
||||
|
||||
if (list_length(so->w) == 0)
|
||||
{
|
||||
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_OFF)
|
||||
if (!hnsw_streaming)
|
||||
break;
|
||||
|
||||
/* Empty index */
|
||||
@@ -229,7 +230,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
break;
|
||||
|
||||
/* Reached max number of additional tuples */
|
||||
if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples)
|
||||
if (hnsw_ef_stream != -1 && so->tuples >= hnsw_ef_search + hnsw_ef_stream)
|
||||
{
|
||||
if (pairingheap_is_empty(so->discarded))
|
||||
break;
|
||||
@@ -242,7 +243,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
{
|
||||
if (pairingheap_is_empty(so->discarded))
|
||||
{
|
||||
ereport(DEBUG1,
|
||||
ereport(NOTICE,
|
||||
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
|
||||
errhint("Increase work_mem to scan more tuples.")));
|
||||
|
||||
@@ -265,7 +266,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
*/
|
||||
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||
|
||||
so->w = ResumeScanItems(scan);
|
||||
HnswBench("scan iteration", so->w = ResumeScanItems(scan));
|
||||
|
||||
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
|
||||
|
||||
@@ -278,19 +279,19 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
break;
|
||||
}
|
||||
|
||||
sc = llast(so->w);
|
||||
element = HnswPtrAccess(base, sc->element);
|
||||
hc = llast(so->w);
|
||||
element = HnswPtrAccess(base, hc->element);
|
||||
|
||||
/* Move to next element if no valid heap TIDs */
|
||||
if (!sc->matches || element->heaptidsLength == 0)
|
||||
if (element->heaptidsLength == 0)
|
||||
{
|
||||
so->w = list_delete_last(so->w);
|
||||
|
||||
/* Mark memory as free for next iteration */
|
||||
if (hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF)
|
||||
if (hnsw_streaming)
|
||||
{
|
||||
pfree(element);
|
||||
pfree(sc);
|
||||
pfree(hc);
|
||||
}
|
||||
|
||||
continue;
|
||||
@@ -298,13 +299,10 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
|
||||
heaptid = &element->heaptids[--element->heaptidsLength];
|
||||
|
||||
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_STRICT)
|
||||
{
|
||||
if (sc->distance < so->previousDistance)
|
||||
continue;
|
||||
if (hc->distance < so->previousDistance)
|
||||
continue;
|
||||
|
||||
so->previousDistance = sc->distance;
|
||||
}
|
||||
so->previousDistance = hc->distance;
|
||||
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
|
||||
|
||||
571
src/hnswutils.c
571
src/hnswutils.c
@@ -100,6 +100,12 @@ hash_offset(Size offset)
|
||||
#define SH_DEFINE
|
||||
#include "lib/simplehash.h"
|
||||
|
||||
typedef union
|
||||
{
|
||||
HnswElement element;
|
||||
ItemPointerData indextid;
|
||||
} HnswUnvisited;
|
||||
|
||||
/*
|
||||
* Get the max number of connections in an upper layer for each element in the index
|
||||
*/
|
||||
@@ -140,45 +146,6 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
|
||||
return index_getprocinfo(index, 1, procnum);
|
||||
}
|
||||
|
||||
/*
|
||||
* Init support functions
|
||||
*/
|
||||
void
|
||||
HnswInitSupport(HnswSupport * support, Relation index)
|
||||
{
|
||||
support->procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
|
||||
|
||||
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
|
||||
support->procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
|
||||
|
||||
support->collation = index->rd_indcollation;
|
||||
support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get element tuple size
|
||||
*/
|
||||
Size
|
||||
HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple)
|
||||
{
|
||||
Size size;
|
||||
|
||||
if (useIndexTuple)
|
||||
{
|
||||
IndexTuple itup = HnswPtrAccess(base, element->itup);
|
||||
|
||||
size = IndexTupleSize(itup);
|
||||
}
|
||||
else
|
||||
{
|
||||
Pointer valuePtr = HnswPtrAccess(base, element->value);
|
||||
|
||||
size = VARSIZE_ANY(valuePtr);
|
||||
}
|
||||
|
||||
return HNSW_ELEMENT_TUPLE_SIZE(size);
|
||||
}
|
||||
|
||||
/*
|
||||
* Normalize value
|
||||
*/
|
||||
@@ -192,40 +159,9 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
|
||||
* Check if non-zero norm
|
||||
*/
|
||||
bool
|
||||
HnswCheckNorm(HnswSupport * support, Datum value)
|
||||
HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value)
|
||||
{
|
||||
return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation[0], value)) > 0;
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if index tuples are equal
|
||||
*/
|
||||
bool
|
||||
HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc)
|
||||
{
|
||||
for (int i = 0; i < tupdesc->natts; i++)
|
||||
{
|
||||
bool nullA;
|
||||
bool nullB;
|
||||
|
||||
Datum datumA = index_getattr(a, i + 1, tupdesc, &nullA);
|
||||
Datum datumB = index_getattr(b, i + 1, tupdesc, &nullB);
|
||||
|
||||
if (nullA || nullB)
|
||||
{
|
||||
if (nullA != nullB)
|
||||
return false;
|
||||
}
|
||||
else
|
||||
{
|
||||
Form_pg_attribute att = TupleDescAttr(tupdesc, i);
|
||||
|
||||
if (!datumIsEqual(datumA, datumB, att->attbyval, att->attlen))
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
return DatumGetFloat8(FunctionCall1Coll(procinfo, collation, value)) > 0;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -254,7 +190,7 @@ HnswInitPage(Buffer buf, Page page)
|
||||
/*
|
||||
* Allocate a neighbor array
|
||||
*/
|
||||
HnswNeighborArray *
|
||||
static HnswNeighborArray *
|
||||
HnswInitNeighborArray(int lm, HnswAllocator * allocator)
|
||||
{
|
||||
HnswNeighborArray *a = HnswAlloc(allocator, HNSW_NEIGHBOR_ARRAY_SIZE(lm));
|
||||
@@ -316,7 +252,6 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
|
||||
HnswInitNeighbors(base, element, m, allocator);
|
||||
|
||||
HnswPtrStore(base, element->value, (Pointer) NULL);
|
||||
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
|
||||
|
||||
return element;
|
||||
}
|
||||
@@ -343,7 +278,6 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
|
||||
element->offno = offno;
|
||||
HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
|
||||
HnswPtrStore(base, element->value, (Pointer) NULL);
|
||||
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
|
||||
return element;
|
||||
}
|
||||
|
||||
@@ -455,45 +389,14 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
|
||||
UnlockReleaseBuffer(buf);
|
||||
}
|
||||
|
||||
/*
|
||||
* Form index tuple
|
||||
*/
|
||||
bool
|
||||
HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc)
|
||||
{
|
||||
Datum newValues[2];
|
||||
|
||||
/* Detoast once for all calls */
|
||||
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
|
||||
/* Check value */
|
||||
if (typeInfo->checkValue != NULL)
|
||||
typeInfo->checkValue(DatumGetPointer(value));
|
||||
|
||||
/* Normalize if needed */
|
||||
if (support->normprocinfo != NULL)
|
||||
{
|
||||
if (!HnswCheckNorm(support, value))
|
||||
return false;
|
||||
|
||||
value = HnswNormValue(typeInfo, support->collation[0], value);
|
||||
}
|
||||
|
||||
newValues[0] = value;
|
||||
for (int i = 1; i < tupdesc->natts; i++)
|
||||
newValues[i] = values[i];
|
||||
|
||||
*out = index_form_tuple(tupdesc, newValues, isnull);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/*
|
||||
* Set element tuple, except for neighbor info
|
||||
*/
|
||||
void
|
||||
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple)
|
||||
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
|
||||
{
|
||||
Pointer valuePtr = HnswPtrAccess(base, element->value);
|
||||
|
||||
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
|
||||
etup->level = element->level;
|
||||
etup->deleted = 0;
|
||||
@@ -505,19 +408,7 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool
|
||||
else
|
||||
ItemPointerSetInvalid(&etup->heaptids[i]);
|
||||
}
|
||||
|
||||
if (useIndexTuple)
|
||||
{
|
||||
IndexTuple itup = HnswPtrAccess(base, element->itup);
|
||||
|
||||
memcpy(&etup->data, itup, IndexTupleSize(itup));
|
||||
}
|
||||
else
|
||||
{
|
||||
Pointer valuePtr = HnswPtrAccess(base, element->value);
|
||||
|
||||
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
|
||||
}
|
||||
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -555,11 +446,74 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
|
||||
ntup->version = e->version;
|
||||
}
|
||||
|
||||
/*
|
||||
* Load neighbors from page
|
||||
*/
|
||||
static void
|
||||
LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
|
||||
{
|
||||
char *base = NULL;
|
||||
|
||||
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
|
||||
int neighborCount = (element->level + 2) * m;
|
||||
|
||||
Assert(HnswIsNeighborTuple(ntup));
|
||||
|
||||
HnswInitNeighbors(base, element, m, NULL);
|
||||
|
||||
/* Ensure expected neighbors */
|
||||
if (ntup->count != neighborCount)
|
||||
return;
|
||||
|
||||
for (int i = 0; i < neighborCount; i++)
|
||||
{
|
||||
HnswElement e;
|
||||
int level;
|
||||
HnswCandidate *hc;
|
||||
ItemPointer indextid;
|
||||
HnswNeighborArray *neighbors;
|
||||
|
||||
indextid = &ntup->indextids[i];
|
||||
|
||||
if (!ItemPointerIsValid(indextid))
|
||||
continue;
|
||||
|
||||
e = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
|
||||
|
||||
/* Calculate level based on offset */
|
||||
level = element->level - i / m;
|
||||
if (level < 0)
|
||||
level = 0;
|
||||
|
||||
neighbors = HnswGetNeighbors(base, element, level);
|
||||
hc = &neighbors->items[neighbors->length++];
|
||||
HnswPtrStore(base, hc->element, e);
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Load neighbors
|
||||
*/
|
||||
void
|
||||
HnswLoadNeighbors(HnswElement element, Relation index, int m)
|
||||
{
|
||||
Buffer buf;
|
||||
Page page;
|
||||
|
||||
buf = ReadBuffer(index, element->neighborPage);
|
||||
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
||||
page = BufferGetPage(buf);
|
||||
|
||||
LoadNeighborsFromPage(element, index, page, m);
|
||||
|
||||
UnlockReleaseBuffer(buf);
|
||||
}
|
||||
|
||||
/*
|
||||
* Load an element from a tuple
|
||||
*/
|
||||
void
|
||||
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index)
|
||||
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
|
||||
{
|
||||
element->level = etup->level;
|
||||
element->deleted = etup->deleted;
|
||||
@@ -583,128 +537,17 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
|
||||
if (loadVec)
|
||||
{
|
||||
char *base = NULL;
|
||||
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
|
||||
|
||||
if (HnswUseIndexTuple(index))
|
||||
{
|
||||
IndexTuple itup = CopyIndexTuple((IndexTuple) &etup->data);
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
bool unused;
|
||||
|
||||
HnswPtrStore(base, element->itup, itup);
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
|
||||
}
|
||||
else
|
||||
{
|
||||
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
|
||||
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(value));
|
||||
}
|
||||
HnswPtrStore(base, element->value, DatumGetPointer(value));
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the attribute distance
|
||||
*/
|
||||
static inline double
|
||||
AttributeDistance(double e)
|
||||
{
|
||||
/* TODO Better bias */
|
||||
/* must be >> max(w * g) + 1 / log10(2) */
|
||||
double bias = 4.32;
|
||||
|
||||
return e > 0 ? bias - 1.0 / log10(e + 1) : 0;
|
||||
}
|
||||
|
||||
/*
|
||||
* Calculate the distance between values
|
||||
*/
|
||||
static double
|
||||
HnswGetDistance(IndexTuple itup, Datum vec, HnswQuery * q, Relation index, HnswSupport * support, bool *matches)
|
||||
{
|
||||
double g;
|
||||
|
||||
if (DatumGetPointer(q->value) == NULL)
|
||||
g = 0;
|
||||
else
|
||||
g = DatumGetFloat8(FunctionCall2Coll(support->procinfo[0], support->collation[0], q->value, vec));
|
||||
|
||||
Assert(PointerIsValid(matches));
|
||||
*matches = true;
|
||||
|
||||
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
|
||||
{
|
||||
double w = 0.25;
|
||||
double e = 0.0;
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
|
||||
if (q->keyData)
|
||||
{
|
||||
/* TODO need to pass length of key data */
|
||||
int keyCount = 1;
|
||||
|
||||
for (int i = 0; i < keyCount; i++)
|
||||
{
|
||||
ScanKey key = &q->keyData[i];
|
||||
bool isnull;
|
||||
Datum value = index_getattr(itup, key->sk_attno, tupdesc, &isnull);
|
||||
bool attnull = key->sk_flags & SK_ISNULL;
|
||||
|
||||
if (isnull || attnull)
|
||||
{
|
||||
if (isnull != attnull)
|
||||
{
|
||||
e += 1000;
|
||||
*matches = false;
|
||||
}
|
||||
}
|
||||
else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
|
||||
{
|
||||
double ei = fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[key->sk_attno - 1], support->collation[key->sk_attno - 1], value, key->sk_argument)));
|
||||
|
||||
if (ei > 0)
|
||||
e += ei;
|
||||
else
|
||||
/* Distance is zero for inequality */
|
||||
e += 1000;
|
||||
|
||||
*matches = false;
|
||||
}
|
||||
}
|
||||
|
||||
return w * g + AttributeDistance(e);
|
||||
}
|
||||
else if (q->itup)
|
||||
{
|
||||
int keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1;
|
||||
|
||||
for (int i = 0; i < keyCount; i++)
|
||||
{
|
||||
bool isnull;
|
||||
bool attnull;
|
||||
Datum value = index_getattr(itup, i + 2, tupdesc, &isnull);
|
||||
Datum value2 = index_getattr(q->itup, i + 2, tupdesc, &attnull);
|
||||
|
||||
if (isnull || attnull)
|
||||
{
|
||||
if (isnull != attnull)
|
||||
e += 1000;
|
||||
}
|
||||
else
|
||||
e += fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[i + 1], support->collation[i + 1], value, value2)));
|
||||
}
|
||||
|
||||
return w * g + AttributeDistance(e);
|
||||
}
|
||||
}
|
||||
|
||||
return g;
|
||||
}
|
||||
|
||||
/*
|
||||
* Load an element and optionally get its distance from q
|
||||
*/
|
||||
static void
|
||||
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
|
||||
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance, HnswElement * element)
|
||||
{
|
||||
Buffer buf;
|
||||
Page page;
|
||||
@@ -722,23 +565,10 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
|
||||
/* Calculate distance */
|
||||
if (distance != NULL)
|
||||
{
|
||||
IndexTuple itup = NULL;
|
||||
Datum value;
|
||||
|
||||
if (HnswUseIndexTuple(index))
|
||||
{
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
bool unused;
|
||||
|
||||
itup = (IndexTuple) &etup->data;
|
||||
value = index_getattr(itup, 1, tupdesc, &unused);
|
||||
}
|
||||
if (DatumGetPointer(*q) == NULL)
|
||||
*distance = 0;
|
||||
else
|
||||
{
|
||||
value = PointerGetDatum(&etup->data);
|
||||
}
|
||||
|
||||
*distance = HnswGetDistance(itup, value, q, index, support, matches);
|
||||
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data)));
|
||||
}
|
||||
|
||||
/* Load element */
|
||||
@@ -747,7 +577,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
|
||||
if (*element == NULL)
|
||||
*element = HnswInitElementFromBlock(blkno, offno);
|
||||
|
||||
HnswLoadElementFromTuple(*element, etup, true, loadVec, index);
|
||||
HnswLoadElementFromTuple(*element, etup, true, loadVec);
|
||||
}
|
||||
|
||||
UnlockReleaseBuffer(buf);
|
||||
@@ -757,37 +587,36 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
|
||||
* Load an element and optionally get its distance from q
|
||||
*/
|
||||
void
|
||||
HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
|
||||
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec, float *maxDistance)
|
||||
{
|
||||
HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, index, support, loadVec, maxDistance, &element);
|
||||
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, procinfo, collation, loadVec, maxDistance, &element);
|
||||
}
|
||||
|
||||
/*
|
||||
* Get the distance for an element
|
||||
*/
|
||||
static double
|
||||
GetElementDistance(char *base, HnswElement element, bool *matches, HnswQuery * q, Relation index, HnswSupport * support)
|
||||
static float
|
||||
GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
Datum value = HnswGetValue(base, element);
|
||||
IndexTuple itup = HnswPtrAccess(base, element->itup);
|
||||
|
||||
return HnswGetDistance(itup, value, q, index, support, matches);
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, value));
|
||||
}
|
||||
|
||||
/*
|
||||
* Create a candidate for the entry point
|
||||
*/
|
||||
HnswSearchCandidate *
|
||||
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, bool inMemory)
|
||||
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
|
||||
{
|
||||
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
|
||||
HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate));
|
||||
|
||||
HnswPtrStore(base, sc->element, entryPoint);
|
||||
if (inMemory)
|
||||
sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, index, support);
|
||||
HnswPtrStore(base, hc->element, entryPoint);
|
||||
if (index == NULL)
|
||||
hc->distance = GetElementDistance(base, entryPoint, q, procinfo, collation);
|
||||
else
|
||||
HnswLoadElement(entryPoint, &sc->distance, &sc->matches, q, index, support, loadVec, NULL);
|
||||
return sc;
|
||||
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec, NULL);
|
||||
return hc;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -839,9 +668,9 @@ CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b,
|
||||
* Init visited
|
||||
*/
|
||||
static inline void
|
||||
InitVisited(char *base, visited_hash * v, bool inMemory, int ef, int m)
|
||||
InitVisited(char *base, visited_hash * v, Relation index, int ef, int m)
|
||||
{
|
||||
if (!inMemory)
|
||||
if (index != NULL)
|
||||
v->tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
|
||||
else if (base != NULL)
|
||||
v->offsets = offsethash_create(CurrentMemoryContext, ef * m * 2, NULL);
|
||||
@@ -853,9 +682,9 @@ InitVisited(char *base, visited_hash * v, bool inMemory, int ef, int m)
|
||||
* Add to visited
|
||||
*/
|
||||
static inline void
|
||||
AddToVisited(char *base, visited_hash * v, HnswElementPtr elementPtr, bool inMemory, bool *found)
|
||||
AddToVisited(char *base, visited_hash * v, HnswElementPtr elementPtr, Relation index, bool *found)
|
||||
{
|
||||
if (!inMemory)
|
||||
if (index != NULL)
|
||||
{
|
||||
HnswElement element = HnswPtrAccess(base, elementPtr);
|
||||
ItemPointerData indextid;
|
||||
@@ -916,7 +745,7 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswUnvisited * unv
|
||||
HnswCandidate *hc = &localNeighborhood->items[i];
|
||||
bool found;
|
||||
|
||||
AddToVisited(base, v, hc->element, true, &found);
|
||||
AddToVisited(base, v, hc->element, NULL, &found);
|
||||
|
||||
if (!found)
|
||||
unvisited[(*unvisitedLength)++].element = HnswPtrAccess(base, hc->element);
|
||||
@@ -924,15 +753,18 @@ HnswLoadUnvisitedFromMemory(char *base, HnswElement element, HnswUnvisited * unv
|
||||
}
|
||||
|
||||
/*
|
||||
* Load neighbor index TIDs
|
||||
* Load unvisited neighbors from disk
|
||||
*/
|
||||
bool
|
||||
HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc)
|
||||
static void
|
||||
HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc)
|
||||
{
|
||||
Buffer buf;
|
||||
Page page;
|
||||
HnswNeighborTuple ntup;
|
||||
int start;
|
||||
ItemPointerData indextids[HNSW_MAX_M * 2];
|
||||
|
||||
*unvisitedLength = 0;
|
||||
|
||||
buf = ReadBuffer(index, element->neighborPage);
|
||||
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
||||
@@ -947,29 +779,14 @@ HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation i
|
||||
if (ntup->version != element->version || ntup->count != (element->level + 2) * m)
|
||||
{
|
||||
UnlockReleaseBuffer(buf);
|
||||
return false;
|
||||
return;
|
||||
}
|
||||
|
||||
/* Copy to minimize lock time */
|
||||
start = (element->level - lc) * m;
|
||||
memcpy(indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
|
||||
memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
|
||||
|
||||
UnlockReleaseBuffer(buf);
|
||||
return true;
|
||||
}
|
||||
|
||||
/*
|
||||
* Load unvisited neighbors from disk
|
||||
*/
|
||||
static void
|
||||
HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *unvisitedLength, visited_hash * v, Relation index, int m, int lm, int lc)
|
||||
{
|
||||
ItemPointerData indextids[HNSW_MAX_M * 2];
|
||||
|
||||
*unvisitedLength = 0;
|
||||
|
||||
if (!HnswLoadNeighborTids(element, indextids, index, m, lm, lc))
|
||||
return;
|
||||
|
||||
for (int i = 0; i < lm; i++)
|
||||
{
|
||||
@@ -990,7 +807,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
|
||||
* Algorithm 2 from paper
|
||||
*/
|
||||
List *
|
||||
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
|
||||
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 *w = NIL;
|
||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
||||
@@ -1003,8 +820,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
int lm = HnswGetLayerM(m, lc);
|
||||
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
|
||||
int unvisitedLength;
|
||||
uint64 additional = 0;
|
||||
uint64 maxAdditional = q->keyData && lc == 0 ? 10000 : 0;
|
||||
|
||||
if (v == NULL)
|
||||
{
|
||||
@@ -1014,14 +829,14 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
|
||||
if (initVisited)
|
||||
{
|
||||
InitVisited(base, v, inMemory, ef, m);
|
||||
InitVisited(base, v, index, ef, m);
|
||||
|
||||
if (discarded != NULL)
|
||||
*discarded = pairingheap_allocate(CompareNearestDiscardedCandidates, NULL);
|
||||
}
|
||||
|
||||
/* Create local memory for neighborhood if needed */
|
||||
if (inMemory)
|
||||
if (index == NULL)
|
||||
{
|
||||
neighborhoodSize = HNSW_NEIGHBOR_ARRAY_SIZE(lm);
|
||||
localNeighborhood = palloc(neighborhoodSize);
|
||||
@@ -1030,30 +845,26 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
/* Add entry points to v, C, and W */
|
||||
foreach(lc2, ep)
|
||||
{
|
||||
HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2);
|
||||
HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2);
|
||||
bool found;
|
||||
|
||||
if (initVisited)
|
||||
{
|
||||
AddToVisited(base, v, sc->element, inMemory, &found);
|
||||
AddToVisited(base, v, hc->element, index, &found);
|
||||
|
||||
if (tuples != NULL)
|
||||
(*tuples)++;
|
||||
}
|
||||
|
||||
pairingheap_add(C, &sc->c_node);
|
||||
pairingheap_add(W, &sc->w_node);
|
||||
|
||||
/* Do not count elements that do not match filter towards ef */
|
||||
if (!sc->matches && ++additional <= maxAdditional)
|
||||
continue;
|
||||
pairingheap_add(C, &hc->c_node);
|
||||
pairingheap_add(W, &hc->w_node);
|
||||
|
||||
/*
|
||||
* Do not count elements being deleted towards ef when vacuuming. It
|
||||
* would be ideal to do this for inserts as well, but this could
|
||||
* affect insert performance.
|
||||
*/
|
||||
if (CountElement(skipElement, HnswPtrAccess(base, sc->element)))
|
||||
if (CountElement(skipElement, HnswPtrAccess(base, hc->element)))
|
||||
wlen++;
|
||||
}
|
||||
|
||||
@@ -1068,7 +879,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
|
||||
cElement = HnswPtrAccess(base, c->element);
|
||||
|
||||
if (inMemory)
|
||||
if (index == NULL)
|
||||
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize);
|
||||
else
|
||||
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
|
||||
@@ -1080,16 +891,15 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
{
|
||||
HnswElement eElement;
|
||||
HnswSearchCandidate *e;
|
||||
double eDistance;
|
||||
bool eMatches;
|
||||
float eDistance;
|
||||
bool alwaysAdd = wlen < ef;
|
||||
|
||||
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
|
||||
|
||||
if (inMemory)
|
||||
if (index == NULL)
|
||||
{
|
||||
eElement = unvisited[i].element;
|
||||
eDistance = GetElementDistance(base, eElement, &eMatches, q, index, support);
|
||||
eDistance = GetElementDistance(base, eElement, q, procinfo, collation);
|
||||
}
|
||||
else
|
||||
{
|
||||
@@ -1099,10 +909,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
|
||||
/* Avoid any allocations if not adding */
|
||||
eElement = NULL;
|
||||
HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
|
||||
|
||||
if (eElement == NULL)
|
||||
continue;
|
||||
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
|
||||
}
|
||||
|
||||
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
|
||||
@@ -1127,7 +934,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
e = palloc(sizeof(HnswSearchCandidate));
|
||||
HnswPtrStore(base, e->element, eElement);
|
||||
e->distance = eDistance;
|
||||
e->matches = eMatches;
|
||||
pairingheap_add(C, &e->c_node);
|
||||
pairingheap_add(W, &e->w_node);
|
||||
|
||||
@@ -1138,10 +944,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
*/
|
||||
if (CountElement(skipElement, eElement))
|
||||
{
|
||||
/* Do not count elements that do not match filter towards ef */
|
||||
if (!e->matches && ++additional <= maxAdditional)
|
||||
continue;
|
||||
|
||||
wlen++;
|
||||
|
||||
/* No need to decrement wlen */
|
||||
@@ -1159,9 +961,9 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
|
||||
/* Add each element of W to w */
|
||||
while (!pairingheap_is_empty(W))
|
||||
{
|
||||
HnswSearchCandidate *sc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
|
||||
HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
|
||||
|
||||
w = lappend(w, sc);
|
||||
w = lappend(w, hc);
|
||||
}
|
||||
|
||||
return w;
|
||||
@@ -1215,28 +1017,32 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
|
||||
return 0;
|
||||
}
|
||||
|
||||
/*
|
||||
* Calculate the distance between elements
|
||||
*/
|
||||
static float
|
||||
HnswGetDistance(char *base, HnswElement a, HnswElement b, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
Datum aValue = HnswGetValue(base, a);
|
||||
Datum bValue = HnswGetValue(base, b);
|
||||
|
||||
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, aValue, bValue));
|
||||
}
|
||||
|
||||
/*
|
||||
* Check if an element is closer to q than any element from R
|
||||
*/
|
||||
static bool
|
||||
CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswSupport * support)
|
||||
CheckElementCloser(char *base, HnswCandidate * e, List *r, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
HnswElement eElement = HnswPtrAccess(base, e->element);
|
||||
HnswQuery q;
|
||||
ListCell *lc2;
|
||||
|
||||
q.value = HnswGetValue(base, eElement);
|
||||
q.itup = HnswPtrAccess(base, eElement->itup);
|
||||
q.keyData = NULL;
|
||||
|
||||
foreach(lc2, r)
|
||||
{
|
||||
HnswCandidate *ri = lfirst(lc2);
|
||||
HnswElement riElement = HnswPtrAccess(base, ri->element);
|
||||
Datum riValue = HnswGetValue(base, riElement);
|
||||
IndexTuple ritup = HnswPtrAccess(base, riElement->itup);
|
||||
bool matches;
|
||||
float distance = HnswGetDistance(ritup, riValue, &q, index, support, &matches);
|
||||
float distance = HnswGetDistance(base, eElement, riElement, procinfo, collation);
|
||||
|
||||
if (distance <= e->distance)
|
||||
return false;
|
||||
@@ -1249,14 +1055,15 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswS
|
||||
* Algorithm 4 from paper
|
||||
*/
|
||||
static List *
|
||||
SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
|
||||
SelectNeighbors(char *base, List *c, int lm, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
|
||||
{
|
||||
List *r = NIL;
|
||||
List *w = list_copy(c);
|
||||
HnswCandidate **wd;
|
||||
int wdlen = 0;
|
||||
int wdoff = 0;
|
||||
bool mustCalculate = !(*closerSet);
|
||||
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e2, lc);
|
||||
bool mustCalculate = !neighbors->closerSet;
|
||||
List *added = NIL;
|
||||
bool removedAny = false;
|
||||
|
||||
@@ -1283,7 +1090,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
|
||||
|
||||
/* Use previous state of r and wd to skip work when possible */
|
||||
if (mustCalculate)
|
||||
e->closer = CheckElementCloser(base, e, r, index, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
else if (list_length(added) > 0)
|
||||
{
|
||||
/* Keep Valgrind happy for in-memory, parallel builds */
|
||||
@@ -1296,7 +1103,8 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
|
||||
*/
|
||||
if (e->closer)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, added, index, support);
|
||||
e->closer = CheckElementCloser(base, e, added, procinfo, collation);
|
||||
|
||||
if (!e->closer)
|
||||
removedAny = true;
|
||||
}
|
||||
@@ -1308,7 +1116,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
|
||||
*/
|
||||
if (removedAny)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, r, index, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
if (e->closer)
|
||||
added = lappend(added, e);
|
||||
}
|
||||
@@ -1316,7 +1124,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
|
||||
}
|
||||
else if (e == newCandidate)
|
||||
{
|
||||
e->closer = CheckElementCloser(base, e, r, index, support);
|
||||
e->closer = CheckElementCloser(base, e, r, procinfo, collation);
|
||||
if (e->closer)
|
||||
added = lappend(added, e);
|
||||
}
|
||||
@@ -1332,7 +1140,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
|
||||
}
|
||||
|
||||
/* Cached value can only be used in future if sorted deterministically */
|
||||
*closerSet = sortCandidates;
|
||||
neighbors->closerSet = sortCandidates;
|
||||
|
||||
/* Keep pruned connections */
|
||||
while (wdoff < wdlen && list_length(r) < lm)
|
||||
@@ -1367,16 +1175,18 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
|
||||
* Update connections
|
||||
*/
|
||||
void
|
||||
HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support)
|
||||
HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
|
||||
{
|
||||
HnswCandidate newHc;
|
||||
HnswElement hce = HnswPtrAccess(base, hc->element);
|
||||
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
|
||||
HnswCandidate hc2;
|
||||
|
||||
HnswPtrStore(base, newHc.element, newElement);
|
||||
newHc.distance = distance;
|
||||
HnswPtrStore(base, hc2.element, element);
|
||||
hc2.distance = hc->distance;
|
||||
|
||||
if (neighbors->length < lm)
|
||||
if (currentNeighbors->length < lm)
|
||||
{
|
||||
neighbors->items[neighbors->length++] = newHc;
|
||||
currentNeighbors->items[currentNeighbors->length++] = hc2;
|
||||
|
||||
/* Track update */
|
||||
if (updateIdx != NULL)
|
||||
@@ -1385,26 +1195,54 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
|
||||
else
|
||||
{
|
||||
/* Shrink connections */
|
||||
List *c = NIL;
|
||||
HnswCandidate *pruned = NULL;
|
||||
|
||||
/* Add candidates */
|
||||
for (int i = 0; i < neighbors->length; i++)
|
||||
c = lappend(c, &neighbors->items[i]);
|
||||
c = lappend(c, &newHc);
|
||||
/* Load elements on insert */
|
||||
if (index != NULL)
|
||||
{
|
||||
Datum q = HnswGetValue(base, hce);
|
||||
|
||||
SelectNeighbors(base, c, lm, index, support, &neighbors->closerSet, &newHc, &pruned, true);
|
||||
for (int i = 0; i < currentNeighbors->length; i++)
|
||||
{
|
||||
HnswCandidate *hc3 = ¤tNeighbors->items[i];
|
||||
HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
|
||||
|
||||
if (HnswPtrIsNull(base, hc3Element->value))
|
||||
HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true, NULL);
|
||||
else
|
||||
hc3->distance = GetElementDistance(base, hc3Element, q, procinfo, collation);
|
||||
|
||||
/* Prune element if being deleted */
|
||||
if (hc3Element->heaptidsLength == 0)
|
||||
{
|
||||
pruned = ¤tNeighbors->items[i];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* Should not happen */
|
||||
if (pruned == NULL)
|
||||
return;
|
||||
{
|
||||
List *c = NIL;
|
||||
|
||||
/* Add candidates */
|
||||
for (int i = 0; i < currentNeighbors->length; i++)
|
||||
c = lappend(c, ¤tNeighbors->items[i]);
|
||||
c = lappend(c, &hc2);
|
||||
|
||||
SelectNeighbors(base, c, lm, lc, procinfo, collation, hce, &hc2, &pruned, true);
|
||||
|
||||
/* Should not happen */
|
||||
if (pruned == NULL)
|
||||
return;
|
||||
}
|
||||
|
||||
/* Find and replace the pruned element */
|
||||
for (int i = 0; i < neighbors->length; i++)
|
||||
for (int i = 0; i < currentNeighbors->length; i++)
|
||||
{
|
||||
if (HnswPtrEqual(base, neighbors->items[i].element, pruned->element))
|
||||
if (HnswPtrEqual(base, currentNeighbors->items[i].element, pruned->element))
|
||||
{
|
||||
neighbors->items[i] = newHc;
|
||||
currentNeighbors->items[i] = hc2;
|
||||
|
||||
/* Track update */
|
||||
if (updateIdx != NULL)
|
||||
@@ -1464,22 +1302,17 @@ PrecomputeHash(char *base, HnswElement element)
|
||||
* Algorithm 1 from paper
|
||||
*/
|
||||
void
|
||||
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, bool inMemory)
|
||||
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
|
||||
{
|
||||
List *ep;
|
||||
List *w;
|
||||
int level = element->level;
|
||||
int entryLevel;
|
||||
HnswQuery q;
|
||||
|
||||
Datum q = HnswGetValue(base, element);
|
||||
HnswElement skipElement = existing ? element : NULL;
|
||||
|
||||
q.value = HnswGetValue(base, element);
|
||||
q.itup = HnswPtrAccess(base, element->itup);
|
||||
q.keyData = NULL;
|
||||
|
||||
/* Precompute hash */
|
||||
if (inMemory)
|
||||
if (index == NULL)
|
||||
PrecomputeHash(base, element);
|
||||
|
||||
/* No neighbors if no entry point */
|
||||
@@ -1487,13 +1320,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
return;
|
||||
|
||||
/* Get entry point and level */
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true, inMemory));
|
||||
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, true));
|
||||
entryLevel = entryPoint->level;
|
||||
|
||||
/* 1st phase: greedy search to insert level */
|
||||
for (int lc = entryLevel; lc >= level + 1; lc--)
|
||||
{
|
||||
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
@@ -1512,7 +1345,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
List *lw = NIL;
|
||||
ListCell *lc2;
|
||||
|
||||
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
|
||||
|
||||
/* Convert search candidates to candidates */
|
||||
foreach(lc2, w)
|
||||
@@ -1528,7 +1361,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
|
||||
/* Elements being deleted or skipped can help with search */
|
||||
/* but should be removed before selecting neighbors */
|
||||
if (!inMemory)
|
||||
if (index != NULL)
|
||||
lw = RemoveElements(base, lw, skipElement);
|
||||
|
||||
/*
|
||||
@@ -1536,7 +1369,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
* sortCandidates to true for in-memory builds to enable closer
|
||||
* caching, but there does not seem to be a difference in performance.
|
||||
*/
|
||||
neighbors = SelectNeighbors(base, lw, lm, index, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
|
||||
neighbors = SelectNeighbors(base, lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
|
||||
|
||||
AddConnections(base, element, neighbors, lc);
|
||||
|
||||
|
||||
@@ -184,12 +184,13 @@ static void
|
||||
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint)
|
||||
{
|
||||
Relation index = vacuumstate->index;
|
||||
HnswSupport *support = &vacuumstate->support;
|
||||
Buffer buf;
|
||||
Page page;
|
||||
GenericXLogState *state;
|
||||
int m = vacuumstate->m;
|
||||
int efConstruction = vacuumstate->efConstruction;
|
||||
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);
|
||||
@@ -204,7 +205,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
||||
element->heaptidsLength = 0;
|
||||
|
||||
/* Find neighbors for element, skipping itself */
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, false);
|
||||
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
|
||||
|
||||
/* Zero memory for each element */
|
||||
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
|
||||
@@ -228,7 +229,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
|
||||
UnlockReleaseBuffer(buf);
|
||||
|
||||
/* Update neighbors */
|
||||
HnswUpdateNeighborsOnDisk(index, support, element, m, true, false);
|
||||
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -238,7 +239,6 @@ static void
|
||||
RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
|
||||
{
|
||||
Relation index = vacuumstate->index;
|
||||
HnswSupport *support = &vacuumstate->support;
|
||||
HnswElement highestPoint = &vacuumstate->highestPoint;
|
||||
HnswElement entryPoint;
|
||||
MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
|
||||
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
|
||||
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
|
||||
|
||||
/* Load element */
|
||||
HnswLoadElement(highestPoint, NULL, NULL, NULL, index, support, true, NULL);
|
||||
HnswLoadElement(highestPoint, 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, NULL, index, support, true, NULL);
|
||||
HnswLoadElement(entryPoint, 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, index);
|
||||
HnswLoadElementFromTuple(element, etup, false, true);
|
||||
|
||||
elements = lappend(elements, element);
|
||||
}
|
||||
@@ -440,7 +440,6 @@ 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
|
||||
@@ -522,14 +521,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
|
||||
|
||||
/* Overwrite element */
|
||||
etup->deleted = 1;
|
||||
if (useIndexTuple)
|
||||
{
|
||||
IndexTuple itup = (IndexTuple) &etup->data;
|
||||
|
||||
MemSet(itup, 0, IndexTupleSize(itup));
|
||||
}
|
||||
else
|
||||
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
|
||||
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
|
||||
|
||||
/* Overwrite neighbors */
|
||||
for (int i = 0; i < ntup->count; i++)
|
||||
@@ -589,13 +581,13 @@ 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];
|
||||
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
|
||||
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
|
||||
"Hnsw vacuum temporary context",
|
||||
ALLOCSET_DEFAULT_SIZES);
|
||||
|
||||
HnswInitSupport(&vacuumstate->support, index);
|
||||
|
||||
/* Get m from metapage */
|
||||
HnswGetMetaPageInfo(index, &vacuumstate->m, NULL);
|
||||
|
||||
|
||||
@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
|
||||
* Add tuple to sort
|
||||
*/
|
||||
static void
|
||||
AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, IvfflatBuildState * buildstate)
|
||||
AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate)
|
||||
{
|
||||
double distance;
|
||||
double minDistance = DBL_MAX;
|
||||
@@ -184,11 +184,6 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, Ivf
|
||||
slot->tts_isnull[1] = false;
|
||||
slot->tts_values[2] = value;
|
||||
slot->tts_isnull[2] = false;
|
||||
for (int i = 1; i < buildstate->tupdesc->natts; i++)
|
||||
{
|
||||
slot->tts_values[2 + i] = values[i];
|
||||
slot->tts_isnull[2 + i] = isnull[i];
|
||||
}
|
||||
ExecStoreVirtualTuple(slot);
|
||||
|
||||
/*
|
||||
@@ -220,7 +215,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
|
||||
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
|
||||
|
||||
/* Add tuple to sort */
|
||||
AddTupleToSort(index, tid, values, isnull, buildstate);
|
||||
AddTupleToSort(index, tid, values, buildstate);
|
||||
|
||||
/* Reset memory context */
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
@@ -231,20 +226,19 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
|
||||
* Get index tuple from sort state
|
||||
*/
|
||||
static inline void
|
||||
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, Datum *values, bool *isnull, IndexTuple *itup, int *list)
|
||||
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list)
|
||||
{
|
||||
Datum value;
|
||||
bool isnull;
|
||||
|
||||
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
|
||||
{
|
||||
bool unused;
|
||||
|
||||
*list = DatumGetInt32(slot_getattr(slot, 1, &unused));
|
||||
|
||||
for (int i = 0; i < tupdesc->natts; i++)
|
||||
values[i] = slot_getattr(slot, 3 + i, &isnull[i]);
|
||||
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
|
||||
value = slot_getattr(slot, 3, &isnull);
|
||||
|
||||
/* Form the index tuple */
|
||||
*itup = index_form_tuple(tupdesc, values, isnull);
|
||||
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused)));
|
||||
*itup = index_form_tuple(tupdesc, &value, &isnull);
|
||||
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
|
||||
}
|
||||
else
|
||||
*list = -1;
|
||||
@@ -260,16 +254,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
|
||||
IndexTuple itup = NULL; /* silence compiler warning */
|
||||
int64 inserted = 0;
|
||||
|
||||
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
|
||||
TupleDesc tupdesc = buildstate->tupdesc;
|
||||
Datum *values = palloc(tupdesc->natts * sizeof(Datum));
|
||||
bool *isnull = palloc(tupdesc->natts * sizeof(bool));
|
||||
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple);
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
|
||||
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
|
||||
|
||||
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples);
|
||||
|
||||
GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
|
||||
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
|
||||
|
||||
for (int i = 0; i < buildstate->centers->length; i++)
|
||||
{
|
||||
@@ -305,7 +297,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
|
||||
|
||||
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted);
|
||||
|
||||
GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list);
|
||||
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
|
||||
}
|
||||
|
||||
insertPage = BufferGetBlockNumber(buf);
|
||||
@@ -315,9 +307,6 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
|
||||
/* Set the start and insert pages */
|
||||
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
|
||||
}
|
||||
|
||||
pfree(values);
|
||||
pfree(isnull);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -330,7 +319,6 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
|
||||
buildstate->index = index;
|
||||
buildstate->indexInfo = indexInfo;
|
||||
buildstate->typeInfo = IvfflatGetTypeInfo(index);
|
||||
buildstate->tupdesc = RelationGetDescr(index);
|
||||
|
||||
buildstate->lists = IvfflatGetLists(index);
|
||||
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
|
||||
@@ -368,13 +356,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
|
||||
errmsg("dimensions must be greater than one for this opclass")));
|
||||
|
||||
/* Create tuple description for sorting */
|
||||
buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts);
|
||||
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
|
||||
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
|
||||
for (int i = 0; i < buildstate->tupdesc->natts; i++)
|
||||
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);
|
||||
buildstate->tupdesc = CreateTemplateTupleDesc(3);
|
||||
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
|
||||
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
|
||||
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
|
||||
|
||||
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
|
||||
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
|
||||
|
||||
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
|
||||
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
|
||||
@@ -646,7 +633,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
|
||||
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
|
||||
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
|
||||
buildstate.centers->length = buildstate.centers->maxlen;
|
||||
ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate);
|
||||
ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate);
|
||||
buildstate.sortstate = ivfspool->sortstate;
|
||||
scan = table_beginscan_parallel(ivfspool->heap,
|
||||
ParallelTableScanFromIvfflatShared(ivfshared));
|
||||
@@ -963,7 +950,7 @@ AssignTuples(IvfflatBuildState * buildstate)
|
||||
}
|
||||
|
||||
/* Begin serial/leader tuplesort */
|
||||
buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);
|
||||
buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate);
|
||||
|
||||
/* Add tuples to sort */
|
||||
if (buildstate->heap != NULL)
|
||||
|
||||
@@ -17,16 +17,8 @@
|
||||
#endif
|
||||
|
||||
int ivfflat_probes;
|
||||
int ivfflat_iterative_search;
|
||||
int ivfflat_iterative_search_max_probes;
|
||||
static relopt_kind ivfflat_relopt_kind;
|
||||
|
||||
static const struct config_enum_entry ivfflat_iterative_search_options[] = {
|
||||
{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
|
||||
{"on", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
|
||||
{NULL, 0, false}
|
||||
};
|
||||
|
||||
/*
|
||||
* Initialize index options and variables
|
||||
*/
|
||||
@@ -41,14 +33,6 @@ IvfflatInit(void)
|
||||
"Valid range is 1..lists.", &ivfflat_probes,
|
||||
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||
|
||||
DefineCustomEnumVariable("ivfflat.iterative_search", "Sets whether to use iterative search",
|
||||
NULL, &ivfflat_iterative_search,
|
||||
IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||
|
||||
DefineCustomIntVariable("ivfflat.iterative_search_max_probes", "Sets the max number of probes for iterative search",
|
||||
"Zero sets to the number of lists", &ivfflat_iterative_search_max_probes,
|
||||
0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||
|
||||
MarkGUCPrefixReserved("ivfflat");
|
||||
}
|
||||
|
||||
@@ -85,8 +69,6 @@ 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;
|
||||
|
||||
@@ -103,8 +85,6 @@ 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);
|
||||
@@ -114,26 +94,34 @@ 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 */
|
||||
startupPages = costs.numIndexPages * ratio;
|
||||
if (startupPages > path->indexinfo->rel->pages && ratio < 0.5)
|
||||
if (costs.numIndexPages > path->indexinfo->rel->pages && ratio < 0.5)
|
||||
{
|
||||
/* Change rest of page cost from random to sequential */
|
||||
costs.indexStartupCost -= (1 - sequentialRatio) * startupPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
/* Change all page cost from random to sequential */
|
||||
costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
|
||||
/* Remove cost of extra pages */
|
||||
costs.indexStartupCost -= (startupPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
|
||||
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);
|
||||
}
|
||||
|
||||
*indexStartupCost = costs.indexStartupCost;
|
||||
/* Use total cost since most work happens before first tuple is returned */
|
||||
*indexStartupCost = costs.indexTotalCost;
|
||||
*indexTotalCost = costs.indexTotalCost;
|
||||
*indexSelectivity = costs.indexSelectivity;
|
||||
*indexCorrelation = costs.indexCorrelation;
|
||||
|
||||
@@ -80,14 +80,6 @@
|
||||
|
||||
/* Variables */
|
||||
extern int ivfflat_probes;
|
||||
extern int ivfflat_iterative_search;
|
||||
extern int ivfflat_iterative_search_max_probes;
|
||||
|
||||
typedef enum IvfflatIterativeSearchType
|
||||
{
|
||||
IVFFLAT_ITERATIVE_SEARCH_OFF,
|
||||
IVFFLAT_ITERATIVE_SEARCH_RELAXED
|
||||
} IvfflatIterativeSearchType;
|
||||
|
||||
typedef struct VectorArrayData
|
||||
{
|
||||
@@ -173,7 +165,6 @@ typedef struct IvfflatBuildState
|
||||
Relation index;
|
||||
IndexInfo *indexInfo;
|
||||
const IvfflatTypeInfo *typeInfo;
|
||||
TupleDesc tupdesc;
|
||||
|
||||
/* Settings */
|
||||
int dimensions;
|
||||
@@ -207,7 +198,7 @@ typedef struct IvfflatBuildState
|
||||
|
||||
/* Sorting */
|
||||
Tuplesortstate *sortstate;
|
||||
TupleDesc sortdesc;
|
||||
TupleDesc tupdesc;
|
||||
TupleTableSlot *slot;
|
||||
|
||||
/* Memory */
|
||||
@@ -256,10 +247,8 @@ typedef struct IvfflatScanOpaqueData
|
||||
{
|
||||
const IvfflatTypeInfo *typeInfo;
|
||||
int probes;
|
||||
int maxProbes;
|
||||
int dimensions;
|
||||
bool first;
|
||||
Datum value;
|
||||
|
||||
/* Sorting */
|
||||
Tuplesortstate *sortstate;
|
||||
@@ -276,8 +265,6 @@ typedef struct IvfflatScanOpaqueData
|
||||
|
||||
/* Lists */
|
||||
pairingheap *listQueue;
|
||||
BlockNumber *listPages;
|
||||
int listIndex;
|
||||
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
|
||||
} IvfflatScanOpaqueData;
|
||||
|
||||
|
||||
@@ -78,8 +78,6 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
|
||||
BlockNumber insertPage = InvalidBlockNumber;
|
||||
ListInfo listInfo;
|
||||
BlockNumber originalInsertPage;
|
||||
TupleDesc tupdesc = RelationGetDescr(index);
|
||||
Datum *newValues = palloc(tupdesc->natts * sizeof(Datum));
|
||||
|
||||
/* Detoast once for all calls */
|
||||
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
|
||||
@@ -100,16 +98,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
|
||||
IvfflatGetMetaPageInfo(index, NULL, NULL);
|
||||
|
||||
/* Find the insert page - sets the page and list info */
|
||||
FindInsertPage(index, &value, &insertPage, &listInfo);
|
||||
FindInsertPage(index, values, &insertPage, &listInfo);
|
||||
Assert(BlockNumberIsValid(insertPage));
|
||||
originalInsertPage = insertPage;
|
||||
|
||||
newValues[0] = value;
|
||||
for (int i = 1; i < tupdesc->natts; i++)
|
||||
newValues[i] = values[i];
|
||||
|
||||
/* Form tuple */
|
||||
itup = index_form_tuple(tupdesc, newValues, isnull);
|
||||
itup = index_form_tuple(RelationGetDescr(index), &value, isnull);
|
||||
itup->t_tid = *heap_tid;
|
||||
|
||||
/* Get tuple size */
|
||||
|
||||
@@ -65,7 +65,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
||||
/* Use procinfo from the index instead of scan key for performance */
|
||||
distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
|
||||
|
||||
if (listCount < so->maxProbes)
|
||||
if (listCount < so->probes)
|
||||
{
|
||||
IvfflatScanList *scanlist;
|
||||
|
||||
@@ -78,7 +78,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
||||
pairingheap_add(so->listQueue, &scanlist->ph_node);
|
||||
|
||||
/* Calculate max distance */
|
||||
if (listCount == so->maxProbes)
|
||||
if (listCount == so->probes)
|
||||
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
|
||||
}
|
||||
else if (distance < maxDistance)
|
||||
@@ -102,11 +102,6 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
||||
|
||||
UnlockReleaseBuffer(cbuf);
|
||||
}
|
||||
|
||||
for (int i = listCount - 1; i >= 0; i--)
|
||||
so->listPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
|
||||
|
||||
Assert(pairingheap_is_empty(so->listQueue));
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -119,14 +114,11 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
|
||||
double tuples = 0;
|
||||
TupleTableSlot *slot = so->vslot;
|
||||
int batchProbes = 0;
|
||||
|
||||
tuplesort_reset(so->sortstate);
|
||||
|
||||
/* Search closest probes lists */
|
||||
while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes)
|
||||
while (!pairingheap_is_empty(so->listQueue))
|
||||
{
|
||||
BlockNumber searchPage = so->listPages[so->listIndex++];
|
||||
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
|
||||
|
||||
/* Search all entry pages for list */
|
||||
while (BlockNumberIsValid(searchPage))
|
||||
@@ -174,17 +166,13 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
}
|
||||
}
|
||||
|
||||
if (tuples < 100 && ivfflat_iterative_search == IVFFLAT_ITERATIVE_SEARCH_OFF)
|
||||
if (tuples < 100)
|
||||
ereport(DEBUG1,
|
||||
(errmsg("index scan found few tuples"),
|
||||
errdetail("Index may have been created with little data."),
|
||||
errhint("Recreate the index and possibly decrease lists.")));
|
||||
|
||||
tuplesort_performsort(so->sortstate);
|
||||
|
||||
#if defined(IVFFLAT_MEMORY)
|
||||
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
|
||||
#endif
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -252,7 +240,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
|
||||
int lists;
|
||||
int dimensions;
|
||||
int probes = ivfflat_probes;
|
||||
int maxProbes;
|
||||
|
||||
scan = RelationGetIndexScan(index, nkeys, norderbys);
|
||||
|
||||
@@ -262,21 +249,10 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
|
||||
if (probes > lists)
|
||||
probes = lists;
|
||||
|
||||
if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF)
|
||||
{
|
||||
if (ivfflat_iterative_search_max_probes == 0)
|
||||
maxProbes = lists;
|
||||
else
|
||||
maxProbes = Min(ivfflat_iterative_search_max_probes, lists);
|
||||
}
|
||||
else
|
||||
maxProbes = probes;
|
||||
|
||||
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
|
||||
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
|
||||
so->typeInfo = IvfflatGetTypeInfo(index);
|
||||
so->first = true;
|
||||
so->probes = probes;
|
||||
so->maxProbes = maxProbes;
|
||||
so->dimensions = dimensions;
|
||||
|
||||
/* Set support functions */
|
||||
@@ -304,8 +280,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
|
||||
so->bas = GetAccessStrategy(BAS_BULKREAD);
|
||||
|
||||
so->listQueue = pairingheap_allocate(CompareLists, scan);
|
||||
so->listPages = palloc(maxProbes * sizeof(BlockNumber));
|
||||
so->listIndex = 0;
|
||||
|
||||
scan->opaque = so;
|
||||
|
||||
@@ -320,9 +294,11 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
|
||||
{
|
||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||
|
||||
if (!so->first)
|
||||
tuplesort_reset(so->sortstate);
|
||||
|
||||
so->first = true;
|
||||
pairingheap_reset(so->listQueue);
|
||||
so->listIndex = 0;
|
||||
|
||||
if (keys && scan->numberOfKeys > 0)
|
||||
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
|
||||
@@ -338,8 +314,6 @@ bool
|
||||
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
{
|
||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||
ItemPointer heaptid;
|
||||
bool isnull;
|
||||
|
||||
/*
|
||||
* Index can be used to scan backward, but Postgres doesn't support
|
||||
@@ -367,25 +341,28 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
IvfflatBench("GetScanLists", GetScanLists(scan, value));
|
||||
IvfflatBench("GetScanItems", GetScanItems(scan, value));
|
||||
so->first = false;
|
||||
so->value = value;
|
||||
|
||||
/* TODO clean up if we allocated a new value */
|
||||
#if defined(IVFFLAT_MEMORY)
|
||||
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
|
||||
#endif
|
||||
|
||||
/* Clean up if we allocated a new value */
|
||||
if (value != scan->orderByData->sk_argument)
|
||||
pfree(DatumGetPointer(value));
|
||||
}
|
||||
|
||||
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
|
||||
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
|
||||
{
|
||||
if (so->listIndex == so->maxProbes)
|
||||
return false;
|
||||
bool isnull;
|
||||
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
|
||||
|
||||
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
|
||||
scan->xs_heaptid = *heaptid;
|
||||
scan->xs_recheck = false;
|
||||
scan->xs_recheckorderby = false;
|
||||
return true;
|
||||
}
|
||||
|
||||
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
|
||||
|
||||
scan->xs_heaptid = *heaptid;
|
||||
scan->xs_recheck = false;
|
||||
scan->xs_recheckorderby = false;
|
||||
return true;
|
||||
return false;
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -397,7 +374,6 @@ ivfflatendscan(IndexScanDesc scan)
|
||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||
|
||||
pairingheap_free(so->listQueue);
|
||||
pfree(so->listPages);
|
||||
tuplesort_end(so->sortstate);
|
||||
FreeAccessStrategy(so->bas);
|
||||
FreeTupleDesc(so->tupdesc);
|
||||
|
||||
@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
|
||||
Page cpage;
|
||||
OffsetNumber coffno;
|
||||
OffsetNumber cmaxoffno;
|
||||
BlockNumber listPages[MaxOffsetNumber];
|
||||
BlockNumber startPages[MaxOffsetNumber];
|
||||
ListInfo listInfo;
|
||||
|
||||
cbuf = ReadBuffer(index, blkno);
|
||||
@@ -40,7 +40,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
|
||||
{
|
||||
IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
|
||||
|
||||
listPages[coffno - FirstOffsetNumber] = list->startPage;
|
||||
startPages[coffno - FirstOffsetNumber] = list->startPage;
|
||||
}
|
||||
|
||||
listInfo.blkno = blkno;
|
||||
@@ -50,7 +50,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
|
||||
|
||||
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
|
||||
{
|
||||
BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
|
||||
BlockNumber searchPage = startPages[coffno - FirstOffsetNumber];
|
||||
BlockNumber insertPage = InvalidBlockNumber;
|
||||
|
||||
/* Iterate over entry pages */
|
||||
|
||||
@@ -6,7 +6,13 @@ use Test::More;
|
||||
|
||||
my $dim = 3;
|
||||
|
||||
my $array_sql = join(",", ('random()') x $dim);
|
||||
my @r = ();
|
||||
for (1 .. $dim)
|
||||
{
|
||||
my $v = int(rand(1000)) + 1;
|
||||
push(@r, "i % $v");
|
||||
}
|
||||
my $array_sql = join(", ", @r);
|
||||
|
||||
# Initialize node
|
||||
my $node = PostgreSQL::Test::Cluster->new('node');
|
||||
@@ -17,20 +23,19 @@ $node->start;
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
|
||||
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
|
||||
|
||||
# Get size
|
||||
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
|
||||
|
||||
# Store values
|
||||
$node->safe_psql("postgres", "CREATE TABLE tmp AS SELECT * FROM tst;");
|
||||
|
||||
# Delete all, vacuum, and insert same data
|
||||
$node->safe_psql("postgres", "DELETE FROM tst;");
|
||||
$node->safe_psql("postgres", "VACUUM tst;");
|
||||
$node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
|
||||
# Check size
|
||||
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
|
||||
|
||||
@@ -41,7 +41,8 @@ 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;
|
||||
));
|
||||
like($explain, qr/Seq Scan/);
|
||||
# TODO Do not use index
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test attribute filtering with few rows removed
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
@@ -59,7 +60,8 @@ 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;
|
||||
));
|
||||
like($explain, qr/Seq Scan/);
|
||||
# TODO Do not use index
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test attribute filtering with few rows removed like
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
|
||||
@@ -17,11 +17,12 @@ $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, 2000) i;"
|
||||
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, $n) i;"
|
||||
);
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
|
||||
$node->safe_psql("postgres", "ANALYZE tst;");
|
||||
@@ -39,21 +40,6 @@ 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;");
|
||||
}
|
||||
|
||||
|
||||
@@ -39,11 +39,6 @@ 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;");
|
||||
}
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ $node->safe_psql("postgres", qq(
|
||||
|
||||
my $count = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.iterative_search = on;
|
||||
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;
|
||||
));
|
||||
@@ -34,16 +34,16 @@ is($count, 10);
|
||||
|
||||
foreach ((30000, 50000, 70000))
|
||||
{
|
||||
my $max_tuples = $_;
|
||||
my $expected = $max_tuples / 10000;
|
||||
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.iterative_search = on;
|
||||
SET hnsw.iterative_search_max_tuples = $max_tuples;
|
||||
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;
|
||||
));
|
||||
@@ -57,8 +57,7 @@ foreach ((30000, 50000, 70000))
|
||||
|
||||
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.iterative_search = on;
|
||||
SET client_min_messages = debug1;
|
||||
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;
|
||||
));
|
||||
@@ -1,54 +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", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
|
||||
|
||||
my $count = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = 10;
|
||||
SET ivfflat.iterative_search = on;
|
||||
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 ((30, 50, 70))
|
||||
{
|
||||
my $max_probes = $_;
|
||||
my $expected = $max_probes / 10;
|
||||
my $sum = 0;
|
||||
|
||||
for my $i (1 .. 20)
|
||||
{
|
||||
$count = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = 10;
|
||||
SET ivfflat.iterative_search = on;
|
||||
SET ivfflat.iterative_search_max_probes = $max_probes;
|
||||
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);
|
||||
}
|
||||
|
||||
done_testing();
|
||||
@@ -21,7 +21,7 @@ sub test_recall
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
SET hnsw.iterative_search = on;
|
||||
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/);
|
||||
@@ -31,7 +31,7 @@ sub test_recall
|
||||
my $actual = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.ef_search = $ef_search;
|
||||
SET hnsw.iterative_search = on;
|
||||
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);
|
||||
@@ -1,125 +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 @cs = (100, 1000);
|
||||
|
||||
sub test_recall
|
||||
{
|
||||
my ($c, $probes, $min, $operator) = @_;
|
||||
my $correct = 0;
|
||||
my $total = 0;
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET ivfflat.probes = $probes;
|
||||
SET ivfflat.iterative_search = 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 ivfflat.probes = $probes;
|
||||
SET ivfflat.iterative_search = 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(3));");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
|
||||
# Generate queries
|
||||
for (1 .. 20)
|
||||
{
|
||||
my $r1 = rand();
|
||||
my $r2 = rand();
|
||||
my $r3 = rand();
|
||||
push(@queries, "[$r1,$r2,$r3]");
|
||||
}
|
||||
|
||||
# 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", "CREATE INDEX idx ON tst USING ivfflat (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, 1, 0.58, $operator);
|
||||
test_recall($c, 10, 0.98, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
if ($operator eq "<->")
|
||||
{
|
||||
test_recall($c, 1, 0.80, $operator);
|
||||
}
|
||||
else
|
||||
{
|
||||
test_recall($c, 1, 0.88, $operator);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
$node->safe_psql("postgres", "DROP INDEX idx;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
@@ -1,113 +0,0 @@
|
||||
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 = 1000;
|
||||
|
||||
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", qq(
|
||||
SET maintenance_work_mem = '256MB';
|
||||
SET max_parallel_maintenance_workers = 2;
|
||||
CREATE INDEX ON tst USING hnsw (v vector_cosine_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_cosine_ops);");
|
||||
like($stderr, qr/first column must be a vector/);
|
||||
|
||||
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_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_cosine_ops, v vector_cosine_ops);");
|
||||
like($stderr, qr/column 2 cannot be a vector/);
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user