Compare commits

..

125 Commits

Author SHA1 Message Date
Andrew Kane
ac79f79637 Allow override 2024-01-20 17:07:36 -08:00
Andrew Kane
57d5d07ef9 Test CTE [skip ci] 2024-01-20 16:24:07 -08:00
Andrew Kane
bab5daf6df Improved order [skip ci] 2024-01-20 15:19:05 -08:00
Andrew Kane
63c4af0454 Do not use index without limit or if limit + offset > expected tuples 2024-01-20 15:16:30 -08:00
Andrew Kane
042ddfdc8a Updated readme [skip ci] 2024-01-20 11:40:49 -08:00
Andrew Kane
56870ce04d Added common paths on Mac [skip ci] 2024-01-19 21:43:55 -08:00
Andrew Kane
4ab77f3d24 Added section on missing SDK [skip ci] 2024-01-19 20:47:23 -08:00
Andrew Kane
cc9e6a6778 Added section on max_parallel_maintenance_workers [skip ci] 2024-01-18 21:57:38 -08:00
Andrew Kane
8f1b669c4f Added IVFFLAT_KMEANS_DEBUG to readme [skip ci] 2024-01-17 22:12:39 -08:00
Andrew Kane
1ff9ab5133 Revert "Improved vector_in"
This reverts commit 4894dc5da1.
2024-01-17 17:14:42 -08:00
Andrew Kane
4894dc5da1 Improved vector_in 2024-01-17 16:57:16 -08:00
Andrew Kane
7390f31261 Updated changelog [skip ci] 2024-01-17 10:50:07 -08:00
Heikki Linnakangas
b7304a3a4a Don't modify input string in vector_in() (#413)
Fixes issue #399
2024-01-17 10:49:04 -08:00
Andrew Kane
018ceb7a46 Updated PG_CONFIG example [skip ci] 2024-01-16 22:26:28 -08:00
Andrew Kane
0b2be00622 Added more defines to contributing docs [skip ci] 2024-01-16 17:21:33 -08:00
Andrew Kane
0ce497a1b1 Updated Homebrew note [skip ci] 2024-01-15 12:12:04 -08:00
Andrew Kane
c7d60346d8 Improved macro [skip ci] 2024-01-13 20:02:41 -08:00
Andrew Kane
597bfdc76b Added HnswGetNeighbors macro 2024-01-13 20:00:34 -08:00
Andrew Kane
cbf3eb4fa5 Improved HNSW build and insert code 2024-01-13 10:07:42 -08:00
Andrew Kane
cacd389f6d Improved pattern for duplicates 2024-01-12 14:30:13 -08:00
Andrew Kane
423cc2b06c Homebrew now adds to postgresql@15 as well [skip ci] 2024-01-11 16:45:50 -08:00
Andrew Kane
85c4ef6a14 Updated Postgres versions in readme [skip ci] 2024-01-11 12:36:24 -08:00
Andrew Kane
c6160a783a Homebrew now adds to postgresql@16 [skip ci] 2024-01-11 12:32:14 -08:00
Andrew Kane
1881b857f9 Simplified code 2024-01-09 18:53:31 -08:00
Andrew Kane
51bde5fb22 Updated readme [skip ci] 2024-01-09 14:38:25 -08:00
Andrew Kane
10e65ce349 Added note about maintenance_work_mem [skip ci] 2024-01-09 14:31:54 -08:00
Andrew Kane
61279f5a59 Updated readme [skip ci] 2024-01-09 14:26:55 -08:00
Andrew Kane
72b3889e26 Updated readme [skip ci] 2024-01-09 14:22:19 -08:00
Andrew Kane
bb21b2decf Updated readme [skip ci] 2024-01-09 14:19:01 -08:00
Andrew Kane
8a65c0e831 Moved section [skip ci] 2024-01-09 13:33:03 -08:00
Andrew Kane
7d75d423e4 Added section on index build time [skip ci] 2024-01-09 13:27:27 -08:00
Andrew Kane
6cad1f5de0 Updated example [skip ci] 2024-01-09 13:04:47 -08:00
Andrew Kane
67eeade63c Moved HNSW first in readme [skip ci] 2024-01-09 13:04:18 -08:00
Andrew Kane
108fb09d7b Improved code [skip ci] 2024-01-08 17:54:49 -08:00
Andrew Kane
65d060ac86 Reverted FlushPages pattern for parallel builds 2024-01-08 10:45:31 -08:00
Andrew Kane
62ee33bb92 Improved locking code 2024-01-08 09:05:12 -08:00
Andrew Kane
520e274dde Improved locking code 2024-01-07 22:34:41 -08:00
Andrew Kane
9e680884bd Moved indtuples to HnswGraph 2024-01-07 22:23:49 -08:00
Andrew Kane
19a0e1b341 Moved graph to separate struct 2024-01-07 20:15:30 -08:00
Andrew Kane
c7fe1571ee Improved code 2024-01-07 18:30:51 -08:00
Andrew Kane
cb4c770df2 Switched to slist for elements to reduce allocations and remove limit 2024-01-07 18:26:19 -08:00
Andrew Kane
85fdecd79b Moved FlushPages before HnswEndParallel 2024-01-07 17:50:46 -08:00
Andrew Kane
6132428914 Improved number of parallel workers for HNSW index builds - closes #397 2024-01-05 19:46:08 -08:00
Andrew Kane
81d13bd40f Improved code [skip ci] 2024-01-03 13:53:23 -05:00
Andrew Kane
8ee37b60a0 Improved memory estimate for HNSW index builds 2024-01-03 13:47:50 -05:00
Andrew Kane
9b73b3d1a6 Reduced memory and allocations for heap TIDs - closes #385 2024-01-03 13:41:34 -05:00
Andrew Kane
cae630784b Improved BuildCallback [skip ci] 2023-12-30 20:55:29 -05:00
Andrew Kane
d87bcd2deb Added comments [skip ci] 2023-12-30 18:29:01 -05:00
Andrew Kane
736576220a Improved BuildCallback 2023-12-30 18:24:03 -05:00
Andrew Kane
a508b120c1 Added IVFFLAT_MEMORY flag to show memory usage [skip ci] 2023-12-24 09:27:09 -05:00
Andrew Kane
9a782d29f8 Use consistent style [skip ci] 2023-12-22 16:41:25 -05:00
Andrew Kane
1e422cd62b Improved readability [skip ci] 2023-12-22 16:39:13 -05:00
Andrew Kane
569c69580a Improved InsertTuple code - #384
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-12-22 15:08:28 -05:00
Andrew Kane
59509c3a17 Added extra 5% to memory estimate 2023-12-22 14:04:05 -05:00
Andrew Kane
61738846af Updated comment [skip ci] 2023-12-22 14:03:33 -05:00
Andrew Kane
e8c3bf0cef Improved memory tracking for HNSW index builds - #384 2023-12-22 13:35:43 -05:00
Andrew Kane
50d1aed3d8 Improved memory usage logging [skip ci] 2023-12-22 13:09:11 -05:00
Andrew Kane
66e14d2434 Updated indentation [skip ci] 2023-12-22 12:59:50 -05:00
Andrew Kane
42cd4c6833 Fixed call to GenerationContextCreate for Postgres < 15 2023-12-22 12:49:07 -05:00
Andrew Kane
dcbe0b6f0d Reduced memory usage for HNSW index builds - #384
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-12-22 12:41:47 -05:00
Andrew Kane
f61d4087b5 Slightly improved memory estimation [skip ci] 2023-12-21 10:31:36 -05:00
Andrew Kane
57554e5b46 Added todo [skip ci] 2023-12-20 17:52:31 -05:00
Andrew Kane
6738fa0bd7 Added HNSW_MEMORY flag to show memory usage - #384 [skip ci] 2023-12-20 16:49:16 -05:00
Andrew Kane
9ab10aa674 Fixed CI 2023-12-20 16:29:13 -05:00
Andrew Kane
ec41dfa1d7 Mark meta buffer contents as dirty when not logging 2023-12-20 16:20:15 -05:00
Andrew Kane
43e0b3d9d4 Mark buffer contents as dirty when not logging 2023-12-20 16:16:25 -05:00
Andrew Kane
2bff7ccaa2 Improved memory estimation - #384 [skip ci] 2023-12-20 11:10:27 -05:00
Andrew Kane
e88a425c9b Reduced WAL generation for HNSW index builds - thanks @hlinnaka 2023-12-19 20:37:32 -05:00
Andrew Kane
921427ee03 Replace dynahash hash table in HNSW with simplehash for speed - #378
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-12-17 11:24:13 -05:00
Andrew Kane
a59aa02dd9 Only show message if flushed due to memory [skip ci] 2023-12-04 15:29:23 -08:00
Andrew Kane
2fef497b7e Fixed check 2023-12-04 15:22:12 -08:00
Andrew Kane
0e19a984fb Updated changelog [skip ci] 2023-12-04 15:15:20 -08:00
Andrew Kane
a156f6c7ae Fixed invalid memory alloc request size error with HNSW - fixes #43 2023-12-04 15:00:32 -08:00
Andrew Kane
c653ac524f Simplified code [skip ci] 2023-12-04 12:19:38 -08:00
Andrew Kane
bf0d56e78e Updated comment [skip ci] 2023-12-03 13:09:57 -08:00
Andrew Kane
4d6739a7af Added Lisp to readme [skip ci] 2023-12-03 10:58:54 -08:00
Andrew Kane
ff744214d0 Added Visual Basic to readme [skip ci] 2023-12-03 00:25:07 -08:00
Andrew Kane
7ca9298163 Updated badge [skip ci] 2023-12-01 15:36:49 -08:00
Andrew Kane
ff3bffd9a8 Moved FAQ [skip ci] 2023-11-29 22:09:09 -08:00
Andrew Kane
014753eb9c Added FAQ about memory [skip ci] 2023-11-29 22:02:27 -08:00
Andrew Kane
6763661d3d Added OCaml to readme [skip ci] 2023-11-28 21:42:30 -08:00
Andrew Kane
d287921d15 Added F# to readme [skip ci] 2023-11-28 01:22:12 -08:00
Andrew Kane
5b12ae8225 Added note about ef_construction [skip ci] 2023-11-16 18:53:48 -08:00
Japin Li
4549e8aeb1 Fix coredump about HnswFreeElement() (#357)
The HnswInitElement() allocate an element with not initialize value
filed, which may has garbage that lead HnswFreeElement() free an
invalid pointer.
2023-11-15 16:19:59 -08:00
Andrew Kane
3263b350f5 Updated HnswLoadElementFromTuple to be less vector-specific 2023-11-11 21:14:12 -08:00
Andrew Kane
dfee5d4045 Added support for on-disk parallel index builds for HNSW 2023-11-11 19:29:45 -08:00
Andrew Kane
69a2ce0d43 Use datumIsEqual to compare 2023-11-10 10:46:48 -08:00
Andrew Kane
c5e8c46b80 Switched from VECTOR_SIZE to VARSIZE_ANY [skip ci] 2023-11-09 19:41:38 -08:00
Andrew Kane
94f7304ccd Keep vector for now to be overly cautious about packing [skip ci] 2023-11-09 18:43:55 -08:00
Andrew Kane
d078db3d25 Switched HnswElementTuple to generic data and zero full section 2023-11-09 18:28:25 -08:00
Andrew Kane
fbb904ae2f Use pointer for VARSIZE_ANY 2023-11-09 17:50:28 -08:00
Andrew Kane
3cf6f62900 Switched to datum for HnswElement 2023-11-09 17:35:39 -08:00
Andrew Kane
2a69e22ca4 Switched from VECTOR_SIZE to VARSIZE_ANY where possible (less vector-specific) 2023-11-09 17:16:43 -08:00
Andrew Kane
84e073888c Removed vector-specific code from HNSW_ELEMENT_TUPLE_SIZE [skip ci] 2023-11-09 16:57:01 -08:00
Andrew Kane
81a62d55d1 Switched from HNSW_ELEMENT_TUPLE_SIZE to ItemIdGetLength where possible (less vector-specific) 2023-11-09 16:32:00 -08:00
Andrew Kane
3f3463bde5 Improved memory calculation for HNSW and removed vector-specific code 2023-11-09 16:21:26 -08:00
Andrew Kane
a01a72d812 Updated comment [skip ci] 2023-11-05 08:42:06 -08:00
Andrew Kane
0c2fc18a80 Updated comment [skip ci] 2023-11-05 08:40:21 -08:00
Andrew Kane
e860042d3c Improved variable name [skip ci] 2023-11-05 08:35:54 -08:00
Andrew Kane
5986862bd2 Added note about check constraint [skip ci] 2023-11-04 15:01:37 -07:00
Andrew Kane
5d24f5d09a Improved header installation on Windows 2023-11-04 11:16:40 -07:00
Andrew Kane
7c43b0d8ee Updated example [skip ci] 2023-11-03 23:54:50 -07:00
Andrew Kane
7be40036f4 Updated readme [skip ci] 2023-11-03 23:46:23 -07:00
Andrew Kane
9b5a1a69db Updated readme [skip ci] 2023-11-03 23:43:47 -07:00
Andrew Kane
04b96506f5 Added info on storing vectors with more precision [skip ci] 2023-11-03 20:14:28 -07:00
Andrew Kane
35cd7b63cb Updated readme [skip ci] 2023-11-03 17:02:30 -07:00
Andrew Kane
b5416d6f10 Updated readme [skip ci] 2023-11-03 16:48:57 -07:00
Andrew Kane
f361bf2704 Improved docs on indexing vectors with different dimensions [skip ci] 2023-11-03 16:42:14 -07:00
Andrew Kane
3d8c1921aa Improved upgrading docs - #339 [skip ci] 2023-11-03 16:15:06 -07:00
Andrew Kane
154207bc17 Added info on columns with different dimensions [skip ci] 2023-11-03 16:02:00 -07:00
Andrew Kane
8e507f3bf5 Free remaining allocation from deconstruct_array - #332 2023-11-02 21:20:21 -07:00
Andrew Kane
e115773a55 Removed unneeded allocation 2023-11-02 21:16:06 -07:00
Andrew Kane
9333bef046 Added link to setup-pgvector [skip ci] 2023-11-02 13:22:19 -07:00
Andrew Kane
4851e47d9f Added Reciprocal Rank Fusion example to readme [skip ci] 2023-11-01 13:20:49 -07:00
Andrew Kane
12aecfb4f5 Added Nim and Zig to readme [skip ci] 2023-10-31 02:26:18 -07:00
Andrew Kane
800697fb14 Updated column alias [skip ci] 2023-10-29 16:47:55 -07:00
Andrew Kane
de1f2b09dd Improved indexing progress queries [skip ci] 2023-10-29 16:41:39 -07:00
Andrew Kane
bcccb7f5a5 Improved docs for indexing progress - closes #320 and closes #321 [skip ci] 2023-10-29 16:13:12 -07:00
Andrew Kane
bec3d30d68 Added TypeScript to readme [skip ci] 2023-10-29 12:49:01 -07:00
Andrew Kane
588de60445 Added Groovy to readme [skip ci] 2023-10-29 12:39:53 -07:00
Andrew Kane
c599f92b52 Updated readme [skip ci] 2023-10-27 13:22:37 -07:00
Andrew Kane
2a17b335da Added Kotlin to readme [skip ci] 2023-10-26 12:25:58 -07:00
Andrew Kane
6ede6ac301 Added link to pgvector-c [skip ci] 2023-10-26 00:30:06 -07:00
Andrew Kane
3f49b95f01 Added Postgres 17 to CI [skip ci] 2023-10-19 00:37:24 -07:00
Andrew Kane
ef1bea7163 Updated checkout action [skip ci] 2023-10-19 00:36:53 -07:00
38 changed files with 1605 additions and 569 deletions

View File

@@ -8,6 +8,8 @@ jobs:
fail-fast: false fail-fast: false
matrix: matrix:
include: include:
- postgres: 17
os: ubuntu-22.04
- postgres: 16 - postgres: 16
os: ubuntu-22.04 os: ubuntu-22.04
- postgres: 15 - postgres: 15
@@ -21,7 +23,7 @@ jobs:
- postgres: 11 - postgres: 11
os: ubuntu-20.04 os: ubuntu-20.04
steps: steps:
- uses: actions/checkout@v3 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: ${{ matrix.postgres }} postgres-version: ${{ matrix.postgres }}
@@ -43,7 +45,7 @@ jobs:
runs-on: macos-latest runs-on: macos-latest
if: ${{ !startsWith(github.ref_name, 'windows') }} if: ${{ !startsWith(github.ref_name, 'windows') }}
steps: steps:
- uses: actions/checkout@v3 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: 14 postgres-version: 14
@@ -65,7 +67,7 @@ jobs:
runs-on: windows-latest runs-on: windows-latest
if: ${{ !startsWith(github.ref_name, 'mac') }} if: ${{ !startsWith(github.ref_name, 'mac') }}
steps: steps:
- uses: actions/checkout@v3 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
with: with:
postgres-version: 14 postgres-version: 14

View File

@@ -1,3 +1,12 @@
## 0.5.2 (unreleased)
- Improved performance of HNSW
- Added support for on-disk parallel index builds for HNSW
- Reduced memory usage for HNSW index builds
- Reduced WAL generation for HNSW index builds
- Fixed error with logical replication
- Fixed `invalid memory alloc request size` error with HNSW index build
## 0.5.1 (2023-10-10) ## 0.5.1 (2023-10-10)
- Improved performance of HNSW index builds - Improved performance of HNSW index builds

View File

@@ -56,7 +56,7 @@ install:
copy $(EXTENSION).control "$(SHAREDIR)\extension" copy $(EXTENSION).control "$(SHAREDIR)\extension"
copy sql\$(EXTENSION)--*.sql "$(SHAREDIR)\extension" copy sql\$(EXTENSION)--*.sql "$(SHAREDIR)\extension"
mkdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)" mkdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
copy $(HEADERS) "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)" for %f in ($(HEADERS)) do copy %f "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
installcheck: installcheck:
"$(BINDIR)\pg_regress" --bindir="$(BINDIR)" $(REGRESS_OPTS) $(REGRESS) "$(BINDIR)\pg_regress" --bindir="$(BINDIR)" $(REGRESS_OPTS) $(REGRESS)

290
README.md
View File

@@ -10,7 +10,7 @@ Store your vectors with the rest of your data. Supports:
Plus [ACID](https://en.wikipedia.org/wiki/ACID) compliance, point-in-time recovery, JOINs, and all of the other [great features](https://www.postgresql.org/about/) of Postgres Plus [ACID](https://en.wikipedia.org/wiki/ACID) compliance, point-in-time recovery, JOINs, and all of the other [great features](https://www.postgresql.org/about/) of Postgres
[![Build Status](https://github.com/pgvector/pgvector/workflows/build/badge.svg?branch=master)](https://github.com/pgvector/pgvector/actions) [![Build Status](https://github.com/pgvector/pgvector/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector/actions)
## Installation ## Installation
@@ -26,7 +26,7 @@ make install # may need sudo
See the [installation notes](#installation-notes) if you run into issues See the [installation notes](#installation-notes) if you run into issues
You can also install it with [Docker](#docker), [Homebrew](#homebrew), [PGXN](#pgxn), [APT](#apt), [Yum](#yum), or [conda-forge](#conda-forge), and it comes preinstalled with [Postgres.app](#postgresapp) and many [hosted providers](#hosted-postgres) You can also install it with [Docker](#docker), [Homebrew](#homebrew), [PGXN](#pgxn), [APT](#apt), [Yum](#yum), or [conda-forge](#conda-forge), and it comes preinstalled with [Postgres.app](#postgresapp) and many [hosted providers](#hosted-postgres). There are also instructions for [GitHub Actions](https://github.com/pgvector/setup-pgvector).
## Getting Started ## Getting Started
@@ -161,8 +161,97 @@ You can add an index to use approximate nearest neighbor search, which trades so
Supported index types are: Supported index types are:
- [IVFFlat](#ivfflat)
- [HNSW](#hnsw) - added in 0.5.0 - [HNSW](#hnsw) - added in 0.5.0
- [IVFFlat](#ivfflat)
## HNSW
An HNSW index creates a multilayer graph. It has better query performance than IVFFlat (in terms of speed-recall tradeoff), but has slower build times and uses more memory. Also, an index can be created without any data in the table since there isnt a training step like IVFFlat.
Add an index for each distance function you want to use.
L2 distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);
```
Inner product
```sql
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops);
```
Cosine distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
```
Vectors with up to 2,000 dimensions can be indexed.
### Index Options
Specify HNSW parameters
- `m` - the max number of connections per layer (16 by default)
- `ef_construction` - the size of the dynamic candidate list for constructing the graph (64 by default)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);
```
A higher value of `ef_construction` provides better recall at the cost of index build time / insert speed.
### Query Options
Specify the size of the dynamic candidate list for search (40 by default)
```sql
SET hnsw.ef_search = 100;
```
A higher value provides better recall at the cost of speed.
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT ...
COMMIT;
```
### Index Build Time
Indexes build significantly faster when the graph fits into `maintenance_work_mem`
```sql
SET maintenance_work_mem = '8GB';
```
A notice is shown when the graph no longer fits
```text
NOTICE: hnsw graph no longer fits into maintenance_work_mem after 100000 tuples
DETAIL: Building will take significantly more time.
HINT: Increase maintenance_work_mem to speed up builds.
```
Note: Do not set `maintenance_work_mem` so high that it exhausts the memory on the server
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for HNSW are:
1. `initializing`
2. `loading tuples`
## IVFFlat ## IVFFlat
@@ -215,78 +304,32 @@ SELECT ...
COMMIT; COMMIT;
``` ```
## HNSW ### Index Build Time
An HNSW index creates a multilayer graph. It has slower build times and uses more memory than IVFFlat, but has better query performance (in terms of speed-recall tradeoff). Theres no training step like IVFFlat, so the index can be created without any data in the table. Speed up index creation on large tables by increasing the number of parallel workers (2 by default)
Add an index for each distance function you want to use.
L2 distance
```sql ```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops); SET max_parallel_maintenance_workers = 7; -- plus leader
``` ```
Inner product For a large number of workers, you may also need to increase `max_parallel_workers` (8 by default)
```sql ### Indexing Progress
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops);
```
Cosine distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
```
Vectors with up to 2,000 dimensions can be indexed.
### Index Options
Specify HNSW parameters
- `m` - the max number of connections per layer (16 by default)
- `ef_construction` - the size of the dynamic candidate list for constructing the graph (64 by default)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);
```
### Query Options
Specify the size of the dynamic candidate list for search (40 by default)
```sql
SET hnsw.ef_search = 100;
```
A higher value provides better recall at the cost of speed.
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT ...
COMMIT;
```
## Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+ Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql ```sql
SELECT phase, tuples_done, tuples_total FROM pg_stat_progress_create_index; SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
``` ```
The phases are: The phases for IVFFlat are:
1. `initializing` 1. `initializing`
2. `performing k-means` - IVFFlat only 2. `performing k-means`
3. `assigning tuples` - IVFFlat only 3. `assigning tuples`
4. `loading tuples` 4. `loading tuples`
Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase Note: `%` is only populated during the `loading tuples` phase
## Filtering ## Filtering
@@ -305,8 +348,7 @@ CREATE INDEX ON items (category_id);
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql ```sql
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100) CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WHERE (category_id = 123);
WHERE (category_id = 123);
``` ```
Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html) for approximate search on many different values of the `WHERE` columns Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html) for approximate search on many different values of the `WHERE` columns
@@ -317,13 +359,15 @@ CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(cate
## Hybrid Search ## Hybrid Search
Use together with Postgres [full-text search](https://www.postgresql.org/docs/current/textsearch-intro.html) for hybrid search ([Python example](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py)). Use together with Postgres [full-text search](https://www.postgresql.org/docs/current/textsearch-intro.html) for hybrid search.
```sql ```sql
SELECT id, content FROM items, plainto_tsquery('hello search') query SELECT id, content FROM items, plainto_tsquery('hello search') query
WHERE textsearch @@ query ORDER BY ts_rank_cd(textsearch, query) DESC LIMIT 5; WHERE textsearch @@ query ORDER BY ts_rank_cd(textsearch, query) DESC LIMIT 5;
``` ```
You can use [Reciprocal Rank Fusion](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search_rrf.py) or a [cross-encoder](https://github.com/pgvector/pgvector-python/blob/master/examples/hybrid_search.py) to combine results.
## Performance ## Performance
Use `EXPLAIN ANALYZE` to debug performance. Use `EXPLAIN ANALYZE` to debug performance.
@@ -360,17 +404,21 @@ Use pgvector from any language with a Postgres client. You can even generate and
Language | Libraries / Examples Language | Libraries / Examples
--- | --- --- | ---
C | [pgvector-c](https://github.com/pgvector/pgvector-c)
C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp) C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp)
C# | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet) C#, F#, Visual Basic | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet)
Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal) Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal)
Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart) Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart)
Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir) Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir)
Go | [pgvector-go](https://github.com/pgvector/pgvector-go) Go | [pgvector-go](https://github.com/pgvector/pgvector-go)
Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell) Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell)
Java, Scala | [pgvector-java](https://github.com/pgvector/pgvector-java) Java, Kotlin, Groovy, Scala | [pgvector-java](https://github.com/pgvector/pgvector-java)
JavaScript, TypeScript | [pgvector-node](https://github.com/pgvector/pgvector-node)
Julia | [pgvector-julia](https://github.com/pgvector/pgvector-julia) Julia | [pgvector-julia](https://github.com/pgvector/pgvector-julia)
Lisp | [pgvector-lisp](https://github.com/pgvector/pgvector-lisp)
Lua | [pgvector-lua](https://github.com/pgvector/pgvector-lua) Lua | [pgvector-lua](https://github.com/pgvector/pgvector-lua)
Node.js | [pgvector-node](https://github.com/pgvector/pgvector-node) Nim | [pgvector-nim](https://github.com/pgvector/pgvector-nim)
OCaml | [pgvector-ocaml](https://github.com/pgvector/pgvector-ocaml)
Perl | [pgvector-perl](https://github.com/pgvector/pgvector-perl) Perl | [pgvector-perl](https://github.com/pgvector/pgvector-perl)
PHP | [pgvector-php](https://github.com/pgvector/pgvector-php) PHP | [pgvector-php](https://github.com/pgvector/pgvector-php)
Python | [pgvector-python](https://github.com/pgvector/pgvector-python) Python | [pgvector-python](https://github.com/pgvector/pgvector-python)
@@ -378,6 +426,7 @@ R | [pgvector-r](https://github.com/pgvector/pgvector-r)
Ruby | [pgvector-ruby](https://github.com/pgvector/pgvector-ruby), [Neighbor](https://github.com/ankane/neighbor) Ruby | [pgvector-ruby](https://github.com/pgvector/pgvector-ruby), [Neighbor](https://github.com/ankane/neighbor)
Rust | [pgvector-rust](https://github.com/pgvector/pgvector-rust) Rust | [pgvector-rust](https://github.com/pgvector/pgvector-rust)
Swift | [pgvector-swift](https://github.com/pgvector/pgvector-swift) Swift | [pgvector-swift](https://github.com/pgvector/pgvector-swift)
Zig | [pgvector-zig](https://github.com/pgvector/pgvector-zig)
## Frequently Asked Questions ## Frequently Asked Questions
@@ -393,6 +442,63 @@ Yes, pgvector uses the write-ahead log (WAL), which allows for replication and p
Youll need to use [dimensionality reduction](https://en.wikipedia.org/wiki/Dimensionality_reduction) at the moment. Youll need to use [dimensionality reduction](https://en.wikipedia.org/wiki/Dimensionality_reduction) at the moment.
#### Can I store vectors with different dimensions in the same column?
You can use `vector` as the type (instead of `vector(3)`).
```sql
CREATE TABLE embeddings (model_id bigint, item_id bigint, embedding vector, PRIMARY KEY (model_id, item_id));
```
However, you can only create indexes on rows with the same number of dimensions (using [expression](https://www.postgresql.org/docs/current/indexes-expressional.html) and [partial](https://www.postgresql.org/docs/current/indexes-partial.html) indexing):
```sql
CREATE INDEX ON embeddings USING hnsw ((embedding::vector(3)) vector_l2_ops) WHERE (model_id = 123);
```
and query with:
```sql
SELECT * FROM embeddings WHERE model_id = 123 ORDER BY embedding::vector(3) <-> '[3,1,2]' LIMIT 5;
```
#### Can I store vectors with more precision?
You can use the `double precision[]` or `numeric[]` type to store vectors with more precision.
```sql
CREATE TABLE items (id bigserial PRIMARY KEY, embedding double precision[]);
-- use {} instead of [] for Postgres arrays
INSERT INTO items (embedding) VALUES ('{1,2,3}'), ('{4,5,6}');
```
Optionally, add a [check constraint](https://www.postgresql.org/docs/current/ddl-constraints.html) to ensure data can be converted to the `vector` type and has the expected dimensions.
```sql
ALTER TABLE items ADD CHECK (vector_dims(embedding::vector) = 3);
```
Use [expression indexing](https://www.postgresql.org/docs/current/indexes-expressional.html) to index (at a lower precision):
```sql
CREATE INDEX ON items USING hnsw ((embedding::vector(3)) vector_l2_ops);
```
and query with:
```sql
SELECT * FROM items ORDER BY embedding::vector(3) <-> '[3,1,2]' LIMIT 5;
```
#### Do indexes need to fit into memory?
No, but like other index types, youll likely see better performance if they do. You can get the size of an index with:
```sql
SELECT pg_size_pretty(pg_relation_size('index_name'));
```
## Troubleshooting ## Troubleshooting
#### Why isnt a query using an index? #### Why isnt a query using an index?
@@ -406,6 +512,8 @@ SELECT ...
COMMIT; COMMIT;
``` ```
Also, if the table is small, a table scan may be faster.
#### Why isnt a query using a parallel table scan? #### Why isnt a query using a parallel table scan?
The planner doesnt consider [out-of-line storage](https://www.postgresql.org/docs/current/storage-toast.html) in cost estimates, which can make a serial scan look cheaper. You can reduce the cost of a parallel scan for a query with: The planner doesnt consider [out-of-line storage](https://www.postgresql.org/docs/current/storage-toast.html) in cost estimates, which can make a serial scan look cheaper. You can reduce the cost of a parallel scan for a query with:
@@ -474,7 +582,7 @@ sum(vector) → vector | sum | 0.5.0
If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with: If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with:
```sh ```sh
export PG_CONFIG=/Applications/Postgres.app/Contents/Versions/latest/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: Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use:
@@ -483,6 +591,14 @@ Then re-run the installation instructions (run `make clean` before `make` if nee
sudo --preserve-env=PG_CONFIG make install sudo --preserve-env=PG_CONFIG make install
``` ```
A few common paths on Mac are:
- EDB installer - `/Library/PostgreSQL/16/bin/pg_config`
- Homebrew (arm64) - `/opt/homebrew/opt/postgresql@16/bin/pg_config`
- Homebrew (x86-64) - `/usr/local/opt/postgresql@16/bin/pg_config`
Note: Replace `16` with your Postgres server version
### Missing Header ### Missing Header
If compilation fails with `fatal error: postgres.h: No such file or directory`, make sure Postgres development files are installed on the server. If compilation fails with `fatal error: postgres.h: No such file or directory`, make sure Postgres development files are installed on the server.
@@ -490,10 +606,14 @@ If compilation fails with `fatal error: postgres.h: No such file or directory`,
For Ubuntu and Debian, use: For Ubuntu and Debian, use:
```sh ```sh
sudo apt install postgresql-server-dev-15 sudo apt install postgresql-server-dev-16
``` ```
Note: Replace `15` with your Postgres server version Note: Replace `16` with your Postgres server version
### Missing SDK
If compilation fails and the output includes `warning: no such sysroot directory` on Mac, reinstall Xcode Command Line Tools.
### Windows ### Windows
@@ -508,7 +628,7 @@ Note: The exact path will vary depending on your Visual Studio version and editi
Then use `nmake` to build: Then use `nmake` to build:
```cmd ```cmd
set "PGROOT=C:\Program Files\PostgreSQL\15" set "PGROOT=C:\Program Files\PostgreSQL\16"
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector cd pgvector
nmake /F Makefile.win nmake /F Makefile.win
@@ -558,22 +678,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: 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 ```sh
sudo apt install postgresql-15-pgvector sudo apt install postgresql-16-pgvector
``` ```
Note: Replace `15` with your Postgres server version Note: Replace `16` with your Postgres server version
### Yum ### 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: 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 ```sh
sudo yum install pgvector_15 sudo yum install pgvector_16
# or # or
sudo dnf install pgvector_15 sudo dnf install pgvector_16
``` ```
Note: Replace `15` with your Postgres server version Note: Replace `16` with your Postgres server version
### conda-forge ### conda-forge
@@ -595,7 +715,7 @@ pgvector is available on [these providers](https://github.com/pgvector/pgvector/
## Upgrading ## Upgrading
Install the latest version. Then in each database you want to upgrade, run: [Install](#installation) the latest version (use the same method as the original installation). Then in each database you want to upgrade, run:
```sql ```sql
ALTER EXTENSION vector UPDATE; ALTER EXTENSION vector UPDATE;
@@ -683,7 +803,25 @@ make prove_installcheck PROVE_TESTS=test/t/001_wal.pl # TAP test
To enable benchmarking: To enable benchmarking:
```sh ```sh
make clean && PG_CFLAGS=-DIVFFLAT_BENCH make && make install make clean && PG_CFLAGS="-DIVFFLAT_BENCH" make && make install
```
To show memory usage:
```sh
make clean && PG_CFLAGS="-DHNSW_MEMORY -DIVFFLAT_MEMORY" make && make install
```
To enable assertions:
```sh
make clean && PG_CFLAGS="-DUSE_ASSERT_CHECKING" make && make install
```
To get k-means metrics:
```sh
make clean && PG_CFLAGS="-DIVFFLAT_KMEANS_DEBUG" make && make install
``` ```
Resources for contributors Resources for contributors

View File

@@ -14,6 +14,7 @@
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
bool hnsw_enable_parallel_build;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
/* /*
@@ -39,6 +40,11 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search", DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* Behind a variable for now since can be slower than building in memory */
DefineCustomBoolVariable("hnsw.enable_parallel_build", "Enables or disables building indexes in parallel",
NULL, &hnsw_enable_parallel_build,
false, PGC_USERSET, 0, NULL, NULL, NULL);
} }
/* /*
@@ -77,8 +83,8 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
List *qinfos; List *qinfos;
#endif #endif
/* Never use index without order */ /* Never use index without order or limit */
if (path->indexorderbys == NULL) if (path->indexorderbys == NULL || root->limit_tuples < 0)
{ {
*indexStartupCost = DBL_MAX; *indexStartupCost = DBL_MAX;
*indexTotalCost = DBL_MAX; *indexTotalCost = DBL_MAX;
@@ -88,6 +94,20 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
return; return;
} }
/*
* Do not use index if limit + offset > ef_search unless enable_seqscan =
* off
*/
if (root->limit_tuples > hnsw_ef_search)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
MemSet(&costs, 0, sizeof(costs)); MemSet(&costs, 0, sizeof(costs));
index = index_open(path->indexinfo->indexoid, NoLock); index = index_open(path->indexinfo->indexoid, NoLock);

View File

@@ -4,7 +4,9 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h" #include "access/generic_xlog.h"
#include "access/parallel.h"
#include "access/reloptions.h" #include "access/reloptions.h"
#include "lib/ilist.h"
#include "nodes/execnodes.h" #include "nodes/execnodes.h"
#include "port.h" /* for random() */ #include "port.h" /* for random() */
#include "utils/sampling.h" #include "utils/sampling.h"
@@ -14,6 +16,10 @@
#error "Requires PostgreSQL 11+" #error "Requires PostgreSQL 11+"
#endif #endif
#if PG_VERSION_NUM < 120000
#include "access/relscan.h"
#endif
#define HNSW_MAX_DIM 2000 #define HNSW_MAX_DIM 2000
/* Support functions */ /* Support functions */
@@ -59,7 +65,7 @@
#define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData)) #define HNSW_MAX_SIZE (BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - sizeof(ItemIdData))
#define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim)) #define HNSW_ELEMENT_TUPLE_SIZE(size) MAXALIGN(offsetof(HnswElementTupleData, data) + (size))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData)) #define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
@@ -67,8 +73,10 @@
#if PG_VERSION_NUM >= 150000 #if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state) #define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
#else #else
#define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE) #define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE)
#define SeedRandom(seed) srandom(seed)
#endif #endif
#if PG_VERSION_NUM < 130000 #if PG_VERSION_NUM < 130000
@@ -86,24 +94,28 @@
#define HnswGetMl(m) (1 / log(m)) #define HnswGetMl(m) (1 / log(m))
/* Ensure fits on page and in uint8 */ /* 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 HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
#define HnswGetNeighbors(element, lc) (AssertMacro((element)->level >= (lc)), &(element)->neighbors[lc])
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build;
typedef struct HnswNeighborArray HnswNeighborArray;
typedef struct HnswElementData typedef struct HnswElementData
{ {
List *heaptids; slist_node next;
ItemPointerData heaptids[HNSW_HEAPTIDS];
uint8 heaptidsLength;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
HnswNeighborArray *neighbors; uint32 hash;
struct HnswNeighborArray *neighbors;
BlockNumber blkno; BlockNumber blkno;
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Vector *vec; Datum value;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
@@ -136,6 +148,59 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */ int efConstruction; /* size of dynamic candidate list */
} HnswOptions; } HnswOptions;
typedef struct HnswGraph
{
slist_head elements;
HnswElement entryPoint;
long memoryUsed;
long memoryTotal;
bool flushed;
double indtuples;
} HnswGraph;
typedef struct HnswSpool
{
Relation heap;
Relation index;
} HnswSpool;
typedef struct HnswShared
{
/* Immutable state */
Oid heaprelid;
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
/* Worker progress */
ConditionVariable workersdonecv;
/* Mutex for mutable state */
slock_t mutex;
/* Mutable state */
int nparticipantsdone;
double reltuples;
HnswGraph graphData;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} HnswShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromHnswShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared)))
#endif
typedef struct HnswLeader
{
ParallelContext *pcxt;
int nparticipanttuplesorts;
HnswShared *hnswshared;
Snapshot snapshot;
} HnswLeader;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -159,16 +224,19 @@ typedef struct HnswBuildState
Oid collation; Oid collation;
/* Variables */ /* Variables */
List *elements; HnswGraph graphData;
HnswElement entryPoint; HnswGraph *graph;
double ml; double ml;
int maxLevel; int maxLevel;
double maxInMemoryElements;
bool flushed;
Vector *normvec; Vector *normvec;
/* Memory */ /* Memory */
MemoryContext graphCtx;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Parallel builds */
HnswLeader *hnswleader;
HnswShared *hnswshared;
} HnswBuildState; } HnswBuildState;
typedef struct HnswMetaPageData typedef struct HnswMetaPageData
@@ -204,7 +272,7 @@ typedef struct HnswElementTupleData
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused2; uint16 unused2;
Vector vec; Vector data;
} HnswElementTupleData; } HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple; typedef HnswElementTupleData * HnswElementTuple;
@@ -250,7 +318,7 @@ typedef struct HnswVacuumState
Oid collation; Oid collation;
/* Variables */ /* Variables */
HTAB *deleted; struct tidhash_hash *deleted;
BufferAccessStrategy bas; BufferAccessStrategy bas;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
HnswElementData highestPoint; HnswElementData highestPoint;
@@ -264,31 +332,28 @@ int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index); int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result); bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
void HnswCommitBuffer(Buffer buf, GenericXLogState *state);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement); List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing); void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswElement HnswFindDuplicate(HnswElement e);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m); void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel); bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building);
void HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting); void HnswUpdateNeighborPages(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 HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec); void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation); void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
@@ -306,4 +371,31 @@ void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys,
bool hnswgettuple(IndexScanDesc scan, ScanDirection dir); bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan); void hnswendscan(IndexScanDesc scan);
/* Hash tables */
typedef struct TidHashEntry
{
ItemPointerData tid;
char status;
} TidHashEntry;
#define SH_PREFIX tidhash
#define SH_ELEMENT_TYPE TidHashEntry
#define SH_KEY_TYPE ItemPointerData
#define SH_SCOPE extern
#define SH_DECLARE
#include "lib/simplehash.h"
typedef struct PointerHashEntry
{
uintptr_t ptr;
char status;
} PointerHashEntry;
#define SH_PREFIX pointerhash
#define SH_ELEMENT_TYPE PointerHashEntry
#define SH_KEY_TYPE uintptr_t
#define SH_SCOPE extern
#define SH_DECLARE
#include "lib/simplehash.h"
#endif #endif

View File

@@ -2,12 +2,16 @@
#include <math.h> #include <math.h>
#include "access/parallel.h"
#include "access/xact.h"
#include "catalog/index.h" #include "catalog/index.h"
#include "hnsw.h" #include "hnsw.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "nodes/pg_list.h" #include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -35,6 +39,27 @@
#define UpdateProgress(index, val) ((void)val) #define UpdateProgress(index, val) ((void)val)
#endif #endif
#if PG_VERSION_NUM >= 140000
#include "utils/backend_status.h"
#include "utils/wait_event.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/table.h"
#include "optimizer/optimizer.h"
#else
#include "access/heapam.h"
#include "optimizer/planner.h"
#include "pgstat.h"
#endif
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000002)
#if PG_VERSION_NUM < 130000
#define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2)
#endif
/* /*
* Create the metapage * Create the metapage
*/ */
@@ -45,11 +70,11 @@ CreateMetaPage(HnswBuildState * buildstate)
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state;
HnswMetaPage metap; HnswMetaPage metap;
buf = HnswNewBuffer(index, forkNum); buf = HnswNewBuffer(index, forkNum);
HnswInitRegisterPage(index, &buf, &page, &state); page = BufferGetPage(buf);
HnswInitPage(buf, page);
/* Set metapage data */ /* Set metapage data */
metap = HnswPageGetMeta(page); metap = HnswPageGetMeta(page);
@@ -65,14 +90,15 @@ CreateMetaPage(HnswBuildState * buildstate)
((PageHeader) page)->pd_lower = ((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page; ((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
HnswCommitBuffer(buf, state); MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
} }
/* /*
* Add a new page * Add a new page
*/ */
static void static void
HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, ForkNumber forkNum) HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, ForkNumber forkNum)
{ {
/* Add a new page */ /* Add a new page */
Buffer newbuf = HnswNewBuffer(index, forkNum); Buffer newbuf = HnswNewBuffer(index, forkNum);
@@ -81,7 +107,7 @@ HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf); HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */ /* Commit */
GenericXLogFinish(*state); MarkBufferDirty(*buf);
UnlockReleaseBuffer(*buf); UnlockReleaseBuffer(*buf);
/* Can take a while, so ensure we can interrupt */ /* Can take a while, so ensure we can interrupt */
@@ -92,8 +118,7 @@ HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **
/* Prepare new page */ /* Prepare new page */
*buf = newbuf; *buf = newbuf;
*state = GenericXLogStart(index); *page = BufferGetPage(*buf);
*page = GenericXLogRegisterBuffer(*state, *buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*buf, *page); HnswInitPage(*buf, *page);
} }
@@ -105,46 +130,52 @@ CreateElementPages(HnswBuildState * buildstate)
{ {
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
int dimensions = buildstate->dimensions; Size etupAllocSize;
Size etupSize;
Size maxSize; Size maxSize;
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
BlockNumber insertPage; BlockNumber insertPage;
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; slist_iter iter;
ListCell *lc;
/* Calculate sizes */ /* Calculate sizes */
etupAllocSize = BLCKSZ;
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */ /* Allocate once */
etup = palloc0(etupSize); etup = palloc0(etupAllocSize);
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
/* Prepare first page */ /* Prepare first page */
buf = HnswNewBuffer(index, forkNum); buf = HnswNewBuffer(index, forkNum);
state = GenericXLogStart(index); page = BufferGetPage(buf);
page = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(buf, page); HnswInitPage(buf, page);
foreach(lc, buildstate->elements) slist_foreach(iter, &buildstate->graph->elements)
{ {
HnswElement element = lfirst(lc); HnswElement element = slist_container(HnswElementData, next, iter.cur);
Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
HnswSetElementTuple(etup, element); /* Zero memory for each element */
MemSet(etup, 0, etupAllocSize);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(element->value)));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
/* Initial size check */
if (etupSize > etupAllocSize)
elog(ERROR, "index tuple too large");
HnswSetElementTuple(etup, element);
/* Keep element and neighbors on the same page if possible */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize)) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
HnswBuildAppendPage(index, &buf, &page, &state, forkNum); HnswBuildAppendPage(index, &buf, &page, forkNum);
/* Calculate offsets */ /* Calculate offsets */
element->blkno = BufferGetBlockNumber(buf); element->blkno = BufferGetBlockNumber(buf);
@@ -168,7 +199,7 @@ CreateElementPages(HnswBuildState * buildstate)
/* Add new page if needed */ /* Add new page if needed */
if (PageGetFreeSpace(page) < ntupSize) if (PageGetFreeSpace(page) < ntupSize)
HnswBuildAppendPage(index, &buf, &page, &state, forkNum); HnswBuildAppendPage(index, &buf, &page, forkNum);
/* Add placeholder for neighbors */ /* Add placeholder for neighbors */
if (PageAddItem(page, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != element->neighborOffno) if (PageAddItem(page, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != element->neighborOffno)
@@ -178,10 +209,10 @@ CreateElementPages(HnswBuildState * buildstate)
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
/* Commit */ /* Commit */
GenericXLogFinish(state); MarkBufferDirty(buf);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum); HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->graph->entryPoint, insertPage, forkNum, true);
pfree(etup); pfree(etup);
pfree(ntup); pfree(ntup);
@@ -196,18 +227,17 @@ CreateNeighborPages(HnswBuildState * buildstate)
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
int m = buildstate->m; int m = buildstate->m;
ListCell *lc; slist_iter iter;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
/* Allocate once */ /* Allocate once */
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
foreach(lc, buildstate->elements) slist_foreach(iter, &buildstate->graph->elements)
{ {
HnswElement e = lfirst(lc); HnswElement e = slist_container(HnswElementData, next, iter.cur);
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
/* Can take a while, so ensure we can interrupt */ /* Can take a while, so ensure we can interrupt */
@@ -216,8 +246,7 @@ CreateNeighborPages(HnswBuildState * buildstate)
buf = ReadBufferExtended(index, forkNum, e->neighborPage, RBM_NORMAL, NULL); buf = ReadBufferExtended(index, forkNum, e->neighborPage, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); page = BufferGetPage(buf);
page = GenericXLogRegisterBuffer(state, buf, 0);
HnswSetNeighborTuple(ntup, e, m); HnswSetNeighborTuple(ntup, e, m);
@@ -225,26 +254,30 @@ CreateNeighborPages(HnswBuildState * buildstate)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
GenericXLogFinish(state); MarkBufferDirty(buf);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
pfree(ntup); pfree(ntup);
} }
#ifdef HNSW_MEMORY
/* /*
* Free elements * Show memory usage
*/ */
static void static void
FreeElements(HnswBuildState * buildstate) ShowMemoryUsage(HnswBuildState * buildstate)
{ {
ListCell *lc; #if PG_VERSION_NUM >= 130000
elog(INFO, "graph memory: %zu MB, total memory: %zu MB",
foreach(lc, buildstate->elements) MemoryContextMemAllocated(buildstate->graphCtx, false) / (1024 * 1024),
HnswFreeElement(lfirst(lc)); MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
list_free(buildstate->elements); MemoryContextStats(CurrentMemoryContext);
elog(INFO, "estimated memory: %zu MB", buildstate->memoryUsed / (1024 * 1024));
#endif
} }
#endif
/* /*
* Flush pages * Flush pages
@@ -252,25 +285,78 @@ FreeElements(HnswBuildState * buildstate)
static void static void
FlushPages(HnswBuildState * buildstate) FlushPages(HnswBuildState * buildstate)
{ {
#ifdef HNSW_MEMORY
ShowMemoryUsage(buildstate);
#endif
CreateMetaPage(buildstate); CreateMetaPage(buildstate);
CreateElementPages(buildstate); CreateElementPages(buildstate);
CreateNeighborPages(buildstate); CreateNeighborPages(buildstate);
buildstate->flushed = true; buildstate->graph->flushed = true;
FreeElements(buildstate); MemoryContextReset(buildstate->graphCtx);
}
#if PG_VERSION_NUM < 130000
/*
* Get the memory used by an element
*/
static long
HnswElementMemory(HnswElement e, int m)
{
long elementSize = sizeof(HnswElementData);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
/* Each allocation has a chunk header */
elementSize += (e->level + 4) * GENERATIONCHUNK_RAWSIZE;
/* Add an extra 5% for alignment and other overhead */
return elementSize * 1.05;
}
#endif
/*
* Find duplicate element
*/
static bool
HnswFindDuplicateInMemory(HnswElement element)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1))
return false;
/* Check for space */
if (neighbor->element->heaptidsLength < HNSW_HEAPTIDS)
{
HnswAddHeapTid(neighbor->element, &element->heaptids[0]);
return true;
}
}
return false;
} }
/* /*
* Insert tuple * Insert tuple into in-memory graph
*/ */
static bool static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup) InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuildState * buildstate)
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint; HnswGraph *graph = buildstate->graph;
HnswElement entryPoint = graph->entryPoint;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
MemoryContext oldCtx;
HnswElement element;
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -282,35 +368,67 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
return false; return false;
} }
/* Copy value to element so accessible outside of memory context */ /* Allocate element in graph memory context */
memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions)); oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
element = HnswInitElement(heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Update memory usage */
#if PG_VERSION_NUM >= 130000
graph->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
#else
graph->memoryUsed += HnswElementMemory(element, buildstate->m);
#endif
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */ /* Look for duplicate */
*dup = HnswFindDuplicate(element); if (HnswFindDuplicateInMemory(element))
/* Update neighbors if needed */
if (*dup == NULL)
{ {
for (int lc = element->level; lc >= 0; lc--) /* No need to free element since memory unlikely to be reallocated */
{ return true;
int lm = HnswGetLayerM(m, lc); }
HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++) /* Add element */
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation); slist_push_head(&graph->elements, &element->next);
}
/* Update neighbors */
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(element, lc);
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
} }
/* Update entry point if needed */ /* Update entry point if needed */
if (*dup == NULL && (entryPoint == NULL || element->level > entryPoint->level)) if (entryPoint == NULL || element->level > entryPoint->level)
buildstate->entryPoint = element; graph->entryPoint = element;
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples); return true;
}
return *dup == NULL; /*
* Acquire a lock if needed
*/
static inline void
HnswLockAcquire(HnswShared * hnswshared)
{
if (hnswshared)
SpinLockAcquire(&hnswshared->mutex);
}
/*
* Release a lock if needed
*/
static inline void
HnswLockRelease(HnswShared * hnswshared)
{
if (hnswshared)
SpinLockRelease(&hnswshared->mutex);
} }
/* /*
@@ -321,9 +439,9 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state) bool *isnull, bool tupleIsAlive, void *state)
{ {
HnswBuildState *buildstate = (HnswBuildState *) state; HnswBuildState *buildstate = (HnswBuildState *) state;
HnswGraph *graph = buildstate->graph;
HnswShared *hnswshared = buildstate->hnswshared;
MemoryContext oldCtx; MemoryContext oldCtx;
HnswElement element;
HnswElement dup = NULL;
bool inserted; bool inserted;
#if PG_VERSION_NUM < 130000 #if PG_VERSION_NUM < 130000
@@ -334,69 +452,50 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0]) if (isnull[0])
return; return;
if (buildstate->indtuples >= buildstate->maxInMemoryElements) /* Flush pages if needed */
if (!graph->flushed && graph->memoryUsed >= graph->memoryTotal)
{ {
if (!buildstate->flushed) ereport(NOTICE,
{ (errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) graph->indtuples),
ereport(NOTICE, errdetail("Building will take significantly more time."),
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) buildstate->indtuples), errhint("Increase maintenance_work_mem to speed up builds.")));
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
FlushPages(buildstate); FlushPages(buildstate);
}
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap))
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
return;
} }
/* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->vec = palloc(VECTOR_SIZE(buildstate->dimensions));
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup); if (graph->flushed)
inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true);
else
inserted = InsertTupleInMemory(index, values, tid, buildstate);
/* Update progress */
if (inserted)
{
HnswLockAcquire(hnswshared);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples);
HnswLockRelease(hnswshared);
}
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx); MemoryContextReset(buildstate->tmpCtx);
/* Add outside memory context */
if (dup != NULL)
HnswAddHeapTid(dup, tid);
/* Add to buildstate or free */
if (inserted)
buildstate->elements = lappend(buildstate->elements, element);
else
HnswFreeElement(element);
} }
/* /*
* Get the max number of elements that fit into maintenance_work_mem * Initialize the graph
*/ */
static double static void
HnswGetMaxInMemoryElements(int m, double ml, int dimensions) InitGraph(HnswGraph * graph)
{ {
Size elementSize = sizeof(HnswElementData); slist_init(&graph->elements);
double avgLevel = -log(0.5) * ml; graph->entryPoint = NULL;
graph->memoryUsed = 0;
elementSize += sizeof(HnswNeighborArray) * (avgLevel + 1); graph->memoryTotal = maintenance_work_mem * 1024L;
elementSize += sizeof(HnswCandidate) * (m * (avgLevel + 2)); graph->flushed = false;
elementSize += sizeof(ItemPointerData); graph->indtuples = 0;
elementSize += VECTOR_SIZE(dimensions);
return (maintenance_work_mem * 1024L) / elementSize;
} }
/* /*
@@ -432,19 +531,26 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0]; buildstate->collation = index->rd_indcollation[0];
buildstate->elements = NIL; InitGraph(&buildstate->graphData);
buildstate->entryPoint = NULL; buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->maxInMemoryElements = HnswGetMaxInMemoryElements(buildstate->m, buildstate->ml, buildstate->dimensions);
buildstate->flushed = false;
/* Reuse for each tuple */ /* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions); buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",
#if PG_VERSION_NUM >= 150000
1024 * 1024, 1024 * 1024,
#endif
1024 * 1024);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context", "Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
buildstate->hnswleader = NULL;
buildstate->hnswshared = NULL;
} }
/* /*
@@ -454,26 +560,427 @@ static void
FreeBuildState(HnswBuildState * buildstate) FreeBuildState(HnswBuildState * buildstate)
{ {
pfree(buildstate->normvec); pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx); MemoryContextDelete(buildstate->tmpCtx);
} }
/*
* Within leader, wait for end of heap scan
*/
static double
ParallelHeapScan(HnswBuildState * buildstate)
{
HnswShared *hnswshared = buildstate->hnswleader->hnswshared;
int nparticipanttuplesorts;
double reltuples;
nparticipanttuplesorts = buildstate->hnswleader->nparticipanttuplesorts;
for (;;)
{
SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->graph = &hnswshared->graphData;
reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex);
break;
}
SpinLockRelease(&hnswshared->mutex);
ConditionVariableSleep(&hnswshared->workersdonecv,
WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN);
}
ConditionVariableCancelSleep();
return reltuples;
}
/*
* Perform a worker's portion of a parallel insert
*/
static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
{
HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples;
IndexInfo *indexInfo;
/* Join parallel scan */
indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
buildstate.graph = &hnswshared->graphData;
buildstate.hnswshared = hnswshared;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
ParallelTableScanFromHnswShared(hnswshared));
reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo,
true, progress, BuildCallback,
(void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Record statistics */
SpinLockAcquire(&hnswshared->mutex);
hnswshared->nparticipantsdone++;
hnswshared->reltuples += reltuples;
SpinLockRelease(&hnswshared->mutex);
/* Log statistics */
if (progress)
ereport(DEBUG1, (errmsg("leader processed " INT64_FORMAT " tuples", (int64) reltuples)));
else
ereport(DEBUG1, (errmsg("worker processed " INT64_FORMAT " tuples", (int64) reltuples)));
/* Notify leader */
ConditionVariableSignal(&hnswshared->workersdonecv);
FreeBuildState(&buildstate);
}
/*
* Perform work within a launched parallel process
*/
void
HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{
char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
LOCKMODE indexLockmode;
/* Set debug_query_string for individual workers first */
sharedquery = shm_toc_lookup(toc, PARALLEL_KEY_QUERY_TEXT, true);
debug_query_string = sharedquery;
/* Report the query string from leader */
pgstat_report_activity(STATE_RUNNING, debug_query_string);
/* Look up shared state */
hnswshared = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_SHARED, false);
/* Open relations using lock modes known to be obtained by index.c */
if (!hnswshared->isconcurrent)
{
heapLockmode = ShareLock;
indexLockmode = AccessExclusiveLock;
}
else
{
heapLockmode = ShareUpdateExclusiveLock;
indexLockmode = RowExclusiveLock;
}
/* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(hnswshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(hnswshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(hnswshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */
hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
/* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false);
/* Close relations within worker */
index_close(indexRel, indexLockmode);
#if PG_VERSION_NUM >= 120000
table_close(heapRel, heapLockmode);
#else
heap_close(heapRel, heapLockmode);
#endif
}
/*
* End parallel build
*/
static void
HnswEndParallel(HnswLeader * hnswleader)
{
/* Shutdown worker processes */
WaitForParallelWorkersToFinish(hnswleader->pcxt);
/* Free last reference to MVCC snapshot, if one was used */
if (IsMVCCSnapshot(hnswleader->snapshot))
UnregisterSnapshot(hnswleader->snapshot);
DestroyParallelContext(hnswleader->pcxt);
ExitParallelMode();
}
/*
* Return size of shared memory required for parallel index build
*/
static Size
ParallelEstimateShared(Relation heap, Snapshot snapshot)
{
#if PG_VERSION_NUM >= 120000
return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(HnswShared);
}
return add_size(offsetof(HnswShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
}
/*
* Within leader, participate as a parallel worker
*/
static void
HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
{
HnswLeader *hnswleader = buildstate->hnswleader;
HnswSpool *leaderworker;
/* Allocate memory and initialize private spool */
leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
}
/*
* Begin parallel build
*/
static void
HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
{
ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot;
Size esthnswshared;
HnswShared *hnswshared;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true;
int querylen;
#ifdef DISABLE_LEADER_PARTICIPATION
leaderparticipates = false;
#endif
/* Enter parallel mode and create context */
EnterParallelMode();
Assert(request > 0);
#if PG_VERSION_NUM >= 120000
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request;
/* Get snapshot for table scan */
if (!isconcurrent)
snapshot = SnapshotAny;
else
snapshot = RegisterSnapshot(GetTransactionSnapshot());
/* Estimate size of workspaces */
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
shm_toc_estimate_keys(&pcxt->estimator, 1);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string)
{
querylen = strlen(debug_query_string);
shm_toc_estimate_chunk(&pcxt->estimator, querylen + 1);
shm_toc_estimate_keys(&pcxt->estimator, 1);
}
else
querylen = 0; /* keep compiler quiet */
/* Everyone's had a chance to ask for space, so now create the DSM */
InitializeParallelDSM(pcxt);
/* If no DSM segment was available, back out (do serial build) */
if (pcxt->seg == NULL)
{
if (IsMVCCSnapshot(snapshot))
UnregisterSnapshot(snapshot);
DestroyParallelContext(pcxt);
ExitParallelMode();
return;
}
/* Store shared build state, for which we reserved space */
hnswshared = (HnswShared *) shm_toc_allocate(pcxt->toc, esthnswshared);
/* Initialize immutable state */
hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
hnswshared->indexrelid = RelationGetRelid(buildstate->index);
hnswshared->isconcurrent = isconcurrent;
hnswshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&hnswshared->workersdonecv);
SpinLockInit(&hnswshared->mutex);
/* Initialize mutable state */
hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0;
InitGraph(&hnswshared->graphData);
/* TODO Support in-memory builds */
hnswshared->graphData.memoryTotal = 0;
hnswshared->graphData.flushed = true;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared),
snapshot);
#else
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
/* Store query string for workers */
if (debug_query_string)
{
char *sharedquery;
sharedquery = (char *) shm_toc_allocate(pcxt->toc, querylen + 1);
memcpy(sharedquery, debug_query_string, querylen + 1);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_QUERY_TEXT, sharedquery);
}
/* Launch workers, saving status for leader/caller */
LaunchParallelWorkers(pcxt);
hnswleader->pcxt = pcxt;
hnswleader->nparticipanttuplesorts = pcxt->nworkers_launched;
if (leaderparticipates)
hnswleader->nparticipanttuplesorts++;
hnswleader->hnswshared = hnswshared;
hnswleader->snapshot = snapshot;
/* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0)
{
HnswEndParallel(hnswleader);
return;
}
/* Log participants */
ereport(DEBUG1, (errmsg("using %d parallel workers", pcxt->nworkers_launched)));
/* Save leader state now that it's clear build will be parallel */
buildstate->hnswleader = hnswleader;
/* Join heap scan ourselves */
if (leaderparticipates)
HnswLeaderParticipateAsWorker(buildstate);
/* Wait for all launched workers */
WaitForParallelWorkersToAttach(pcxt);
}
/*
* Compute parallel workers
*/
static int
ComputeParallelWorkers(Relation heap, Relation index)
{
int parallel_workers;
/* Make sure it's safe to use parallel workers */
parallel_workers = plan_create_index_workers(RelationGetRelid(heap), RelationGetRelid(index));
if (parallel_workers == 0)
return 0;
/* Use parallel_workers storage parameter on table if set */
parallel_workers = RelationGetParallelWorkers(heap, -1);
if (parallel_workers != -1)
return Min(parallel_workers, max_parallel_maintenance_workers);
return max_parallel_maintenance_workers;
}
/* /*
* Build graph * Build graph
*/ */
static void static void
BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum) BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
{ {
int parallel_workers = 0;
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD); UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */
if (buildstate->heap != NULL && hnsw_enable_parallel_build)
parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to graph */
if (buildstate->heap != NULL)
{
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL); true, true, BuildCallback, (void *) buildstate, NULL);
#else #else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL); true, BuildCallback, (void *) buildstate, NULL);
#endif #endif
}
buildstate->indtuples = buildstate->graph->indtuples;
}
/* Flush pages */
if (!buildstate->graph->flushed)
FlushPages(buildstate);
/* End parallel build */
if (buildstate->hnswleader)
HnswEndParallel(buildstate->hnswleader);
} }
#if PG_VERSION_NUM < 110008
void
log_newpage_range(Relation rel, ForkNumber forkNum, BlockNumber startblk, BlockNumber endblk, bool page_std)
{
for (BlockNumber blkno = startblk; blkno < endblk; blkno++)
{
Buffer buf = ReadBufferExtended(rel, forkNum, blkno, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
MarkBufferDirty(buf);
log_newpage_buffer(buf, page_std);
UnlockReleaseBuffer(buf);
}
}
#endif
/* /*
* Build the index * Build the index
*/ */
@@ -481,13 +988,16 @@ static void
BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo, BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
HnswBuildState * buildstate, ForkNumber forkNum) HnswBuildState * buildstate, ForkNumber forkNum)
{ {
#ifdef HNSW_MEMORY
SeedRandom(42);
#endif
InitBuildState(buildstate, heap, index, indexInfo, forkNum); InitBuildState(buildstate, heap, index, indexInfo, forkNum);
if (buildstate->heap != NULL) BuildGraph(buildstate, forkNum);
BuildGraph(buildstate, forkNum);
if (!buildstate->flushed) if (RelationNeedsWAL(index))
FlushPages(buildstate); log_newpage_range(index, forkNum, 0, RelationGetNumberOfBlocks(index), true);
FreeBuildState(buildstate); FreeBuildState(buildstate);
} }

View File

@@ -5,6 +5,7 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
/* /*
@@ -92,7 +93,7 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
* Add a new page * Add a new page
*/ */
static void static void
HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState *state, Page page) HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState *state, Page page, bool building)
{ {
/* Add a new page */ /* Add a new page */
LockRelationForExtension(index, ExclusiveLock); LockRelationForExtension(index, ExclusiveLock);
@@ -100,7 +101,11 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
UnlockRelationForExtension(index, ExclusiveLock); UnlockRelationForExtension(index, ExclusiveLock);
/* Init new page */ /* Init new page */
*npage = GenericXLogRegisterBuffer(state, *nbuf, GENERIC_XLOG_FULL_IMAGE); if (building)
*npage = BufferGetPage(*nbuf);
else
*npage = GenericXLogRegisterBuffer(state, *nbuf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*nbuf, *npage); HnswInitPage(*nbuf, *npage);
/* Update previous buffer */ /* Update previous buffer */
@@ -111,7 +116,7 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
* Add to element and neighbor pages * Add to element and neighbor pages
*/ */
static void static void
WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage) WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -123,7 +128,6 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Size minCombinedSize; Size minCombinedSize;
HnswElementTuple etup; HnswElementTuple etup;
BlockNumber currentPage = insertPage; BlockNumber currentPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Buffer nbuf; Buffer nbuf;
Page npage; Page npage;
@@ -132,7 +136,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions); etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(e->value)));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -152,8 +156,16 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
buf = ReadBuffer(index, currentPage); buf = ReadBuffer(index, currentPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); if (building)
page = GenericXLogRegisterBuffer(state, buf, 0); {
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Keep track of first page where element at level 0 can fit */ /* Keep track of first page where element at level 0 can fit */
if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize) if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize)
@@ -173,7 +185,12 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage)) if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
{ {
if (nbuf != buf) if (nbuf != buf)
npage = GenericXLogRegisterBuffer(state, nbuf, 0); {
if (building)
npage = BufferGetPage(nbuf);
else
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
break; break;
} }
@@ -182,7 +199,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
/* Skip if both tuples can fit on the same page */ /* Skip if both tuples can fit on the same page */
if (combinedSize > maxSize && PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno)) if (combinedSize > maxSize && PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
{ {
HnswInsertAppendPage(index, &nbuf, &npage, state, page); HnswInsertAppendPage(index, &nbuf, &npage, state, page, building);
break; break;
} }
@@ -191,7 +208,8 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
if (BlockNumberIsValid(currentPage)) if (BlockNumberIsValid(currentPage))
{ {
/* Move to next page */ /* Move to next page */
GenericXLogAbort(state); if (!building)
GenericXLogAbort(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
else else
@@ -199,22 +217,33 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
Buffer newbuf; Buffer newbuf;
Page newpage; Page newpage;
HnswInsertAppendPage(index, &newbuf, &newpage, state, page); HnswInsertAppendPage(index, &newbuf, &newpage, state, page, building);
/* Commit */ /* Commit */
GenericXLogFinish(state); if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
/* Unlock previous buffer */ /* Unlock previous buffer */
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Prepare new buffer */ /* Prepare new buffer */
state = GenericXLogStart(index);
buf = newbuf; buf = newbuf;
page = GenericXLogRegisterBuffer(state, buf, 0); if (building)
{
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Create new page for neighbors if needed */ /* Create new page for neighbors if needed */
if (PageGetFreeSpace(page) < combinedSize) if (PageGetFreeSpace(page) < combinedSize)
HnswInsertAppendPage(index, &nbuf, &npage, state, page); HnswInsertAppendPage(index, &nbuf, &npage, state, page, building);
else else
{ {
nbuf = buf; nbuf = buf;
@@ -268,7 +297,14 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
} }
/* Commit */ /* Commit */
GenericXLogFinish(state); if (building)
{
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
}
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
if (nbuf != buf) if (nbuf != buf)
UnlockReleaseBuffer(nbuf); UnlockReleaseBuffer(nbuf);
@@ -302,12 +338,12 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
* Update neighbors * Update neighbors
*/ */
void void
HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting) HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc]; HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
@@ -343,8 +379,16 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
/* Register page */ /* Register page */
buf = ReadBuffer(index, hc->element->neighborPage); buf = ReadBuffer(index, hc->element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); if (building)
page = GenericXLogRegisterBuffer(state, buf, 0); {
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Get tuple */ /* Get tuple */
itemid = PageGetItemId(page, offno); itemid = PageGetItemId(page, offno);
@@ -386,9 +430,12 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
GenericXLogFinish(state); if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
} }
else else if (!building)
GenericXLogAbort(state); GenericXLogAbort(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -400,23 +447,34 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
* Add a heap TID to an existing element * Add a heap TID to an existing element
*/ */
static bool static bool
HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup) HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool building)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
GenericXLogState *state; GenericXLogState *state;
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim); ItemId itemid;
HnswElementTuple etup; HnswElementTuple etup;
Size etupSize;
int i; int i;
/* Read page */ /* Read page */
buf = ReadBuffer(index, dup->blkno); buf = ReadBuffer(index, dup->blkno);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); if (building)
page = GenericXLogRegisterBuffer(state, buf, 0); {
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
/* Find space */ /* Find space */
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, dup->offno)); itemid = PageGetItemId(page, dup->offno);
etup = (HnswElementTuple) PageGetItem(page, itemid);
etupSize = ItemIdGetLength(itemid);
for (i = 0; i < HNSW_HEAPTIDS; i++) for (i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (!ItemPointerIsValid(&etup->heaptids[i])) if (!ItemPointerIsValid(&etup->heaptids[i]))
@@ -426,60 +484,84 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
/* Either being deleted or we lost our chance to another backend */ /* Either being deleted or we lost our chance to another backend */
if (i == 0 || i == HNSW_HEAPTIDS) if (i == 0 || i == HNSW_HEAPTIDS)
{ {
GenericXLogAbort(state); if (!building)
GenericXLogAbort(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
return false; return false;
} }
/* Add heap TID */ /* Add heap TID */
etup->heaptids[i] = *((ItemPointer) linitial(element->heaptids)); etup->heaptids[i] = element->heaptids[0];
/* Overwrite tuple */ /* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, dup->offno, (Item) etup, etupSize)) if (!PageIndexTupleOverwrite(page, dup->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */ /* Commit */
GenericXLogFinish(state); if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
return true; return true;
} }
/*
* Find duplicate element
*/
static bool
HnswFindDuplicate(Relation index, HnswElement element, bool building)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1))
return false;
if (HnswAddDuplicate(index, element, neighbor->element, building))
return true;
}
return false;
}
/* /*
* Write changes to disk * Write changes to disk
*/ */
static void static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint) WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Try to add to existing page */ /* Look for duplicate */
if (dup != NULL) if (HnswFindDuplicate(index, element, building))
{ return;
if (HnswAddDuplicate(index, element, dup))
return;
}
/* Write element and neighbor tuples */ /* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage); WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage, building);
/* Update insert page if needed */ /* Update insert page if needed */
if (BlockNumberIsValid(newInsertPage)) if (BlockNumberIsValid(newInsertPage))
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM); HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false); HnswUpdateNeighborPages(index, procinfo, collation, element, m, false, building);
/* Update metapage if needed */ /* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM); HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM, building);
} }
/* /*
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel) HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building)
{ {
Datum value; Datum value;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -489,7 +571,6 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
HnswElement dup;
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
/* Detoast once for all calls */ /* Detoast once for all calls */
@@ -515,7 +596,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
element->vec = DatumGetVector(value); element->value = value;
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -534,11 +615,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Insert element in graph */ /* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false); HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
dup = HnswFindDuplicate(element);
/* Write to disk */ /* Write to disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint); WriteElement(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -572,7 +650,7 @@ hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid,
oldCtx = MemoryContextSwitchTo(insertCtx); oldCtx = MemoryContextSwitchTo(insertCtx);
/* Insert tuple */ /* Insert tuple */
HnswInsertTuple(index, values, isnull, heap_tid, heap); HnswInsertTuple(index, values, isnull, heap_tid, heap, false);
/* Delete memory context */ /* Delete memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);

View File

@@ -188,15 +188,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
ItemPointer heaptid; ItemPointer heaptid;
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (list_length(hc->element->heaptids) == 0) if (hc->element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);
continue; continue;
} }
heaptid = llast(hc->element->heaptids); heaptid = &hc->element->heaptids[--hc->element->heaptidsLength];
hc->element->heaptids = list_delete_last(hc->element->heaptids);
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);

View File

@@ -4,8 +4,92 @@
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h"
#include "vector.h" #include "vector.h"
#if PG_VERSION_NUM >= 130000
#include "common/hashfn.h"
#else
#include "utils/hashutils.h"
#endif
#if PG_VERSION_NUM < 170000
static inline uint64
murmurhash64(uint64 data)
{
uint64 h = data;
h ^= h >> 33;
h *= 0xff51afd7ed558ccd;
h ^= h >> 33;
h *= 0xc4ceb9fe1a85ec53;
h ^= h >> 33;
return h;
}
#endif
/* TID hash table */
static uint32
hash_tid(ItemPointerData tid)
{
union
{
uint64 i;
ItemPointerData tid;
} x;
/* Initialize unused bytes */
x.i = 0;
x.tid = tid;
return murmurhash64(x.i);
}
#define SH_PREFIX tidhash
#define SH_ELEMENT_TYPE TidHashEntry
#define SH_KEY_TYPE ItemPointerData
#define SH_KEY tid
#define SH_HASH_KEY(tb, key) hash_tid(key)
#define SH_EQUAL(tb, a, b) ItemPointerEquals(&a, &b)
#define SH_SCOPE extern
#define SH_DEFINE
#include "lib/simplehash.h"
/* Needed to include simplehash.h twice */
#if PG_VERSION_NUM < 120000
#undef SH_EQUAL
#define sh_log2 pointerhash_sh_log2
#define sh_pow2 pointerhash_sh_pow2
#endif
/* Pointer hash table */
static uint32
hash_pointer(uintptr_t ptr)
{
#if SIZEOF_VOID_P == 8
return murmurhash64((uint64) ptr);
#else
return murmurhash32((uint32) ptr);
#endif
}
#define SH_PREFIX pointerhash
#define SH_ELEMENT_TYPE PointerHashEntry
#define SH_KEY_TYPE uintptr_t
#define SH_KEY ptr
#define SH_HASH_KEY(tb, key) hash_pointer(key)
#define SH_EQUAL(tb, a, b) (a == b)
#define SH_SCOPE extern
#define SH_DEFINE
#include "lib/simplehash.h"
typedef union
{
pointerhash_hash *pointers;
tidhash_hash *tids;
} visited_hash;
/* /*
* Get the max number of connections in an upper layer for each element in the index * Get the max number of connections in an upper layer for each element in the index
*/ */
@@ -100,27 +184,6 @@ HnswInitPage(Buffer buf, Page page)
HnswPageGetOpaque(page)->page_id = HNSW_PAGE_ID; HnswPageGetOpaque(page)->page_id = HNSW_PAGE_ID;
} }
/*
* Init and register page
*/
void
HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state)
{
*state = GenericXLogStart(index);
*page = GenericXLogRegisterBuffer(*state, *buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*buf, *page);
}
/*
* Commit buffer
*/
void
HnswCommitBuffer(Buffer buf, GenericXLogState *state)
{
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
}
/* /*
* Allocate neighbors * Allocate neighbors
*/ */
@@ -143,17 +206,6 @@ HnswInitNeighbors(HnswElement element, int m)
} }
} }
/*
* Free neighbors
*/
static void
HnswFreeNeighbors(HnswElement element)
{
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
}
/* /*
* Allocate an element * Allocate an element
*/ */
@@ -168,7 +220,7 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
if (level > maxLevel) if (level > maxLevel)
level = maxLevel; level = maxLevel;
element->heaptids = NIL; element->heaptidsLength = 0;
HnswAddHeapTid(element, heaptid); HnswAddHeapTid(element, heaptid);
element->level = level; element->level = level;
@@ -176,19 +228,9 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
HnswInitNeighbors(element, m); HnswInitNeighbors(element, m);
return element; element->value = PointerGetDatum(NULL);
}
/* return element;
* Free an element
*/
void
HnswFreeElement(HnswElement element)
{
HnswFreeNeighbors(element);
list_free_deep(element->heaptids);
pfree(element->vec);
pfree(element);
} }
/* /*
@@ -197,10 +239,7 @@ HnswFreeElement(HnswElement element)
void void
HnswAddHeapTid(HnswElement element, ItemPointer heaptid) HnswAddHeapTid(HnswElement element, ItemPointer heaptid)
{ {
ItemPointer copy = palloc(sizeof(ItemPointerData)); element->heaptids[element->heaptidsLength++] = *heaptid;
ItemPointerCopy(heaptid, copy);
element->heaptids = lappend(element->heaptids, copy);
} }
/* /*
@@ -214,7 +253,7 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; element->neighbors = NULL;
element->vec = NULL; element->value = PointerGetDatum(NULL);
return element; return element;
} }
@@ -292,7 +331,7 @@ HnswUpdateMetaPageInfo(Page page, int updateEntry, HnswElement entryPoint, Block
* Update the metapage * Update the metapage
*/ */
void void
HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum) HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -300,12 +339,24 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
buf = ReadBufferExtended(index, forkNum, HNSW_METAPAGE_BLKNO, RBM_NORMAL, NULL); buf = ReadBufferExtended(index, forkNum, HNSW_METAPAGE_BLKNO, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index); if (building)
page = GenericXLogRegisterBuffer(state, buf, 0); {
state = NULL;
page = BufferGetPage(buf);
}
else
{
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage); HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage);
HnswCommitBuffer(buf, state); if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
} }
/* /*
@@ -319,12 +370,12 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
etup->deleted = 0; etup->deleted = 0;
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (i < list_length(element->heaptids)) if (i < element->heaptidsLength)
etup->heaptids[i] = *((ItemPointer) list_nth(element->heaptids, i)); etup->heaptids[i] = element->heaptids[i];
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim)); memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value)));
} }
/* /*
@@ -339,7 +390,7 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
HnswNeighborArray *neighbors = &e->neighbors[lc]; HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc);
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
for (int i = 0; i < lm; i++) for (int i = 0; i < lm; i++)
@@ -397,7 +448,7 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
if (level < 0) if (level < 0)
level = 0; level = 0;
neighbors = &element->neighbors[level]; neighbors = HnswGetNeighbors(element, level);
hc = &neighbors->items[neighbors->length++]; hc = &neighbors->items[neighbors->length++];
hc->element = e; hc->element = e;
} }
@@ -431,7 +482,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
element->deleted = etup->deleted; element->deleted = etup->deleted;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid); element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid); element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptids = NIL; element->heaptidsLength = 0;
if (loadHeaptids) if (loadHeaptids)
{ {
@@ -446,10 +497,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
} }
if (loadVec) if (loadVec)
{ element->value = datumCopy(PointerGetDatum(&etup->data), false, -1);
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
} }
/* /*
@@ -476,7 +524,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec))); *distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->data)));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
} }
@@ -487,7 +535,7 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
static float static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation) GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{ {
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec))); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value));
} }
/* /*
@@ -552,16 +600,22 @@ CreatePairingHeapNode(HnswCandidate * c)
* Add to visited * Add to visited
*/ */
static inline void static inline void
AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found) AddToVisited(visited_hash v, HnswCandidate * hc, Relation index, bool *found)
{ {
if (index == NULL) if (index == NULL)
hash_search(v, &hc->element, HASH_ENTER, found); {
#if PG_VERSION_NUM >= 130000
pointerhash_insert_hash(v.pointers, (uintptr_t) hc->element, hc->element->hash, found);
#else
pointerhash_insert(v.pointers, (uintptr_t) hc->element, found);
#endif
}
else else
{ {
ItemPointerData indextid; ItemPointerData indextid;
ItemPointerSet(&indextid, hc->element->blkno, hc->element->offno); ItemPointerSet(&indextid, hc->element->blkno, hc->element->offno);
hash_search(v, &indextid, HASH_ENTER, found); tidhash_insert(v.tids, indextid, found);
} }
} }
@@ -571,36 +625,26 @@ AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
List * List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement) HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
{ {
ListCell *lc2;
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL); pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0; int wlen = 0;
HASHCTL hash_ctl; visited_hash v;
HTAB *v; ListCell *lc2;
/* Create hash table */ /* Create hash table */
if (index == NULL) if (index == NULL)
{ v.pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL);
hash_ctl.keysize = sizeof(HnswElement *);
hash_ctl.entrysize = sizeof(HnswElement *);
}
else else
{ v.tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
hash_ctl.keysize = sizeof(ItemPointerData);
hash_ctl.entrysize = sizeof(ItemPointerData);
}
hash_ctl.hcxt = CurrentMemoryContext;
v = hash_create("hnsw visited", 256, &hash_ctl, HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
/* Add entry points to v, C, and W */ /* Add entry points to v, C, and W */
foreach(lc2, ep) foreach(lc2, ep)
{ {
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
bool found;
AddToVisited(v, hc, index, NULL); AddToVisited(v, hc, index, &found);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
@@ -610,7 +654,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* would be ideal to do this for inserts as well, but this could * would be ideal to do this for inserts as well, but this could
* affect insert performance. * affect insert performance.
*/ */
if (skipElement == NULL || list_length(hc->element->heaptids) != 0) if (skipElement == NULL || hc->element->heaptidsLength != 0)
wlen++; wlen++;
} }
@@ -627,7 +671,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswLoadNeighbors(c->element, index, m); HnswLoadNeighbors(c->element, index, m);
/* Get the neighborhood at layer lc */ /* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc]; neighborhood = HnswGetNeighbors(c->element, lc);
for (int i = 0; i < neighborhood->length; i++) for (int i = 0; i < neighborhood->length; i++)
{ {
@@ -669,7 +713,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* vacuuming. It would be ideal to do this for inserts as * vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance. * well, but this could affect insert performance.
*/ */
if (skipElement == NULL || list_length(e->element->heaptids) != 0) if (skipElement == NULL || e->element->heaptidsLength != 0)
{ {
wlen++; wlen++;
@@ -730,27 +774,27 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
/* Look for cached distance */ /* Look for cached distance */
if (a->neighbors != NULL) if (a->neighbors != NULL)
{ {
Assert(a->level >= lc); HnswNeighborArray *neighbors = HnswGetNeighbors(a, lc);
for (int i = 0; i < a->neighbors[lc].length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
if (a->neighbors[lc].items[i].element == b) if (neighbors->items[i].element == b)
return a->neighbors[lc].items[i].distance; return neighbors->items[i].distance;
} }
} }
if (b->neighbors != NULL) if (b->neighbors != NULL)
{ {
Assert(b->level >= lc); HnswNeighborArray *neighbors = HnswGetNeighbors(b, lc);
for (int i = 0; i < b->neighbors[lc].length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
if (b->neighbors[lc].items[i].element == a) if (neighbors->items[i].element == a)
return b->neighbors[lc].items[i].distance; return neighbors->items[i].distance;
} }
} }
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec))); return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value));
} }
/* /*
@@ -782,7 +826,8 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
pairingheap *wd; pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet; HnswNeighborArray *neighbors = HnswGetNeighbors(e2, lc);
bool mustCalculate = !neighbors->closerSet;
List *added = NIL; List *added = NIL;
bool removedAny = false; bool removedAny = false;
@@ -846,7 +891,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
} }
/* Cached value can only be used in future if sorted deterministically */ /* Cached value can only be used in future if sorted deterministically */
e2->neighbors[lc].closerSet = sortCandidates; neighbors->closerSet = sortCandidates;
/* Keep pruned connections */ /* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m) while (!pairingheap_is_empty(wd) && list_length(r) < m)
@@ -864,30 +909,6 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
return r; return r;
} }
/*
* Find duplicate element
*/
HnswElement
HnswFindDuplicate(HnswElement e)
{
HnswNeighborArray *neighbors = &e->neighbors[0];
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (vector_cmp_internal(e->vec, neighbor->element->vec) != 0)
break;
/* Check for space */
if (list_length(neighbor->element->heaptids) < HNSW_HEAPTIDS)
return neighbor->element;
}
return NULL;
}
/* /*
* Add connections * Add connections
*/ */
@@ -895,7 +916,7 @@ static void
AddConnections(HnswElement element, List *neighbors, int m, int lc) AddConnections(HnswElement element, List *neighbors, int m, int lc)
{ {
ListCell *lc2; ListCell *lc2;
HnswNeighborArray *a = &element->neighbors[lc]; HnswNeighborArray *a = HnswGetNeighbors(element, lc);
foreach(lc2, neighbors) foreach(lc2, neighbors)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2)); a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
@@ -907,7 +928,7 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
void void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation) HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{ {
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc]; HnswNeighborArray *currentNeighbors = HnswGetNeighbors(hc->element, lc);
HnswCandidate hc2; HnswCandidate hc2;
@@ -930,19 +951,19 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = PointerGetDatum(hc->element->vec); Datum q = hc->element->value;
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
if (hc3->element->vec == NULL) if (DatumGetPointer(hc3->element->value) == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */ /* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0) if (hc3->element->heaptidsLength == 0)
{ {
pruned = &currentNeighbors->items[i]; pruned = &currentNeighbors->items[i];
break; break;
@@ -1000,7 +1021,7 @@ RemoveElements(List *w, HnswElement skipElement)
if (skipElement != NULL && hc->element->blkno == skipElement->blkno && hc->element->offno == skipElement->offno) if (skipElement != NULL && hc->element->blkno == skipElement->blkno && hc->element->offno == skipElement->offno)
continue; continue;
if (list_length(hc->element->heaptids) != 0) if (hc->element->heaptidsLength != 0)
w2 = lappend(w2, hc); w2 = lappend(w2, hc);
} }
@@ -1017,9 +1038,15 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = PointerGetDatum(element->vec); Datum q = element->value;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
#if PG_VERSION_NUM >= 130000
/* Precompute hash */
if (index == NULL)
element->hash = hash_pointer((uintptr_t) element);
#endif
/* No neighbors if no entry point */ /* No neighbors if no entry point */
if (entryPoint == NULL) if (entryPoint == NULL)
return; return;

View File

@@ -12,12 +12,9 @@
* Check if deleted list contains an index TID * Check if deleted list contains an index TID
*/ */
static bool static bool
DeletedContains(HTAB *deleted, ItemPointer indextid) DeletedContains(tidhash_hash * deleted, ItemPointer indextid)
{ {
bool found; return tidhash_lookup(deleted, *indextid) != NULL;
hash_search(deleted, indextid, HASH_FIND, &found);
return found;
} }
/* /*
@@ -62,7 +59,8 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
/* Iterate over nodes */ /* Iterate over nodes */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno)); ItemId itemid = PageGetItemId(page, offno);
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
int idx = 0; int idx = 0;
bool itemUpdated = false; bool itemUpdated = false;
@@ -93,7 +91,7 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (itemUpdated) if (itemUpdated)
{ {
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim); Size etupSize = ItemIdGetLength(itemid);
/* Mark rest as invalid */ /* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++) for (int i = idx; i < HNSW_HEAPTIDS; i++)
@@ -109,11 +107,13 @@ RemoveHeapTids(HnswVacuumState * vacuumstate)
if (!ItemPointerIsValid(&etup->heaptids[0])) if (!ItemPointerIsValid(&etup->heaptids[0]))
{ {
ItemPointerData ip; ItemPointerData ip;
bool found;
/* Add to deleted list */ /* Add to deleted list */
ItemPointerSet(&ip, blkno, offno); ItemPointerSet(&ip, blkno, offno);
(void) hash_search(vacuumstate->deleted, &ip, HASH_ENTER, NULL); tidhash_insert(vacuumstate->deleted, ip, &found);
Assert(!found);
} }
else if (etup->level > highestLevel && !(entryPoint != NULL && blkno == entryPoint->blkno && offno == entryPoint->offno)) else if (etup->level > highestLevel && !(entryPoint != NULL && blkno == entryPoint->blkno && offno == entryPoint->offno))
{ {
@@ -206,7 +206,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
/* Init fields */ /* Init fields */
HnswInitNeighbors(element, m); HnswInitNeighbors(element, m);
element->heaptids = NIL; element->heaptidsLength = 0;
/* Add element to graph, skipping itself */ /* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true); HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
@@ -230,7 +230,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, true); HnswUpdateNeighborPages(index, procinfo, collation, element, m, true, false);
} }
/* /*
@@ -286,7 +286,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* point is outdated and empty, the entry point will be empty * point is outdated and empty, the entry point will be empty
* until an element is repaired. * until an element is repaired.
*/ */
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, highestPoint, InvalidBlockNumber, MAIN_FORKNUM); HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, highestPoint, InvalidBlockNumber, MAIN_FORKNUM, false);
} }
else else
{ {
@@ -419,7 +419,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
* was replaced and highest point was outdated. * was replaced and highest point was outdated.
*/ */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM); HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM, false);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -477,7 +477,8 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Update element and neighbors together */ /* Update element and neighbors together */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno)) for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{ {
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno)); ItemId itemid = PageGetItemId(page, offno);
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, itemid);
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
@@ -505,7 +506,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
continue; continue;
/* Calculate sizes */ /* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim); etupSize = ItemIdGetLength(itemid);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
/* Get neighbor page */ /* Get neighbor page */
@@ -528,7 +529,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float)); MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)
@@ -563,7 +564,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
} }
/* Update insert page last, after everything has been marked as deleted */ /* Update insert page last, after everything has been marked as deleted */
HnswUpdateMetaPage(index, 0, NULL, insertPage, MAIN_FORKNUM); HnswUpdateMetaPage(index, 0, NULL, insertPage, MAIN_FORKNUM, false);
} }
/* /*
@@ -573,7 +574,6 @@ static void
InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkDeleteResult *stats, IndexBulkDeleteCallback callback, void *callback_state) InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkDeleteResult *stats, IndexBulkDeleteCallback callback, void *callback_state)
{ {
Relation index = info->index; Relation index = info->index;
HASHCTL hash_ctl;
if (stats == NULL) if (stats == NULL)
stats = (IndexBulkDeleteResult *) palloc0(sizeof(IndexBulkDeleteResult)); stats = (IndexBulkDeleteResult *) palloc0(sizeof(IndexBulkDeleteResult));
@@ -595,10 +595,7 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
HnswGetMetaPageInfo(index, &vacuumstate->m, NULL); HnswGetMetaPageInfo(index, &vacuumstate->m, NULL);
/* Create hash table */ /* Create hash table */
hash_ctl.keysize = sizeof(ItemPointerData); vacuumstate->deleted = tidhash_create(CurrentMemoryContext, 256, NULL);
hash_ctl.entrysize = sizeof(ItemPointerData);
hash_ctl.hcxt = CurrentMemoryContext;
vacuumstate->deleted = hash_create("hnswbulkdelete indextids", 256, &hash_ctl, HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
} }
/* /*
@@ -607,7 +604,7 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
static void static void
FreeVacuumState(HnswVacuumState * vacuumstate) FreeVacuumState(HnswVacuumState * vacuumstate)
{ {
hash_destroy(vacuumstate->deleted); tidhash_destroy(vacuumstate->deleted);
FreeAccessStrategy(vacuumstate->bas); FreeAccessStrategy(vacuumstate->bas);
pfree(vacuumstate->ntup); pfree(vacuumstate->ntup);
MemoryContextDelete(vacuumstate->tmpCtx); MemoryContextDelete(vacuumstate->tmpCtx);

View File

@@ -11,7 +11,6 @@
#include "miscadmin.h" #include "miscadmin.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h" #include "tcop/tcopprot.h"
#include "utils/datum.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
@@ -66,18 +65,11 @@
static void static void
AddSample(Datum *values, IvfflatBuildState * buildstate) AddSample(Datum *values, IvfflatBuildState * buildstate)
{ {
MemoryContext oldCtx; VectorArray samples = buildstate->samples;
Datum value; int targsamples = samples->maxlen;
int targsamples = buildstate->targsamples;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Detoast once for all calls */ /* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Restore memory context */
MemoryContextSwitchTo(oldCtx);
/* /*
* Normalize with KMEANS_NORM_PROC since spherical distance function * Normalize with KMEANS_NORM_PROC since spherical distance function
@@ -89,23 +81,18 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
return; return;
} }
/* Copy datum */ if (samples->length < targsamples)
value = datumCopy(value, false, -1); {
VectorArraySet(samples, samples->length, DatumGetVector(value));
/* Reset memory context */ samples->length++;
MemoryContextReset(buildstate->tmpCtx); }
if (list_length(buildstate->samples) < targsamples)
buildstate->samples = lappend(buildstate->samples, DatumGetVector(value));
else else
{ {
if (buildstate->rowstoskip < 0) if (buildstate->rowstoskip < 0)
buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, list_length(buildstate->samples), targsamples); buildstate->rowstoskip = reservoir_get_next_S(&buildstate->rstate, samples->length, targsamples);
if (buildstate->rowstoskip <= 0) if (buildstate->rowstoskip <= 0)
{ {
ListCell *lc;
#if PG_VERSION_NUM >= 150000 #if PG_VERSION_NUM >= 150000
int k = (int) (targsamples * sampler_random_fract(&buildstate->rstate.randstate)); int k = (int) (targsamples * sampler_random_fract(&buildstate->rstate.randstate));
#else #else
@@ -113,8 +100,7 @@ AddSample(Datum *values, IvfflatBuildState * buildstate)
#endif #endif
Assert(k >= 0 && k < targsamples); Assert(k >= 0 && k < targsamples);
lc = list_nth_cell(buildstate->samples, k); VectorArraySet(samples, k, DatumGetVector(value));
lfirst(lc) = DatumGetVector(value);
} }
buildstate->rowstoskip -= 1; buildstate->rowstoskip -= 1;
@@ -129,13 +115,21 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state) bool *isnull, bool tupleIsAlive, void *state)
{ {
IvfflatBuildState *buildstate = (IvfflatBuildState *) state; IvfflatBuildState *buildstate = (IvfflatBuildState *) state;
MemoryContext oldCtx;
/* Skip nulls */ /* Skip nulls */
if (isnull[0]) if (isnull[0])
return; return;
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add sample */ /* Add sample */
AddSample(values, buildstate); AddSample(values, state);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
} }
/* /*
@@ -144,7 +138,7 @@ SampleCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
static void static void
SampleRows(IvfflatBuildState * buildstate) SampleRows(IvfflatBuildState * buildstate)
{ {
int targsamples = buildstate->targsamples; int targsamples = buildstate->samples->maxlen;
BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap); BlockNumber totalblocks = RelationGetNumberOfBlocks(buildstate->heap);
buildstate->rowstoskip = -1; buildstate->rowstoskip = -1;
@@ -455,13 +449,12 @@ ComputeCenters(IvfflatBuildState * buildstate)
/* Sample rows */ /* Sample rows */
/* TODO Ensure within maintenance_work_mem */ /* TODO Ensure within maintenance_work_mem */
buildstate->samples = NIL; buildstate->samples = VectorArrayInit(numSamples, buildstate->dimensions);
buildstate->targsamples = numSamples;
if (buildstate->heap != NULL) if (buildstate->heap != NULL)
{ {
SampleRows(buildstate); SampleRows(buildstate);
if (list_length(buildstate->samples) < buildstate->lists) if (buildstate->samples->length < buildstate->lists)
{ {
ereport(NOTICE, ereport(NOTICE,
(errmsg("ivfflat index created with little data"), (errmsg("ivfflat index created with little data"),
@@ -474,7 +467,7 @@ ComputeCenters(IvfflatBuildState * buildstate)
IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers)); IvfflatBench("k-means", IvfflatKmeans(buildstate->index, buildstate->samples, buildstate->centers));
/* Free samples before we allocate more memory */ /* Free samples before we allocate more memory */
list_free_deep(buildstate->samples); VectorArrayFree(buildstate->samples);
} }
/* /*
@@ -550,10 +543,10 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
pfree(list); pfree(list);
} }
#ifdef IVFFLAT_KMEANS_DEBUG
/* /*
* Print k-means metrics * Print k-means metrics
*/ */
#ifdef IVFFLAT_KMEANS_DEBUG
static void static void
PrintKmeansMetrics(IvfflatBuildState * buildstate) PrintKmeansMetrics(IvfflatBuildState * buildstate)
{ {

View File

@@ -76,8 +76,8 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
List *qinfos; List *qinfos;
#endif #endif
/* Never use index without order */ /* Never use index without order or limit */
if (path->indexorderbys == NULL) if (path->indexorderbys == NULL || root->limit_tuples < 0)
{ {
*indexStartupCost = DBL_MAX; *indexStartupCost = DBL_MAX;
*indexTotalCost = DBL_MAX; *indexTotalCost = DBL_MAX;
@@ -105,6 +105,20 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*/ */
costs.numIndexTuples = path->indexinfo->tuples * ratio; costs.numIndexTuples = path->indexinfo->tuples * ratio;
/*
* Do not use index if limit + offset > expected tuples unless
* enable_seqscan = off
*/
if (root->limit_tuples > costs.numIndexTuples)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM >= 120000
genericcostestimate(root, path, loop_count, &costs); genericcostestimate(root, path, loop_count, &costs);
#else #else

View File

@@ -80,10 +80,6 @@
#define RandomInt() random() #define RandomInt() random()
#endif #endif
#if PG_VERSION_NUM < 130000
#define list_sort(list, cmp) list_qsort(list, cmp)
#endif
/* Variables */ /* Variables */
extern int ivfflat_probes; extern int ivfflat_probes;
@@ -182,8 +178,7 @@ typedef struct IvfflatBuildState
Oid collation; Oid collation;
/* Variables */ /* Variables */
List *samples; VectorArray samples;
int targsamples;
VectorArray centers; VectorArray centers;
ListInfo *listInfo; ListInfo *listInfo;
Vector *normvec; Vector *normvec;
@@ -279,7 +274,7 @@ typedef IvfflatScanOpaqueData * IvfflatScanOpaque;
VectorArray VectorArrayInit(int maxlen, int dimensions); VectorArray VectorArrayInit(int maxlen, int dimensions);
void VectorArrayFree(VectorArray arr); void VectorArrayFree(VectorArray arr);
void PrintVectorArray(char *msg, VectorArray arr); void PrintVectorArray(char *msg, VectorArray arr);
void IvfflatKmeans(Relation index, List *samples, VectorArray centers); void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum); FmgrInfo *IvfflatOptionalProcInfo(Relation index, uint16 procnum);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result); bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
int IvfflatGetLists(Relation index); int IvfflatGetLists(Relation index);

View File

@@ -6,26 +6,30 @@
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h"
#endif
/* /*
* Initialize with kmeans++ * Initialize with kmeans++
* *
* https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf * https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf
*/ */
static void static void
InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBound) InitCenters(Relation index, VectorArray samples, VectorArray centers, float *lowerBound)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
Oid collation; Oid collation;
int64 j; int64 j;
float *weight = palloc(list_length(samples) * sizeof(float)); float *weight = palloc(samples->length * sizeof(float));
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = list_length(samples); int numSamples = samples->length;
procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC); procinfo = index_getprocinfo(index, 1, IVFFLAT_KMEANS_DISTANCE_PROC);
collation = index->rd_indcollation[0]; collation = index->rd_indcollation[0];
/* Choose an initial center uniformly at random */ /* Choose an initial center uniformly at random */
VectorArraySet(centers, 0, list_nth(samples, RandomInt() % list_length(samples))); VectorArraySet(centers, 0, VectorArrayGet(samples, RandomInt() % samples->length));
centers->length++; centers->length++;
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
@@ -42,7 +46,7 @@ InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBoun
for (j = 0; j < numSamples; j++) for (j = 0; j < numSamples; j++)
{ {
Vector *vec = list_nth(samples, j); Vector *vec = VectorArrayGet(samples, j);
double distance; double distance;
/* Only need to compute distance for new center */ /* Only need to compute distance for new center */
@@ -74,7 +78,7 @@ InitCenters(Relation index, List *samples, VectorArray centers, float *lowerBoun
break; break;
} }
VectorArraySet(centers, i + 1, list_nth(samples, j)); VectorArraySet(centers, i + 1, VectorArrayGet(samples, j));
centers->length++; centers->length++;
} }
@@ -106,41 +110,25 @@ CompareVectors(const void *a, const void *b)
return vector_cmp_internal((Vector *) a, (Vector *) b); return vector_cmp_internal((Vector *) a, (Vector *) b);
} }
/*
* Compare list vectors
*/
static int
#if PG_VERSION_NUM >= 130000
CompareListVectors(const ListCell *a, const ListCell *b)
#else
CompareListVectors(const void *a, const void *b)
#endif
{
Vector *va = lfirst((ListCell *) a);
Vector *vb = lfirst((ListCell *) b);
return CompareVectors(va, vb);
}
/* /*
* Quick approach if we have little data * Quick approach if we have little data
*/ */
static void static void
QuickCenters(Relation index, List *samples, VectorArray centers) QuickCenters(Relation index, VectorArray samples, VectorArray centers)
{ {
int dimensions = centers->dim; int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC); FmgrInfo *normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_KMEANS_NORM_PROC);
/* Copy existing vectors while avoiding duplicates */ /* Copy existing vectors while avoiding duplicates */
if (list_length(samples) > 0) if (samples->length > 0)
{ {
list_sort(samples, CompareListVectors); qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
for (int i = 0; i < list_length(samples); i++) for (int i = 0; i < samples->length; i++)
{ {
Vector *vec = list_nth(samples, i); Vector *vec = VectorArrayGet(samples, i);
if (i == 0 || CompareVectors(vec, list_nth(samples, i - 1)) != 0) if (i == 0 || CompareVectors(vec, VectorArrayGet(samples, i - 1)) != 0)
{ {
VectorArraySet(centers, centers->length, vec); VectorArraySet(centers, centers->length, vec);
centers->length++; centers->length++;
@@ -167,6 +155,23 @@ QuickCenters(Relation index, List *samples, VectorArray centers)
} }
} }
#ifdef IVFFLAT_MEMORY
/*
* Show memory usage
*/
static void
ShowMemoryUsage(Size estimatedSize)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "total memory: %zu MB",
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
#endif
elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
}
#endif
/* /*
* Use Elkan for performance. This requires distance function to satisfy triangle inequality. * Use Elkan for performance. This requires distance function to satisfy triangle inequality.
* *
@@ -176,7 +181,7 @@ QuickCenters(Relation index, List *samples, VectorArray centers)
* https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf * https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf
*/ */
static void static void
ElkanKmeans(Relation index, List *samples, VectorArray centers) ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
{ {
FmgrInfo *procinfo; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
@@ -187,7 +192,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
int64 k; int64 k;
int dimensions = centers->dim; int dimensions = centers->dim;
int numCenters = centers->maxlen; int numCenters = centers->maxlen;
int numSamples = list_length(samples); int numSamples = samples->length;
VectorArray newCenters; VectorArray newCenters;
int *centerCounts; int *centerCounts;
int *closestCenters; int *closestCenters;
@@ -198,7 +203,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
float *newcdist; float *newcdist;
/* Calculate allocation sizes */ /* Calculate allocation sizes */
Size samplesSize = 0; Size samplesSize = VECTOR_ARRAY_SIZE(samples->maxlen, samples->dim);
Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim); Size centersSize = VECTOR_ARRAY_SIZE(centers->maxlen, centers->dim);
Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions); Size newCentersSize = VECTOR_ARRAY_SIZE(numCenters, dimensions);
Size centerCountsSize = sizeof(int) * numCenters; Size centerCountsSize = sizeof(int) * numCenters;
@@ -247,6 +252,10 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
vec->dim = dimensions; vec->dim = dimensions;
} }
#ifdef IVFFLAT_MEMORY
ShowMemoryUsage(totalSize);
#endif
/* Pick initial centers */ /* Pick initial centers */
InitCenters(index, samples, centers, lowerBound); InitCenters(index, samples, centers, lowerBound);
@@ -342,7 +351,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k]) if (upperBound[j] <= halfcdist[closestCenters[j] * numCenters + k])
continue; continue;
vec = list_nth(samples, j); vec = VectorArrayGet(samples, j);
/* Step 3a */ /* Step 3a */
if (rj) if (rj)
@@ -393,7 +402,7 @@ ElkanKmeans(Relation index, List *samples, VectorArray centers)
{ {
int closestCenter; int closestCenter;
vec = list_nth(samples, j); vec = VectorArrayGet(samples, j);
closestCenter = closestCenters[j]; closestCenter = closestCenters[j];
/* Increment sum and count of closest center */ /* Increment sum and count of closest center */
@@ -530,9 +539,9 @@ CheckCenters(Relation index, VectorArray centers)
* We use spherical k-means for inner product and cosine * We use spherical k-means for inner product and cosine
*/ */
void void
IvfflatKmeans(Relation index, List *samples, VectorArray centers) IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers)
{ {
if (list_length(samples) <= centers->maxlen) if (samples->length <= centers->maxlen)
QuickCenters(index, samples, centers); QuickCenters(index, samples, centers);
else else
ElkanKmeans(index, samples, centers); ElkanKmeans(index, samples, centers);

View File

@@ -89,7 +89,7 @@ CheckDim(int dim)
} }
/* /*
* Ensure finite elements * Ensure finite element
*/ */
static inline void static inline void
CheckElement(float value) CheckElement(float value)
@@ -177,14 +177,15 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_in);
Datum Datum
vector_in(PG_FUNCTION_ARGS) vector_in(PG_FUNCTION_ARGS)
{ {
char *str = PG_GETARG_CSTRING(0); char *lit = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2); int32 typmod = PG_GETARG_INT32(2);
float x[VECTOR_MAX_DIM]; float x[VECTOR_MAX_DIM];
int dim = 0; int dim = 0;
char *pt; char *pt;
char *stringEnd; char *stringEnd;
Vector *result; Vector *result;
char *lit = pstrdup(str); char *litcopy = pstrdup(lit);
char *str = litcopy;
while (vector_isspace(*str)) while (vector_isspace(*str))
str++; str++;
@@ -268,7 +269,7 @@ vector_in(PG_FUNCTION_ARGS)
(errcode(ERRCODE_DATA_EXCEPTION), (errcode(ERRCODE_DATA_EXCEPTION),
errmsg("vector must have at least 1 dimension"))); errmsg("vector must have at least 1 dimension")));
pfree(lit); pfree(litcopy);
CheckExpectedDim(typmod, dim); CheckExpectedDim(typmod, dim);
@@ -437,17 +438,18 @@ vector_send(PG_FUNCTION_ARGS)
/* /*
* Convert vector to vector * Convert vector to vector
* This is needed to check the type modifier
*/ */
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector); PGDLLEXPORT PG_FUNCTION_INFO_V1(vector);
Datum Datum
vector(PG_FUNCTION_ARGS) vector(PG_FUNCTION_ARGS)
{ {
Vector *arg = PG_GETARG_VECTOR_P(0); Vector *vec = PG_GETARG_VECTOR_P(0);
int32 typmod = PG_GETARG_INT32(1); int32 typmod = PG_GETARG_INT32(1);
CheckExpectedDim(typmod, arg->dim); CheckExpectedDim(typmod, vec->dim);
PG_RETURN_POINTER(arg); PG_RETURN_POINTER(vec);
} }
/* /*
@@ -464,7 +466,6 @@ array_to_vector(PG_FUNCTION_ARGS)
bool typbyval; bool typbyval;
char typalign; char typalign;
Datum *elemsp; Datum *elemsp;
bool *nullsp;
int nelemsp; int nelemsp;
if (ARR_NDIM(array) > 1) if (ARR_NDIM(array) > 1)
@@ -478,7 +479,7 @@ array_to_vector(PG_FUNCTION_ARGS)
errmsg("array must not contain nulls"))); errmsg("array must not contain nulls")));
get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign); get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign);
deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, &nullsp, &nelemsp); deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, NULL, &nelemsp);
CheckDim(nelemsp); CheckDim(nelemsp);
CheckExpectedDim(typmod, nelemsp); CheckExpectedDim(typmod, nelemsp);
@@ -512,6 +513,12 @@ array_to_vector(PG_FUNCTION_ARGS)
errmsg("unsupported array type"))); errmsg("unsupported array type")));
} }
/*
* Free allocation from deconstruct_array. Do not free individual elements
* when pass-by-reference since they point to original array.
*/
pfree(elemsp);
/* Check elements */ /* Check elements */
for (int i = 0; i < result->dim; i++) for (int i = 0; i < result->dim; i++)
CheckElement(result->x[i]); CheckElement(result->x[i]);
@@ -718,7 +725,7 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
} }
/* /*
* Get the L1 distance between vectors * Get the L1 distance between two vectors
*/ */
PGDLLEXPORT PG_FUNCTION_INFO_V1(l1_distance); PGDLLEXPORT PG_FUNCTION_INFO_V1(l1_distance);
Datum Datum
@@ -897,8 +904,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_lt);
Datum Datum
vector_lt(PG_FUNCTION_ARGS) vector_lt(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) < 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) < 0);
} }
@@ -910,8 +917,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_le);
Datum Datum
vector_le(PG_FUNCTION_ARGS) vector_le(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) <= 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) <= 0);
} }
@@ -923,8 +930,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_eq);
Datum Datum
vector_eq(PG_FUNCTION_ARGS) vector_eq(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) == 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) == 0);
} }
@@ -936,8 +943,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_ne);
Datum Datum
vector_ne(PG_FUNCTION_ARGS) vector_ne(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) != 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) != 0);
} }
@@ -949,8 +956,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_ge);
Datum Datum
vector_ge(PG_FUNCTION_ARGS) vector_ge(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) >= 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) >= 0);
} }
@@ -962,8 +969,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_gt);
Datum Datum
vector_gt(PG_FUNCTION_ARGS) vector_gt(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_BOOL(vector_cmp_internal(a, b) > 0); PG_RETURN_BOOL(vector_cmp_internal(a, b) > 0);
} }
@@ -975,8 +982,8 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_cmp);
Datum Datum
vector_cmp(PG_FUNCTION_ARGS) vector_cmp(PG_FUNCTION_ARGS)
{ {
Vector *a = (Vector *) PG_GETARG_VECTOR_P(0); Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = (Vector *) PG_GETARG_VECTOR_P(1); Vector *b = PG_GETARG_VECTOR_P(1);
PG_RETURN_INT32(vector_cmp_internal(a, b)); PG_RETURN_INT32(vector_cmp_internal(a, b));
} }

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_cosine_ops); CREATE INDEX ON t USING hnsw (val vector_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]'; SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
val val
--------- ---------
[1,1,1] [1,1,1]
@@ -11,13 +11,13 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
[1,2,4] [1,2,4]
(3 rows) (3 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
count count
------- -------
3 3
(1 row) (1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
count count
------- -------
3 3

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_ip_ops); CREATE INDEX ON t USING hnsw (val vector_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]'; SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
val val
--------- ---------
[1,2,4] [1,2,4]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <#> '[3,3,3]';
[0,0,0] [0,0,0]
(4 rows) (4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
count count
------- -------
4 4

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops); CREATE INDEX ON t USING hnsw (val vector_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]'; SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
val val
--------- ---------
[1,2,3] [1,2,3]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]';
[0,0,0] [0,0,0]
(4 rows) (4 rows)
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector); SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
val val
--------- ---------
[0,0,0] [0,0,0]
@@ -28,7 +28,7 @@ SELECT COUNT(*) FROM t;
(1 row) (1 row)
TRUNCATE t; TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]'; SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
val val
----- -----
(0 rows) (0 rows)

View File

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

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_cosine_ops) WITH (lists = 1); CREATE INDEX ON t USING ivfflat (val vector_cosine_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]'; SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
val val
--------- ---------
[1,1,1] [1,1,1]
@@ -11,13 +11,13 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]';
[1,2,4] [1,2,4]
(3 rows) (3 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
count count
------- -------
3 3
(1 row) (1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
count count
------- -------
3 3

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 1); CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]'; SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
val val
--------- ---------
[1,2,4] [1,2,4]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <#> '[3,3,3]';
[0,0,0] [0,0,0]
(4 rows) (4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
count count
------- -------
4 4

View File

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

View File

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

View File

@@ -6,8 +6,8 @@ CREATE INDEX ON t USING hnsw (val vector_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]'; SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
DROP TABLE t; DROP TABLE t;

View File

@@ -6,7 +6,7 @@ CREATE INDEX ON t USING hnsw (val vector_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]'; SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
DROP TABLE t; DROP TABLE t;

View File

@@ -6,11 +6,11 @@ CREATE INDEX ON t USING hnsw (val vector_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]'; SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector); SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
SELECT COUNT(*) FROM t; SELECT COUNT(*) FROM t;
TRUNCATE t; TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]'; SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
DROP TABLE t; DROP TABLE t;

View File

@@ -4,6 +4,6 @@ CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL); INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops); CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SELECT * FROM t ORDER BY val <-> '[3,3,3]'; SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
DROP TABLE t; DROP TABLE t;

View File

@@ -6,8 +6,8 @@ CREATE INDEX ON t USING ivfflat (val vector_cosine_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]'; SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
DROP TABLE t; DROP TABLE t;

View File

@@ -6,7 +6,7 @@ CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]'); INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]'; SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2; SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
DROP TABLE t; DROP TABLE t;

View File

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

View File

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

View File

@@ -49,7 +49,7 @@ is(idx_scan(), 0);
$count = $node->safe_psql("postgres", qq( $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET ivfflat.probes = 100; SET ivfflat.probes = 100;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t; SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 20000) t;
)); ));
is($count, $expected); is($count, $expected);
is(idx_scan(), 1); is(idx_scan(), 1);

View File

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

View File

@@ -42,7 +42,7 @@ for my $i (1 .. 20)
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t; SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 20) t;
)); ));
is($count, 10); is($count, 10);
@@ -63,7 +63,7 @@ $node->pgbench(
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = 1000; SET hnsw.ef_search = 1000;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t; SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 1000) t;
)); ));
# Elements may lose all incoming connections with the HNSW algorithm # Elements may lose all incoming connections with the HNSW algorithm
# Vacuuming can fix this if one of the elements neighbors is deleted # Vacuuming can fix this if one of the elements neighbors is deleted

View File

@@ -26,7 +26,7 @@ sub test_duplicates
my $res = $node->safe_psql("postgres", qq( my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = 1; SET hnsw.ef_search = 1;
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t; SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]' LIMIT 20) t;
)); ));
is($res, 10); is($res, 10);
} }

View File

@@ -0,0 +1,64 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 1000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 10);");
# Test limit
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 100;
));
like($explain, qr/Index Scan/);
# Test limit with probes
$explain = $node->safe_psql("postgres", qq(
SET ivfflat.probes = 2;
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 200;
));
like($explain, qr/Index Scan/);
# Test limit + offset
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 90 OFFSET 10;
));
like($explain, qr/Index Scan/);
# Test limit > expected tuples
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 101;
));
like($explain, qr/Seq Scan/);
# Test limit > expected tuples with probes
$explain = $node->safe_psql("postgres", qq(
SET ivfflat.probes = 2;
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 201;
));
like($explain, qr/Seq Scan/);
# Test limit + offset > expected tuples
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 91 OFFSET 10;
));
like($explain, qr/Seq Scan/);
# Test no limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]';
));
like($explain, qr/Seq Scan/);
done_testing();

62
test/t/020_hnsw_limit.pl Normal file
View File

@@ -0,0 +1,62 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 1000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
# Test limit
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 40;
));
like($explain, qr/Index Scan/);
# Test limit with CTE
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE WITH cte AS (SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 40) SELECT * FROM cte;
));
like($explain, qr/Index Scan/);
# Test limit + offset
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 30 OFFSET 10;
));
like($explain, qr/Index Scan/);
# Test limit > ef_search
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 41;
));
like($explain, qr/Seq Scan/);
# Test limit > ef_search with CTE
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE WITH cte AS (SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 41) SELECT * FROM cte;
));
like($explain, qr/Seq Scan/);
# Test limit + offset > ef_search
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 31 OFFSET 10;
));
like($explain, qr/Seq Scan/);
# Test no limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]';
));
like($explain, qr/Seq Scan/);
done_testing();