Compare commits

...

188 Commits

Author SHA1 Message Date
Andrew Kane
3d866844d3 Added PrintGraph function for HNSW [skip ci] 2023-08-26 13:08:27 -07:00
Andrew Kane
bb1e5ed98f Improved code [skip ci] 2023-08-25 21:05:18 -07:00
Andrew Kane
8801832946 Fixed check in SelectNeighbors 2023-08-25 20:46:05 -07:00
Andrew Kane
552c64d492 Improved HNSW code 2023-08-25 20:39:07 -07:00
Andrew Kane
86c29b3bf0 Improved param code 2023-08-23 21:10:00 -07:00
Andrew Kane
e406b7f5ea Added comments [skip ci] 2023-08-23 21:03:07 -07:00
Andrew Kane
6d88a9e1d2 Updated HNSW_PAGE_ID [skip ci] 2023-08-21 22:59:53 -07:00
Andrew Kane
1e851c12c0 Updated comment [skip ci] 2023-08-21 22:55:24 -07:00
Andrew Kane
8ed3cc5f0b Improved macro [skip ci] 2023-08-21 22:52:58 -07:00
Andrew Kane
bace0891bd Updated comments [skip ci] 2023-08-21 22:51:24 -07:00
Andrew Kane
4600979504 Updated comments [skip ci] 2023-08-21 22:38:09 -07:00
Andrew Kane
69addf05d0 Updated comments [skip ci] 2023-08-21 22:21:53 -07:00
Andrew Kane
4a5ae8a8df Updated comment [skip ci] 2023-08-21 21:59:47 -07:00
Andrew Kane
ce7da66ca2 Updated comment [skip ci] 2023-08-21 20:57:54 -07:00
Andrew Kane
c6be1c5e13 Fixed test 2023-08-21 20:45:56 -07:00
Andrew Kane
9b3d1a32ff Updated comment [skip ci] 2023-08-21 16:27:03 -07:00
Andrew Kane
8420734350 Updated comments [skip ci] 2023-08-21 16:26:32 -07:00
Andrew Kane
8b03267267 Improved locking for HNSW vacuum [skip ci] 2023-08-21 16:24:55 -07:00
Andrew Kane
641ddf5413 Improved locking for HNSW vacuum 2023-08-21 16:12:29 -07:00
Andrew Kane
782a1051e3 Improved locking for HNSW vacuum 2023-08-21 16:06:32 -07:00
Andrew Kane
ca2be5be6e Updated comment [skip ci] 2023-08-21 16:00:44 -07:00
Andrew Kane
0e1de45463 Improved locking code [skip ci] 2023-08-21 15:42:59 -07:00
Andrew Kane
3f3b3ca8e3 Made function static [skip ci] 2023-08-21 03:27:42 -07:00
Andrew Kane
d4fe67e8ee Simplified locking for entry point 2023-08-21 03:22:23 -07:00
Andrew Kane
02f4e0ec8b Revert "Added version to reduce stale reads and writes and prepare for optimistic locking"
This reverts commit ef1209eaf4.
2023-08-21 02:47:27 -07:00
Andrew Kane
1301706d30 Increased concurrency in HNSW insert test 2023-08-21 02:46:16 -07:00
Andrew Kane
0d58683011 Improved HNSW insert test 2023-08-21 02:44:20 -07:00
Andrew Kane
90a042e5aa Wait for selects to complete 2023-08-21 02:24:53 -07:00
Andrew Kane
bbd57dfebf Moved wait [skip ci] 2023-08-21 01:02:05 -07:00
Andrew Kane
40a013a622 Wait for inserts to complete when vacuuming HNSW 2023-08-21 00:36:32 -07:00
Andrew Kane
ef1209eaf4 Added version to reduce stale reads and writes and prepare for optimistic locking 2023-08-20 17:08:20 -07:00
Andrew Kane
687263ccd4 DRY HNSW vacuum code 2023-08-20 14:52:31 -07:00
Andrew Kane
a62c045c93 Update metapage if needed for HNSW vacuum 2023-08-19 12:19:42 -07:00
Andrew Kane
651e4343c1 Made entryPoint argument for RepairGraphElement 2023-08-19 11:47:13 -07:00
Andrew Kane
2ca2ef94e6 Use rmdir [skip ci] 2023-08-19 10:15:50 -07:00
Andrew Kane
6ba95691d3 Added headers for Windows 2023-08-19 10:03:15 -07:00
Andrew Kane
e9de248e4f Skip i386 tests for windows branches [skip ci] 2023-08-19 10:01:55 -07:00
Andrew Kane
fe79d435c0 Use HEADERS for simplicity [skip ci] 2023-08-19 09:58:25 -07:00
Florents Tselai
b34525fbc2 Make pgvector headers available to others. (#233) 2023-08-19 09:56:02 -07:00
Andrew Kane
206c42e170 Cleaned up header 2023-08-19 01:22:28 -07:00
Andrew Kane
4f8d824280 Use variable for uninstall [skip ci] 2023-08-19 01:21:05 -07:00
Andrew Kane
ca847d02b0 Fixed highest point in HNSW vacuum [skip ci] 2023-08-18 22:23:19 -07:00
Andrew Kane
ca726052ae Fixed flaky test [skip ci] 2023-08-18 20:54:31 -07:00
Andrew Kane
2b25127f3d Fixed flaky test [skip ci] 2023-08-18 20:53:47 -07:00
Andrew Kane
fecb8c06c9 Simplified code 2023-08-18 20:49:00 -07:00
Andrew Kane
a03f6ae4bc Added prefix to function name [skip ci] 2023-08-18 00:54:09 -07:00
Andrew Kane
ed20d98777 Check if item pointer is valid [skip ci] 2023-08-16 17:29:59 -07:00
Andrew Kane
b72a22b3c0 Fixed duplicate connections when repairing graph 2023-08-16 17:07:19 -07:00
Andrew Kane
8c5c0f791e Improved HNSW insert code [skip ci] 2023-08-16 16:31:11 -07:00
Andrew Kane
f9f5ae61df Added test for select with no data for ivfflat 2023-08-15 23:44:18 -07:00
Andrew Kane
2b138d6cb5 Fixed select with no data 2023-08-15 23:25:28 -07:00
Andrew Kane
e8f36aee92 Improved HNSW vacuum code [skip ci] 2023-08-15 22:36:47 -07:00
Andrew Kane
508a8a9ac7 Updated comments [skip ci] 2023-08-15 20:32:15 -07:00
Andrew Kane
d0b0298cda Improved HNSW vacuum code [skip ci] 2023-08-15 20:02:21 -07:00
Andrew Kane
c3bafc76e8 Only update entry point on inserts if level is greater 2023-08-15 12:31:06 -07:00
Andrew Kane
86a062e504 Set entry level in CreateMetaPage [skip ci] 2023-08-15 11:47:05 -07:00
Andrew Kane
b421e76f29 Updated comment [skip ci] 2023-08-13 17:23:52 -07:00
Andrew Kane
ef480ff0b9 Added comment [skip ci] 2023-08-13 17:22:25 -07:00
Andrew Kane
e1d6654063 Revert "Improved HNSW vacuum performance"
This reverts commit c5b2f3ac8b.
2023-08-13 17:19:22 -07:00
Andrew Kane
c5b2f3ac8b Improved HNSW vacuum performance 2023-08-13 16:50:36 -07:00
Andrew Kane
0f238b1fa9 Update neighbors when vacuuming 2023-08-13 16:17:30 -07:00
Andrew Kane
6fc7d06313 Improved HNSW vacuuming 2023-08-13 15:53:07 -07:00
Andrew Kane
304e262a41 Fixed initialization 2023-08-13 15:50:05 -07:00
Andrew Kane
23cb79c1dc Fixed recall for HNSW after vacuuming 2023-08-13 15:37:47 -07:00
Andrew Kane
dc8d467ba5 Improved HNSW vacuum recall test 2023-08-13 13:46:06 -07:00
Andrew Kane
29a0a00731 Improved test [skip ci] 2023-08-13 13:42:51 -07:00
Andrew Kane
eda07ccf7c Improved test 2023-08-13 13:35:43 -07:00
Andrew Kane
c694dbb776 Added test for HNSW vacuum recall 2023-08-13 13:09:51 -07:00
Andrew Kane
3a857ef9f9 Updated readme [skip ci] 2023-08-13 12:16:36 -07:00
Andrew Kane
b60e2483a9 Improved HNSW vacuum stats [skip ci] 2023-08-12 14:09:13 -07:00
Andrew Kane
4295ee3b3a Updated todos [skip ci] 2023-08-12 13:47:20 -07:00
Andrew Kane
8b1ea42d42 Improved HNSW vacuum test 2023-08-12 12:41:27 -07:00
Andrew Kane
3c291fea41 Added todo [skip ci] 2023-08-12 12:41:02 -07:00
Andrew Kane
3d568ef6a1 Added comments [skip ci] 2023-08-11 21:12:18 -07:00
Andrew Kane
7db3b08ef4 Simplified tests [skip ci] 2023-08-11 19:41:55 -07:00
Andrew Kane
4196c88dbd Added opclass to tests [skip ci] 2023-08-11 19:40:04 -07:00
Andrew Kane
3e7ee6ea04 Simplified opclasses in TAP tests [skip ci] 2023-08-11 19:20:03 -07:00
Andrew Kane
5616c1d7a7 Simplified TAP tests [skip ci] 2023-08-11 19:16:08 -07:00
Andrew Kane
0778a507de Use consistent style in TAP tests, part 2 [skip ci] 2023-08-11 19:13:22 -07:00
Andrew Kane
aabbeac30c Use consistent style in TAP tests [skip ci] 2023-08-11 19:08:15 -07:00
Andrew Kane
746dce72ef Use concurrent inserts for insert recall test 2023-08-11 18:41:00 -07:00
Andrew Kane
098813c483 Improved concurrent inserts test for HNSW [skip ci] 2023-08-11 18:33:01 -07:00
Andrew Kane
9d689afd7e Updated test name [skip ci] 2023-08-11 17:17:28 -07:00
Andrew Kane
03930e15a4 Updated comment [skip ci] 2023-08-11 17:09:09 -07:00
Andrew Kane
8a807c01f6 Test fallback path for duplicates [skip ci] 2023-08-11 17:07:38 -07:00
Andrew Kane
85ebb7714e Updated test name [skip ci] 2023-08-11 17:07:17 -07:00
Andrew Kane
34e99593ab Added comment [skip ci] 2023-08-11 17:02:50 -07:00
Andrew Kane
fc45231ef0 Improving naming [skip ci] 2023-08-11 14:56:46 -07:00
Andrew Kane
0191a90f9f Updated comment [skip ci] 2023-08-11 14:37:56 -07:00
Andrew Kane
82e83bd3b1 Added comment [skip ci] 2023-08-11 14:29:24 -07:00
Andrew Kane
d8a14c658a Fixed flaky test [skip ci] 2023-08-11 13:38:42 -07:00
Andrew Kane
ae5de26893 Improved logic for updating HNSW insert page 2023-08-11 13:24:56 -07:00
Andrew Kane
daef83a112 Added comment [skip ci] 2023-08-11 13:03:38 -07:00
Andrew Kane
c9d82c6bdc Revert "Improved insert logic [skip ci]"
This reverts commit 573a336f53.
2023-08-11 12:52:58 -07:00
Andrew Kane
573a336f53 Improved insert logic [skip ci] 2023-08-11 12:39:22 -07:00
Andrew Kane
e7a913f361 Updated comment [skip ci] 2023-08-11 10:11:01 -07:00
Andrew Kane
33438c3cf9 Use neighbor page for insert page [skip ci] 2023-08-11 09:51:41 -07:00
Andrew Kane
8ec1821f1e Added comment [skip ci] 2023-08-11 09:45:08 -07:00
Andrew Kane
b288de719d Consider previously deleted tuples for insert page [skip ci] 2023-08-11 09:41:31 -07:00
Andrew Kane
c13f713f29 Improved test [skip ci] 2023-08-11 09:36:29 -07:00
Andrew Kane
6d15435003 Updated comment [skip ci] 2023-08-11 09:26:54 -07:00
Andrew Kane
12d15a9c41 Improved inserts for HNSW 2023-08-11 09:23:56 -07:00
Andrew Kane
c64288377b Updated min ef_construction to 4 [skip ci] 2023-08-10 21:11:10 -07:00
Andrew Kane
451e36cee7 Added check for ef_construction [skip ci] 2023-08-10 21:09:38 -07:00
Andrew Kane
3ff2e34d7f Updated min M to 2 [skip ci] 2023-08-10 20:57:50 -07:00
Andrew Kane
94a684c9e9 Added test for duplicates 2023-08-10 20:52:58 -07:00
Andrew Kane
1a0d7bccc7 Updated min ef_search to 1 [skip ci] 2023-08-10 20:47:15 -07:00
Andrew Kane
ded8bda72a Fixed flaky test [skip ci] 2023-08-10 18:59:11 -07:00
Andrew Kane
9b967d592f Improved concurrent inserts with empty entry point, part 2 2023-08-10 18:31:58 -07:00
Andrew Kane
ed513e62c1 Improved code for skipping element [skip ci] 2023-08-10 15:29:08 -07:00
Andrew Kane
27ccb5fa7a Improved code [skip ci] 2023-08-10 15:23:26 -07:00
Andrew Kane
92e25e7733 Fixed warning 2023-08-10 15:10:07 -07:00
Andrew Kane
a18bea24b8 Improved concurrent inserts with low number of elements 2023-08-10 15:07:28 -07:00
Andrew Kane
483173460b Improved concurrent inserts with empty entry point 2023-08-10 14:35:24 -07:00
Andrew Kane
da8a914106 Moved FindDuplicate [skip ci] 2023-08-10 13:44:21 -07:00
Andrew Kane
4d6da72b08 Simplified code 2023-08-10 09:56:17 -07:00
Andrew Kane
443c7a5dab Updated variable name [skip ci] 2023-08-10 09:02:01 -07:00
Andrew Kane
9287fe8bcc Updated comment [skip ci] 2023-08-10 08:59:38 -07:00
Andrew Kane
309fa94c05 Improved performance for duplicates 2023-08-10 08:52:45 -07:00
Andrew Kane
6f15dd266c Improved construction code 2023-08-10 08:38:31 -07:00
Andrew Kane
df68eb4570 Look for duplicates last, since may need to fallback 2023-08-10 00:00:27 -07:00
Andrew Kane
7c0d94c99c Improved concurrent inserts 2023-08-09 23:51:35 -07:00
Andrew Kane
d63d430af8 Fixed flaky test [skip ci] 2023-08-09 20:06:28 -07:00
Andrew Kane
5d62e4d080 Moved SelectNeighbors after duplicate check 2023-08-09 19:19:41 -07:00
Andrew Kane
dab8f25d1c Fixed overflow with vector_norm 2023-08-09 16:33:54 -07:00
Andrew Kane
4b887a98ae Moved define [skip ci] 2023-08-09 10:21:01 -07:00
Andrew Kane
d253bafee6 Added HNSW paper to thanks [skip ci] 2023-08-08 23:34:57 -07:00
Andrew Kane
600ca5a797 Improved logic for pruning elements 2023-08-08 18:22:55 -07:00
Andrew Kane
c17d51588a Removed distance from neighbor tuples 2023-08-08 18:11:11 -07:00
Andrew Kane
51d292c93d Added HNSW index type - #181 2023-08-08 16:42:47 -07:00
Andrew Kane
19a6c81367 Fixed link [skip ci] 2023-08-08 01:48:27 -07:00
Andrew Kane
6f212d7cc1 Improved Windows instructions, part 2 [skip ci] 2023-08-07 14:01:24 -07:00
Andrew Kane
25ecabf1ea Improved Windows instructions - closes #218 [skip ci] 2023-08-07 13:56:00 -07:00
Andrew Kane
03457f41ba Added link to pgvector-dart - closes #215 [skip ci] 2023-08-06 22:37:54 -07:00
Andrew Kane
e45c84fd46 Added space before comment [skip ci] 2023-08-05 10:20:30 -07:00
Andrew Kane
37b49b0a37 Fixed results for NULL and NaN distances - fixes #205
Co-authored-by: Xiaoran Wang <wxiaoran@vmware.com>
2023-08-05 09:36:58 -07:00
Andrew Kane
4df01af8b6 Fixed Docker build instructions - #197 [skip ci] 2023-07-27 09:34:11 -07:00
Andrew Kane
237a6df51f Fixed Docker build - fixes #197 2023-07-27 09:27:47 -07:00
Andrew Kane
bcc1366d86 Added tests for large distances 2023-07-25 16:49:21 -07:00
Andrew Kane
047877f495 Improved test 2023-07-25 16:04:32 -07:00
Andrew Kane
e36cf20bce Fixed CI 2023-07-25 15:44:04 -07:00
Andrew Kane
47e5a86b63 Fixed out of range results for cosine distance - fixes #196 2023-07-25 14:54:13 -07:00
Andrew Kane
f210791846 Updated readme [skip ci] 2023-07-19 22:21:12 -07:00
Andrew Kane
280ec74cbd Updated readme [skip ci] 2023-07-19 16:37:32 -07:00
Andrew Kane
67e6b9ee6d Updated readme [skip ci] 2023-07-19 15:58:20 -07:00
Andrew Kane
0160f8f2f1 Added troubleshooting section - closes #186 [skip ci] 2023-07-19 15:48:34 -07:00
Andrew Kane
b710dc68a0 Use fused multiply-add for cosine distance 2023-07-18 16:10:17 -07:00
Andrew Kane
8a05debda1 Ran pgindent [skip ci] 2023-07-18 16:08:50 -07:00
Andrew Kane
1a81b21029 Updated changelog [skip ci] 2023-07-18 16:03:42 -07:00
Andrew Kane
cf1f151cff Now available on DigitalOcean [skip ci] 2023-07-18 13:20:54 -07:00
Pavel Borisov
3950bc3dc6 Speed up ivfflat build: use float instead of double for dot product (#180)
calculation

On ARM this makes CPU using vector multiply-add instruction (fmadd)
instead of vector multiplication + conversion to double + addition
(fmul + fcvt + fadd) at each vector dimension.

Output of distance functions and calculations that are are done once
per vector pair are left double as this don't make speed difference
and for compatibility.
2023-07-18 12:55:38 -07:00
Andrew Kane
f4c28b1c06 Improved includes 2023-07-17 13:17:48 -07:00
Andrew Kane
d13eb8563e Improved includes 2023-07-17 13:15:10 -07:00
Andrew Kane
f6fc033622 Removed unneeded comments [skip ci] 2023-07-17 13:08:27 -07:00
Andrew Kane
518a35784d Improved variable name [skip ci] 2023-07-17 00:18:14 -07:00
Andrew Kane
c3394ace03 Removed unused variable from IvfflatUpdateList 2023-07-17 00:15:20 -07:00
Andrew Kane
98ba795d83 Improved variable scoping 2023-07-17 00:11:21 -07:00
Andrew Kane
6e8746277e Moved _PG_init 2023-07-16 20:03:34 -07:00
Andrew Kane
ad8df27fbb Changed sorting tuples indexing phase to assigning tuples [skip ci] 2023-07-16 18:46:45 -07:00
Andrew Kane
b77451f19e Renamed tests 2023-07-16 16:07:33 -07:00
Andrew Kane
f1d3aa2ba8 Updated comment [skip ci] 2023-07-16 16:02:47 -07:00
Andrew Kane
a722df9cac Test sum aggregate 2023-07-16 15:49:24 -07:00
Andrew Kane
a0c7f618ff Simplified sum aggregate 2023-07-16 15:30:03 -07:00
Andrew Kane
f9d9d64a4e Updated header [skip ci] 2023-07-16 00:55:42 -07:00
Andrew Kane
1fd9794d19 Updated readme [skip ci] 2023-07-15 23:15:49 -07:00
Andrew Kane
13ecd271ba Removed unneeded lists check [skip ci] 2023-07-15 23:06:34 -07:00
Andrew Kane
fd3c3ce83d Added test for avg with large values 2023-07-15 23:03:15 -07:00
Andrew Kane
5796f38ad2 Added test for 2-d array 2023-07-15 22:58:15 -07:00
Andrew Kane
1d47d7218d Improved array_to_vector [skip ci] 2023-07-15 22:56:29 -07:00
Andrew Kane
42007b41ea Improved null check and fixed message [skip ci] 2023-07-15 22:51:40 -07:00
Andrew Kane
21f56374b5 Added test for empty string [skip ci] 2023-07-15 22:45:19 -07:00
Andrew Kane
890f189495 Removed todo [skip ci] 2023-07-15 22:37:52 -07:00
Andrew Kane
ed1bc3e856 Added sum aggregate 2023-07-15 20:25:30 -07:00
Andrew Kane
08e7209810 Added element-wise multiplication for vectors 2023-07-15 20:19:51 -07:00
Andrew Kane
cd4ac17f9f Added l1_distance function - #166 2023-07-15 20:13:45 -07:00
Andrew Kane
b6a822918f Added support for parallel index builds 2023-07-15 19:52:25 -07:00
Andrew Kane
21d5d7e934 Improved variable scoping 2023-07-15 19:01:40 -07:00
Andrew Kane
ea47342870 Improved variable scoping 2023-07-15 18:53:41 -07:00
Andrew Kane
22e9be6528 Use LockRelationForExtension 2023-07-15 18:25:26 -07:00
Andrew Kane
b56971febe Updated link [skip ci] 2023-07-11 21:47:08 -07:00
Jonathan S. Katz
de6502ab6e Simplify TAP test structure to align with upstream (#169)
postgres/postgres@549ec20 moved to using "done_testing()" from
the Perl testing framework, which removed the need to include test
counts.
2023-06-29 11:36:08 -04:00
Andrew Kane
7aaba14440 Added comment [skip ci] 2023-06-29 10:25:06 -04:00
Andrew Kane
ee0bf10d7d Fixed CI for equal distances 2023-06-29 10:24:02 -04:00
Andrew Kane
88094fc39f Fixed tests 2023-06-29 10:17:15 -04:00
Andrew Kane
209394faab Now available on Google Cloud SQL [skip ci] 2023-06-26 11:17:10 -04:00
Andrew Kane
673aca97fc Fixed CI for branches for i386 2023-06-16 17:38:51 -04:00
Andrew Kane
0e2dc0e6d4 Updated readme [skip ci] 2023-06-16 10:15:22 -04:00
Andrew Kane
3c681b94fd Updated example [skip ci] 2023-06-16 10:03:07 -04:00
68 changed files with 5665 additions and 355 deletions

View File

@@ -78,6 +78,7 @@ jobs:
nmake /NOLOGO /F Makefile.win uninstall
shell: cmd
i386:
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
runs-on: ubuntu-latest
container:
image: debian:11
@@ -88,7 +89,8 @@ jobs:
- run: |
git clone https://github.com/${{ github.repository }}.git pgvector
cd pgvector
git checkout ${{ github.ref }}
git fetch origin ${{ github.ref }}
git reset --hard FETCH_HEAD
make
make install
chown -R postgres .

View File

@@ -1,3 +1,14 @@
## 0.5.0 (unreleased)
- Added HNSW index type
- Added support for parallel index builds
- Added `l1_distance` function
- Added element-wise multiplication for vectors
- Added `sum` aggregate
- Improved performance of distance functions
- Fixed out of range results for cosine distance
- Fixed results for NULL and NaN distances
## 0.4.4 (2023-06-12)
- Improved error message for malformed vector literal

View File

@@ -5,6 +5,7 @@ ARG PG_MAJOR
COPY . /tmp/pgvector
RUN apt-get update && \
apt-mark hold locales && \
apt-get install -y --no-install-recommends build-essential postgresql-server-dev-$PG_MAJOR && \
cd /tmp/pgvector && \
make clean && \
@@ -15,4 +16,5 @@ RUN apt-get update && \
rm -r /tmp/pgvector && \
apt-get remove -y build-essential postgresql-server-dev-$PG_MAJOR && \
apt-get autoremove -y && \
apt-mark unhold locales && \
rm -rf /var/lib/apt/lists/*

View File

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

View File

@@ -1,7 +1,8 @@
EXTENSION = vector
EXTVERSION = 0.4.4
OBJS = src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
OBJS = src\hnsw.obj src\hnswbuild.obj src\hnswinsert.obj src\hnswscan.obj src\hnswutils.obj src\hnswvacuum.obj src\ivfbuild.obj src\ivfflat.obj src\ivfinsert.obj src\ivfkmeans.obj src\ivfscan.obj src\ivfutils.obj src\ivfvacuum.obj src\vector.obj
HEADERS = src\vector.h
REGRESS = btree cast copy functions input ivfflat_cosine ivfflat_ip ivfflat_l2 ivfflat_options ivfflat_unlogged
REGRESS_OPTS = --inputdir=test --load-extension=vector
@@ -54,6 +55,8 @@ install:
copy $(SHLIB) "$(PKGLIBDIR)"
copy $(EXTENSION).control "$(SHAREDIR)\extension"
copy sql\$(EXTENSION)--*.sql "$(SHAREDIR)\extension"
mkdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
copy $(HEADERS) "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
installcheck:
"$(BINDIR)\pg_regress" --bindir="$(BINDIR)" $(REGRESS_OPTS) $(REGRESS)
@@ -61,7 +64,9 @@ installcheck:
uninstall:
del /f "$(PKGLIBDIR)\$(SHLIB)"
del /f "$(SHAREDIR)\extension\$(EXTENSION).control"
del /f "$(SHAREDIR)\extension\vector--*.sql"
del /f "$(SHAREDIR)\extension\$(EXTENSION)--*.sql"
del /f "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)\*.h"
rmdir "$(INCLUDEDIR_SERVER)\extension\$(EXTENSION)"
clean:
del /f $(SHLIB) $(EXTENSION).lib $(EXTENSION).exp

View File

@@ -2,7 +2,7 @@
Open-source vector similarity search for Postgres
Supports
Store your vectors with the rest of your data. Supports:
- exact and approximate nearest neighbor search
- L2 distance, inner product, and cosine distance
@@ -223,7 +223,7 @@ The phases are:
Note: `tuples_done` and `tuples_total` are only populated during the `loading tuples` phase
### Filtering
## Filtering
There are a few ways to index nearest neighbor queries with a `WHERE` clause
@@ -255,7 +255,7 @@ CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(cate
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)).
```sql
SELECT id, content FROM items, to_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;
```
@@ -298,6 +298,7 @@ Language | Libraries / Examples
C++ | [pgvector-cpp](https://github.com/pgvector/pgvector-cpp)
C# | [pgvector-dotnet](https://github.com/pgvector/pgvector-dotnet)
Crystal | [pgvector-crystal](https://github.com/pgvector/pgvector-crystal)
Dart | [pgvector-dart](https://github.com/pgvector/pgvector-dart)
Elixir | [pgvector-elixir](https://github.com/pgvector/pgvector-elixir)
Go | [pgvector-go](https://github.com/pgvector/pgvector-go)
Haskell | [pgvector-haskell](https://github.com/pgvector/pgvector-haskell)
@@ -327,10 +328,45 @@ 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.
#### Why am I seeing less results after adding an index?
## Troubleshooting
#### Why isnt a query using an index?
The cost estimation in pgvector < 0.4.3 does not always work well with the planner. You can encourage the planner to use an index for a query with:
```sql
BEGIN;
SET LOCAL enable_seqscan = off;
SELECT ...
COMMIT;
```
#### 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:
```sql
BEGIN;
SET LOCAL min_parallel_table_scan_size = 1;
SET LOCAL parallel_setup_cost = 1;
SELECT ...
COMMIT;
```
or choose to store vectors inline:
```sql
ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN;
```
#### Why are there less results for a query after adding an index?
The index was likely created with too little data for the number of lists. Drop the index until the table has more data.
```sql
DROP INDEX index_name;
```
## Reference
### Vector Type
@@ -343,6 +379,7 @@ Operator | Description
--- | ---
\+ | element-wise addition
\- | element-wise subtraction
\* | element-wise multiplication [unreleased]
<-> | Euclidean distance
<#> | negative inner product
<=> | cosine distance
@@ -354,6 +391,7 @@ Function | Description
cosine_distance(vector, vector) → double precision | cosine distance
inner_product(vector, vector) → double precision | inner product
l2_distance(vector, vector) → double precision | Euclidean distance
l1_distance(vector, vector) → double precision | taxicab distance [unreleased]
vector_dims(vector) → integer | number of dimensions
vector_norm(vector) → double precision | Euclidean norm
@@ -362,6 +400,7 @@ vector_norm(vector) → double precision | Euclidean norm
Function | Description
--- | ---
avg(vector) → vector | arithmetic mean
sum(vector) → vector | sum [unreleased]
## Installation Notes
@@ -393,7 +432,15 @@ Note: Replace `15` with your Postgres server version
### Windows
Support for Windows is currently experimental. Use `nmake` to build:
Support for Windows is currently experimental. Ensure [C++ support in Visual Studio](https://learn.microsoft.com/en-us/cpp/build/building-on-the-command-line?view=msvc-170#download-and-install-the-tools) is installed, and run:
```cmd
call "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat"
```
Note: The exact path will vary depending on your Visual Studio version and edition
Then use `nmake` to build:
```cmd
set "PGROOT=C:\Program Files\PostgreSQL\15"
@@ -420,6 +467,7 @@ You can also build the image manually:
```sh
git clone --branch v0.4.4 https://github.com/pgvector/pgvector.git
cd pgvector
git cherry-pick 237a6df
docker build --build-arg PG_MAJOR=15 -t myuser/pgvector .
```
@@ -481,12 +529,6 @@ Download the [latest release](https://postgresapp.com/downloads.html) with Postg
pgvector is available on [these providers](https://github.com/pgvector/pgvector/issues/54).
To request a new extension on other providers:
- Google Cloud SQL - vote or comment on [this page](https://issuetracker.google.com/issues/265172065)
- DigitalOcean Managed Databases - vote or comment on [this page](https://ideas.digitalocean.com/managed-database/p/pgvector-extension-for-postgresql)
- Heroku Postgres - vote or comment on [this page](https://github.com/heroku/roadmap/issues/156)
## Upgrading
Install the latest version and run:
@@ -527,9 +569,10 @@ Thanks to:
- [PASE: PostgreSQL Ultra-High-Dimensional Approximate Nearest Neighbor Search Extension](https://dl.acm.org/doi/pdf/10.1145/3318464.3386131)
- [Faiss: A Library for Efficient Similarity Search and Clustering of Dense Vectors](https://github.com/facebookresearch/faiss)
- [Using the Triangle Inequality to Accelerate k-means](https://www.aaai.org/Papers/ICML/2003/ICML03-022.pdf)
- [Using the Triangle Inequality to Accelerate k-means](https://cdn.aaai.org/ICML/2003/ICML03-022.pdf)
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
- [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
## History
@@ -577,4 +620,4 @@ Resources for contributors
- [Extension Building Infrastructure](https://www.postgresql.org/docs/current/extend-pgxs.html)
- [Index Access Method Interface Definition](https://www.postgresql.org/docs/current/indexam.html)
- [Generic WAL Records](https://www.postgresql.org/docs/13/generic-wal.html)
- [Generic WAL Records](https://www.postgresql.org/docs/current/generic-wal.html)

View File

@@ -0,0 +1,43 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.5.0'" to load this file. \quit
CREATE FUNCTION l1_distance(vector, vector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR * (
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_mul,
COMMUTATOR = *
);
CREATE AGGREGATE sum(vector) (
SFUNC = vector_add,
STYPE = vector,
COMBINEFUNC = vector_add,
PARALLEL = SAFE
);
CREATE FUNCTION hnswhandler(internal) RETURNS index_am_handler
AS 'MODULE_PATHNAME' LANGUAGE C;
CREATE ACCESS METHOD hnsw TYPE INDEX HANDLER hnswhandler;
COMMENT ON ACCESS METHOD hnsw IS 'hnsw index access method';
CREATE OPERATOR CLASS vector_l2_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_l2_squared_distance(vector, vector);
CREATE OPERATOR CLASS vector_ip_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector);
CREATE OPERATOR CLASS vector_cosine_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector);

View File

@@ -40,6 +40,9 @@ CREATE FUNCTION inner_product(vector, vector) RETURNS float8
CREATE FUNCTION cosine_distance(vector, vector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION l1_distance(vector, vector) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_dims(vector) RETURNS integer
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
@@ -52,6 +55,9 @@ CREATE FUNCTION vector_add(vector, vector) RETURNS vector
CREATE FUNCTION vector_sub(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE FUNCTION vector_mul(vector, vector) RETURNS vector
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
-- private functions
CREATE FUNCTION vector_lt(vector, vector) RETURNS bool
@@ -104,6 +110,13 @@ CREATE AGGREGATE avg(vector) (
PARALLEL = SAFE
);
CREATE AGGREGATE sum(vector) (
SFUNC = vector_add,
STYPE = vector,
COMBINEFUNC = vector_add,
PARALLEL = SAFE
);
-- cast functions
CREATE FUNCTION vector(vector, integer, boolean) RETURNS vector
@@ -171,6 +184,11 @@ CREATE OPERATOR - (
COMMUTATOR = -
);
CREATE OPERATOR * (
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_mul,
COMMUTATOR = *
);
CREATE OPERATOR < (
LEFTARG = vector, RIGHTARG = vector, PROCEDURE = vector_lt,
COMMUTATOR = > , NEGATOR = >= ,
@@ -209,7 +227,7 @@ CREATE OPERATOR > (
RESTRICT = scalargtsel, JOIN = scalargtjoinsel
);
-- access method
-- access methods
CREATE FUNCTION ivfflathandler(internal) RETURNS index_am_handler
AS 'MODULE_PATHNAME' LANGUAGE C;
@@ -218,6 +236,13 @@ CREATE ACCESS METHOD ivfflat TYPE INDEX HANDLER ivfflathandler;
COMMENT ON ACCESS METHOD ivfflat IS 'ivfflat index access method';
CREATE FUNCTION hnswhandler(internal) RETURNS index_am_handler
AS 'MODULE_PATHNAME' LANGUAGE C;
CREATE ACCESS METHOD hnsw TYPE INDEX HANDLER hnswhandler;
COMMENT ON ACCESS METHOD hnsw IS 'hnsw index access method';
-- opclasses
CREATE OPERATOR CLASS vector_ops
@@ -249,3 +274,19 @@ CREATE OPERATOR CLASS vector_cosine_ops
FUNCTION 2 vector_norm(vector),
FUNCTION 3 vector_spherical_distance(vector, vector),
FUNCTION 4 vector_norm(vector);
CREATE OPERATOR CLASS vector_l2_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_l2_squared_distance(vector, vector);
CREATE OPERATOR CLASS vector_ip_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector);
CREATE OPERATOR CLASS vector_cosine_ops
FOR TYPE vector USING hnsw AS
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector);

224
src/hnsw.c Normal file
View File

@@ -0,0 +1,224 @@
#include "postgres.h"
#include <float.h>
#include <math.h>
#include "access/amapi.h"
#include "commands/vacuum.h"
#include "hnsw.h"
#include "utils/guc.h"
#include "utils/selfuncs.h"
#if PG_VERSION_NUM >= 120000
#include "commands/progress.h"
#endif
int hnsw_ef_search;
static relopt_kind hnsw_relopt_kind;
/*
* Initialize index options and variables
*/
void
HnswInit(void)
{
hnsw_relopt_kind = add_reloption_kind();
add_int_reloption(hnsw_relopt_kind, "m", "Max number of connections",
HNSW_DEFAULT_M, HNSW_MIN_M, HNSW_MAX_M
#if PG_VERSION_NUM >= 130000
,AccessExclusiveLock
#endif
);
add_int_reloption(hnsw_relopt_kind, "ef_construction", "Size of the dynamic candidate list for construction",
HNSW_DEFAULT_EF_CONSTRUCTION, HNSW_MIN_EF_CONSTRUCTION, HNSW_MAX_EF_CONSTRUCTION
#if PG_VERSION_NUM >= 130000
,AccessExclusiveLock
#endif
);
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for 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);
}
/*
* Get the name of index build phase
*/
#if PG_VERSION_NUM >= 120000
static char *
hnswbuildphasename(int64 phasenum)
{
switch (phasenum)
{
case PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE:
return "initializing";
case PROGRESS_HNSW_PHASE_LOAD:
return "loading tuples";
default:
return NULL;
}
}
#endif
/*
* Estimate the cost of an index scan
*/
static void
hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
Cost *indexStartupCost, Cost *indexTotalCost,
Selectivity *indexSelectivity, double *indexCorrelation,
double *indexPages)
{
GenericCosts costs;
int m;
int entryLevel;
Relation index;
#if PG_VERSION_NUM < 120000
List *qinfos;
#endif
/* Never use index without order */
if (path->indexorderbys == NULL)
{
*indexStartupCost = DBL_MAX;
*indexTotalCost = DBL_MAX;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
MemSet(&costs, 0, sizeof(costs));
index = index_open(path->indexinfo->indexoid, NoLock);
m = HnswGetM(index);
index_close(index, NoLock);
/* Approximate entry level */
entryLevel = (int) -log(1.0 / path->indexinfo->tuples) * HnswGetMl(m);
/* TODO Improve estimate of visited tuples (currently underestimates) */
/* Account for number of tuples (or entry level), m, and ef_search */
costs.numIndexTuples = (entryLevel + 2) * m;
#if PG_VERSION_NUM >= 120000
genericcostestimate(root, path, loop_count, &costs);
#else
qinfos = deconstruct_indexquals(path);
genericcostestimate(root, path, loop_count, qinfos, &costs);
#endif
/* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost;
*indexTotalCost = costs.indexTotalCost;
*indexSelectivity = costs.indexSelectivity;
*indexCorrelation = costs.indexCorrelation;
*indexPages = costs.numIndexPages;
}
/*
* Parse and validate the reloptions
*/
static bytea *
hnswoptions(Datum reloptions, bool validate)
{
static const relopt_parse_elt tab[] = {
{"m", RELOPT_TYPE_INT, offsetof(HnswOptions, m)},
{"ef_construction", RELOPT_TYPE_INT, offsetof(HnswOptions, efConstruction)},
};
#if PG_VERSION_NUM >= 130000
return (bytea *) build_reloptions(reloptions, validate,
hnsw_relopt_kind,
sizeof(HnswOptions),
tab, lengthof(tab));
#else
relopt_value *options;
int numoptions;
HnswOptions *rdopts;
options = parseRelOptions(reloptions, validate, hnsw_relopt_kind, &numoptions);
rdopts = allocateReloptStruct(sizeof(HnswOptions), options, numoptions);
fillRelOptions((void *) rdopts, sizeof(HnswOptions), options, numoptions,
validate, tab, lengthof(tab));
return (bytea *) rdopts;
#endif
}
/*
* Validate catalog entries for the specified operator class
*/
static bool
hnswvalidate(Oid opclassoid)
{
return true;
}
/*
* Define index handler
*
* See https://www.postgresql.org/docs/current/index-api.html
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hnswhandler);
Datum
hnswhandler(PG_FUNCTION_ARGS)
{
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 2;
#if PG_VERSION_NUM >= 130000
amroutine->amoptsprocnum = 0;
#endif
amroutine->amcanorder = false;
amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false;
amroutine->amcanmulticol = false;
amroutine->amoptionalkey = true;
amroutine->amsearcharray = false;
amroutine->amsearchnulls = false;
amroutine->amstorage = false;
amroutine->amclusterable = false;
amroutine->ampredlocks = false;
amroutine->amcanparallel = false;
amroutine->amcaninclude = false;
#if PG_VERSION_NUM >= 130000
amroutine->amusemaintenanceworkmem = false; /* not used during VACUUM */
amroutine->amparallelvacuumoptions = VACUUM_OPTION_PARALLEL_BULKDEL;
#endif
amroutine->amkeytype = InvalidOid;
/* Interface functions */
amroutine->ambuild = hnswbuild;
amroutine->ambuildempty = hnswbuildempty;
amroutine->aminsert = hnswinsert;
amroutine->ambulkdelete = hnswbulkdelete;
amroutine->amvacuumcleanup = hnswvacuumcleanup;
amroutine->amcanreturn = NULL; /* tuple not included in heapsort */
amroutine->amcostestimate = hnswcostestimate;
amroutine->amoptions = hnswoptions;
amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */
#if PG_VERSION_NUM >= 120000
amroutine->ambuildphasename = hnswbuildphasename;
#endif
amroutine->amvalidate = hnswvalidate;
#if PG_VERSION_NUM >= 140000
amroutine->amadjustmembers = NULL;
#endif
amroutine->ambeginscan = hnswbeginscan;
amroutine->amrescan = hnswrescan;
amroutine->amgettuple = hnswgettuple;
amroutine->amgetbitmap = NULL;
amroutine->amendscan = hnswendscan;
amroutine->ammarkpos = NULL;
amroutine->amrestrpos = NULL;
/* Interface functions to support parallel index scans */
amroutine->amestimateparallelscan = NULL;
amroutine->aminitparallelscan = NULL;
amroutine->amparallelrescan = NULL;
PG_RETURN_POINTER(amroutine);
}

305
src/hnsw.h Normal file
View File

@@ -0,0 +1,305 @@
#ifndef HNSW_H
#define HNSW_H
#include "postgres.h"
#include "access/generic_xlog.h"
#include "access/reloptions.h"
#include "nodes/execnodes.h"
#include "port.h" /* for random() */
#include "utils/sampling.h"
#include "vector.h"
#if PG_VERSION_NUM < 110000
#error "Requires PostgreSQL 11+"
#endif
#define HNSW_MAX_DIM 2000
/* Support functions */
#define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2
#define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953
#define HNSW_PAGE_ID 0xFF90
/* Preserved page numbers */
#define HNSW_METAPAGE_BLKNO 0
#define HNSW_HEAD_BLKNO 1 /* first element page */
/* Must correspond to page numbers since page lock is used */
#define HNSW_UPDATE_LOCK 0
#define HNSW_SCAN_LOCK 1
/* HNSW parameters */
#define HNSW_DEFAULT_M 16
#define HNSW_MIN_M 2
#define HNSW_MAX_M 100
#define HNSW_DEFAULT_EF_CONSTRUCTION 40
#define HNSW_MIN_EF_CONSTRUCTION 4
#define HNSW_MAX_EF_CONSTRUCTION 1000
#define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000
/* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1
#define HNSW_NEIGHBOR_TUPLE_TYPE 2
/* Make graph robust against non-HOT updates */
#define HNSW_HEAPTIDS 10
#define HNSW_UPDATE_ENTRY_GREATER 1
#define HNSW_UPDATE_ENTRY_ALWAYS 2
/* Build phases */
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_HNSW_PHASE_LOAD 2
#define HNSW_ELEMENT_TUPLE_SIZE(_dim) MAXALIGN(offsetof(HnswElementTupleData, vec) + VECTOR_SIZE(_dim))
#define HNSW_NEIGHBOR_TUPLE_SIZE(level, m) MAXALIGN(offsetof(HnswNeighborTupleData, indextids) + ((level) + 2) * (m) * sizeof(ItemPointerData))
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page))
#if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
#else
#define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE)
#endif
#if PG_VERSION_NUM < 130000
#define list_delete_last(list) list_truncate(list, list_length(list) - 1)
#define list_sort(list, cmp) list_qsort(list, cmp)
#endif
#define HnswIsElementTuple(tup) ((tup)->type == HNSW_ELEMENT_TUPLE_TYPE)
#define HnswIsNeighborTuple(tup) ((tup)->type == HNSW_NEIGHBOR_TUPLE_TYPE)
/* 2 * M connections for ground layer */
#define HnswGetLayerM(m, layer) (layer == 0 ? (m) * 2 : (m))
/* Optimal ML from paper */
#define HnswGetMl(m) (1 / log(m))
/* 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)
/* Variables */
extern int hnsw_ef_search;
typedef struct HnswNeighborArray HnswNeighborArray;
typedef struct HnswElementData
{
List *heaptids;
uint8 level;
uint8 deleted;
HnswNeighborArray *neighbors;
BlockNumber blkno;
OffsetNumber offno;
OffsetNumber neighborOffno;
BlockNumber neighborPage;
Vector *vec;
} HnswElementData;
typedef HnswElementData * HnswElement;
typedef struct HnswCandidate
{
HnswElement element;
float distance;
} HnswCandidate;
typedef struct HnswNeighborArray
{
int length;
HnswCandidate *items;
} HnswNeighborArray;
typedef struct HnswPairingHeapNode
{
pairingheap_node ph_node;
HnswCandidate *inner;
} HnswPairingHeapNode;
/* HNSW index options */
typedef struct HnswOptions
{
int32 vl_len_; /* varlena header (do not touch directly!) */
int m; /* number of connections */
int efConstruction; /* size of dynamic candidate list */
} HnswOptions;
typedef struct HnswBuildState
{
/* Info */
Relation heap;
Relation index;
IndexInfo *indexInfo;
ForkNumber forkNum;
/* Settings */
int dimensions;
int m;
int efConstruction;
/* Statistics */
double indtuples;
double reltuples;
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
Oid collation;
/* Variables */
List *elements;
HnswElement entryPoint;
double ml;
int maxLevel;
double maxInMemoryElements;
bool flushed;
Vector *normvec;
/* Memory */
MemoryContext tmpCtx;
} HnswBuildState;
typedef struct HnswMetaPageData
{
uint32 magicNumber;
uint32 version;
uint32 dimensions;
uint16 m;
uint16 efConstruction;
BlockNumber entryBlkno;
OffsetNumber entryOffno;
int16 entryLevel;
BlockNumber insertPage;
} HnswMetaPageData;
typedef HnswMetaPageData * HnswMetaPage;
typedef struct HnswPageOpaqueData
{
BlockNumber nextblkno;
uint16 unused;
uint16 page_id; /* for identification of HNSW indexes */
} HnswPageOpaqueData;
typedef HnswPageOpaqueData * HnswPageOpaque;
typedef struct HnswElementTupleData
{
uint8 type;
uint8 level;
uint8 deleted;
uint8 unused;
ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid;
uint16 unused2;
Vector vec;
} HnswElementTupleData;
typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData
{
uint8 type;
uint8 unused;
uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData;
typedef HnswNeighborTupleData * HnswNeighborTuple;
typedef struct HnswScanOpaqueData
{
bool first;
Buffer buf;
List *w;
MemoryContext tmpCtx;
/* Support functions */
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
Oid collation;
} HnswScanOpaqueData;
typedef HnswScanOpaqueData * HnswScanOpaque;
typedef struct HnswVacuumState
{
/* Info */
Relation index;
IndexBulkDeleteResult *stats;
IndexBulkDeleteCallback callback;
void *callback_state;
/* Settings */
int m;
int efConstruction;
/* Support functions */
FmgrInfo *procinfo;
Oid collation;
/* Variables */
HTAB *deleted;
BufferAccessStrategy bas;
HnswNeighborTuple ntup;
HnswElementData highestPoint;
/* Memory */
MemoryContext tmpCtx;
} HnswVacuumState;
/* Methods */
int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation rel, uint16 procnum);
bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
void HnswCommitBuffer(Buffer buf, GenericXLogState *state);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void HnswInit(void);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement);
HnswElement HnswGetEntryPoint(Relation index);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswElement HnswFindDuplicate(HnswElement e);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel);
void HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index);
/* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);
void hnswbuildempty(Relation index);
bool hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heap, IndexUniqueCheck checkUnique
#if PG_VERSION_NUM >= 140000
,bool indexUnchanged
#endif
,IndexInfo *indexInfo
);
IndexBulkDeleteResult *hnswbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats, IndexBulkDeleteCallback callback, void *callback_state);
IndexBulkDeleteResult *hnswvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats);
IndexScanDesc hnswbeginscan(Relation index, int nkeys, int norderbys);
void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys);
bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan);
#endif

526
src/hnswbuild.c Normal file
View File

@@ -0,0 +1,526 @@
#include "postgres.h"
#include <math.h>
#include "catalog/index.h"
#include "hnsw.h"
#include "miscadmin.h"
#include "lib/pairingheap.h"
#include "nodes/pg_list.h"
#include "storage/bufmgr.h"
#include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000
#include "utils/backend_progress.h"
#elif PG_VERSION_NUM >= 120000
#include "pgstat.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/tableam.h"
#include "commands/progress.h"
#else
#define PROGRESS_CREATEIDX_TUPLES_DONE 0
#endif
#if PG_VERSION_NUM >= 130000
#define CALLBACK_ITEM_POINTER ItemPointer tid
#else
#define CALLBACK_ITEM_POINTER HeapTuple hup
#endif
#if PG_VERSION_NUM >= 120000
#define UpdateProgress(index, val) pgstat_progress_update_param(index, val)
#else
#define UpdateProgress(index, val) ((void)val)
#endif
/*
* Create the metapage
*/
static void
CreateMetaPage(HnswBuildState * buildstate)
{
Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum;
Buffer buf;
Page page;
GenericXLogState *state;
HnswMetaPage metap;
buf = HnswNewBuffer(index, forkNum);
HnswInitRegisterPage(index, &buf, &page, &state);
/* Set metapage data */
metap = HnswPageGetMeta(page);
metap->magicNumber = HNSW_MAGIC_NUMBER;
metap->version = HNSW_VERSION;
metap->dimensions = buildstate->dimensions;
metap->m = buildstate->m;
metap->efConstruction = buildstate->efConstruction;
metap->entryBlkno = InvalidBlockNumber;
metap->entryOffno = InvalidOffsetNumber;
metap->entryLevel = -1;
metap->insertPage = InvalidBlockNumber;
((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
HnswCommitBuffer(buf, state);
}
/*
* Add a new page
*/
static void
HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, ForkNumber forkNum)
{
/* Add a new page */
Buffer newbuf = HnswNewBuffer(index, forkNum);
/* Update previous page */
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */
MarkBufferDirty(*buf);
GenericXLogFinish(*state);
UnlockReleaseBuffer(*buf);
/* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */
LockBuffer(newbuf, BUFFER_LOCK_UNLOCK);
CHECK_FOR_INTERRUPTS();
LockBuffer(newbuf, BUFFER_LOCK_EXCLUSIVE);
/* Prepare new page */
*buf = newbuf;
*state = GenericXLogStart(index);
*page = GenericXLogRegisterBuffer(*state, *buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*buf, *page);
}
/*
* Create element pages
*/
static void
CreateElementPages(HnswBuildState * buildstate)
{
Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum;
int dimensions = buildstate->dimensions;
Size etupSize;
Size maxSize;
HnswElementTuple etup;
HnswNeighborTuple ntup;
BlockNumber insertPage;
Buffer buf;
Page page;
GenericXLogState *state;
ListCell *lc;
/* Calculate sizes */
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
/* Allocate once */
etup = palloc0(etupSize);
ntup = palloc0(maxSize);
/* Prepare first page */
buf = HnswNewBuffer(index, forkNum);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(buf, page);
foreach(lc, buildstate->elements)
{
HnswElement element = lfirst(lc);
Size ntupSize;
Size combinedSize;
HnswSetElementTuple(etup, element);
/* Calculate sizes */
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
/* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
/* Calculate offsets */
element->blkno = BufferGetBlockNumber(buf);
element->offno = OffsetNumberNext(PageGetMaxOffsetNumber(page));
if (combinedSize <= maxSize)
{
element->neighborPage = element->blkno;
element->neighborOffno = OffsetNumberNext(element->offno);
}
else
{
element->neighborPage = element->blkno + 1;
element->neighborOffno = FirstOffsetNumber;
}
ItemPointerSet(&etup->neighbortid, element->neighborPage, element->neighborOffno);
/* Add element */
if (PageAddItem(page, (Item) etup, etupSize, InvalidOffsetNumber, false, false) != element->offno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Add new page if needed */
if (PageGetFreeSpace(page) < ntupSize)
HnswBuildAppendPage(index, &buf, &page, &state, forkNum);
/* Add placeholder for neighbors */
if (PageAddItem(page, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != element->neighborOffno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
}
insertPage = BufferGetBlockNumber(buf);
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum);
pfree(etup);
pfree(ntup);
}
/*
* Create neighbor pages
*/
static void
CreateNeighborPages(HnswBuildState * buildstate)
{
Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum;
int m = buildstate->m;
ListCell *lc;
HnswNeighborTuple ntup;
/* Allocate once */
ntup = palloc0(BLCKSZ);
foreach(lc, buildstate->elements)
{
HnswElement e = lfirst(lc);
Buffer buf;
Page page;
GenericXLogState *state;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
/* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */
CHECK_FOR_INTERRUPTS();
buf = ReadBufferExtended(index, forkNum, e->neighborPage, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
HnswSetNeighborTuple(ntup, e, m);
if (!PageIndexTupleOverwrite(page, e->neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
}
pfree(ntup);
}
/*
* Free elements
*/
static void
FreeElements(HnswBuildState * buildstate)
{
ListCell *lc;
foreach(lc, buildstate->elements)
HnswFreeElement(lfirst(lc));
list_free(buildstate->elements);
}
/*
* Flush pages
*/
static void
FlushPages(HnswBuildState * buildstate)
{
CreateMetaPage(buildstate);
CreateElementPages(buildstate);
CreateNeighborPages(buildstate);
buildstate->flushed = true;
FreeElements(buildstate);
}
/*
* Insert tuple
*/
static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup)
{
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction;
int m = buildstate->m;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
return false;
}
/* Copy value to element so accessible outside of memory context */
memcpy(element->vec, DatumGetVector(value), VECTOR_SIZE(buildstate->dimensions));
/* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
*dup = HnswFindDuplicate(element);
/* Update neighbors if needed */
if (*dup == NULL)
{
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
}
}
/* Update entry point if needed */
if (*dup == NULL && (entryPoint == NULL || element->level > entryPoint->level))
buildstate->entryPoint = element;
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
return *dup == NULL;
}
/*
* Callback for table_index_build_scan
*/
static void
BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state)
{
HnswBuildState *buildstate = (HnswBuildState *) state;
MemoryContext oldCtx;
HnswElement element;
HnswElement dup = NULL;
bool inserted;
#if PG_VERSION_NUM < 130000
ItemPointer tid = &hup->t_self;
#endif
/* Skip nulls */
if (isnull[0])
return;
if (buildstate->indtuples >= buildstate->maxInMemoryElements)
{
if (!buildstate->flushed)
{
ereport(NOTICE,
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) buildstate->indtuples),
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
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);
/* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
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
*/
static double
HnswGetMaxInMemoryElements(int m, double ml, int dimensions)
{
Size elementSize = sizeof(HnswElementData);
double avgLevel = -log(0.5) * ml;
elementSize += sizeof(HnswNeighborArray) * (avgLevel + 1);
elementSize += sizeof(HnswCandidate) * (m * (avgLevel + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VECTOR_SIZE(dimensions);
return (maintenance_work_mem * 1024L) / elementSize;
}
/*
* Initialize the build state
*/
static void
InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, IndexInfo *indexInfo, ForkNumber forkNum)
{
buildstate->heap = heap;
buildstate->index = index;
buildstate->indexInfo = indexInfo;
buildstate->forkNum = forkNum;
buildstate->m = HnswGetM(index);
buildstate->efConstruction = HnswGetEfConstruction(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
/* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0)
elog(ERROR, "column does not have dimensions");
if (buildstate->dimensions > HNSW_MAX_DIM)
elog(ERROR, "column cannot have more than %d dimensions for hnsw index", HNSW_MAX_DIM);
if (buildstate->efConstruction < 2 * buildstate->m)
elog(ERROR, "ef_construction must be greater than or equal to 2 * m");
buildstate->reltuples = 0;
buildstate->indtuples = 0;
/* Get support functions */
buildstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0];
buildstate->elements = NIL;
buildstate->entryPoint = NULL;
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->maxInMemoryElements = HnswGetMaxInMemoryElements(buildstate->m, buildstate->ml, buildstate->dimensions);
buildstate->flushed = false;
/* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES);
}
/*
* Free resources
*/
static void
FreeBuildState(HnswBuildState * buildstate)
{
pfree(buildstate->normvec);
MemoryContextDelete(buildstate->tmpCtx);
}
/*
* Build graph
*/
static void
BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
{
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
#if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
/*
* Build the index
*/
static void
BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
HnswBuildState * buildstate, ForkNumber forkNum)
{
InitBuildState(buildstate, heap, index, indexInfo, forkNum);
if (buildstate->heap != NULL)
BuildGraph(buildstate, forkNum);
if (!buildstate->flushed)
FlushPages(buildstate);
FreeBuildState(buildstate);
}
/*
* Build the index for a logged table
*/
IndexBuildResult *
hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo)
{
IndexBuildResult *result;
HnswBuildState buildstate;
BuildIndex(heap, index, indexInfo, &buildstate, MAIN_FORKNUM);
result = (IndexBuildResult *) palloc(sizeof(IndexBuildResult));
result->heap_tuples = buildstate.reltuples;
result->index_tuples = buildstate.indtuples;
return result;
}
/*
* Build the index for an unlogged table
*/
void
hnswbuildempty(Relation index)
{
IndexInfo *indexInfo = BuildIndexInfo(index);
HnswBuildState buildstate;
BuildIndex(NULL, index, indexInfo, &buildstate, INIT_FORKNUM);
}

590
src/hnswinsert.c Normal file
View File

@@ -0,0 +1,590 @@
#include "postgres.h"
#include <math.h>
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
* Get the insert page
*/
static BlockNumber
GetInsertPage(Relation index)
{
Buffer buf;
Page page;
HnswMetaPage metap;
BlockNumber insertPage;
buf = ReadBuffer(index, HNSW_METAPAGE_BLKNO);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
metap = HnswPageGetMeta(page);
insertPage = metap->insertPage;
UnlockReleaseBuffer(buf);
return insertPage;
}
/*
* Check for a free offset
*/
static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
{
OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
if (etup->deleted)
{
BlockNumber elementPage = BufferGetBlockNumber(buf);
BlockNumber neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
OffsetNumber neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
ItemId itemid;
if (!BlockNumberIsValid(*newInsertPage))
*newInsertPage = elementPage;
if (neighborPage == elementPage)
{
*nbuf = buf;
*npage = page;
}
else
{
*nbuf = ReadBuffer(index, neighborPage);
LockBuffer(*nbuf, BUFFER_LOCK_EXCLUSIVE);
/* Skip WAL for now */
*npage = BufferGetPage(*nbuf);
}
itemid = PageGetItemId(*npage, neighborOffno);
/* Check for space on neighbor tuple page */
if (PageGetFreeSpace(*npage) + ItemIdGetLength(itemid) - sizeof(ItemIdData) >= ntupSize)
{
*freeOffno = offno;
*freeNeighborOffno = neighborOffno;
return true;
}
else if (*nbuf != buf)
UnlockReleaseBuffer(*nbuf);
}
}
return false;
}
/*
* Add a new page
*/
static void
HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState *state, Page page)
{
/* Add a new page */
LockRelationForExtension(index, ExclusiveLock);
*nbuf = HnswNewBuffer(index, MAIN_FORKNUM);
UnlockRelationForExtension(index, ExclusiveLock);
/* Init new page */
*npage = GenericXLogRegisterBuffer(state, *nbuf, GENERIC_XLOG_FULL_IMAGE);
HnswInitPage(*nbuf, *npage);
/* Update previous buffer */
HnswPageGetOpaque(page)->nextblkno = BufferGetBlockNumber(*nbuf);
}
/*
* Add to element and neighbor pages
*/
static void
WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage)
{
Buffer buf;
Page page;
GenericXLogState *state;
Size etupSize;
Size ntupSize;
Size combinedSize;
Size maxSize;
Size minCombinedSize;
HnswElementTuple etup;
BlockNumber currentPage = insertPage;
int dimensions = e->vec->dim;
HnswNeighborTuple ntup;
Buffer nbuf;
Page npage;
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(dimensions);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData));
minCombinedSize = etupSize + HNSW_NEIGHBOR_TUPLE_SIZE(0, m) + sizeof(ItemIdData);
/* Prepare element tuple */
etup = palloc0(etupSize);
HnswSetElementTuple(etup, e);
/* Prepare neighbor tuple */
ntup = palloc0(ntupSize);
HnswSetNeighborTuple(ntup, e, m);
/* Find a page (or two if needed) to insert the tuples */
for (;;)
{
buf = ReadBuffer(index, currentPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Keep track of first page where element at level 0 can fit */
if (!BlockNumberIsValid(newInsertPage) && PageGetFreeSpace(page) >= minCombinedSize)
newInsertPage = currentPage;
/* First, try the fastest path */
/* Space for both tuples on the current page */
/* This can split existing tuples in rare cases */
if (PageGetFreeSpace(page) >= combinedSize)
{
nbuf = buf;
npage = page;
break;
}
/* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
{
if (nbuf != buf)
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
break;
}
/* Finally, try space for element only if last page */
/* Skip if both tuples can fit on the same page */
if (combinedSize > maxSize && PageGetFreeSpace(page) >= etupSize && !BlockNumberIsValid(HnswPageGetOpaque(page)->nextblkno))
{
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
break;
}
currentPage = HnswPageGetOpaque(page)->nextblkno;
if (BlockNumberIsValid(currentPage))
{
/* Move to next page */
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
else
{
Buffer newbuf;
Page newpage;
HnswInsertAppendPage(index, &newbuf, &newpage, state, page);
/* Commit */
MarkBufferDirty(newbuf);
MarkBufferDirty(buf);
GenericXLogFinish(state);
/* Unlock previous buffer */
UnlockReleaseBuffer(buf);
/* Prepare new buffer */
state = GenericXLogStart(index);
buf = newbuf;
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Create new page for neighbors if needed */
if (PageGetFreeSpace(page) < combinedSize)
HnswInsertAppendPage(index, &nbuf, &npage, state, page);
else
{
nbuf = buf;
npage = page;
}
break;
}
}
e->blkno = BufferGetBlockNumber(buf);
e->neighborPage = BufferGetBlockNumber(nbuf);
/* Added tuple to new page if newInsertPage is not set */
/* So can set to neighbor page instead of element page */
if (!BlockNumberIsValid(newInsertPage))
newInsertPage = e->neighborPage;
if (OffsetNumberIsValid(freeOffno))
{
e->offno = freeOffno;
e->neighborOffno = freeNeighborOffno;
}
else
{
e->offno = OffsetNumberNext(PageGetMaxOffsetNumber(page));
if (nbuf == buf)
e->neighborOffno = OffsetNumberNext(e->offno);
else
e->neighborOffno = FirstOffsetNumber;
}
ItemPointerSet(&etup->neighbortid, e->neighborPage, e->neighborOffno);
/* Add element and neighbors */
if (OffsetNumberIsValid(freeOffno))
{
if (!PageIndexTupleOverwrite(page, e->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
if (!PageIndexTupleOverwrite(npage, e->neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
}
else
{
if (PageAddItem(page, (Item) etup, etupSize, InvalidOffsetNumber, false, false) != e->offno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
if (PageAddItem(npage, (Item) ntup, ntupSize, InvalidOffsetNumber, false, false) != e->neighborOffno)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
}
/* Commit */
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
if (nbuf != buf)
UnlockReleaseBuffer(nbuf);
/* Update the insert page */
if (BlockNumberIsValid(newInsertPage) && newInsertPage != insertPage)
*updatedInsertPage = newInsertPage;
}
/*
* Check if connection already exists
*/
static bool
ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
{
for (int i = 0; i < lm; i++)
{
ItemPointer indextid = &ntup->indextids[startIdx + i];
if (!ItemPointerIsValid(indextid))
break;
if (ItemPointerGetBlockNumber(indextid) == e->blkno && ItemPointerGetOffsetNumber(indextid) == e->offno)
return true;
}
return false;
}
/*
* Update neighbors
*/
void
HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting)
{
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
Buffer buf;
Page page;
GenericXLogState *state;
ItemId itemid;
HnswNeighborTuple ntup;
Size ntupSize;
int idx = -1;
int startIdx;
OffsetNumber offno = hc->element->neighborOffno;
/* Get latest neighbors since they may have changed */
/* Do not lock yet since selecting neighbors can take time */
HnswLoadNeighbors(hc->element, index);
/*
* Could improve performance for vacuuming by checking neighbors
* against list of elements being deleted to find index. It's
* important to exclude already deleted elements for this since
* they can be replaced at any time.
*/
/* Select neighbors */
HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
/* New element was not selected as a neighbor */
if (idx == -1)
continue;
/* Register page */
buf = ReadBuffer(index, hc->element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Get tuple */
itemid = PageGetItemId(page, offno);
ntup = (HnswNeighborTuple) PageGetItem(page, itemid);
ntupSize = ItemIdGetLength(itemid);
/* Calculate index for update */
startIdx = (hc->element->level - lc) * m;
/* Check for existing connection */
if (checkExisting && ConnectionExists(e, ntup, startIdx, lm))
idx = -1;
else if (idx == -2)
{
/* Find free offset if still exists */
/* TODO Retry updating connections if not */
for (int j = 0; j < lm; j++)
{
if (!ItemPointerIsValid(&ntup->indextids[startIdx + j]))
{
idx = startIdx + j;
break;
}
}
}
else
idx += startIdx;
/* Make robust to issues */
if (idx >= 0 && idx < ntup->count)
{
ItemPointer indextid = &ntup->indextids[idx];
/* Update neighbor */
ItemPointerSet(indextid, e->blkno, e->offno);
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, offno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
}
else
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
}
}
/*
* Add a heap TID to an existing element
*/
static bool
HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup)
{
Buffer buf;
Page page;
GenericXLogState *state;
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(dup->vec->dim);
HnswElementTuple etup;
int i;
/* Read page */
buf = ReadBuffer(index, dup->blkno);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Find space */
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, dup->offno));
for (i = 0; i < HNSW_HEAPTIDS; i++)
{
if (!ItemPointerIsValid(&etup->heaptids[i]))
break;
}
/* Either being deleted or we lost our chance to another backend */
if (i == 0 || i == HNSW_HEAPTIDS)
{
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
return false;
}
/* Add heap TID */
etup->heaptids[i] = *((ItemPointer) linitial(element->heaptids));
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, dup->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
return true;
}
/*
* Write changes to disk
*/
static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint)
{
BlockNumber newInsertPage = InvalidBlockNumber;
/* Try to add to existing page */
if (dup != NULL)
{
if (HnswAddDuplicate(index, element, dup))
return;
}
/* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage);
/* Update insert page if needed */
if (BlockNumberIsValid(newInsertPage))
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false);
/* Update metapage if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM);
}
/*
* Insert a tuple into the index
*/
bool
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel)
{
Datum value;
FmgrInfo *normprocinfo;
HnswElement entryPoint;
HnswElement element;
int m = HnswGetM(index);
int efConstruction = HnswGetEfConstruction(index);
double ml = HnswGetMl(m);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
HnswElement dup;
LOCKMODE lockmode = ShareLock;
/* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Normalize if needed */
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswNormValue(normprocinfo, collation, &value, NULL))
return false;
}
/* Create an element */
element = HnswInitElement(heap_tid, m, ml, HnswGetMaxLevel(m));
element->vec = DatumGetVector(value);
/*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
* before repairing graph. Use a page lock so it does not interfere with
* buffer lock (or reads when vacuuming).
*/
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get entry point */
entryPoint = HnswGetEntryPoint(index);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get exclusive lock */
lockmode = ExclusiveLock;
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get latest entry point after lock is acquired */
entryPoint = HnswGetEntryPoint(index);
}
/* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
dup = HnswFindDuplicate(element);
/* Write to disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
return true;
}
/*
* Insert a tuple into the index
*/
bool
hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid,
Relation heap, IndexUniqueCheck checkUnique
#if PG_VERSION_NUM >= 140000
,bool indexUnchanged
#endif
,IndexInfo *indexInfo
)
{
MemoryContext oldCtx;
MemoryContext insertCtx;
/* Skip nulls */
if (isnull[0])
return false;
/* Create memory context */
insertCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw insert temporary context",
ALLOCSET_DEFAULT_SIZES);
oldCtx = MemoryContextSwitchTo(insertCtx);
/* Insert tuple */
HnswInsertTuple(index, values, isnull, heap_tid, heap);
/* Delete memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextDelete(insertCtx);
return false;
}

223
src/hnswscan.c Normal file
View File

@@ -0,0 +1,223 @@
#include "postgres.h"
#include "access/relscan.h"
#include "hnsw.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
* Algorithm 5 from paper
*/
static List *
GetScanItems(IndexScanDesc scan, Datum q)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep;
List *w;
HnswElement entryPoint = HnswGetEntryPoint(index);
if (entryPoint == NULL)
return NIL;
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, false, NULL);
ep = w;
}
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, false, NULL);
}
/*
* Get dimensions from metapage
*/
static int
GetDimensions(Relation index)
{
Buffer buf;
Page page;
HnswMetaPage metap;
int dimensions;
buf = ReadBuffer(index, HNSW_METAPAGE_BLKNO);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
metap = HnswPageGetMeta(page);
dimensions = metap->dimensions;
UnlockReleaseBuffer(buf);
return dimensions;
}
/*
* Prepare for an index scan
*/
IndexScanDesc
hnswbeginscan(Relation index, int nkeys, int norderbys)
{
IndexScanDesc scan;
HnswScanOpaque so;
scan = RelationGetIndexScan(index, nkeys, norderbys);
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->buf = InvalidBuffer;
so->first = true;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */
so->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->collation = index->rd_indcollation[0];
scan->opaque = so;
return scan;
}
/*
* Start or restart an index scan
*/
void
hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int norderbys)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
so->first = true;
MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
if (orderbys && scan->numberOfOrderBys > 0)
memmove(scan->orderByData, orderbys, scan->numberOfOrderBys * sizeof(ScanKeyData));
}
/*
* Fetch the next tuple in the given scan
*/
bool
hnswgettuple(IndexScanDesc scan, ScanDirection dir)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/*
* Index can be used to scan backward, but Postgres doesn't support
* backward scan on operators
*/
Assert(ScanDirectionIsForward(dir));
if (so->first)
{
Datum value;
/* Count index scan for stats */
pgstat_count_index_scan(scan->indexRelation);
/* Safety check */
if (scan->orderByData == NULL)
elog(ERROR, "cannot scan hnsw index without order");
if (scan->orderByData->sk_flags & SK_ISNULL)
value = PointerGetDatum(InitVector(GetDimensions(scan->indexRelation)));
else
{
value = scan->orderByData->sk_argument;
/* Value should not be compressed or toasted */
Assert(!VARATT_IS_COMPRESSED(DatumGetPointer(value)));
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
HnswNormValue(so->normprocinfo, so->collation, &value, NULL);
}
/*
* Get a shared lock. This allows vacuum to ensure no in-flight scans
* before marking tuples as deleted.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->w = GetScanItems(scan, value);
/* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->first = false;
}
while (list_length(so->w) > 0)
{
HnswCandidate *hc = llast(so->w);
ItemPointer tid;
BlockNumber indexblkno;
/* Move to next element if no valid heap tids */
if (list_length(hc->element->heaptids) == 0)
{
so->w = list_delete_last(so->w);
continue;
}
tid = llast(hc->element->heaptids);
indexblkno = hc->element->blkno;
hc->element->heaptids = list_delete_last(hc->element->heaptids);
MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *tid;
#else
scan->xs_ctup.t_self = *tid;
#endif
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
/*
* An index scan must maintain a pin on the index page holding the
* item last returned by amgettuple
*
* https://www.postgresql.org/docs/current/index-locking.html
*/
so->buf = ReadBuffer(scan->indexRelation, indexblkno);
scan->xs_recheckorderby = false;
return true;
}
MemoryContextSwitchTo(oldCtx);
return false;
}
/*
* End a scan and release resources
*/
void
hnswendscan(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
/* Release pin */
if (BufferIsValid(so->buf))
ReleaseBuffer(so->buf);
MemoryContextDelete(so->tmpCtx);
pfree(so);
scan->opaque = NULL;
}

985
src/hnswutils.c Normal file
View File

@@ -0,0 +1,985 @@
#include "postgres.h"
#include <math.h>
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "vector.h"
/*
* Get the max number of connections in an upper layer for each element in the index
*/
int
HnswGetM(Relation index)
{
HnswOptions *opts = (HnswOptions *) index->rd_options;
if (opts)
return opts->m;
return HNSW_DEFAULT_M;
}
/*
* Get the size of the dynamic candidate list in the index
*/
int
HnswGetEfConstruction(Relation index)
{
HnswOptions *opts = (HnswOptions *) index->rd_options;
if (opts)
return opts->efConstruction;
return HNSW_DEFAULT_EF_CONSTRUCTION;
}
/*
* Get proc
*/
FmgrInfo *
HnswOptionalProcInfo(Relation rel, uint16 procnum)
{
if (!OidIsValid(index_getprocid(rel, 1, procnum)))
return NULL;
return index_getprocinfo(rel, 1, procnum);
}
/*
* Divide by the norm
*
* Returns false if value should not be indexed
*
* The caller needs to free the pointer stored in value
* if it's different than the original value
*/
bool
HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
{
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
Vector *v = DatumGetVector(*value);
if (result == NULL)
result = InitVector(v->dim);
for (int i = 0; i < v->dim; i++)
result->x[i] = v->x[i] / norm;
*value = PointerGetDatum(result);
return true;
}
return false;
}
/*
* New buffer
*/
Buffer
HnswNewBuffer(Relation index, ForkNumber forkNum)
{
Buffer buf = ReadBufferExtended(index, forkNum, P_NEW, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
return buf;
}
/*
* Init page
*/
void
HnswInitPage(Buffer buf, Page page)
{
PageInit(page, BufferGetPageSize(buf), sizeof(HnswPageOpaqueData));
HnswPageGetOpaque(page)->nextblkno = InvalidBlockNumber;
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)
{
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
}
/*
* Allocate neighbors
*/
void
HnswInitNeighbors(HnswElement element, int m)
{
int level = element->level;
element->neighbors = palloc(sizeof(HnswNeighborArray) * (level + 1));
for (int lc = 0; lc <= level; lc++)
{
HnswNeighborArray *a;
int lm = HnswGetLayerM(m, lc);
a = &element->neighbors[lc];
a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm);
}
}
/*
* Allocate an element
*/
HnswElement
HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
{
HnswElement element = palloc(sizeof(HnswElementData));
int level = (int) (-log(RandomDouble()) * ml);
/* Cap level */
if (level > maxLevel)
level = maxLevel;
element->heaptids = NIL;
HnswAddHeapTid(element, heaptid);
element->level = level;
element->deleted = 0;
HnswInitNeighbors(element, m);
return element;
}
/*
* Free an element
*/
void
HnswFreeElement(HnswElement element)
{
list_free_deep(element->heaptids);
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
pfree(element->vec);
pfree(element);
}
/*
* Add a heap TID to an element
*/
void
HnswAddHeapTid(HnswElement element, ItemPointer heaptid)
{
ItemPointer copy = palloc(sizeof(ItemPointerData));
ItemPointerCopy(heaptid, copy);
element->heaptids = lappend(element->heaptids, copy);
}
/*
* Allocate an element from block and offset numbers
*/
HnswElement
HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
{
HnswElement element = palloc(sizeof(HnswElementData));
element->blkno = blkno;
element->offno = offno;
element->neighbors = NULL;
element->vec = NULL;
return element;
}
/*
* Get the entry point
*/
HnswElement
HnswGetEntryPoint(Relation index)
{
Buffer buf;
Page page;
HnswMetaPage metap;
HnswElement entryPoint = NULL;
buf = ReadBuffer(index, HNSW_METAPAGE_BLKNO);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
metap = HnswPageGetMeta(page);
if (BlockNumberIsValid(metap->entryBlkno))
entryPoint = HnswInitElementFromBlock(metap->entryBlkno, metap->entryOffno);
UnlockReleaseBuffer(buf);
return entryPoint;
}
/*
* Update the metapage info
*/
static void
HnswUpdateMetaPageInfo(Page page, int updateEntry, HnswElement entryPoint, BlockNumber insertPage)
{
HnswMetaPage metap = HnswPageGetMeta(page);
if (updateEntry)
{
if (entryPoint == NULL)
{
metap->entryBlkno = InvalidBlockNumber;
metap->entryOffno = InvalidOffsetNumber;
metap->entryLevel = -1;
}
else if (entryPoint->level > metap->entryLevel || updateEntry == HNSW_UPDATE_ENTRY_ALWAYS)
{
metap->entryBlkno = entryPoint->blkno;
metap->entryOffno = entryPoint->offno;
metap->entryLevel = entryPoint->level;
}
}
if (BlockNumberIsValid(insertPage))
metap->insertPage = insertPage;
}
/*
* Update the metapage
*/
void
HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum)
{
Buffer buf;
Page page;
GenericXLogState *state;
buf = ReadBufferExtended(index, forkNum, HNSW_METAPAGE_BLKNO, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage);
HnswCommitBuffer(buf, state);
}
/*
* Set element tuple, except for neighbor info
*/
void
HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
{
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level;
etup->deleted = 0;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
if (i < list_length(element->heaptids))
etup->heaptids[i] = *((ItemPointer) list_nth(element->heaptids, i));
else
ItemPointerSetInvalid(&etup->heaptids[i]);
}
memcpy(&etup->vec, element->vec, VECTOR_SIZE(element->vec->dim));
}
/*
* Set neighbor tuple
*/
void
HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
{
int idx = 0;
ntup->type = HNSW_NEIGHBOR_TUPLE_TYPE;
for (int lc = e->level; lc >= 0; lc--)
{
HnswNeighborArray *neighbors = &e->neighbors[lc];
int lm = HnswGetLayerM(m, lc);
for (int i = 0; i < lm; i++)
{
ItemPointer indextid = &ntup->indextids[idx++];
if (i < neighbors->length)
{
HnswCandidate *hc = &neighbors->items[i];
ItemPointerSet(indextid, hc->element->blkno, hc->element->offno);
}
else
ItemPointerSetInvalid(indextid);
}
}
ntup->count = idx;
}
/*
* Load neighbors from page
*/
static void
LoadNeighborsFromPage(HnswElement element, Relation index, Page page)
{
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
int m = HnswGetM(index);
int neighborCount = (element->level + 2) * m;
Assert(HnswIsNeighborTuple(ntup));
HnswInitNeighbors(element, m);
/* Ensure expected neighbors */
if (ntup->count != neighborCount)
return;
for (int i = 0; i < neighborCount; i++)
{
HnswElement e;
int level;
HnswCandidate *hc;
ItemPointer indextid;
HnswNeighborArray *neighbors;
indextid = &ntup->indextids[i];
if (!ItemPointerIsValid(indextid))
continue;
e = HnswInitElementFromBlock(ItemPointerGetBlockNumber(indextid), ItemPointerGetOffsetNumber(indextid));
/* Calculate level based on offset */
level = element->level - i / m;
if (level < 0)
level = 0;
neighbors = &element->neighbors[level];
hc = &neighbors->items[neighbors->length++];
hc->element = e;
}
}
/*
* Load neighbors
*/
void
HnswLoadNeighbors(HnswElement element, Relation index)
{
Buffer buf;
Page page;
buf = ReadBuffer(index, element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
LoadNeighborsFromPage(element, index, page);
UnlockReleaseBuffer(buf);
}
/*
* Load an element from a tuple
*/
void
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
{
element->level = etup->level;
element->deleted = etup->deleted;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptids = NIL;
if (loadHeaptids)
{
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
/* Can stop at first invalid */
if (!ItemPointerIsValid(&etup->heaptids[i]))
break;
HnswAddHeapTid(element, &etup->heaptids[i]);
}
}
if (loadVec)
{
element->vec = palloc(VECTOR_SIZE(etup->vec.dim));
memcpy(element->vec, &etup->vec, VECTOR_SIZE(etup->vec.dim));
}
}
/*
* Load an element and optionally get its distance from q
*/
void
HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{
Buffer buf;
Page page;
HnswElementTuple etup;
/* Read vector */
buf = ReadBuffer(index, element->blkno);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, element->offno));
Assert(HnswIsElementTuple(etup));
/* Load element */
HnswLoadElementFromTuple(element, etup, true, loadVec);
/* Calculate distance */
if (distance != NULL)
*distance = (float) DatumGetFloat8(FunctionCall2Coll(procinfo, collation, *q, PointerGetDatum(&etup->vec)));
UnlockReleaseBuffer(buf);
}
/*
* Get the distance for a candidate
*/
static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, PointerGetDatum(hc->element->vec)));
}
/*
* Create a candidate for the entry point
*/
HnswCandidate *
HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = entryPoint;
if (index == NULL)
hc->distance = GetCandidateDistance(hc, q, procinfo, collation);
else
HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadVec);
return hc;
}
/*
* Compare candidate distances
*/
static int
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (((const HnswPairingHeapNode *) a)->inner->distance < ((const HnswPairingHeapNode *) b)->inner->distance)
return 1;
if (((const HnswPairingHeapNode *) a)->inner->distance > ((const HnswPairingHeapNode *) b)->inner->distance)
return -1;
return 0;
}
/*
* Compare candidate distances
*/
static int
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (((const HnswPairingHeapNode *) a)->inner->distance < ((const HnswPairingHeapNode *) b)->inner->distance)
return -1;
if (((const HnswPairingHeapNode *) a)->inner->distance > ((const HnswPairingHeapNode *) b)->inner->distance)
return 1;
return 0;
}
/*
* Create a pairing heap node for a candidate
*/
static HnswPairingHeapNode *
CreatePairingHeapNode(HnswCandidate * c)
{
HnswPairingHeapNode *node = palloc(sizeof(HnswPairingHeapNode));
node->inner = c;
return node;
}
/*
* Add to visited
*/
static inline void
AddToVisited(HTAB *v, HnswCandidate * hc, Relation index, bool *found)
{
if (index == NULL)
hash_search(v, &hc->element, HASH_ENTER, found);
else
{
ItemPointerData indextid;
ItemPointerSet(&indextid, hc->element->blkno, hc->element->offno);
hash_search(v, &indextid, HASH_ENTER, found);
}
}
/*
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, bool inserting, HnswElement skipElement)
{
ListCell *lc2;
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
HASHCTL hash_ctl;
HTAB *v;
/* Create hash table */
if (index == NULL)
{
hash_ctl.keysize = sizeof(HnswElement *);
hash_ctl.entrysize = sizeof(HnswElement *);
}
else
{
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 */
foreach(lc2, ep)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
AddToVisited(v, hc, index, NULL);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
/*
* Do not count elements being deleted towards ef when vacuuming. It
* would be ideal to do this for inserts as well, but this could
* affect insert performance.
*/
if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
wlen++;
}
while (!pairingheap_is_empty(C))
{
HnswNeighborArray *neighborhood;
HnswCandidate *c = ((HnswPairingHeapNode *) pairingheap_remove_first(C))->inner;
HnswCandidate *f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (c->distance > f->distance)
break;
if (c->element->neighbors == NULL)
HnswLoadNeighbors(c->element, index);
/* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc];
for (int i = 0; i < neighborhood->length; i++)
{
HnswCandidate *e = &neighborhood->items[i];
bool visited;
AddToVisited(v, e, index, &visited);
if (!visited)
{
float eDistance;
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (index == NULL)
eDistance = GetCandidateDistance(e, q, procinfo, collation);
else
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
Assert(!e->element->deleted);
/* Skip self for vacuuming update */
if (skipElement != NULL && e->element->blkno == skipElement->blkno && e->element->offno == skipElement->offno)
continue;
/* Make robust to issues */
if (e->element->level < lc)
continue;
if (eDistance < f->distance || wlen < ef)
{
/* Copy e */
HnswCandidate *ec = palloc(sizeof(HnswCandidate));
ec->element = e->element;
ec->distance = eDistance;
pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(ec)->ph_node));
/*
* Do not count elements being deleted towards ef when
* vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance.
*/
if (skipElement == NULL || list_length(e->element->heaptids) != 0)
{
wlen++;
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
}
}
}
}
}
/* Add each element of W to w */
while (!pairingheap_is_empty(W))
{
HnswCandidate *hc = ((HnswPairingHeapNode *) pairingheap_remove_first(W))->inner;
w = lappend(w, hc);
}
return w;
}
/*
* Calculate the distance between elements
*/
static float
HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid collation)
{
/* Look for cached distance */
if (a->neighbors != NULL)
{
Assert(a->level >= lc);
for (int i = 0; i < a->neighbors[lc].length; i++)
{
if (a->neighbors[lc].items[i].element == b)
return a->neighbors[lc].items[i].distance;
}
}
if (b->neighbors != NULL)
{
Assert(b->level >= lc);
for (int i = 0; i < b->neighbors[lc].length; i++)
{
if (b->neighbors[lc].items[i].element == a)
return b->neighbors[lc].items[i].distance;
}
}
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, PointerGetDatum(a->vec), PointerGetDatum(b->vec)));
}
/*
* Check if an element is closer to q than any element from R
*/
static bool
CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid collation)
{
ListCell *lc2;
foreach(lc2, r)
{
HnswCandidate *ri = lfirst(lc2);
float distance = HnswGetDistance(e->element, ri->element, lc, procinfo, collation);
if (distance <= e->distance)
return false;
}
return true;
}
/*
* Algorithm 4 from paper
*/
static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswCandidate * *pruned)
{
List *r = NIL;
List *w = list_copy(c);
pairingheap *wd;
if (list_length(w) <= m)
return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL);
while (list_length(w) > 0 && list_length(r) < m)
{
/* Assumes w is already ordered desc */
HnswCandidate *e = llast(w);
bool closer;
w = list_delete_last(w);
closer = CheckElementCloser(e, r, lc, procinfo, collation);
if (closer)
r = lappend(r, e);
else
pairingheap_add(wd, &(CreatePairingHeapNode(e)->ph_node));
}
/* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
/* Return pruned for update connections */
if (pruned != NULL)
{
if (!pairingheap_is_empty(wd))
*pruned = ((HnswPairingHeapNode *) pairingheap_first(wd))->inner;
else
*pruned = linitial(w);
}
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
*/
static void
AddConnections(HnswElement element, List *neighbors, int m, int lc)
{
ListCell *lc2;
HnswNeighborArray *a = &element->neighbors[lc];
foreach(lc2, neighbors)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
}
/*
* Compare candidate distances
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistances(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistances(const void *a, const void *b)
#endif
{
HnswCandidate *hca = lfirst((ListCell *) a);
HnswCandidate *hcb = lfirst((ListCell *) b);
if (hca->distance < hcb->distance)
return 1;
if (hca->distance > hcb->distance)
return -1;
return 0;
}
/*
* Update connections
*/
void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
HnswCandidate hc2;
hc2.element = element;
hc2.distance = hc->distance;
if (currentNeighbors->length < m)
{
currentNeighbors->items[currentNeighbors->length++] = hc2;
/* Track update */
if (updateIdx != NULL)
*updateIdx = -2;
}
else
{
/* Shrink connections */
HnswCandidate *pruned = NULL;
/* Load elements on insert */
if (index != NULL)
{
Datum q = PointerGetDatum(hc->element->vec);
for (int i = 0; i < currentNeighbors->length; i++)
{
HnswCandidate *hc3 = &currentNeighbors->items[i];
if (hc3->element->vec == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0)
{
pruned = &currentNeighbors->items[i];
break;
}
}
}
if (pruned == NULL)
{
List *c = NIL;
/* Add and sort candidates */
for (int i = 0; i < currentNeighbors->length; i++)
c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2);
list_sort(c, CompareCandidateDistances);
SelectNeighbors(c, m, lc, procinfo, collation, &pruned);
/* Should not happen */
if (pruned == NULL)
return;
}
/* Find and replace the pruned element */
for (int i = 0; i < currentNeighbors->length; i++)
{
if (currentNeighbors->items[i].element == pruned->element)
{
currentNeighbors->items[i] = hc2;
/* Track update */
if (updateIdx != NULL)
*updateIdx = i;
break;
}
}
}
}
/*
* Remove elements being deleted
*/
static List *
RemoveElementsBeingDeleted(List *w)
{
ListCell *lc2;
List *w2 = NIL;
foreach(lc2, w)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
if (list_length(hc->element->heaptids) != 0)
w2 = lappend(w2, hc);
}
return w2;
}
/*
* Algorithm 1 from paper
*/
void
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
{
List *ep;
List *w;
int level = element->level;
int entryLevel;
Datum q = PointerGetDatum(element->vec);
HnswElement skipElement = existing ? element : NULL;
/* No neighbors if no entry point */
if (entryPoint == NULL)
return;
/* Get entry point and level */
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, true));
entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, true, skipElement);
ep = w;
}
if (level > entryLevel)
level = entryLevel;
/* 2nd phase */
for (int lc = level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
List *neighbors;
w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, true, skipElement);
/* Elements being deleted can help with search */
/* but should be removed before selecting neighbors */
if (index != NULL)
w = RemoveElementsBeingDeleted(w);
neighbors = SelectNeighbors(w, lm, lc, procinfo, collation, NULL);
AddConnections(element, neighbors, lm, lc);
ep = w;
}
}

723
src/hnswvacuum.c Normal file
View File

@@ -0,0 +1,723 @@
#include "postgres.h"
#include <math.h>
#include "commands/vacuum.h"
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
* Check if deleted list contains an index TID
*/
static bool
DeletedContains(HTAB *deleted, ItemPointer indextid)
{
bool found;
hash_search(deleted, indextid, HASH_FIND, &found);
return found;
}
/*
* Remove deleted heap TIDs
*
* OK to remove for entry point, since always considered for searches and inserts
*/
static void
RemoveHeapTids(HnswVacuumState * vacuumstate)
{
BlockNumber blkno = HNSW_HEAD_BLKNO;
HnswElement highestPoint = &vacuumstate->highestPoint;
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
HnswElement entryPoint = HnswGetEntryPoint(vacuumstate->index);
IndexBulkDeleteResult *stats = vacuumstate->stats;
/* Store separately since highestPoint.level is uint8 */
int highestLevel = -1;
/* Initialize highest point */
highestPoint->blkno = InvalidBlockNumber;
highestPoint->offno = InvalidOffsetNumber;
while (BlockNumberIsValid(blkno))
{
Buffer buf;
Page page;
GenericXLogState *state;
OffsetNumber offno;
OffsetNumber maxoffno;
bool updated = false;
vacuum_delay_point();
buf = ReadBufferExtended(index, MAIN_FORKNUM, blkno, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
maxoffno = PageGetMaxOffsetNumber(page);
/* Iterate over nodes */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
int idx = 0;
bool itemUpdated = false;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
if (ItemPointerIsValid(&etup->heaptids[0]))
{
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
/* Stop at first unused */
if (!ItemPointerIsValid(&etup->heaptids[i]))
break;
if (vacuumstate->callback(&etup->heaptids[i], vacuumstate->callback_state))
{
itemUpdated = true;
stats->tuples_removed++;
}
else
{
/* Move to front of list */
etup->heaptids[idx++] = etup->heaptids[i];
stats->num_index_tuples++;
}
}
if (itemUpdated)
{
Size etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
/* Mark rest as invalid */
for (int i = idx; i < HNSW_HEAPTIDS; i++)
ItemPointerSetInvalid(&etup->heaptids[i]);
if (!PageIndexTupleOverwrite(page, offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
updated = true;
}
}
if (!ItemPointerIsValid(&etup->heaptids[0]))
{
ItemPointerData ip;
/* Add to deleted list */
ItemPointerSet(&ip, blkno, offno);
(void) hash_search(vacuumstate->deleted, &ip, HASH_ENTER, NULL);
}
else if (etup->level > highestLevel && !(entryPoint != NULL && blkno == entryPoint->blkno && offno == entryPoint->offno))
{
/* Keep track of highest non-entry point */
highestPoint->blkno = blkno;
highestPoint->offno = offno;
highestPoint->level = etup->level;
highestLevel = etup->level;
}
}
blkno = HnswPageGetOpaque(page)->nextblkno;
if (updated)
{
MarkBufferDirty(buf);
GenericXLogFinish(state);
}
else
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
}
/*
* Check for deleted neighbors
*/
static bool
NeedsUpdated(HnswVacuumState * vacuumstate, HnswElement element)
{
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
Buffer buf;
Page page;
HnswNeighborTuple ntup;
bool needsUpdated = false;
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
Assert(HnswIsNeighborTuple(ntup));
/* Check neighbors */
for (int i = 0; i < ntup->count; i++)
{
ItemPointer indextid = &ntup->indextids[i];
if (!ItemPointerIsValid(indextid))
continue;
/* Check if in deleted list */
if (DeletedContains(vacuumstate->deleted, indextid))
{
needsUpdated = true;
break;
}
}
/* Also update if layer 0 is not full */
/* This could indicate too many candidates being deleted during insert */
if (!needsUpdated)
needsUpdated = !ItemPointerIsValid(&ntup->indextids[ntup->count - 1]);
UnlockReleaseBuffer(buf);
return needsUpdated;
}
/*
* Repair graph for a single element
*/
static void
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint)
{
Relation index = vacuumstate->index;
Buffer buf;
Page page;
GenericXLogState *state;
int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction;
FmgrInfo *procinfo = vacuumstate->procinfo;
Oid collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
/* Skip if element is entry point */
if (entryPoint != NULL && element->blkno == entryPoint->blkno && element->offno == entryPoint->offno)
return;
/* Init fields */
HnswInitNeighbors(element, m);
element->heaptids = NIL;
/* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Update neighbor tuple */
/* Do this before getting page to minimize locking */
HnswSetNeighborTuple(ntup, element, m);
/* Get neighbor page */
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, element->neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, true);
}
/*
* Repair graph entry point
*/
static void
RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
{
Relation index = vacuumstate->index;
HnswElement highestPoint = &vacuumstate->highestPoint;
HnswElement entryPoint;
MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
if (!BlockNumberIsValid(highestPoint->blkno))
highestPoint = NULL;
/*
* Repair graph for highest non-entry point. Highest point may be outdated
* due to inserts that happen during and after RemoveHeapTids.
*/
if (highestPoint != NULL)
{
/* Get a shared lock */
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */
HnswLoadElement(highestPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
/* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint))
RepairGraphElement(vacuumstate, highestPoint, HnswGetEntryPoint(index));
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, ShareLock);
}
/* Prevent concurrent inserts when possibly updating entry point */
LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Get latest entry point */
entryPoint = HnswGetEntryPoint(index);
if (entryPoint != NULL)
{
ItemPointerData epData;
ItemPointerSet(&epData, entryPoint->blkno, entryPoint->offno);
if (DeletedContains(vacuumstate->deleted, &epData))
{
/*
* Replace the entry point with the highest point. If highest
* point is outdated and empty, the entry point will be empty
* until an element is repaired.
*/
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, highestPoint, InvalidBlockNumber, MAIN_FORKNUM);
}
else
{
/*
* Repair the entry point with the highest point. If highest point
* is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine.
*/
HnswLoadElement(entryPoint, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true);
if (NeedsUpdated(vacuumstate, entryPoint))
{
/* Reset neighbors from previous update */
if (highestPoint != NULL)
highestPoint->neighbors = NULL;
RepairGraphElement(vacuumstate, entryPoint, highestPoint);
}
}
}
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(vacuumstate->tmpCtx);
}
/*
* Repair graph for all elements
*/
static void
RepairGraph(HnswVacuumState * vacuumstate)
{
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
BlockNumber blkno = HNSW_HEAD_BLKNO;
/* Wait for inserts to complete */
LockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
UnlockPage(index, HNSW_UPDATE_LOCK, ExclusiveLock);
/* Repair entry point first */
RepairGraphEntryPoint(vacuumstate);
while (BlockNumberIsValid(blkno))
{
Buffer buf;
Page page;
OffsetNumber offno;
OffsetNumber maxoffno;
List *elements = NIL;
ListCell *lc2;
MemoryContext oldCtx;
vacuum_delay_point();
oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
buf = ReadBufferExtended(index, MAIN_FORKNUM, blkno, RBM_NORMAL, bas);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
/* Load items into memory to minimize locking */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElement element;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
/* Skip updating neighbors if being deleted */
if (!ItemPointerIsValid(&etup->heaptids[0]))
continue;
/* Create an element */
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true);
elements = lappend(elements, element);
}
blkno = HnswPageGetOpaque(page)->nextblkno;
UnlockReleaseBuffer(buf);
/* Update neighbor pages */
foreach(lc2, elements)
{
HnswElement element = (HnswElement) lfirst(lc2);
HnswElement entryPoint;
LOCKMODE lockmode = ShareLock;
/* Check if any neighbors point to deleted values */
if (!NeedsUpdated(vacuumstate, element))
continue;
/* Get a shared lock */
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Refresh entry point for each element */
entryPoint = HnswGetEntryPoint(index);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get exclusive lock */
lockmode = ExclusiveLock;
LockPage(index, HNSW_UPDATE_LOCK, lockmode);
/* Get latest entry point after lock is acquired */
entryPoint = HnswGetEntryPoint(index);
}
/* Repair connections */
RepairGraphElement(vacuumstate, element, entryPoint);
/*
* Update metapage if needed. Should only happen if entry point
* was replaced and highest point was outdated.
*/
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
}
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(vacuumstate->tmpCtx);
}
}
/*
* Mark items as deleted
*/
static void
MarkDeleted(HnswVacuumState * vacuumstate)
{
BlockNumber blkno = HNSW_HEAD_BLKNO;
BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
/* Wait for selects to complete */
LockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
UnlockPage(index, HNSW_SCAN_LOCK, ExclusiveLock);
while (BlockNumberIsValid(blkno))
{
Buffer buf;
Page page;
GenericXLogState *state;
OffsetNumber offno;
OffsetNumber maxoffno;
vacuum_delay_point();
buf = ReadBufferExtended(index, MAIN_FORKNUM, blkno, RBM_NORMAL, bas);
/*
* ambulkdelete cannot delete entries from pages that are pinned by
* other backends
*
* https://www.postgresql.org/docs/current/index-locking.html
*/
LockBufferForCleanup(buf);
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
maxoffno = PageGetMaxOffsetNumber(page);
/* Update element and neighbors together */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswNeighborTuple ntup;
Size etupSize;
Size ntupSize;
Buffer nbuf;
Page npage;
BlockNumber neighborPage;
OffsetNumber neighborOffno;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
/* Skip deleted tuples */
if (etup->deleted)
{
/* Set to first free page */
if (!BlockNumberIsValid(insertPage))
insertPage = blkno;
continue;
}
/* Skip live tuples */
if (ItemPointerIsValid(&etup->heaptids[0]))
continue;
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(etup->vec.dim);
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(etup->level, vacuumstate->m);
/* Get neighbor page */
neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
if (neighborPage == blkno)
{
nbuf = buf;
npage = page;
}
else
{
nbuf = ReadBufferExtended(index, MAIN_FORKNUM, neighborPage, RBM_NORMAL, bas);
LockBuffer(nbuf, BUFFER_LOCK_EXCLUSIVE);
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
ntup = (HnswNeighborTuple) PageGetItem(npage, PageGetItemId(npage, neighborOffno));
/* Overwrite element */
etup->deleted = 1;
MemSet(&etup->vec.x, 0, etup->vec.dim * sizeof(float));
/* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++)
ItemPointerSetInvalid(&ntup->indextids[i]);
/* Overwrite element tuple */
if (!PageIndexTupleOverwrite(page, offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Overwrite neighbor tuple */
if (!PageIndexTupleOverwrite(npage, neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
GenericXLogFinish(state);
if (nbuf != buf)
UnlockReleaseBuffer(nbuf);
/* Set to first free page */
if (!BlockNumberIsValid(insertPage))
insertPage = blkno;
/* Prepare new xlog */
state = GenericXLogStart(index);
page = GenericXLogRegisterBuffer(state, buf, 0);
}
blkno = HnswPageGetOpaque(page)->nextblkno;
GenericXLogAbort(state);
UnlockReleaseBuffer(buf);
}
/* Update insert page last, after everything has been marked as deleted */
HnswUpdateMetaPage(index, 0, NULL, insertPage, MAIN_FORKNUM);
}
/*
* Initialize the vacuum state
*/
static void
InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkDeleteResult *stats, IndexBulkDeleteCallback callback, void *callback_state)
{
Relation index = info->index;
HASHCTL hash_ctl;
if (stats == NULL)
stats = (IndexBulkDeleteResult *) palloc0(sizeof(IndexBulkDeleteResult));
vacuumstate->index = index;
vacuumstate->stats = stats;
vacuumstate->callback = callback;
vacuumstate->callback_state = callback_state;
vacuumstate->m = HnswGetM(index);
vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
vacuumstate->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
vacuumstate->collation = index->rd_indcollation[0];
vacuumstate->ntup = palloc0(BLCKSZ);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Create hash table */
hash_ctl.keysize = sizeof(ItemPointerData);
hash_ctl.entrysize = sizeof(ItemPointerData);
hash_ctl.hcxt = CurrentMemoryContext;
vacuumstate->deleted = hash_create("hnswbulkdelete indextids", 256, &hash_ctl, HASH_ELEM | HASH_BLOBS | HASH_CONTEXT);
}
/*
* Free resources
*/
static void
FreeVacuumState(HnswVacuumState * vacuumstate)
{
hash_destroy(vacuumstate->deleted);
FreeAccessStrategy(vacuumstate->bas);
pfree(vacuumstate->ntup);
MemoryContextDelete(vacuumstate->tmpCtx);
}
/*
* Print graph
*/
#ifdef HNSW_DEBUG
static void
PrintGraph(HnswVacuumState * vacuumstate)
{
BlockNumber blkno = HNSW_HEAD_BLKNO;
Relation index = vacuumstate->index;
while (BlockNumberIsValid(blkno))
{
Buffer buf;
Page page;
OffsetNumber offno;
OffsetNumber maxoffno;
buf = ReadBuffer(index, blkno);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page);
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
HnswElementTuple etup = (HnswElementTuple) PageGetItem(page, PageGetItemId(page, offno));
HnswElement element;
/* Skip neighbor tuples */
if (!HnswIsElementTuple(etup))
continue;
/* Skip deleted tuples */
if (etup->deleted)
continue;
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true);
HnswLoadNeighbors(element, index);
elog(INFO, "element (%d,%d)", element->blkno, element->offno);
for (int lc = element->level; lc >= 0; lc--)
{
HnswNeighborArray *neighbors = &element->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
{
HnswElement e = neighbors->items[i].element;
elog(INFO, "%d: (%d,%d)", lc, e->blkno, e->offno);
}
}
}
blkno = HnswPageGetOpaque(page)->nextblkno;
UnlockReleaseBuffer(buf);
}
}
#endif
/*
* Bulk delete tuples from the index
*/
IndexBulkDeleteResult *
hnswbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
IndexBulkDeleteCallback callback, void *callback_state)
{
HnswVacuumState vacuumstate;
InitVacuumState(&vacuumstate, info, stats, callback, callback_state);
/* Pass 1: Remove heap TIDs */
RemoveHeapTids(&vacuumstate);
/* Pass 2: Repair graph */
RepairGraph(&vacuumstate);
/* Pass 3: Mark as deleted */
MarkDeleted(&vacuumstate);
#ifdef HNSW_DEBUG
PrintGraph(&vacuumstate);
#endif
FreeVacuumState(&vacuumstate);
return vacuumstate.stats;
}
/*
* Clean up after a VACUUM operation
*/
IndexBulkDeleteResult *
hnswvacuumcleanup(IndexVacuumInfo *info, IndexBulkDeleteResult *stats)
{
Relation rel = info->index;
if (info->analyze_only)
return stats;
/* stats is NULL if ambulkdelete not called */
/* OK to return NULL if index not changed */
if (stats == NULL)
return NULL;
stats->num_pages = RelationGetNumberOfBlocks(rel);
return stats;
}

View File

@@ -2,11 +2,16 @@
#include <float.h>
#include "access/parallel.h"
#include "access/xact.h"
#include "catalog/index.h"
#include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h"
#include "ivfflat.h"
#include "miscadmin.h"
#include "storage/bufmgr.h"
#include "utils/memutils.h"
#include "tcop/tcopprot.h"
#if PG_VERSION_NUM >= 140000
#include "utils/backend_progress.h"
@@ -23,9 +28,6 @@
#define PROGRESS_CREATEIDX_TUPLES_DONE 0
#endif
#include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h"
#if PG_VERSION_NUM >= 130000
#define CALLBACK_ITEM_POINTER ItemPointer tid
#else
@@ -38,6 +40,25 @@
#define UpdateProgress(index, val) ((void)val)
#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_IVFFLAT_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_TUPLESORT UINT64CONST(0xA000000000000002)
#define PARALLEL_KEY_IVFFLAT_CENTERS UINT64CONST(0xA000000000000003)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000004)
/*
* Add sample
*/
@@ -150,7 +171,6 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
int closestCenter = 0;
VectorArray centers = buildstate->centers;
TupleTableSlot *slot = buildstate->slot;
int i;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -163,7 +183,7 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState
}
/* Find the list that minimizes the distance */
for (i = 0; i < centers->length; i++)
for (int i = 0; i < centers->length; i++)
{
distance = DatumGetFloat8(FunctionCall2Coll(buildstate->procinfo, buildstate->collation, value, PointerGetDatum(VectorArrayGet(centers, i))));
@@ -258,15 +278,8 @@ GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot,
static void
InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
{
Buffer buf;
Page page;
GenericXLogState *state;
int list;
IndexTuple itup = NULL; /* silence compiler warning */
BlockNumber startPage;
BlockNumber insertPage;
Size itemsz;
int i;
int64 inserted = 0;
#if PG_VERSION_NUM >= 120000
@@ -282,8 +295,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
for (i = 0; i < buildstate->centers->length; i++)
for (int i = 0; i < buildstate->centers->length; i++)
{
Buffer buf;
Page page;
GenericXLogState *state;
BlockNumber startPage;
BlockNumber insertPage;
/* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */
CHECK_FOR_INTERRUPTS();
@@ -297,7 +316,8 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
while (list == i)
{
/* Check for free space */
itemsz = MAXALIGN(IndexTupleSize(itup));
Size itemsz = MAXALIGN(IndexTupleSize(itup));
if (PageGetFreeSpace(page) < itemsz)
IvfflatAppendPage(index, &buf, &page, &state, forkNum);
@@ -317,7 +337,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IvfflatCommitBuffer(buf, state);
/* Set the start and insert pages */
IvfflatUpdateList(index, state, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
}
}
@@ -351,9 +371,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->collation = index->rd_indcollation[0];
/* Require more than one dimension for spherical k-means */
/* Lists check for backwards compatibility */
/* TODO Remove lists check in 0.3.0 */
if (buildstate->kmeansnormprocinfo != NULL && buildstate->dimensions == 1 && buildstate->lists > 1)
if (buildstate->kmeansnormprocinfo != NULL && buildstate->dimensions == 1)
elog(ERROR, "dimensions must be greater than one for this opclass");
/* Create tuple description for sorting */
@@ -387,6 +405,8 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->listSums = palloc0(sizeof(double) * buildstate->lists);
buildstate->listCounts = palloc0(sizeof(int) * buildstate->lists);
#endif
buildstate->ivfleader = NULL;
}
/*
@@ -483,7 +503,6 @@ static void
CreateListPages(Relation index, VectorArray centers, int dimensions,
int lists, ForkNumber forkNum, ListInfo * *listInfo)
{
int i;
Buffer buf;
Page page;
GenericXLogState *state;
@@ -497,7 +516,7 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
buf = IvfflatNewBuffer(index, forkNum);
IvfflatInitRegisterPage(index, &buf, &page, &state);
for (i = 0; i < lists; i++)
for (int i = 0; i < lists; i++)
{
/* Load list */
list->startPage = InvalidBlockNumber;
@@ -533,7 +552,7 @@ PrintKmeansMetrics(IvfflatBuildState * buildstate)
elog(INFO, "inertia: %.3e", buildstate->inertia);
/* Calculate Davies-Bouldin index */
if (buildstate->lists > 1)
if (buildstate->lists > 1 && !buildstate->ivfleader)
{
double db = 0.0;
@@ -568,19 +587,455 @@ PrintKmeansMetrics(IvfflatBuildState * buildstate)
}
#endif
/*
* Within leader, wait for end of heap scan
*/
static double
ParallelHeapScan(IvfflatBuildState * buildstate)
{
IvfflatShared *ivfshared = buildstate->ivfleader->ivfshared;
int nparticipanttuplesorts;
double reltuples;
nparticipanttuplesorts = buildstate->ivfleader->nparticipanttuplesorts;
for (;;)
{
SpinLockAcquire(&ivfshared->mutex);
if (ivfshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->indtuples = ivfshared->indtuples;
reltuples = ivfshared->reltuples;
#ifdef IVFFLAT_KMEANS_DEBUG
buildstate->inertia = ivfshared->inertia;
#endif
SpinLockRelease(&ivfshared->mutex);
break;
}
SpinLockRelease(&ivfshared->mutex);
ConditionVariableSleep(&ivfshared->workersdonecv,
WAIT_EVENT_PARALLEL_CREATE_INDEX_SCAN);
}
ConditionVariableCancelSleep();
return reltuples;
}
/*
* Perform a worker's portion of a parallel sort
*/
static void
IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, Sharedsort *sharedsort, Vector * ivfcenters, int sortmem, bool progress)
{
SortCoordinate coordinate;
IvfflatBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples;
IndexInfo *indexInfo;
/* Sort options, which must match AssignTuples */
AttrNumber attNums[] = {1};
Oid sortOperators[] = {Int4LessOperator};
Oid sortCollations[] = {InvalidOid};
bool nullsFirstFlags[] = {false};
/* Initialize local tuplesort coordination state */
coordinate = palloc0(sizeof(SortCoordinateData));
coordinate->isWorker = true;
coordinate->nParticipants = -1;
coordinate->sharedsort = sharedsort;
/* Join parallel scan */
indexInfo = BuildIndexInfo(ivfspool->index);
indexInfo->ii_Concurrent = ivfshared->isconcurrent;
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, VECTOR_SIZE(buildstate.centers->dim) * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = tuplesort_begin_heap(buildstate.tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, sortmem, coordinate, false);
buildstate.sortstate = ivfspool->sortstate;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared));
reltuples = table_index_build_scan(ivfspool->heap, ivfspool->index, indexInfo,
true, progress, BuildCallback,
(void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(ivfspool->heap, &ivfshared->heapdesc);
reltuples = IndexBuildHeapScan(ivfspool->heap, ivfspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Execute this worker's part of the sort */
tuplesort_performsort(ivfspool->sortstate);
/* Record statistics */
SpinLockAcquire(&ivfshared->mutex);
ivfshared->nparticipantsdone++;
ivfshared->reltuples += reltuples;
ivfshared->indtuples += buildstate.indtuples;
#ifdef IVFFLAT_KMEANS_DEBUG
ivfshared->inertia += buildstate.inertia;
#endif
SpinLockRelease(&ivfshared->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(&ivfshared->workersdonecv);
/* We can end tuplesorts immediately */
tuplesort_end(ivfspool->sortstate);
FreeBuildState(&buildstate);
}
/*
* Perform work within a launched parallel process
*/
void
IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{
char *sharedquery;
IvfflatSpool *ivfspool;
IvfflatShared *ivfshared;
Sharedsort *sharedsort;
Vector *ivfcenters;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
LOCKMODE indexLockmode;
int sortmem;
/* 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 */
ivfshared = shm_toc_lookup(toc, PARALLEL_KEY_IVFFLAT_SHARED, false);
/* Open relations using lock modes known to be obtained by index.c */
if (!ivfshared->isconcurrent)
{
heapLockmode = ShareLock;
indexLockmode = AccessExclusiveLock;
}
else
{
heapLockmode = ShareUpdateExclusiveLock;
indexLockmode = RowExclusiveLock;
}
/* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(ivfshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(ivfshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(ivfshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */
ivfspool = (IvfflatSpool *) palloc0(sizeof(IvfflatSpool));
ivfspool->heap = heapRel;
ivfspool->index = indexRel;
/* Look up shared state private to tuplesort.c */
sharedsort = shm_toc_lookup(toc, PARALLEL_KEY_TUPLESORT, false);
tuplesort_attach_shared(sharedsort, seg);
ivfcenters = shm_toc_lookup(toc, PARALLEL_KEY_IVFFLAT_CENTERS, false);
/* Perform sorting */
sortmem = maintenance_work_mem / ivfshared->scantuplesortstates;
IvfflatParallelScanAndSort(ivfspool, ivfshared, sharedsort, ivfcenters, sortmem, 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
IvfflatEndParallel(IvfflatLeader * ivfleader)
{
/* Shutdown worker processes */
WaitForParallelWorkersToFinish(ivfleader->pcxt);
/* Free last reference to MVCC snapshot, if one was used */
if (IsMVCCSnapshot(ivfleader->snapshot))
UnregisterSnapshot(ivfleader->snapshot);
DestroyParallelContext(ivfleader->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(IvfflatShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(IvfflatShared);
}
return add_size(offsetof(IvfflatShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
}
/*
* Within leader, participate as a parallel worker
*/
static void
IvfflatLeaderParticipateAsWorker(IvfflatBuildState * buildstate)
{
IvfflatLeader *ivfleader = buildstate->ivfleader;
IvfflatSpool *leaderworker;
int sortmem;
/* Allocate memory and initialize private spool */
leaderworker = (IvfflatSpool *) palloc0(sizeof(IvfflatSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
sortmem = maintenance_work_mem / ivfleader->nparticipanttuplesorts;
IvfflatParallelScanAndSort(leaderworker, ivfleader->ivfshared,
ivfleader->sharedsort, ivfleader->ivfcenters,
sortmem, true);
}
/*
* Begin parallel build
*/
static void
IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int request)
{
ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot;
Size estivfshared;
Size estsort;
Size estcenters;
IvfflatShared *ivfshared;
Sharedsort *sharedsort;
Vector *ivfcenters;
IvfflatLeader *ivfleader = (IvfflatLeader *) palloc0(sizeof(IvfflatLeader));
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", "IvfflatParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "IvfflatParallelBuildMain", 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 */
estivfshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, estivfshared);
estsort = tuplesort_estimate_shared(scantuplesortstates);
shm_toc_estimate_chunk(&pcxt->estimator, estsort);
estcenters = VECTOR_SIZE(buildstate->dimensions) * buildstate->lists;
shm_toc_estimate_chunk(&pcxt->estimator, estcenters);
shm_toc_estimate_keys(&pcxt->estimator, 3);
/* 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 */
ivfshared = (IvfflatShared *) shm_toc_allocate(pcxt->toc, estivfshared);
/* Initialize immutable state */
ivfshared->heaprelid = RelationGetRelid(buildstate->heap);
ivfshared->indexrelid = RelationGetRelid(buildstate->index);
ivfshared->isconcurrent = isconcurrent;
ivfshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&ivfshared->workersdonecv);
SpinLockInit(&ivfshared->mutex);
/* Initialize mutable state */
ivfshared->nparticipantsdone = 0;
ivfshared->reltuples = 0;
ivfshared->indtuples = 0;
#ifdef IVFFLAT_KMEANS_DEBUG
ivfshared->inertia = 0;
#endif
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromIvfflatShared(ivfshared),
snapshot);
#else
heap_parallelscan_initialize(&ivfshared->heapdesc, buildstate->heap, snapshot);
#endif
/* Store shared tuplesort-private state, for which we reserved space */
sharedsort = (Sharedsort *) shm_toc_allocate(pcxt->toc, estsort);
tuplesort_initialize_shared(sharedsort, scantuplesortstates,
pcxt->seg);
ivfcenters = (Vector *) shm_toc_allocate(pcxt->toc, estcenters);
memcpy(ivfcenters, buildstate->centers->items, estcenters);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_IVFFLAT_SHARED, ivfshared);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_TUPLESORT, sharedsort);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_IVFFLAT_CENTERS, ivfcenters);
/* 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);
ivfleader->pcxt = pcxt;
ivfleader->nparticipanttuplesorts = pcxt->nworkers_launched;
if (leaderparticipates)
ivfleader->nparticipanttuplesorts++;
ivfleader->ivfshared = ivfshared;
ivfleader->sharedsort = sharedsort;
ivfleader->snapshot = snapshot;
ivfleader->ivfcenters = ivfcenters;
/* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0)
{
IvfflatEndParallel(ivfleader);
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->ivfleader = ivfleader;
/* Join heap scan ourselves */
if (leaderparticipates)
IvfflatLeaderParticipateAsWorker(buildstate);
/* Wait for all launched workers */
WaitForParallelWorkersToAttach(pcxt);
}
/*
* Scan table for tuples to index
*/
static void
ScanTable(IvfflatBuildState * buildstate)
AssignTuples(IvfflatBuildState * buildstate)
{
int parallel_workers = 0;
SortCoordinate coordinate = NULL;
/* Sort options, which must match IvfflatParallelScanAndSort */
AttrNumber attNums[] = {1};
Oid sortOperators[] = {Int4LessOperator};
Oid sortCollations[] = {InvalidOid};
bool nullsFirstFlags[] = {false};
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_ASSIGN);
/* Calculate parallel workers */
if (buildstate->heap != NULL)
parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
IvfflatBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
/* Set up coordination state if at least one worker launched */
if (buildstate->ivfleader)
{
coordinate = (SortCoordinate) palloc0(sizeof(SortCoordinateData));
coordinate->isWorker = false;
coordinate->nParticipants = buildstate->ivfleader->nparticipanttuplesorts;
coordinate->sharedsort = buildstate->ivfleader->sharedsort;
}
/* Begin serial/leader tuplesort */
buildstate->sortstate = tuplesort_begin_heap(buildstate->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, maintenance_work_mem, coordinate, false);
/* Add tuples to sort */
if (buildstate->heap != NULL)
{
if (buildstate->ivfleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL);
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
#ifdef IVFFLAT_KMEANS_DEBUG
PrintKmeansMetrics(buildstate);
#endif
}
}
/*
@@ -589,29 +1044,21 @@ ScanTable(IvfflatBuildState * buildstate)
static void
CreateEntryPages(IvfflatBuildState * buildstate, ForkNumber forkNum)
{
AttrNumber attNums[] = {1};
Oid sortOperators[] = {Int4LessOperator};
Oid sortCollations[] = {InvalidOid};
bool nullsFirstFlags[] = {false};
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_SORT);
buildstate->sortstate = tuplesort_begin_heap(buildstate->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, maintenance_work_mem, NULL, false);
/* Add tuples to sort */
if (buildstate->heap != NULL)
IvfflatBench("assign tuples", ScanTable(buildstate));
/* Assign */
IvfflatBench("assign tuples", AssignTuples(buildstate));
/* Sort */
IvfflatBench("sort tuples", tuplesort_performsort(buildstate->sortstate));
#ifdef IVFFLAT_KMEANS_DEBUG
PrintKmeansMetrics(buildstate);
#endif
/* Insert */
/* Load */
IvfflatBench("load tuples", InsertTuples(buildstate->index, buildstate, forkNum));
/* End sort */
tuplesort_end(buildstate->sortstate);
/* End parallel build */
if (buildstate->ivfleader)
IvfflatEndParallel(buildstate->ivfleader);
}
/*

View File

@@ -20,11 +20,11 @@ static relopt_kind ivfflat_relopt_kind;
* Initialize index options and variables
*/
void
_PG_init(void)
IvfflatInit(void)
{
ivfflat_relopt_kind = add_reloption_kind();
add_int_reloption(ivfflat_relopt_kind, "lists", "Number of inverted lists",
IVFFLAT_DEFAULT_LISTS, 1, IVFFLAT_MAX_LISTS
IVFFLAT_DEFAULT_LISTS, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS
#if PG_VERSION_NUM >= 130000
,AccessExclusiveLock
#endif
@@ -32,7 +32,7 @@ _PG_init(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
"Valid range is 1..lists.", &ivfflat_probes,
1, 1, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
}
/*
@@ -48,8 +48,8 @@ ivfflatbuildphasename(int64 phasenum)
return "initializing";
case PROGRESS_IVFFLAT_PHASE_KMEANS:
return "performing k-means";
case PROGRESS_IVFFLAT_PHASE_SORT:
return "sorting tuples";
case PROGRESS_IVFFLAT_PHASE_ASSIGN:
return "assigning tuples";
case PROGRESS_IVFFLAT_PHASE_LOAD:
return "loading tuples";
default:

View File

@@ -3,14 +3,11 @@
#include "postgres.h"
#if PG_VERSION_NUM < 110000
#error "Requires PostgreSQL 11+"
#endif
#include "access/generic_xlog.h"
#include "access/parallel.h"
#include "access/reloptions.h"
#include "nodes/execnodes.h"
#include "port.h" /* for strtof() and random() */
#include "port.h" /* for random() */
#include "utils/sampling.h"
#include "utils/tuplesort.h"
#include "vector.h"
@@ -19,6 +16,10 @@
#include "common/pg_prng.h"
#endif
#if PG_VERSION_NUM < 120000
#include "access/relscan.h"
#endif
#ifdef IVFFLAT_BENCH
#include "portability/instr_time.h"
#endif
@@ -39,13 +40,16 @@
#define IVFFLAT_METAPAGE_BLKNO 0
#define IVFFLAT_HEAD_BLKNO 1 /* first list page */
/* IVFFlat parameters */
#define IVFFLAT_DEFAULT_LISTS 100
#define IVFFLAT_MIN_LISTS 1
#define IVFFLAT_MAX_LISTS 32768
#define IVFFLAT_DEFAULT_PROBES 1
/* Build phases */
/* PROGRESS_CREATEIDX_SUBPHASE_INITIALIZE is 1 */
#define PROGRESS_IVFFLAT_PHASE_KMEANS 2
#define PROGRESS_IVFFLAT_PHASE_SORT 3
#define PROGRESS_IVFFLAT_PHASE_ASSIGN 3
#define PROGRESS_IVFFLAT_PHASE_LOAD 4
#define IVFFLAT_LIST_SIZE(_dim) (offsetof(IvfflatListData, center) + VECTOR_SIZE(_dim))
@@ -79,9 +83,6 @@
/* Variables */
extern int ivfflat_probes;
/* Exported functions */
PGDLLEXPORT void _PG_init(void);
typedef struct VectorArrayData
{
int length;
@@ -105,6 +106,56 @@ typedef struct IvfflatOptions
int lists; /* number of lists */
} IvfflatOptions;
typedef struct IvfflatSpool
{
Tuplesortstate *sortstate;
Relation heap;
Relation index;
} IvfflatSpool;
typedef struct IvfflatShared
{
/* Immutable state */
Oid heaprelid;
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
/* Worker progress */
ConditionVariable workersdonecv;
/* Mutex for mutable state */
slock_t mutex;
/* Mutable state */
int nparticipantsdone;
double reltuples;
double indtuples;
#ifdef IVFFLAT_KMEANS_DEBUG
double inertia;
#endif
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} IvfflatShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromIvfflatShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(IvfflatShared)))
#endif
typedef struct IvfflatLeader
{
ParallelContext *pcxt;
int nparticipanttuplesorts;
IvfflatShared *ivfshared;
Sharedsort *sharedsort;
Snapshot snapshot;
Vector *ivfcenters;
} IvfflatLeader;
typedef struct IvfflatBuildState
{
/* Info */
@@ -150,6 +201,9 @@ typedef struct IvfflatBuildState
/* Memory */
MemoryContext tmpCtx;
/* Parallel builds */
IvfflatLeader *ivfleader;
} IvfflatBuildState;
typedef struct IvfflatMetaPageData
@@ -224,12 +278,14 @@ void IvfflatKmeans(Relation index, VectorArray samples, VectorArray centers);
FmgrInfo *IvfflatOptionalProcInfo(Relation rel, uint16 procnum);
bool IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result);
int IvfflatGetLists(Relation index);
void IvfflatUpdateList(Relation index, GenericXLogState *state, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);
void IvfflatUpdateList(Relation index, ListInfo listInfo, BlockNumber insertPage, BlockNumber originalInsertPage, BlockNumber startPage, ForkNumber forkNum);
void IvfflatCommitBuffer(Buffer buf, GenericXLogState *state);
void IvfflatAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state, ForkNumber forkNum);
Buffer IvfflatNewBuffer(Relation index, ForkNumber forkNum);
void IvfflatInitPage(Buffer buf, Page page);
void IvfflatInitRegisterPage(Relation index, Buffer *buf, Page *page, GenericXLogState **state);
void IvfflatInit(void);
PGDLLEXPORT void IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc);
/* Index access methods */
IndexBuildResult *ivfflatbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -4,6 +4,7 @@
#include "ivfflat.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/memutils.h"
/*
@@ -121,23 +122,16 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
}
else
{
Buffer metabuf;
Buffer newbuf;
Page newpage;
/*
* From ReadBufferExtended: Caller is responsible for ensuring
* that only one backend tries to extend a relation at the same
* time!
*/
metabuf = ReadBuffer(rel, IVFFLAT_METAPAGE_BLKNO);
LockBuffer(metabuf, BUFFER_LOCK_EXCLUSIVE);
/* Add a new page */
LockRelationForExtension(rel, ExclusiveLock);
newbuf = IvfflatNewBuffer(rel, MAIN_FORKNUM);
newpage = GenericXLogRegisterBuffer(state, newbuf, GENERIC_XLOG_FULL_IMAGE);
UnlockRelationForExtension(rel, ExclusiveLock);
/* Init new page */
newpage = GenericXLogRegisterBuffer(state, newbuf, GENERIC_XLOG_FULL_IMAGE);
IvfflatInitPage(newbuf, newpage);
/* Update insert page */
@@ -151,9 +145,6 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
MarkBufferDirty(buf);
GenericXLogFinish(state);
/* Unlock extend relation lock as early as possible */
UnlockReleaseBuffer(metabuf);
/* Unlock previous buffer */
UnlockReleaseBuffer(buf);
@@ -173,7 +164,7 @@ InsertTuple(Relation rel, Datum *values, bool *isnull, ItemPointer heap_tid, Rel
/* Update the insert page */
if (insertPage != originalInsertPage)
IvfflatUpdateList(rel, state, listInfo, insertPage, originalInsertPage, InvalidBlockNumber, MAIN_FORKNUM);
IvfflatUpdateList(rel, listInfo, insertPage, originalInsertPage, InvalidBlockNumber, MAIN_FORKNUM);
}
/*

View File

@@ -16,7 +16,6 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
{
FmgrInfo *procinfo;
Oid collation;
int i;
int64 j;
double distance;
double sum;
@@ -36,7 +35,7 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
for (j = 0; j < numSamples; j++)
weight[j] = DBL_MAX;
for (i = 0; i < numCenters; i++)
for (int i = 0; i < numCenters; i++)
{
CHECK_FOR_INTERRUPTS();
@@ -88,13 +87,12 @@ InitCenters(Relation index, VectorArray samples, VectorArray centers, float *low
static inline void
ApplyNorm(FmgrInfo *normprocinfo, Oid collation, Vector * vec)
{
int i;
double norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, PointerGetDatum(vec)));
/* TODO Handle zero norm */
if (norm > 0)
{
for (i = 0; i < vec->dim; i++)
for (int i = 0; i < vec->dim; i++)
vec->x[i] /= norm;
}
}
@@ -114,8 +112,6 @@ CompareVectors(const void *a, const void *b)
static void
QuickCenters(Relation index, VectorArray samples, VectorArray centers)
{
int i;
int j;
Vector *vec;
int dimensions = centers->dim;
Oid collation = index->rd_indcollation[0];
@@ -125,7 +121,7 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
if (samples->length > 0)
{
qsort(samples->items, samples->length, VECTOR_SIZE(samples->dim), CompareVectors);
for (i = 0; i < samples->length; i++)
for (int i = 0; i < samples->length; i++)
{
vec = VectorArrayGet(samples, i);
@@ -145,7 +141,7 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
SET_VARSIZE(vec, VECTOR_SIZE(dimensions));
vec->dim = dimensions;
for (j = 0; j < dimensions; j++)
for (int j = 0; j < dimensions; j++)
vec->x[j] = RandomDouble();
/* Normalize if needed (only needed for random centers) */
@@ -471,19 +467,17 @@ CheckCenters(Relation index, VectorArray centers)
FmgrInfo *normprocinfo;
Oid collation;
Vector *vec;
int i;
int j;
double norm;
if (centers->length != centers->maxlen)
elog(ERROR, "Not enough centers. Please report a bug.");
/* Ensure no NaN or infinite values */
for (i = 0; i < centers->length; i++)
for (int i = 0; i < centers->length; i++)
{
vec = VectorArrayGet(centers, i);
for (j = 0; j < vec->dim; j++)
for (int j = 0; j < vec->dim; j++)
{
if (isnan(vec->x[j]))
elog(ERROR, "NaN detected. Please report a bug.");
@@ -496,7 +490,7 @@ CheckCenters(Relation index, VectorArray centers)
/* Ensure no duplicate centers */
/* Fine to sort in-place */
qsort(centers->items, centers->length, VECTOR_SIZE(centers->dim), CompareVectors);
for (i = 1; i < centers->length; i++)
for (int i = 1; i < centers->length; i++)
{
if (CompareVectors(VectorArrayGet(centers, i), VectorArrayGet(centers, i - 1)) == 0)
elog(ERROR, "Duplicate centers detected. Please report a bug.");
@@ -509,7 +503,7 @@ CheckCenters(Relation index, VectorArray centers)
{
collation = index->rd_indcollation[0];
for (i = 0; i < centers->length; i++)
for (int i = 0; i < centers->length; i++)
{
norm = DatumGetFloat8(FunctionCall1Coll(normprocinfo, collation, PointerGetDatum(VectorArrayGet(centers, i))));
if (norm == 0)

View File

@@ -3,14 +3,13 @@
#include <float.h>
#include "access/relscan.h"
#include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h"
#include "ivfflat.h"
#include "miscadmin.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
#include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h"
/*
* Compare list distances
*/
@@ -172,7 +171,6 @@ GetScanItems(IndexScanDesc scan, Datum value)
FreeAccessStrategy(bas);
/* TODO Scan more lists */
if (tuples < 100)
ereport(DEBUG1,
(errmsg("index scan found few tuples"),
@@ -182,6 +180,29 @@ GetScanItems(IndexScanDesc scan, Datum value)
tuplesort_performsort(so->sortstate);
}
/*
* Get dimensions from metapage
*/
static int
GetDimensions(Relation index)
{
Buffer buf;
Page page;
IvfflatMetaPage metap;
int dimensions;
buf = ReadBuffer(index, IVFFLAT_METAPAGE_BLKNO);
LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf);
metap = IvfflatPageGetMeta(page);
dimensions = metap->dimensions;
UnlockReleaseBuffer(buf);
return dimensions;
}
/*
* Prepare for an index scan
*/
@@ -287,21 +308,19 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
if (scan->orderByData == NULL)
elog(ERROR, "cannot scan ivfflat index without order");
/* No items will match if null */
if (scan->orderByData->sk_flags & SK_ISNULL)
return false;
value = scan->orderByData->sk_argument;
/* Value should not be compressed or toasted */
Assert(!VARATT_IS_COMPRESSED(DatumGetPointer(value)));
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
if (so->normprocinfo != NULL)
value = PointerGetDatum(InitVector(GetDimensions(scan->indexRelation)));
else
{
/* No items will match if normalization fails */
if (!IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL))
return false;
value = scan->orderByData->sk_argument;
/* Value should not be compressed or toasted */
Assert(!VARATT_IS_COMPRESSED(DatumGetPointer(value)));
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Fine if normalization fails */
if (so->normprocinfo != NULL)
IvfflatNormValue(so->normprocinfo, so->collation, &value, NULL);
}
IvfflatBench("GetScanLists", GetScanLists(scan, value));

View File

@@ -35,9 +35,7 @@ VectorArrayFree(VectorArray arr)
void
PrintVectorArray(char *msg, VectorArray arr)
{
int i;
for (i = 0; i < arr->length; i++)
for (int i = 0; i < arr->length; i++)
PrintVector(msg, VectorArrayGet(arr, i));
}
@@ -78,20 +76,16 @@ IvfflatOptionalProcInfo(Relation rel, uint16 procnum)
bool
IvfflatNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * result)
{
Vector *v;
int i;
double norm;
norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
double norm = DatumGetFloat8(FunctionCall1Coll(procinfo, collation, *value));
if (norm > 0)
{
v = DatumGetVector(*value);
Vector *v = DatumGetVector(*value);
if (result == NULL)
result = InitVector(v->dim);
for (i = 0; i < v->dim; i++)
for (int i = 0; i < v->dim; i++)
result->x[i] = v->x[i] / norm;
*value = PointerGetDatum(result);
@@ -182,12 +176,13 @@ IvfflatAppendPage(Relation index, Buffer *buf, Page *page, GenericXLogState **st
* Update the start or insert page of a list
*/
void
IvfflatUpdateList(Relation index, GenericXLogState *state, ListInfo listInfo,
IvfflatUpdateList(Relation index, ListInfo listInfo,
BlockNumber insertPage, BlockNumber originalInsertPage,
BlockNumber startPage, ForkNumber forkNum)
{
Buffer buf;
Page page;
GenericXLogState *state;
IvfflatList list;
bool changed = false;

View File

@@ -12,34 +12,23 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
IndexBulkDeleteCallback callback, void *callback_state)
{
Relation index = info->index;
Buffer cbuf;
Page cpage;
Buffer buf;
Page page;
IvfflatList list;
IndexTuple itup;
ItemPointer htup;
OffsetNumber deletable[MaxOffsetNumber];
int ndeletable;
BlockNumber startPages[MaxOffsetNumber];
BlockNumber nextblkno = IVFFLAT_HEAD_BLKNO;
BlockNumber searchPage;
BlockNumber insertPage;
GenericXLogState *state;
OffsetNumber coffno;
OffsetNumber cmaxoffno;
OffsetNumber offno;
OffsetNumber maxoffno;
ListInfo listInfo;
BlockNumber blkno = IVFFLAT_HEAD_BLKNO;
BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
if (stats == NULL)
stats = (IndexBulkDeleteResult *) palloc0(sizeof(IndexBulkDeleteResult));
/* Iterate over list pages */
while (BlockNumberIsValid(nextblkno))
while (BlockNumberIsValid(blkno))
{
cbuf = ReadBuffer(index, nextblkno);
Buffer cbuf;
Page cpage;
OffsetNumber coffno;
OffsetNumber cmaxoffno;
BlockNumber startPages[MaxOffsetNumber];
ListInfo listInfo;
cbuf = ReadBuffer(index, blkno);
LockBuffer(cbuf, BUFFER_LOCK_SHARE);
cpage = BufferGetPage(cbuf);
@@ -48,23 +37,32 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
/* Iterate over lists */
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
{
list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
startPages[coffno - FirstOffsetNumber] = list->startPage;
}
listInfo.blkno = nextblkno;
nextblkno = IvfflatPageGetOpaque(cpage)->nextblkno;
listInfo.blkno = blkno;
blkno = IvfflatPageGetOpaque(cpage)->nextblkno;
UnlockReleaseBuffer(cbuf);
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
{
searchPage = startPages[coffno - FirstOffsetNumber];
insertPage = InvalidBlockNumber;
BlockNumber searchPage = startPages[coffno - FirstOffsetNumber];
BlockNumber insertPage = InvalidBlockNumber;
/* Iterate over entry pages */
while (BlockNumberIsValid(searchPage))
{
Buffer buf;
Page page;
GenericXLogState *state;
OffsetNumber offno;
OffsetNumber maxoffno;
OffsetNumber deletable[MaxOffsetNumber];
int ndeletable;
vacuum_delay_point();
buf = ReadBufferExtended(index, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
@@ -86,8 +84,8 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
/* Find deleted tuples */
for (offno = FirstOffsetNumber; offno <= maxoffno; offno = OffsetNumberNext(offno))
{
itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
htup = &(itup->t_tid);
IndexTuple itup = (IndexTuple) PageGetItem(page, PageGetItemId(page, offno));
ItemPointer htup = &(itup->t_tid);
if (callback(htup, callback_state))
{
@@ -127,7 +125,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
if (BlockNumberIsValid(insertPage))
{
listInfo.offno = coffno;
IvfflatUpdateList(index, state, listInfo, insertPage, InvalidBlockNumber, InvalidBlockNumber, MAIN_FORKNUM);
IvfflatUpdateList(index, listInfo, insertPage, InvalidBlockNumber, InvalidBlockNumber, MAIN_FORKNUM);
}
}
}

View File

@@ -2,15 +2,22 @@
#include <math.h>
#include "vector.h"
#include "fmgr.h"
#include "catalog/pg_type.h"
#include "fmgr.h"
#include "hnsw.h"
#include "ivfflat.h"
#include "lib/stringinfo.h"
#include "libpq/pqformat.h"
#include "port.h" /* for strtof() */
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/lsyscache.h"
#include "utils/numeric.h"
#include "vector.h"
#if PG_VERSION_NUM >= 160000
#include "varatt.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "common/shortest_dec.h"
@@ -29,6 +36,17 @@
PG_MODULE_MAGIC;
/*
* Initialize index options and variables
*/
PGDLLEXPORT void _PG_init(void);
void
_PG_init(void)
{
HnswInit();
IvfflatInit();
}
/*
* Ensure same dimensions
*/
@@ -87,6 +105,23 @@ CheckElement(float value)
errmsg("infinite value not allowed in vector")));
}
/*
* Allocate and initialize a new vector
*/
Vector *
InitVector(int dim)
{
Vector *result;
int size;
size = VECTOR_SIZE(dim);
result = (Vector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
/*
* Check for whitespace, since array_isspace() is static
*/
@@ -125,6 +160,14 @@ float_overflow_error(void)
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value out of range: overflow")));
}
static pg_noinline void
float_underflow_error(void)
{
ereport(ERROR,
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
errmsg("value out of range: underflow")));
}
#endif
/*
@@ -136,7 +179,6 @@ vector_in(PG_FUNCTION_ARGS)
{
char *str = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2);
int i;
float x[VECTOR_MAX_DIM];
int dim = 0;
char *pt;
@@ -231,7 +273,7 @@ vector_in(PG_FUNCTION_ARGS)
CheckExpectedDim(typmod, dim);
result = InitVector(dim);
for (i = 0; i < dim; i++)
for (int i = 0; i < dim; i++)
result->x[i] = x[i];
PG_RETURN_POINTER(result);
@@ -248,7 +290,6 @@ vector_out(PG_FUNCTION_ARGS)
int dim = vector->dim;
char *buf;
char *ptr;
int i;
int n;
#if PG_VERSION_NUM < 120000
@@ -275,7 +316,7 @@ vector_out(PG_FUNCTION_ARGS)
*ptr = '[';
ptr++;
for (i = 0; i < dim; i++)
for (int i = 0; i < dim; i++)
{
if (i > 0)
{
@@ -353,7 +394,6 @@ vector_recv(PG_FUNCTION_ARGS)
Vector *result;
int16 dim;
int16 unused;
int i;
dim = pq_getmsgint(buf, sizeof(int16));
unused = pq_getmsgint(buf, sizeof(int16));
@@ -367,7 +407,7 @@ vector_recv(PG_FUNCTION_ARGS)
errmsg("expected unused to be 0, not %d", unused)));
result = InitVector(dim);
for (i = 0; i < dim; i++)
for (int i = 0; i < dim; i++)
{
result->x[i] = pq_getmsgfloat4(buf);
CheckElement(result->x[i]);
@@ -385,12 +425,11 @@ vector_send(PG_FUNCTION_ARGS)
{
Vector *vec = PG_GETARG_VECTOR_P(0);
StringInfoData buf;
int i;
pq_begintypsend(&buf);
pq_sendint(&buf, vec->dim, sizeof(int16));
pq_sendint(&buf, vec->unused, sizeof(int16));
for (i = 0; i < vec->dim; i++)
for (int i = 0; i < vec->dim; i++)
pq_sendfloat4(&buf, vec->x[i]);
PG_RETURN_BYTEA_P(pq_endtypsend(&buf));
@@ -420,7 +459,6 @@ array_to_vector(PG_FUNCTION_ARGS)
{
ArrayType *array = PG_GETARG_ARRAYTYPE_P(0);
int32 typmod = PG_GETARG_INT32(1);
int i;
Vector *result;
int16 typlen;
bool typbyval;
@@ -434,6 +472,11 @@ array_to_vector(PG_FUNCTION_ARGS)
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("array must be 1-D")));
if (ARR_HASNULL(array) && array_contains_nulls(array))
ereport(ERROR,
(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
errmsg("array must not contain nulls")));
get_typlenbyvalalign(ARR_ELEMTYPE(array), &typlen, &typbyval, &typalign);
deconstruct_array(array, ARR_ELEMTYPE(array), typlen, typbyval, typalign, &elemsp, &nullsp, &nelemsp);
@@ -441,29 +484,37 @@ array_to_vector(PG_FUNCTION_ARGS)
CheckExpectedDim(typmod, nelemsp);
result = InitVector(nelemsp);
for (i = 0; i < nelemsp; i++)
if (ARR_ELEMTYPE(array) == INT4OID)
{
if (nullsp[i])
ereport(ERROR,
(errcode(ERRCODE_NULL_VALUE_NOT_ALLOWED),
errmsg("array must not containing NULLs")));
/* TODO Move outside loop in 0.5.0 */
if (ARR_ELEMTYPE(array) == INT4OID)
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetInt32(elemsp[i]);
else if (ARR_ELEMTYPE(array) == FLOAT8OID)
result->x[i] = DatumGetFloat8(elemsp[i]);
else if (ARR_ELEMTYPE(array) == FLOAT4OID)
result->x[i] = DatumGetFloat4(elemsp[i]);
else if (ARR_ELEMTYPE(array) == NUMERICOID)
result->x[i] = DatumGetFloat4(DirectFunctionCall1(numeric_float4, elemsp[i]));
else
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("unsupported array type")));
CheckElement(result->x[i]);
}
else if (ARR_ELEMTYPE(array) == FLOAT8OID)
{
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat8(elemsp[i]);
}
else if (ARR_ELEMTYPE(array) == FLOAT4OID)
{
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat4(elemsp[i]);
}
else if (ARR_ELEMTYPE(array) == NUMERICOID)
{
for (int i = 0; i < nelemsp; i++)
result->x[i] = DatumGetFloat4(DirectFunctionCall1(numeric_float4, elemsp[i]));
}
else
{
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("unsupported array type")));
}
/* Check elements */
for (int i = 0; i < result->dim; i++)
CheckElement(result->x[i]);
PG_RETURN_POINTER(result);
}
@@ -478,11 +529,10 @@ vector_to_float4(PG_FUNCTION_ARGS)
Vector *vec = PG_GETARG_VECTOR_P(0);
Datum *datums;
ArrayType *result;
int i;
datums = (Datum *) palloc(sizeof(Datum) * vec->dim);
for (i = 0; i < vec->dim; i++)
for (int i = 0; i < vec->dim; i++)
datums[i] = Float4GetDatum(vec->x[i]);
/* Use TYPALIGN_INT for float4 */
@@ -504,8 +554,8 @@ l2_distance(PG_FUNCTION_ARGS)
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
double distance = 0.0;
double diff;
float distance = 0.0;
float diff;
CheckDims(a, b);
@@ -516,7 +566,7 @@ l2_distance(PG_FUNCTION_ARGS)
distance += diff * diff;
}
PG_RETURN_FLOAT8(sqrt(distance));
PG_RETURN_FLOAT8(sqrt((double) distance));
}
/*
@@ -531,8 +581,8 @@ vector_l2_squared_distance(PG_FUNCTION_ARGS)
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
double distance = 0.0;
double diff;
float distance = 0.0;
float diff;
CheckDims(a, b);
@@ -543,7 +593,7 @@ vector_l2_squared_distance(PG_FUNCTION_ARGS)
distance += diff * diff;
}
PG_RETURN_FLOAT8(distance);
PG_RETURN_FLOAT8((double) distance);
}
/*
@@ -557,7 +607,7 @@ inner_product(PG_FUNCTION_ARGS)
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
double distance = 0.0;
float distance = 0.0;
CheckDims(a, b);
@@ -565,7 +615,7 @@ inner_product(PG_FUNCTION_ARGS)
for (int i = 0; i < a->dim; i++)
distance += ax[i] * bx[i];
PG_RETURN_FLOAT8(distance);
PG_RETURN_FLOAT8((double) distance);
}
/*
@@ -579,7 +629,7 @@ vector_negative_inner_product(PG_FUNCTION_ARGS)
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
double distance = 0.0;
float distance = 0.0;
CheckDims(a, b);
@@ -587,7 +637,7 @@ vector_negative_inner_product(PG_FUNCTION_ARGS)
for (int i = 0; i < a->dim; i++)
distance += ax[i] * bx[i];
PG_RETURN_FLOAT8(distance * -1);
PG_RETURN_FLOAT8((double) distance * -1);
}
/*
@@ -601,9 +651,10 @@ cosine_distance(PG_FUNCTION_ARGS)
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
double distance = 0.0;
double norma = 0.0;
double normb = 0.0;
float distance = 0.0;
float norma = 0.0;
float normb = 0.0;
double similarity;
CheckDims(a, b);
@@ -616,7 +667,21 @@ cosine_distance(PG_FUNCTION_ARGS)
}
/* Use sqrt(a * b) over sqrt(a) * sqrt(b) */
PG_RETURN_FLOAT8(1 - (distance / sqrt(norma * normb)));
similarity = (double) distance / sqrt((double) norma * (double) normb);
#ifdef _MSC_VER
/* /fp:fast may not propagate NaN */
if (isnan(similarity))
PG_RETURN_FLOAT8(NAN);
#endif
/* Keep in range */
if (similarity > 1)
similarity = 1.0;
else if (similarity < -1)
similarity = -1.0;
PG_RETURN_FLOAT8(1.0 - similarity);
}
/*
@@ -630,13 +695,16 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
double distance = 0.0;
float dp = 0.0;
double distance;
CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
distance += a->x[i] * b->x[i];
dp += a->x[i] * b->x[i];
distance = (double) dp;
/* Prevent NaN with acos with loss of precision */
if (distance > 1)
@@ -647,6 +715,28 @@ vector_spherical_distance(PG_FUNCTION_ARGS)
PG_RETURN_FLOAT8(acos(distance) / M_PI);
}
/*
* Get the L1 distance between vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(l1_distance);
Datum
l1_distance(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
float distance = 0.0;
CheckDims(a, b);
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
distance += fabsf(ax[i] - bx[i]);
PG_RETURN_FLOAT8((double) distance);
}
/*
* Get the dimensions of a vector
*/
@@ -672,7 +762,7 @@ vector_norm(PG_FUNCTION_ARGS)
/* Auto-vectorized */
for (int i = 0; i < a->dim; i++)
norm += ax[i] * ax[i];
norm += (double) ax[i] * (double) ax[i];
PG_RETURN_FLOAT8(sqrt(norm));
}
@@ -743,17 +833,51 @@ vector_sub(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result);
}
/*
* Multiply vectors
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_mul);
Datum
vector_mul(PG_FUNCTION_ARGS)
{
Vector *a = PG_GETARG_VECTOR_P(0);
Vector *b = PG_GETARG_VECTOR_P(1);
float *ax = a->x;
float *bx = b->x;
Vector *result;
float *rx;
CheckDims(a, b);
result = InitVector(a->dim);
rx = result->x;
/* Auto-vectorized */
for (int i = 0, imax = a->dim; i < imax; i++)
rx[i] = ax[i] * bx[i];
/* Check for overflow and underflow */
for (int i = 0, imax = a->dim; i < imax; i++)
{
if (isinf(rx[i]))
float_overflow_error();
if (rx[i] == 0 && !(ax[i] == 0 || bx[i] == 0))
float_underflow_error();
}
PG_RETURN_POINTER(result);
}
/*
* Internal helper to compare vectors
*/
int
vector_cmp_internal(Vector * a, Vector * b)
{
int i;
CheckDims(a, b);
for (i = 0; i < a->dim; i++)
for (int i = 0; i < a->dim; i++)
{
if (a->x[i] < b->x[i])
return -1;

View File

@@ -1,12 +1,6 @@
#ifndef VECTOR_H
#define VECTOR_H
#include "postgres.h"
#if PG_VERSION_NUM >= 160000
#include "varatt.h"
#endif
#define VECTOR_MAX_DIM 16000
#define VECTOR_SIZE(_dim) (offsetof(Vector, x) + sizeof(float)*(_dim))
@@ -22,24 +16,8 @@ typedef struct Vector
float x[FLEXIBLE_ARRAY_MEMBER];
} Vector;
Vector *InitVector(int dim);
void PrintVector(char *msg, Vector * vector);
int vector_cmp_internal(Vector * a, Vector * b);
/*
* Allocate and initialize a new vector
*/
static inline Vector *
InitVector(int dim)
{
Vector *result;
int size;
size = VECTOR_SIZE(dim);
result = (Vector *) palloc0(size);
SET_VARSIZE(result, size);
result->dim = dim;
return result;
}
#endif

View File

@@ -29,7 +29,7 @@ SELECT ARRAY[1,2,3]::numeric[]::vector;
(1 row)
SELECT '{NULL}'::real[]::vector;
ERROR: array must not containing NULLs
ERROR: array must not contain nulls
SELECT '{NaN}'::real[]::vector;
ERROR: NaN not allowed in vector
SELECT '{Infinity}'::real[]::vector;
@@ -38,6 +38,8 @@ SELECT '{-Infinity}'::real[]::vector;
ERROR: infinite value not allowed in vector
SELECT '{}'::real[]::vector;
ERROR: vector must have at least 1 dimension
SELECT '{{1}}'::real[]::vector;
ERROR: array must be 1-D
SELECT '[1,2,3]'::vector::real[];
float4
---------

View File

@@ -14,6 +14,16 @@ SELECT '[1,2,3]'::vector - '[4,5,6]';
SELECT '[-3e38]'::vector - '[3e38]';
ERROR: value out of range: overflow
SELECT '[1,2,3]'::vector * '[4,5,6]';
?column?
-----------
[4,10,18]
(1 row)
SELECT '[1e37]'::vector * '[1e37]';
ERROR: value out of range: overflow
SELECT '[1e-37]'::vector * '[1e-37]';
ERROR: value out of range: underflow
SELECT vector_dims('[1,2,3]');
vector_dims
-------------
@@ -38,6 +48,12 @@ SELECT vector_norm('[0,1]');
1
(1 row)
SELECT vector_norm('[3e37,4e37]')::real;
vector_norm
-------------
5e+37
(1 row)
SELECT l2_distance('[0,0]', '[3,4]');
l2_distance
-------------
@@ -52,6 +68,12 @@ SELECT l2_distance('[0,0]', '[0,1]');
SELECT l2_distance('[1,2]', '[3]');
ERROR: different vector dimensions 2 and 1
SELECT l2_distance('[3e38]', '[-3e38]');
l2_distance
-------------
Infinity
(1 row)
SELECT inner_product('[1,2]', '[3,4]');
inner_product
---------------
@@ -60,6 +82,12 @@ SELECT inner_product('[1,2]', '[3,4]');
SELECT inner_product('[1,2]', '[3]');
ERROR: different vector dimensions 2 and 1
SELECT inner_product('[3e38]', '[3e38]');
inner_product
---------------
Infinity
(1 row)
SELECT cosine_distance('[1,2]', '[2,4]');
cosine_distance
-----------------
@@ -86,6 +114,44 @@ SELECT cosine_distance('[1,1]', '[-1,-1]');
SELECT cosine_distance('[1,2]', '[3]');
ERROR: different vector dimensions 2 and 1
SELECT cosine_distance('[1,1]', '[1.1,1.1]');
cosine_distance
-----------------
0
(1 row)
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]');
cosine_distance
-----------------
2
(1 row)
SELECT cosine_distance('[3e38]', '[3e38]');
cosine_distance
-----------------
NaN
(1 row)
SELECT l1_distance('[0,0]', '[3,4]');
l1_distance
-------------
7
(1 row)
SELECT l1_distance('[0,0]', '[0,1]');
l1_distance
-------------
1
(1 row)
SELECT l1_distance('[1,2]', '[3]');
ERROR: different vector dimensions 2 and 1
SELECT l1_distance('[3e38]', '[-3e38]');
l1_distance
-------------
Infinity
(1 row)
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
avg
-----------
@@ -106,5 +172,33 @@ SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2]'::vector, '[3]']) v;
ERROR: expected 2 dimensions, not 1
SELECT avg(v) FROM unnest(ARRAY['[3e38]'::vector, '[3e38]']) v;
avg
---------
[3e+38]
(1 row)
SELECT vector_avg(array_agg(n)) FROM generate_series(1, 16002) n;
ERROR: vector cannot have more than 16000 dimensions
SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
sum
----------
[4,7,10]
(1 row)
SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
sum
----------
[4,7,10]
(1 row)
SELECT sum(v) FROM unnest(ARRAY[]::vector[]) v;
sum
-----
(1 row)
SELECT sum(v) FROM unnest(ARRAY['[1,2]'::vector, '[3]']) v;
ERROR: different vector dimensions 2 and 1
SELECT sum(v) FROM unnest(ARRAY['[3e38]'::vector, '[3e38]']) v;
ERROR: value out of range: overflow

View File

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

21
test/expected/hnsw_ip.out Normal file
View File

@@ -0,0 +1,21 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
val
---------
[1,2,4]
[1,2,3]
[1,1,1]
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2;
count
-------
4
(1 row)
DROP TABLE t;

36
test/expected/hnsw_l2.out Normal file
View File

@@ -0,0 +1,36 @@
SET enable_seqscan = off;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,2,4]
[1,1,1]
[0,0,0]
(4 rows)
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
val
---------
[0,0,0]
[1,1,1]
[1,2,3]
[1,2,4]
(4 rows)
SELECT COUNT(*) FROM t;
count
-------
5
(1 row)
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)
DROP TABLE t;

View File

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

View File

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

View File

@@ -81,6 +81,11 @@ ERROR: malformed vector literal: "1,2,3"
LINE 1: SELECT '1,2,3'::vector;
^
DETAIL: Vector contents must start with "[".
SELECT ''::vector;
ERROR: malformed vector literal: ""
LINE 1: SELECT ''::vector;
^
DETAIL: Vector contents must start with "[".
SELECT '['::vector;
ERROR: malformed vector literal: "["
LINE 1: SELECT '['::vector;

View File

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

View File

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

View File

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

View File

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

View File

@@ -1,7 +1,7 @@
SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val) 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]';
val
---------

View File

@@ -8,6 +8,7 @@ SELECT '{NaN}'::real[]::vector;
SELECT '{Infinity}'::real[]::vector;
SELECT '{-Infinity}'::real[]::vector;
SELECT '{}'::real[]::vector;
SELECT '{{1}}'::real[]::vector;
SELECT '[1,2,3]'::vector::real[];
SELECT array_agg(n)::vector FROM generate_series(1, 16001) n;
SELECT array_to_vector(array_agg(n), 16001, false) FROM generate_series(1, 16001) n;

View File

@@ -2,28 +2,49 @@ SELECT '[1,2,3]'::vector + '[4,5,6]';
SELECT '[3e38]'::vector + '[3e38]';
SELECT '[1,2,3]'::vector - '[4,5,6]';
SELECT '[-3e38]'::vector - '[3e38]';
SELECT '[1,2,3]'::vector * '[4,5,6]';
SELECT '[1e37]'::vector * '[1e37]';
SELECT '[1e-37]'::vector * '[1e-37]';
SELECT vector_dims('[1,2,3]');
SELECT round(vector_norm('[1,1]')::numeric, 5);
SELECT vector_norm('[3,4]');
SELECT vector_norm('[0,1]');
SELECT vector_norm('[3e37,4e37]')::real;
SELECT l2_distance('[0,0]', '[3,4]');
SELECT l2_distance('[0,0]', '[0,1]');
SELECT l2_distance('[1,2]', '[3]');
SELECT l2_distance('[3e38]', '[-3e38]');
SELECT inner_product('[1,2]', '[3,4]');
SELECT inner_product('[1,2]', '[3]');
SELECT inner_product('[3e38]', '[3e38]');
SELECT cosine_distance('[1,2]', '[2,4]');
SELECT cosine_distance('[1,2]', '[0,0]');
SELECT cosine_distance('[1,1]', '[1,1]');
SELECT cosine_distance('[1,1]', '[-1,-1]');
SELECT cosine_distance('[1,2]', '[3]');
SELECT cosine_distance('[1,1]', '[1.1,1.1]');
SELECT cosine_distance('[1,1]', '[-1.1,-1.1]');
SELECT cosine_distance('[3e38]', '[3e38]');
SELECT l1_distance('[0,0]', '[3,4]');
SELECT l1_distance('[0,0]', '[0,1]');
SELECT l1_distance('[1,2]', '[3]');
SELECT l1_distance('[3e38]', '[-3e38]');
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2]'::vector, '[3]']) v;
SELECT avg(v) FROM unnest(ARRAY['[3e38]'::vector, '[3e38]']) v;
SELECT vector_avg(array_agg(n)) FROM generate_series(1, 16002) n;
SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT sum(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
SELECT sum(v) FROM unnest(ARRAY[]::vector[]) v;
SELECT sum(v) FROM unnest(ARRAY['[1,2]'::vector, '[3]']) v;
SELECT sum(v) FROM unnest(ARRAY['[3e38]'::vector, '[3e38]']) v;

13
test/sql/hnsw_cosine.sql Normal file
View File

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

12
test/sql/hnsw_ip.sql Normal file
View File

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

16
test/sql/hnsw_l2.sql Normal file
View File

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

13
test/sql/hnsw_options.sql Normal file
View File

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

View File

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

View File

@@ -14,6 +14,7 @@ SELECT '[4e38,1]'::vector;
SELECT '[1,2,3'::vector;
SELECT '[1,2,3]9'::vector;
SELECT '1,2,3'::vector;
SELECT ''::vector;
SELECT '['::vector;
SELECT '[,'::vector;
SELECT '[]'::vector;

View File

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

View File

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

View File

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

View File

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

View File

@@ -2,7 +2,7 @@ SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val) 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]';

View File

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

View File

@@ -2,12 +2,13 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 1;
use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim) {
for (1 .. $dim)
{
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
@@ -24,7 +25,7 @@ $node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
# Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
@@ -39,3 +40,5 @@ $node->safe_psql("postgres",
# Check size
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
is($size, $new_size, "size does not change");
done_testing();

View File

@@ -0,0 +1,128 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
sub test_recall
{
my ($probes, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
push(@expected, $res);
}
# Build index serially
$node->safe_psql("postgres", qq(
SET max_parallel_maintenance_workers = 0;
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
));
# Test approximate results
if ($operator ne "<#>")
{
# TODO fix test
test_recall(1, 0.71, $operator);
test_recall(10, 0.95, $operator);
}
# Account for equal distances
test_recall(100, 0.9925, $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;
CREATE INDEX idx ON tst USING ivfflat (v $opclass);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
# Test approximate results
if ($operator ne "<#>")
{
# TODO fix test
test_recall(1, 0.71, $operator);
test_recall(10, 0.95, $operator);
}
# Account for equal distances
test_recall(100, 0.9925, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();

View File

@@ -2,7 +2,7 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 3;
use Test::More;
# Initialize node
my $node = get_new_node('node');
@@ -20,7 +20,7 @@ sub test_centers
{
my ($lists, $min) = @_;
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v) WITH (lists = $lists);");
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops) WITH (lists = $lists);");
is($ret, 0, $stderr);
}
@@ -34,3 +34,5 @@ $node->safe_psql("postgres",
# Test no error for duplicate centers
test_centers(10);
done_testing();

View File

@@ -2,7 +2,7 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 60;
use Test::More;
# Initialize node
my $node = get_new_node('node');
@@ -18,24 +18,21 @@ $node->safe_psql("postgres",
# Check each index type
my @operators = ("<->", "<#>", "<=>");
foreach (@operators) {
my $operator = $_;
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
# Add index
my $opclass;
if ($operator == "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator == "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
# Test 100% recall
for (1..20) {
my $i = int(rand() * 100000);
my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $i;");
for (1 .. 20)
{
my $id = int(rand() * 100000);
my $query = $node->safe_psql("postgres", "SELECT v FROM tst WHERE i = $id;");
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT v FROM tst ORDER BY v <-> '$query' LIMIT 1;
@@ -43,3 +40,5 @@ foreach (@operators) {
is($res, $query);
}
}
done_testing();

View File

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

View File

@@ -2,7 +2,7 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 7;
use Test::More;
my $dim = 768;
@@ -19,7 +19,7 @@ $node->safe_psql("postgres", "CREATE TABLE tst (v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v);");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
$node->pgbench(
"--no-vacuum --client=5 --transactions=100",
@@ -28,7 +28,7 @@ $node->pgbench(
[qr{^$}],
"concurrent INSERTs",
{
"007_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
"007_ivfflat_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
}
);
@@ -53,3 +53,5 @@ $count = $node->safe_psql("postgres", qq(
));
is($count, $expected);
is(idx_scan(), 1);
done_testing();

49
test/t/008_aggregates.pl Normal file
View File

@@ -0,0 +1,49 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (r1 real, r2 real, r3 real, v vector(3));");
$node->safe_psql("postgres", qq(
INSERT INTO tst SELECT r1, r2, r3, ARRAY[r1, r2, r3] FROM (
SELECT random() + 1.01 AS r1, random() + 2.01 AS r2, random() + 3.01 AS r3 FROM generate_series(1, 1000000) t
) i;
));
sub test_aggregate
{
my ($agg) = @_;
# Test value
my $res = $node->safe_psql("postgres", "SELECT $agg(v) FROM tst;");
like($res, qr/\[1\.5/);
like($res, qr/,2\.5/);
like($res, qr/,3\.5/);
# Test matches real for avg
# Cannot test sum since sum(real) varies between calls
if ($agg eq 'avg')
{
my $r1 = $node->safe_psql("postgres", "SELECT $agg(r1)::float4 FROM tst;");
my $r2 = $node->safe_psql("postgres", "SELECT $agg(r2)::float4 FROM tst;");
my $r3 = $node->safe_psql("postgres", "SELECT $agg(r3)::float4 FROM tst;");
is($res, "[$r1,$r2,$r3]");
}
# Test explain
my $explain = $node->safe_psql("postgres", "EXPLAIN SELECT $agg(v) FROM tst;");
like($explain, qr/Partial Aggregate/);
}
test_aggregate('avg');
test_aggregate('sum');
done_testing();

View File

@@ -1,35 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 5;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (r1 real, r2 real, r3 real, v vector(3));");
$node->safe_psql("postgres", qq(
INSERT INTO tst SELECT r1, r2, r3, ARRAY[r1, r2, r3] FROM (
SELECT random() + 1.01 AS r1, random() + 2.01 AS r2, random() + 3.01 AS r3 FROM generate_series(1, 1000000) t
) i;
));
# Test avg
my $avg = $node->safe_psql("postgres", "SELECT AVG(v) FROM tst;");
like($avg, qr/\[1\.5/);
like($avg, qr/,2\.5/);
like($avg, qr/,3\.5/);
# Test matches real
my $r1 = $node->safe_psql("postgres", "SELECT AVG(r1)::float4 FROM tst;");
my $r2 = $node->safe_psql("postgres", "SELECT AVG(r2)::float4 FROM tst;");
my $r3 = $node->safe_psql("postgres", "SELECT AVG(r3)::float4 FROM tst;");
is($avg, "[$r1,$r2,$r3]");
# Test explain
my $explain = $node->safe_psql("postgres", "EXPLAIN SELECT AVG(v) FROM tst;");
like($explain, qr/Partial Aggregate/);

View File

@@ -2,7 +2,7 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 1;
use Test::More;
my $dim = 1024;
@@ -30,3 +30,5 @@ my ($ret, $stdout, $stderr) = $node->psql("postgres",
"INSERT INTO tst SELECT array_agg(n), array_agg(n), array_agg(n) FROM generate_series(1, $dim) n"
);
like($stderr, qr/row is too big/);
done_testing();

101
test/t/010_hnsw_wal.pl Normal file
View File

@@ -0,0 +1,101 @@
# Based on postgres/contrib/bloom/t/001_wal.pl
# Test generic xlog record work for hnsw index replication.
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 32;
my $node_primary;
my $node_replica;
# Run few queries on both primary and replica and check their results match.
sub test_index_replay
{
my ($test_name) = @_;
# Wait for replica to catch up
my $applname = $node_replica->name;
my $server_version_num = $node_primary->safe_psql("postgres", "SHOW server_version_num");
my $caughtup_query = "SELECT pg_current_wal_lsn() <= replay_lsn FROM pg_stat_replication WHERE application_name = '$applname';";
$node_primary->poll_query_until('postgres', $caughtup_query)
or die "Timed out while waiting for replica 1 to catch up";
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $sql = join(",", @r);
my $queries = qq(
SET enable_seqscan = off;
SELECT * FROM tst ORDER BY v <-> '[$sql]' LIMIT 10;
);
# Run test queries and compare their result
my $primary_result = $node_primary->safe_psql("postgres", $queries);
my $replica_result = $node_replica->safe_psql("postgres", $queries);
is($primary_result, $replica_result, "$test_name: query result matches");
return;
}
# Use ARRAY[random(), random(), random(), ...] over
# SELECT array_agg(random()) FROM generate_series(1, $dim)
# to generate different values for each row
my $array_sql = join(",", ('random()') x $dim);
# Initialize primary node
$node_primary = get_new_node('primary');
$node_primary->init(allows_streaming => 1);
if ($dim > 32)
{
# TODO use wal_keep_segments for Postgres < 13
$node_primary->append_conf('postgresql.conf', qq(wal_keep_size = 1GB));
}
if ($dim > 1500)
{
$node_primary->append_conf('postgresql.conf', qq(maintenance_work_mem = 128MB));
}
$node_primary->start;
my $backup_name = 'my_backup';
# Take backup
$node_primary->backup($backup_name);
# Create streaming replica linking to primary
$node_replica = get_new_node('replica');
$node_replica->init_from_backup($node_primary, $backup_name, has_streaming => 1);
$node_replica->start;
# Create hnsw index on primary
$node_primary->safe_psql("postgres", "CREATE EXTENSION vector;");
$node_primary->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node_primary->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 1000) i;"
);
$node_primary->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
# Test that queries give same result
test_index_replay('initial');
# Run 10 cycles of table modification. Run test queries after each modification.
for my $i (1 .. 10)
{
$node_primary->safe_psql("postgres", "DELETE FROM tst WHERE i = $i;");
test_index_replay("delete $i");
$node_primary->safe_psql("postgres", "VACUUM tst;");
test_index_replay("vacuum $i");
my ($start, $end) = (1001 + ($i - 1) * 100, 1000 + $i * 100);
$node_primary->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series($start, $end) i;"
);
test_index_replay("insert $i");
}
done_testing();

54
test/t/011_hnsw_vacuum.pl Normal file
View File

@@ -0,0 +1,54 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim)
{
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
my $array_sql = join(", ", @r);
# 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 (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
# Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
# Delete all, vacuum, and insert same data
$node->safe_psql("postgres", "DELETE FROM tst;");
$node->safe_psql("postgres", "VACUUM tst;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
# Check size
# May increase some due to different levels
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
cmp_ok($new_size, "<=", $size * 1.02, "size does not increase too much");
# Delete all but one
$node->safe_psql("postgres", "DELETE FROM tst WHERE i != 123;");
$node->safe_psql("postgres", "VACUUM tst;");
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst ORDER BY v <-> '[0,0,0]' LIMIT 10;
));
is($res, 123);
done_testing();

View File

@@ -2,7 +2,7 @@ use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More tests => 9;
use Test::More;
my $node;
my @queries = ();
@@ -11,14 +11,20 @@ my $limit = 20;
sub test_recall
{
my ($probes, $min, $operator) = @_;
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
for my $i (0 .. $#queries) {
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
@@ -26,8 +32,10 @@ sub test_recall
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids) {
if (exists($actual_set{$_})) {
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
@@ -46,11 +54,12 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 10000) i;"
);
# Generate queries
for (1..20) {
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
@@ -59,30 +68,26 @@ for (1..20) {
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
foreach (@operators) {
my $operator = $_;
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
# Get exact results
@expected = ();
foreach (@queries) {
foreach (@queries)
{
my $res = $node->safe_psql("postgres", "SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;");
push(@expected, $res);
}
# Add index
my $opclass;
if ($operator == "<->") {
$opclass = "vector_l2_ops";
} elsif ($operator == "<#>") {
$opclass = "vector_ip_ops";
} else {
$opclass = "vector_cosine_ops";
}
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v $opclass);");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v $opclass);");
# Test approximate results
test_recall(1, 0.75, $operator);
test_recall(10, 0.95, $operator);
test_recall(100, 1.0, $operator);
my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator);
}
done_testing();

View File

@@ -0,0 +1,108 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = get_new_node('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i serial, v vector(3));");
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# Check each index type
my @operators = ("<->", "<#>", "<=>");
my @opclasses = ("vector_l2_ops", "vector_ip_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
# Add index
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v $opclass);");
# Use concurrent inserts
$node->pgbench(
"--no-vacuum --client=10 --transactions=1000",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"013_hnsw_insert_recall_$opclass" => "INSERT INTO tst (v) VALUES (ARRAY[random(), random(), random()]);"
}
);
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v $operator '$_' LIMIT $limit;
));
push(@expected, $res);
}
my $min = $operator eq "<#>" ? 0.80 : 0.99;
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
$node->safe_psql("postgres", "TRUNCATE tst;");
}
done_testing();

View File

@@ -0,0 +1,74 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Ensures elements and neighbors on both same and different pages
my $dim = 1900;
my $array_sql = join(",", ('random()') x $dim);
# 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($dim));");
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
sub idx_scan
{
# Stats do not update instantaneously
# https://www.postgresql.org/docs/current/monitoring-stats.html#MONITORING-STATS-VIEWS
sleep(1);
$node->safe_psql("postgres", "SELECT idx_scan FROM pg_stat_user_indexes WHERE indexrelid = 'tst_v_idx'::regclass;");
}
for my $i (1 .. 20)
{
$node->pgbench(
"--no-vacuum --client=10 --transactions=1",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"014_hnsw_inserts_$i" => "INSERT INTO tst VALUES (ARRAY[$array_sql]);"
}
);
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
is($count, 10);
$node->safe_psql("postgres", "TRUNCATE tst;");
}
$node->pgbench(
"--no-vacuum --client=20 --transactions=5",
0,
[qr{actually processed}],
[qr{^$}],
"concurrent INSERTs",
{
"014_hnsw_inserts" => "INSERT INTO tst SELECT ARRAY[$array_sql] FROM generate_series(1, 10) i;"
}
);
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1000;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
# Elements may lose all incoming connections with the HNSW algorithm
# Vacuuming can fix this if one of the elements neighbors is deleted
cmp_ok($count, ">=", 997);
is(idx_scan(), 21);
done_testing();

View File

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

View File

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