Compare commits

..

82 Commits

Author SHA1 Message Date
Andrew Kane
13d3bae99b Less space [skip ci] 2024-01-22 18:43:26 -08:00
Andrew Kane
3b9c4c55ee Use consistent order [skip ci] 2024-01-22 18:39:52 -08:00
Andrew Kane
06e2052e40 Added InitVisited [skip ci] 2024-01-22 18:37:37 -08:00
Andrew Kane
891072df68 Added base [skip ci] 2024-01-22 18:02:24 -08:00
Andrew Kane
e9a96d341f Removed prefix for static functions 2024-01-22 17:58:42 -08:00
Andrew Kane
a45ca414f4 Keep logic consistent with insert [skip ci] 2024-01-22 17:43:04 -08:00
Andrew Kane
b4f3cc3e13 Moved outside loop [skip ci] 2024-01-22 17:26:26 -08:00
Andrew Kane
c01cf8a315 Renamed ApplyChanges to UpdateGraph [skip ci] 2024-01-22 17:02:02 -08:00
Andrew Kane
3ecb9a3cb2 Renamed HnswInsertElement to HnswFindElementNeighbors [skip ci] 2024-01-22 16:59:08 -08:00
Andrew Kane
a069e18fe4 Improved function names [skip ci] 2024-01-22 16:48:50 -08:00
Andrew Kane
5174a23094 Updated comment [skip ci] 2024-01-22 16:45:42 -08:00
Andrew Kane
16d7de79f6 Improved function names [skip ci] 2024-01-22 16:43:50 -08:00
Andrew Kane
e54ec4d637 Improved code [skip ci] 2024-01-22 16:37:32 -08:00
Andrew Kane
cc641002d3 Updated comments [skip ci] 2024-01-22 10:59:46 -08:00
Heikki Linnakangas
2f9b1e2893 xAdd overview comment on how HNSW build works (#419)
And rewrite some of the comments in InsertTuple(), to also give more
of a high-level overview of the flow.
2024-01-22 10:50:12 -08:00
Andrew Kane
a3e4fbf6aa Use shared lock for copying neighbors to local memory 2024-01-19 13:44:25 -08:00
Andrew Kane
09a4ec29a0 Added InsertTupleInMemory 2024-01-19 01:27:45 -08:00
Heikki Linnakangas
ca3b4cd029 Remove HnswSpool
It was just used to pass heap/index relations to
HnswParallelScanAndInsert. I think it was copied from nbtsort.c, which
is more complicated. I don't think we need a struct like this.

(That said, I actually think that we should have a state object that
would hold fields like 'heap', 'index', 'procinfo', 'collation'
etc. Passing that object around would simplify the signatures of many
functions. But that's a different story).
2024-01-19 00:26:47 -08:00
Heikki Linnakangas
d96e486274 Remove unused 'scantuplesortstates' field 2024-01-19 00:26:47 -08:00
Heikki Linnakangas
88213186a5 Remove unused argument 2024-01-19 00:26:47 -08:00
Andrew Kane
7dd9534894 Use same locking as insert 2024-01-19 00:18:29 -08:00
Andrew Kane
d801a843f4 Removed HnswPtrSetNull to avoid setting relptr_off directly 2024-01-16 17:08:13 -08:00
Andrew Kane
1458c7bb2a Improved code [skip ci] 2024-01-16 14:03:28 -08:00
Andrew Kane
cad48d9203 Improved locking 2024-01-16 13:34:55 -08:00
Heikki Linnakangas
719b4b7436 Use LWLocks instead of SpinLocks (#410)
Spinlocks should be held only for a few instructions, for multiple
reasons:

- You have to be very careful not to elog() out while holding a
  spinlock, because there is no mechanism to release the spinlock on
  error.

- Waiters can waste a lot of cycles spinning if the lock is
  contended. I you wait on a spinlock for too long, the PostgreSQL
  implementation will actually PANIC, see s_lock_stuck().

The flushLock is particularly problematic. It is held in exclusive
mode, which means it holds a spinlock, over the call to
FlushPages(). FlushPages() performs lots of I/O so it can take a very
long time (>= minutes), and can also easily error out for various
reasons.

allocatorLock would perhaps be OK as a spinlocks, but even that feels
a bit heavy, so I converted that to an LWLock, too.

entryLock is usually held for a very short time, in shared mode, so
that would be fine as a spinlock. However, in the rare case that the
entry point is updated, it's held for a very long time. An LWLock used
in shared mode is about as fast a spinlock, that path is pretty
heavily optimized.

I think we have some problems with the per-element spinlocks too. In
HnswUpdateNeighborPagesInMemory(), it's held over a call to
HnswUpdateConnection(), but HnswUpdateConnection() can error out at
least in case of an out-of-memory error (it uses lappend(), which
calls palloc()). It also calls the distance function, and I don't
think they are guaranteed to be ereport-free either. However, I didn't
address that in this PR, it needs a bit more thinking.
2024-01-16 13:25:03 -08:00
Andrew Kane
fa0acbf62d Fixed CI 2024-01-15 19:55:46 -08:00
Andrew Kane
1612b84069 Fixed error on Windows [skip ci] 2024-01-15 19:33:16 -08:00
Andrew Kane
2f9371516d Leave space for other objects in shared memory 2024-01-15 19:17:50 -08:00
Andrew Kane
9d3e4e74df Added support for in-memory parallel index builds for HNSW 2024-01-15 15:07:31 -08:00
Andrew Kane
0ce497a1b1 Updated Homebrew note [skip ci] 2024-01-15 12:12:04 -08:00
Andrew Kane
c7d60346d8 Improved macro [skip ci] 2024-01-13 20:02:41 -08:00
Andrew Kane
597bfdc76b Added HnswGetNeighbors macro 2024-01-13 20:00:34 -08:00
Andrew Kane
cbf3eb4fa5 Improved HNSW build and insert code 2024-01-13 10:07:42 -08:00
Andrew Kane
cacd389f6d Improved pattern for duplicates 2024-01-12 14:30:13 -08:00
Andrew Kane
423cc2b06c Homebrew now adds to postgresql@15 as well [skip ci] 2024-01-11 16:45:50 -08:00
Andrew Kane
85c4ef6a14 Updated Postgres versions in readme [skip ci] 2024-01-11 12:36:24 -08:00
Andrew Kane
c6160a783a Homebrew now adds to postgresql@16 [skip ci] 2024-01-11 12:32:14 -08:00
Andrew Kane
1881b857f9 Simplified code 2024-01-09 18:53:31 -08:00
Andrew Kane
51bde5fb22 Updated readme [skip ci] 2024-01-09 14:38:25 -08:00
Andrew Kane
10e65ce349 Added note about maintenance_work_mem [skip ci] 2024-01-09 14:31:54 -08:00
Andrew Kane
61279f5a59 Updated readme [skip ci] 2024-01-09 14:26:55 -08:00
Andrew Kane
72b3889e26 Updated readme [skip ci] 2024-01-09 14:22:19 -08:00
Andrew Kane
bb21b2decf Updated readme [skip ci] 2024-01-09 14:19:01 -08:00
Andrew Kane
8a65c0e831 Moved section [skip ci] 2024-01-09 13:33:03 -08:00
Andrew Kane
7d75d423e4 Added section on index build time [skip ci] 2024-01-09 13:27:27 -08:00
Andrew Kane
6cad1f5de0 Updated example [skip ci] 2024-01-09 13:04:47 -08:00
Andrew Kane
67eeade63c Moved HNSW first in readme [skip ci] 2024-01-09 13:04:18 -08:00
Andrew Kane
108fb09d7b Improved code [skip ci] 2024-01-08 17:54:49 -08:00
Andrew Kane
65d060ac86 Reverted FlushPages pattern for parallel builds 2024-01-08 10:45:31 -08:00
Andrew Kane
62ee33bb92 Improved locking code 2024-01-08 09:05:12 -08:00
Andrew Kane
520e274dde Improved locking code 2024-01-07 22:34:41 -08:00
Andrew Kane
9e680884bd Moved indtuples to HnswGraph 2024-01-07 22:23:49 -08:00
Andrew Kane
19a0e1b341 Moved graph to separate struct 2024-01-07 20:15:30 -08:00
Andrew Kane
c7fe1571ee Improved code 2024-01-07 18:30:51 -08:00
Andrew Kane
cb4c770df2 Switched to slist for elements to reduce allocations and remove limit 2024-01-07 18:26:19 -08:00
Andrew Kane
85fdecd79b Moved FlushPages before HnswEndParallel 2024-01-07 17:50:46 -08:00
Andrew Kane
6132428914 Improved number of parallel workers for HNSW index builds - closes #397 2024-01-05 19:46:08 -08:00
Andrew Kane
81d13bd40f Improved code [skip ci] 2024-01-03 13:53:23 -05:00
Andrew Kane
8ee37b60a0 Improved memory estimate for HNSW index builds 2024-01-03 13:47:50 -05:00
Andrew Kane
9b73b3d1a6 Reduced memory and allocations for heap TIDs - closes #385 2024-01-03 13:41:34 -05:00
Andrew Kane
cae630784b Improved BuildCallback [skip ci] 2023-12-30 20:55:29 -05:00
Andrew Kane
d87bcd2deb Added comments [skip ci] 2023-12-30 18:29:01 -05:00
Andrew Kane
736576220a Improved BuildCallback 2023-12-30 18:24:03 -05:00
Andrew Kane
a508b120c1 Added IVFFLAT_MEMORY flag to show memory usage [skip ci] 2023-12-24 09:27:09 -05:00
Andrew Kane
9a782d29f8 Use consistent style [skip ci] 2023-12-22 16:41:25 -05:00
Andrew Kane
1e422cd62b Improved readability [skip ci] 2023-12-22 16:39:13 -05:00
Andrew Kane
569c69580a Improved InsertTuple code - #384
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-12-22 15:08:28 -05:00
Andrew Kane
59509c3a17 Added extra 5% to memory estimate 2023-12-22 14:04:05 -05:00
Andrew Kane
61738846af Updated comment [skip ci] 2023-12-22 14:03:33 -05:00
Andrew Kane
e8c3bf0cef Improved memory tracking for HNSW index builds - #384 2023-12-22 13:35:43 -05:00
Andrew Kane
50d1aed3d8 Improved memory usage logging [skip ci] 2023-12-22 13:09:11 -05:00
Andrew Kane
66e14d2434 Updated indentation [skip ci] 2023-12-22 12:59:50 -05:00
Andrew Kane
42cd4c6833 Fixed call to GenerationContextCreate for Postgres < 15 2023-12-22 12:49:07 -05:00
Andrew Kane
dcbe0b6f0d Reduced memory usage for HNSW index builds - #384
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2023-12-22 12:41:47 -05:00
Andrew Kane
f61d4087b5 Slightly improved memory estimation [skip ci] 2023-12-21 10:31:36 -05:00
Andrew Kane
57554e5b46 Added todo [skip ci] 2023-12-20 17:52:31 -05:00
Andrew Kane
6738fa0bd7 Added HNSW_MEMORY flag to show memory usage - #384 [skip ci] 2023-12-20 16:49:16 -05:00
Andrew Kane
9ab10aa674 Fixed CI 2023-12-20 16:29:13 -05:00
Andrew Kane
ec41dfa1d7 Mark meta buffer contents as dirty when not logging 2023-12-20 16:20:15 -05:00
Andrew Kane
43e0b3d9d4 Mark buffer contents as dirty when not logging 2023-12-20 16:16:25 -05:00
Andrew Kane
2bff7ccaa2 Improved memory estimation - #384 [skip ci] 2023-12-20 11:10:27 -05:00
Andrew Kane
e88a425c9b Reduced WAL generation for HNSW index builds - thanks @hlinnaka 2023-12-19 20:37:32 -05:00
13 changed files with 1154 additions and 531 deletions

View File

@@ -62,7 +62,7 @@ jobs:
wget -q https://github.com/postgres/postgres/archive/refs/tags/REL_14_5.tar.gz
tar xf REL_14_5.tar.gz
- run: make prove_installcheck PROVE_FLAGS="-I ./postgres-REL_14_5/src/test/perl" PERL5LIB="/Users/runner/perl5/lib/perl5"
- run: make clean && /usr/local/opt/llvm@15/bin/scan-build --status-bugs make
- run: make clean && /usr/local/opt/llvm@15/bin/scan-build --status-bugs make PG_CFLAGS="-DUSE_ASSERT_CHECKING"
windows:
runs-on: windows-latest
if: ${{ !startsWith(github.ref_name, 'mac') }}

View File

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

181
README.md
View File

@@ -161,8 +161,97 @@ You can add an index to use approximate nearest neighbor search, which trades so
Supported index types are:
- [IVFFlat](#ivfflat)
- [HNSW](#hnsw) - added in 0.5.0
- [IVFFlat](#ivfflat)
## HNSW
An HNSW index creates a multilayer graph. It has better query performance than IVFFlat (in terms of speed-recall tradeoff), but has slower build times and uses more memory. Also, an index can be created without any data in the table since there isnt a training step like IVFFlat.
Add an index for each distance function you want to use.
L2 distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);
```
Inner product
```sql
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops);
```
Cosine distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
```
Vectors with up to 2,000 dimensions can be indexed.
### Index Options
Specify HNSW parameters
- `m` - the max number of connections per layer (16 by default)
- `ef_construction` - the size of the dynamic candidate list for constructing the graph (64 by default)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);
```
A higher value of `ef_construction` provides better recall at the cost of index build time / insert speed.
### Query Options
Specify the size of the dynamic candidate list for search (40 by default)
```sql
SET hnsw.ef_search = 100;
```
A higher value provides better recall at the cost of speed.
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT ...
COMMIT;
```
### Index Build Time
Indexes build significantly faster when the graph fits into `maintenance_work_mem`
```sql
SET maintenance_work_mem = '8GB';
```
A notice is shown when the graph no longer fits
```text
NOTICE: hnsw graph no longer fits into maintenance_work_mem after 100000 tuples
DETAIL: Building will take significantly more time.
HINT: Increase maintenance_work_mem to speed up builds.
```
Note: Do not set `maintenance_work_mem` so high that it exhausts the memory on the server
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for HNSW are:
1. `initializing`
2. `loading tuples`
## IVFFlat
@@ -232,77 +321,6 @@ The phases for IVFFlat are:
Note: `%` is only populated during the `loading tuples` phase
## HNSW
An HNSW index creates a multilayer graph. It has slower build times and uses more memory than IVFFlat, but has better query performance (in terms of speed-recall tradeoff). Theres no training step like IVFFlat, so the index can be created without any data in the table.
Add an index for each distance function you want to use.
L2 distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops);
```
Inner product
```sql
CREATE INDEX ON items USING hnsw (embedding vector_ip_ops);
```
Cosine distance
```sql
CREATE INDEX ON items USING hnsw (embedding vector_cosine_ops);
```
Vectors with up to 2,000 dimensions can be indexed.
### Index Options
Specify HNSW parameters
- `m` - the max number of connections per layer (16 by default)
- `ef_construction` - the size of the dynamic candidate list for constructing the graph (64 by default)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WITH (m = 16, ef_construction = 64);
```
A higher value of `ef_construction` provides better recall at the cost of index build time / insert speed.
### Query Options
Specify the size of the dynamic candidate list for search (40 by default)
```sql
SET hnsw.ef_search = 100;
```
A higher value provides better recall at the cost of speed.
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL hnsw.ef_search = 100;
SELECT ...
COMMIT;
```
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * blocks_done / nullif(blocks_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for HNSW are:
1. `initializing`
2. `loading tuples`
## Filtering
There are a few ways to index nearest neighbor queries with a `WHERE` clause
@@ -320,8 +338,7 @@ CREATE INDEX ON items (category_id);
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)
WHERE (category_id = 123);
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WHERE (category_id = 123);
```
Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html) for approximate search on many different values of the `WHERE` columns
@@ -571,10 +588,10 @@ If compilation fails with `fatal error: postgres.h: No such file or directory`,
For Ubuntu and Debian, use:
```sh
sudo apt install postgresql-server-dev-15
sudo apt install postgresql-server-dev-16
```
Note: Replace `15` with your Postgres server version
Note: Replace `16` with your Postgres server version
### Windows
@@ -589,7 +606,7 @@ Note: The exact path will vary depending on your Visual Studio version and editi
Then use `nmake` to build:
```cmd
set "PGROOT=C:\Program Files\PostgreSQL\15"
set "PGROOT=C:\Program Files\PostgreSQL\16"
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector
nmake /F Makefile.win
@@ -639,22 +656,22 @@ pgxn install vector
Debian and Ubuntu packages are available from the [PostgreSQL APT Repository](https://wiki.postgresql.org/wiki/Apt). Follow the [setup instructions](https://wiki.postgresql.org/wiki/Apt#Quickstart) and run:
```sh
sudo apt install postgresql-15-pgvector
sudo apt install postgresql-16-pgvector
```
Note: Replace `15` with your Postgres server version
Note: Replace `16` with your Postgres server version
### Yum
RPM packages are available from the [PostgreSQL Yum Repository](https://yum.postgresql.org/). Follow the [setup instructions](https://www.postgresql.org/download/linux/redhat/) for your distribution and run:
```sh
sudo yum install pgvector_15
sudo yum install pgvector_16
# or
sudo dnf install pgvector_15
sudo dnf install pgvector_16
```
Note: Replace `15` with your Postgres server version
Note: Replace `16` with your Postgres server version
### conda-forge

View File

@@ -14,15 +14,45 @@
#endif
int hnsw_ef_search;
bool hnsw_enable_parallel_build;
int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind;
/*
* Assign a tranche ID for our LWLocks. This only needs to be done by one
* backend, as the tranche ID is remembered in shared memory.
*
* This shared memory area is very small, so we just allocate it from the
* "slop" that PostgreSQL reserves for small allocations like this. If
* this grows bigger, we should use a shmem_request_hook and
* RequestAddinShmemSpace() to pre-reserve space for this.
*/
static void
HnswInitLockTranche(void)
{
int *tranche_ids;
bool found;
LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
tranche_ids = ShmemInitStruct("hnsw LWLock ids",
sizeof(int) * 1,
&found);
if (!found)
tranche_ids[0] = LWLockNewTrancheId();
hnsw_lock_tranche_id = tranche_ids[0];
LWLockRelease(AddinShmemInitLock);
/* Per-backend registration of the tranche ID */
LWLockRegisterTranche(hnsw_lock_tranche_id, "HnswBuild");
}
/*
* Initialize index options and variables
*/
void
HnswInit(void)
{
HnswInitLockTranche();
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
@@ -40,11 +70,6 @@ HnswInit(void)
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);
/* Behind a variable for now since can be slower than building in memory */
DefineCustomBoolVariable("hnsw.enable_parallel_build", "Enables or disables building indexes in parallel",
NULL, &hnsw_enable_parallel_build,
false, PGC_USERSET, 0, NULL, NULL, NULL);
}
/*

View File

@@ -8,6 +8,7 @@
#include "access/reloptions.h"
#include "nodes/execnodes.h"
#include "port.h" /* for random() */
#include "utils/relptr.h"
#include "utils/sampling.h"
#include "vector.h"
@@ -72,8 +73,10 @@
#if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
#else
#define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE)
#define SeedRandom(seed) srandom(seed)
#endif
#if PG_VERSION_NUM < 130000
@@ -91,33 +94,64 @@
#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)
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
#define HnswGetValue(base, element) PointerGetDatum(HnswPtrAccess(base, (element)->value))
#if PG_VERSION_NUM < 140005
#define relptr_offset(rp) ((rp).relptr_off - 1)
#endif
/* Pointer macros */
#define HnswPtrAccess(base, hp) ((base) == NULL ? (hp).ptr : relptr_access(base, (hp).relptr))
#define HnswPtrStore(base, hp, value) ((base) == NULL ? (void) ((hp).ptr = (value)) : (void) relptr_store(base, (hp).relptr, value))
#define HnswPtrIsNull(base, hp) ((base) == NULL ? (hp).ptr == NULL : relptr_is_null((hp).relptr))
#define HnswPtrEqual(base, hp1, hp2) ((base) == NULL ? (hp1).ptr == (hp2).ptr : relptr_offset((hp1).relptr) == relptr_offset((hp2).relptr))
/* For code paths dedicated to each type */
#define HnswPtrPointer(hp) (hp).ptr
#define HnswPtrOffset(hp) relptr_offset((hp).relptr)
/* Variables */
extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build;
extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray;
#define HnswPtrDeclare(type, relptrtype, ptrtype) \
relptr_declare(type, relptrtype); \
typedef union { type *ptr; relptrtype relptr; } ptrtype;
/* Pointers that can be absolute or relative */
/* Use char for DatumPtr so works with Pointer */
HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr);
typedef struct HnswElementData
{
List *heaptids;
HnswElementPtr next;
ItemPointerData heaptids[HNSW_HEAPTIDS];
uint8 heaptidsLength;
uint8 level;
uint8 deleted;
uint32 hash;
HnswNeighborArray *neighbors;
HnswNeighborsPtr neighbors;
BlockNumber blkno;
OffsetNumber offno;
OffsetNumber neighborOffno;
BlockNumber neighborPage;
Datum value;
DatumPtr value;
LWLock lock;
} HnswElementData;
typedef HnswElementData * HnswElement;
typedef struct HnswCandidate
{
HnswElement element;
HnswElementPtr element;
float distance;
bool closer;
} HnswCandidate;
@@ -126,7 +160,7 @@ typedef struct HnswNeighborArray
{
int length;
bool closerSet;
HnswCandidate *items;
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborArray;
typedef struct HnswPairingHeapNode
@@ -143,11 +177,26 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */
} HnswOptions;
typedef struct HnswSpool
typedef struct HnswGraph
{
Relation heap;
Relation index;
} HnswSpool;
/* Graph state */
slock_t lock;
HnswElementPtr head;
double indtuples;
/* Entry state */
LWLock entryLock;
HnswElementPtr entryPoint;
/* Allocations state */
LWLock allocatorLock;
long memoryUsed;
long memoryTotal;
/* Flushed state */
LWLock flushLock;
bool flushed;
} HnswGraph;
typedef struct HnswShared
{
@@ -155,7 +204,6 @@ typedef struct HnswShared
Oid heaprelid;
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
/* Worker progress */
ConditionVariable workersdonecv;
@@ -166,7 +214,7 @@ typedef struct HnswShared
/* Mutable state */
int nparticipantsdone;
double reltuples;
double indtuples;
HnswGraph graphData;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
@@ -184,8 +232,15 @@ typedef struct HnswLeader
int nparticipanttuplesorts;
HnswShared *hnswshared;
Snapshot snapshot;
char *hnswarea;
} HnswLeader;
typedef struct HnswAllocator
{
void *(*alloc) (Size size, void *state);
void *state;
} HnswAllocator;
typedef struct HnswBuildState
{
/* Info */
@@ -209,20 +264,21 @@ typedef struct HnswBuildState
Oid collation;
/* Variables */
List *elements;
HnswElement entryPoint;
HnswGraph graphData;
HnswGraph *graph;
double ml;
int maxLevel;
long memoryLeft;
bool flushed;
Vector *normvec;
/* Memory */
MemoryContext graphCtx;
MemoryContext tmpCtx;
HnswAllocator allocator;
/* Parallel builds */
HnswLeader *hnswleader;
HnswShared *hnswshared;
char *hnswarea;
} HnswBuildState;
typedef struct HnswMetaPageData
@@ -321,25 +377,24 @@ bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * re
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement);
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element);
void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void 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 HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);
void HnswSetNeighborTuple(char *base, 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, bool building);
void HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building);
void HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
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 HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
@@ -359,6 +414,16 @@ void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys,
bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan);
static inline HnswNeighborArray *
HnswGetNeighbors(char *base, HnswElement element, int lc)
{
HnswNeighborArrayPtr *neighborList = HnswPtrAccess(base, element->neighbors);
Assert(element->level >= lc);
return HnswPtrAccess(base, neighborList[lc]);
}
/* Hash tables */
typedef struct TidHashEntry
{
@@ -386,4 +451,17 @@ typedef struct PointerHashEntry
#define SH_DECLARE
#include "lib/simplehash.h"
typedef struct OffsetHashEntry
{
Size offset;
char status;
} OffsetHashEntry;
#define SH_PREFIX offsethash
#define SH_ELEMENT_TYPE OffsetHashEntry
#define SH_KEY_TYPE Size
#define SH_SCOPE extern
#define SH_DECLARE
#include "lib/simplehash.h"
#endif

View File

@@ -1,3 +1,39 @@
/*
* The HNSW build happens in two phases:
*
* 1. In-memory phase
*
* In this first phase, the graph is held completely in memory. When the graph
* is fully built, or we run out of memory reserved for the build (determined
* by maintenance_work_mem), we materialize the graph to disk (see
* FlushPages()), and switch to the on-disk phase.
*
* In a parallel build, a large contiguous chunk of shared memory is allocated
* to hold the graph. Each worker process has its own HnswBuildState struct in
* private memory, which contains information that doesn't change throughout
* the build, and pointers to the shared structs in shared memory. The shared
* memory area is mapped to a different address in each worker process, and
* 'HnswBuildState.hnswarea' points to the beginning of the shared area in the
* worker process's address space. All pointers used in the graph are
* "relative pointers", stored as an offset from 'hnswarea'.
*
* Each element is protected by an LWLock. It must be held when reading or
* modifying the element's neighbors or 'heaptids'.
*
* In a non-parallel build, the graph is held in backend-private memory. All
* the elements are allocated in a dedicated memory context, 'graphCtx', and
* the pointers used in the graph are regular pointers.
*
* 2. On-disk phase
*
* In the on-disk phase, the index is built by inserting each vector to the
* index one by one, just like on INSERT. The only difference is that we don't
* WAL-log the individual inserts. If the graph fit completely in memory and
* was fully built in the in-memory phase, the on-disk phase is skipped.
*
* After we have finished building the graph, we perform one more scan through
* the index and write all the pages to the WAL.
*/
#include "postgres.h"
#include <math.h>
@@ -54,9 +90,12 @@
#endif
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000002)
#define PARALLEL_KEY_HNSW_AREA UINT64CONST(0xA000000000000002)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000003)
#define LIST_MAX_LENGTH ((1 << 26) - 1)
#if PG_VERSION_NUM < 130000
#define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2)
#endif
/*
* Create the metapage
@@ -88,6 +127,7 @@ CreateMetaPage(HnswBuildState * buildstate)
((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
}
@@ -104,6 +144,7 @@ HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, ForkNumber forkNum)
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */
MarkBufferDirty(*buf);
UnlockReleaseBuffer(*buf);
/* Can take a while, so ensure we can interrupt */
@@ -131,9 +172,11 @@ CreateElementPages(HnswBuildState * buildstate)
HnswElementTuple etup;
HnswNeighborTuple ntup;
BlockNumber insertPage;
HnswElement entryPoint;
Buffer buf;
Page page;
ListCell *lc;
HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
/* Calculate sizes */
etupAllocSize = BLCKSZ;
@@ -148,18 +191,22 @@ CreateElementPages(HnswBuildState * buildstate)
page = BufferGetPage(buf);
HnswInitPage(buf, page);
foreach(lc, buildstate->elements)
while (!HnswPtrIsNull(base, iter))
{
HnswElement element = lfirst(lc);
HnswElement element = HnswPtrAccess(base, iter);
Size etupSize;
Size ntupSize;
Size combinedSize;
void *valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */
iter = element->next;
/* Zero memory for each element */
MemSet(etup, 0, etupAllocSize);
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(element->value)));
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -167,7 +214,7 @@ CreateElementPages(HnswBuildState * buildstate)
if (etupSize > etupAllocSize)
elog(ERROR, "index tuple too large");
HnswSetElementTuple(etup, element);
HnswSetElementTuple(base, etup, element);
/* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -205,9 +252,11 @@ CreateElementPages(HnswBuildState * buildstate)
insertPage = BufferGetBlockNumber(buf);
/* Commit */
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum, true);
entryPoint = HnswPtrAccess(base, buildstate->graph->entryPoint);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, entryPoint, insertPage, forkNum, true);
pfree(etup);
pfree(ntup);
@@ -222,19 +271,23 @@ CreateNeighborPages(HnswBuildState * buildstate)
Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum;
int m = buildstate->m;
ListCell *lc;
HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
HnswNeighborTuple ntup;
/* Allocate once */
ntup = palloc0(BLCKSZ);
foreach(lc, buildstate->elements)
while (!HnswPtrIsNull(base, iter))
{
HnswElement e = lfirst(lc);
HnswElement e = HnswPtrAccess(base, iter);
Buffer buf;
Page page;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
/* Update iterator */
iter = e->next;
/* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */
CHECK_FOR_INTERRUPTS();
@@ -243,58 +296,203 @@ CreateNeighborPages(HnswBuildState * buildstate)
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
page = BufferGetPage(buf);
HnswSetNeighborTuple(ntup, e, m);
HnswSetNeighborTuple(base, 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);
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)
{
#ifdef HNSW_MEMORY
elog(INFO, "memory: %zu MB", buildstate->graph->memoryUsed / (1024 * 1024));
#endif
CreateMetaPage(buildstate);
CreateElementPages(buildstate);
CreateNeighborPages(buildstate);
buildstate->flushed = true;
FreeElements(buildstate);
buildstate->graph->flushed = true;
MemoryContextReset(buildstate->graphCtx);
}
/*
* Add a heap TID to an existing element
*/
static bool
AddDuplicateInMemory(HnswElement element, HnswElement dup)
{
LWLockAcquire(&dup->lock, LW_EXCLUSIVE);
if (dup->heaptidsLength == HNSW_HEAPTIDS)
{
LWLockRelease(&dup->lock);
return false;
}
HnswAddHeapTid(dup, &element->heaptids[0]);
LWLockRelease(&dup->lock);
return true;
}
/*
* Find duplicate element
*/
static bool
FindDuplicateInMemory(char *base, HnswElement element)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
/* Check for space */
if (AddDuplicateInMemory(element, neighborElement))
return true;
}
return false;
}
/*
* Add to element list
*/
static void
AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
{
SpinLockAcquire(&graph->lock);
element->next = graph->head;
HnswPtrStore(base, graph->head, element);
SpinLockRelease(&graph->lock);
}
/*
* Update neighbors
*/
static void
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
{
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, hc->element);
/* Keep scan-build happy on Mac x86-64 */
Assert(neighborElement);
/* Use element for lock instead of hc since hc can be replaced */
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
LWLockRelease(&neighborElement->lock);
}
}
}
/*
* Update graph in memory
*/
static void
UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
{
HnswGraph *graph = buildstate->graph;
char *base = buildstate->hnswarea;
/* Look for duplicate */
if (FindDuplicateInMemory(base, element))
return;
/* Add element */
AddElementInMemory(base, graph, element);
/* Update neighbors */
UpdateNeighborsInMemory(base, procinfo, collation, element, m);
/* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswPtrStore(base, graph->entryPoint, element);
}
/*
* Insert tuple in memory
*/
static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
HnswGraph *graph = buildstate->graph;
HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock;
int efConstruction = buildstate->efConstruction;
int m = buildstate->m;
char *base = buildstate->hnswarea;
/* Get entry point */
LWLockAcquire(entryLock, LW_SHARED);
entryPoint = HnswPtrAccess(base, graph->entryPoint);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
LWLockRelease(entryLock);
/* Get exclusive lock */
LWLockAcquire(entryLock, LW_EXCLUSIVE);
/* Get latest entry point after lock is acquired */
entryPoint = HnswPtrAccess(base, graph->entryPoint);
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
/* Release entry lock */
LWLockRelease(entryLock);
}
/*
* Insert tuple
*/
static bool
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx)
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
{
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction;
int m = buildstate->m;
MemoryContext oldCtx;
HnswGraph *graph = buildstate->graph;
HnswElement element;
HnswAllocator *allocator = &buildstate->allocator;
Size valueSize;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -302,56 +500,81 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec))
if (!HnswNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
return false;
}
/* Copy value to element so accessible outside of memory context */
oldCtx = MemoryContextSwitchTo(outerCtx);
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Get datum size */
valueSize = VARSIZE_ANY(DatumGetPointer(value));
/* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Ensure graph not flushed when inserting */
LWLockAcquire(flushLock, LW_SHARED);
/* Look for duplicate */
*dup = HnswFindDuplicate(element);
/* Update neighbors if needed */
if (*dup == NULL)
/* Are we in the on-disk phase? */
if (graph->flushed)
{
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &element->neighbors[lc];
LWLockRelease(flushLock);
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
}
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
}
/* Update entry point if needed */
if (*dup == NULL && (entryPoint == NULL || element->level > entryPoint->level))
buildstate->entryPoint = element;
/*
* In a parallel build, the HnswElement is allocated from the shared
* memory area, so we need to coordinate with other processes.
*/
LWLockAcquire(&graph->allocatorLock, LW_EXCLUSIVE);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
/*
* Check that we have enough memory available for the new element now that
* we have the allocator lock, and flush pages if needed.
*/
if (graph->memoryUsed >= graph->memoryTotal)
{
LWLockRelease(&graph->allocatorLock);
return *dup == NULL;
}
LWLockRelease(flushLock);
LWLockAcquire(flushLock, LW_EXCLUSIVE);
/*
* Get the memory used by an element
*/
static long
HnswElementMemory(HnswElement e, int m)
{
long elementSize = sizeof(HnswElementData);
if (!graph->flushed)
{
ereport(NOTICE,
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) graph->indtuples),
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
return elementSize;
FlushPages(buildstate);
}
LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
}
/* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
valuePtr = HnswAlloc(allocator, valueSize);
/*
* We have now allocated the space needed for the element, so we don't
* need the allocator lock anymore. Release it and initialize the rest of
* the element.
*/
LWLockRelease(&graph->allocatorLock);
/* Copy the datum */
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
/* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
/* Insert tuple */
InsertTupleInMemory(buildstate, element);
/* Release flush lock */
LWLockRelease(flushLock);
return true;
}
/*
@@ -362,10 +585,8 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state)
{
HnswBuildState *buildstate = (HnswBuildState *) state;
HnswGraph *graph = buildstate->graph;
MemoryContext oldCtx;
HnswElement element;
HnswElement dup = NULL;
bool inserted;
#if PG_VERSION_NUM < 130000
ItemPointer tid = &hup->t_self;
@@ -375,70 +596,80 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0])
return;
if (buildstate->flushed || buildstate->memoryLeft <= 0 || list_length(buildstate->elements) == LIST_MAX_LENGTH)
{
if (!buildstate->flushed)
{
if (buildstate->memoryLeft <= 0)
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, true))
{
if (buildstate->hnswshared)
{
HnswShared *hnswshared = buildstate->hnswshared;
SpinLockAcquire(&hnswshared->mutex);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++hnswshared->indtuples);
SpinLockRelease(&hnswshared->mutex);
}
else
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
}
/* 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);
/* Use memory context since detoast can allocate */
/* Use memory context */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */
inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx);
if (InsertTuple(index, values, isnull, tid, buildstate))
{
/* Update progress */
SpinLockAcquire(&graph->lock);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples);
SpinLockRelease(&graph->lock);
}
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
}
/* Add outside memory context */
if (dup != NULL)
{
HnswAddHeapTid(dup, tid);
buildstate->memoryLeft -= sizeof(ItemPointerData);
}
/*
* Initialize the graph
*/
static void
InitGraph(HnswGraph * graph, char *base, long memoryTotal)
{
HnswPtrStore(base, graph->head, (HnswElement) NULL);
HnswPtrStore(base, graph->entryPoint, (HnswElement) NULL);
graph->memoryUsed = 0;
graph->memoryTotal = memoryTotal;
graph->flushed = false;
graph->indtuples = 0;
SpinLockInit(&graph->lock);
LWLockInitialize(&graph->entryLock, hnsw_lock_tranche_id);
LWLockInitialize(&graph->allocatorLock, hnsw_lock_tranche_id);
LWLockInitialize(&graph->flushLock, hnsw_lock_tranche_id);
}
/* Add to buildstate or free */
if (inserted)
{
buildstate->elements = lappend(buildstate->elements, element);
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
}
else
HnswFreeElement(element);
/*
* Initialize an allocator
*/
static void
InitAllocator(HnswAllocator * allocator, void *(*alloc) (Size size, void *state), void *state)
{
allocator->alloc = alloc;
allocator->state = state;
}
/*
* Memory context allocator
*/
static void *
HnswMemoryContextAlloc(Size size, void *state)
{
HnswBuildState *buildstate = (HnswBuildState *) state;
void *chunk = MemoryContextAlloc(buildstate->graphCtx, size);
#if PG_VERSION_NUM >= 130000
buildstate->graphData.memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
#else
buildstate->graphData.memoryUsed += MAXALIGN(size);
#endif
return chunk;
}
/*
* Shared memory allocator
*/
static void *
HnswSharedMemoryAlloc(Size size, void *state)
{
HnswBuildState *buildstate = (HnswBuildState *) state;
void *chunk = buildstate->hnswarea + buildstate->graph->memoryUsed;
buildstate->graph->memoryUsed += MAXALIGN(size);
return chunk;
}
/*
@@ -474,22 +705,29 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0];
buildstate->elements = NIL;
buildstate->entryPoint = NULL;
InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->memoryLeft = maintenance_work_mem * 1024L;
buildstate->flushed = false;
/* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",
#if PG_VERSION_NUM >= 150000
1024 * 1024, 1024 * 1024,
#endif
1024 * 1024);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES);
InitAllocator(&buildstate->allocator, &HnswMemoryContextAlloc, buildstate);
buildstate->hnswleader = NULL;
buildstate->hnswshared = NULL;
buildstate->hnswarea = NULL;
}
/*
@@ -499,6 +737,7 @@ static void
FreeBuildState(HnswBuildState * buildstate)
{
pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx);
}
@@ -518,7 +757,8 @@ ParallelHeapScan(HnswBuildState * buildstate)
SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->indtuples = hnswshared->indtuples;
buildstate->graph = &hnswshared->graphData;
buildstate->hnswarea = buildstate->hnswleader->hnswarea;
reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex);
break;
@@ -538,7 +778,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
* Perform a worker's portion of a parallel insert
*/
static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
HnswParallelScanAndInsert(Relation heapRel, Relation indexRel, HnswShared * hnswshared, char *hnswarea, bool progress)
{
HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
@@ -550,22 +790,21 @@ HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool p
IndexInfo *indexInfo;
/* Join parallel scan */
indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo = BuildIndexInfo(indexRel);
indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
/* TODO Support in-memory builds */
buildstate.memoryLeft = 0;
buildstate.flushed = true;
buildstate.hnswshared = hnswshared;
InitBuildState(&buildstate, heapRel, indexRel, indexInfo, MAIN_FORKNUM);
buildstate.graph = &hnswshared->graphData;
buildstate.hnswarea = hnswarea;
InitAllocator(&buildstate.allocator, &HnswSharedMemoryAlloc, &buildstate);
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
scan = table_beginscan_parallel(heapRel,
ParallelTableScanFromHnswShared(hnswshared));
reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo,
reltuples = table_index_build_scan(heapRel, indexRel, indexInfo,
true, progress, BuildCallback,
(void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
scan = heap_beginscan_parallel(heapRel, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(heapRel, indexRel, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
@@ -595,8 +834,8 @@ void
HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{
char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared;
char *hnswarea;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
@@ -632,13 +871,10 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
#endif
indexRel = index_open(hnswshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */
hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
hnswarea = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false);
/* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false);
HnswParallelScanAndInsert(heapRel, indexRel, hnswshared, hnswarea, false);
/* Close relations within worker */
index_close(indexRel, indexLockmode);
@@ -693,15 +929,9 @@ static void
HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
{
HnswLeader *hnswleader = buildstate->hnswleader;
HnswSpool *leaderworker;
/* Allocate memory and initialize private spool */
leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
HnswParallelScanAndInsert(buildstate->heap, buildstate->index, hnswleader->hnswshared, hnswleader->hnswarea, true);
}
/*
@@ -711,10 +941,12 @@ static void
HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
{
ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot;
Size esthnswshared;
Size esthnswarea;
Size estother;
HnswShared *hnswshared;
char *hnswarea;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true;
int querylen;
@@ -732,8 +964,6 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request;
/* Get snapshot for table scan */
if (!isconcurrent)
snapshot = SnapshotAny;
@@ -743,7 +973,17 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
/* Estimate size of workspaces */
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
shm_toc_estimate_keys(&pcxt->estimator, 1);
/* Leave space for other objects in shared memory */
/* Docker has a default limit of 64 MB for shm_size */
/* which happens to be the default value of maintenance_work_mem */
esthnswarea = maintenance_work_mem * 1024L;
estother = 2 * 1024 * 1024;
if (esthnswarea > estother)
esthnswarea -= estother;
shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea);
shm_toc_estimate_keys(&pcxt->estimator, 2);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string)
@@ -774,13 +1014,11 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
hnswshared->indexrelid = RelationGetRelid(buildstate->index);
hnswshared->isconcurrent = isconcurrent;
hnswshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&hnswshared->workersdonecv);
SpinLockInit(&hnswshared->mutex);
/* Initialize mutable state */
hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0;
hnswshared->indtuples = 0;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared),
@@ -789,7 +1027,12 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif
hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea);
/* Report less than allocated so never fails */
InitGraph(&hnswshared->graphData, hnswarea, esthnswarea - 1024 * 1024);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_SHARED, hnswshared);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_AREA, hnswarea);
/* Store query string for workers */
if (debug_query_string)
@@ -809,6 +1052,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
hnswleader->nparticipanttuplesorts++;
hnswleader->hnswshared = hnswshared;
hnswleader->snapshot = snapshot;
hnswleader->hnswarea = hnswarea;
/* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0)
@@ -831,6 +1075,27 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
WaitForParallelWorkersToAttach(pcxt);
}
/*
* Compute parallel workers
*/
static int
ComputeParallelWorkers(Relation heap, Relation index)
{
int parallel_workers;
/* Make sure it's safe to use parallel workers */
parallel_workers = plan_create_index_workers(RelationGetRelid(heap), RelationGetRelid(index));
if (parallel_workers == 0)
return 0;
/* Use parallel_workers storage parameter on table if set */
parallel_workers = RelationGetParallelWorkers(heap, -1);
if (parallel_workers != -1)
return Min(parallel_workers, max_parallel_maintenance_workers);
return max_parallel_maintenance_workers;
}
/*
* Build graph
*/
@@ -842,30 +1107,35 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */
if (hnsw_enable_parallel_build)
parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
if (buildstate->heap != NULL)
parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
/* Add tuples to graph */
if (buildstate->heap != NULL)
{
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000
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
}
buildstate->indtuples = buildstate->graph->indtuples;
}
/* Add tuples to sort */
if (buildstate->hnswleader)
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);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
/* Flush pages */
if (!buildstate->graph->flushed)
FlushPages(buildstate);
/* End parallel build */
if (buildstate->hnswleader)
@@ -895,13 +1165,13 @@ static void
BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
HnswBuildState * buildstate, ForkNumber forkNum)
{
#ifdef HNSW_MEMORY
SeedRandom(42);
#endif
InitBuildState(buildstate, heap, index, indexInfo, forkNum);
if (buildstate->heap != NULL)
BuildGraph(buildstate, forkNum);
if (!buildstate->flushed)
FlushPages(buildstate);
BuildGraph(buildstate, forkNum);
if (RelationNeedsWAL(index))
log_newpage_range(index, forkNum, 0, RelationGetNumberOfBlocks(index), true);

View File

@@ -5,6 +5,7 @@
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h"
/*
@@ -115,7 +116,7 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
* Add to element and neighbor pages
*/
static void
WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building)
AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building)
{
Buffer buf;
Page page;
@@ -133,9 +134,10 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
char *base = NULL;
/* Calculate sizes */
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(DatumGetPointer(e->value)));
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE;
@@ -143,11 +145,11 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
/* Prepare element tuple */
etup = palloc0(etupSize);
HnswSetElementTuple(etup, e);
HnswSetElementTuple(base, etup, e);
/* Prepare neighbor tuple */
ntup = palloc0(ntupSize);
HnswSetNeighborTuple(ntup, e, m);
HnswSetNeighborTuple(base, ntup, e, m);
/* Find a page (or two if needed) to insert the tuples */
for (;;)
@@ -219,7 +221,9 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
HnswInsertAppendPage(index, &newbuf, &newpage, state, page, building);
/* Commit */
if (!building)
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
/* Unlock previous buffer */
@@ -294,7 +298,13 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
}
/* Commit */
if (!building)
if (building)
{
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
}
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
if (nbuf != buf)
@@ -329,12 +339,14 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
* Update neighbors
*/
void
HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
{
char *base = NULL;
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc];
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
for (int i = 0; i < neighbors->length; i++)
{
@@ -347,11 +359,12 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
Size ntupSize;
int idx = -1;
int startIdx;
OffsetNumber offno = hc->element->neighborOffno;
HnswElement neighborElement = HnswPtrAccess(base, hc->element);
OffsetNumber offno = neighborElement->neighborOffno;
/* Get latest neighbors since they may have changed */
/* Do not lock yet since selecting neighbors can take time */
HnswLoadNeighbors(hc->element, index, m);
HnswLoadNeighbors(neighborElement, index, m);
/*
* Could improve performance for vacuuming by checking neighbors
@@ -361,14 +374,14 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
*/
/* Select neighbors */
HnswUpdateConnection(e, hc, lm, lc, &idx, index, procinfo, collation);
HnswUpdateConnection(NULL, 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);
buf = ReadBuffer(index, neighborElement->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
if (building)
{
@@ -387,7 +400,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
ntupSize = ItemIdGetLength(itemid);
/* Calculate index for update */
startIdx = (hc->element->level - lc) * m;
startIdx = (neighborElement->level - lc) * m;
/* Check for existing connection */
if (checkExisting && ConnectionExists(e, ntup, startIdx, lm))
@@ -421,7 +434,9 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
if (!building)
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
}
else if (!building)
@@ -436,7 +451,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
* Add a heap TID to an existing element
*/
static bool
HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool building)
AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool building)
{
Buffer buf;
Page page;
@@ -480,14 +495,16 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool buil
}
/* Add heap TID */
etup->heaptids[i] = *((ItemPointer) linitial(element->heaptids));
etup->heaptids[i] = element->heaptids[0];
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, dup->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
if (!building)
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
@@ -495,31 +512,55 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool buil
}
/*
* Write changes to disk
* Find duplicate element
*/
static bool
FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
{
char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true;
}
return false;
}
/*
* Update graph on disk
*/
static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint, bool building)
UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{
BlockNumber newInsertPage = InvalidBlockNumber;
/* Try to add to existing page */
if (dup != NULL)
{
if (HnswAddDuplicate(index, element, dup, building))
return;
}
/* Look for duplicate */
if (FindDuplicateOnDisk(index, element, building))
return;
/* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage, building);
/* Add element */
AddElementOnDisk(index, element, m, GetInsertPage(index), &newInsertPage, building);
/* Update insert page if needed */
if (BlockNumberIsValid(newInsertPage))
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false, building);
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building);
/* Update metapage if needed */
/* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM, building);
}
@@ -528,29 +569,16 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
* Insert a tuple into the index
*/
bool
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building)
HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building)
{
Datum value;
FmgrInfo *normprocinfo;
HnswElement entryPoint;
HnswElement element;
int m;
int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
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;
}
char *base = NULL;
/*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -563,8 +591,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
HnswGetMetaPageInfo(index, &m, &entryPoint);
/* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m));
element->value = value;
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -580,14 +608,11 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
entryPoint = HnswGetEntryPoint(index);
}
/* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Find neighbors for element */
HnswFindElementNeighbors(base, 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, building);
/* Update graph on disk */
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -595,6 +620,30 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
return true;
}
/*
* Insert a tuple into the index
*/
static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel)
{
Datum value;
FmgrInfo *normprocinfo;
Oid collation = index->rd_indcollation[0];
/* 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;
}
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, heapRel, false);
}
/*
* Insert a tuple into the index
*/
@@ -621,7 +670,7 @@ hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid,
oldCtx = MemoryContextSwitchTo(insertCtx);
/* Insert tuple */
HnswInsertTuple(index, values, isnull, heap_tid, heap, false);
HnswInsertTuple(index, values, isnull, heap_tid, heap);
/* Delete memory context */
MemoryContextSwitchTo(oldCtx);

View File

@@ -21,6 +21,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
List *w;
int m;
HnswElement entryPoint;
char *base = NULL;
/* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -28,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL)
return NIL;
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, false));
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, false));
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL);
ep = w;
}
return HnswSearchLayer(q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL);
}
/*
@@ -184,19 +185,19 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
while (list_length(so->w) > 0)
{
char *base = NULL;
HnswCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid;
/* Move to next element if no valid heap TIDs */
if (list_length(hc->element->heaptids) == 0)
if (element->heaptidsLength == 0)
{
so->w = list_delete_last(so->w);
continue;
}
heaptid = llast(hc->element->heaptids);
hc->element->heaptids = list_delete_last(hc->element->heaptids);
heaptid = &element->heaptids[--element->heaptidsLength];
MemoryContextSwitchTo(oldCtx);

View File

@@ -84,9 +84,40 @@ hash_pointer(uintptr_t ptr)
#define SH_DEFINE
#include "lib/simplehash.h"
/* Needed to include simplehash.h again */
#if PG_VERSION_NUM < 120000
#undef SH_EQUAL
#undef sh_log2
#undef sh_pow2
#define sh_log2 offsethash_sh_log2
#define sh_pow2 offsethash_sh_pow2
#endif
/* Offset hash table */
static uint32
hash_offset(Size offset)
{
#if SIZEOF_SIZE_T == 8
return murmurhash64((uint64) offset);
#else
return murmurhash32((uint32) offset);
#endif
}
#define SH_PREFIX offsethash
#define SH_ELEMENT_TYPE OffsetHashEntry
#define SH_KEY_TYPE Size
#define SH_KEY offset
#define SH_HASH_KEY(tb, key) hash_offset(key)
#define SH_EQUAL(tb, a, b) (a == b)
#define SH_SCOPE extern
#define SH_DEFINE
#include "lib/simplehash.h"
typedef union
{
pointerhash_hash *pointers;
offsethash_hash *offsets;
tidhash_hash *tids;
} visited_hash;
@@ -188,42 +219,46 @@ HnswInitPage(Buffer buf, Page page)
* Allocate neighbors
*/
void
HnswInitNeighbors(HnswElement element, int m)
HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * allocator)
{
int level = element->level;
element->neighbors = palloc(sizeof(HnswNeighborArray) * (level + 1));
HnswNeighborArrayPtr *neighborList = (HnswNeighborArrayPtr *) HnswAlloc(allocator, sizeof(HnswNeighborArrayPtr) * (level + 1));
HnswPtrStore(base, element->neighbors, neighborList);
for (int lc = 0; lc <= level; lc++)
{
HnswNeighborArray *a;
int lm = HnswGetLayerM(m, lc);
a = &element->neighbors[lc];
HnswPtrStore(base, neighborList[lc], (HnswNeighborArray *) HnswAlloc(allocator, offsetof(HnswNeighborArray, items) + sizeof(HnswCandidate) * lm));
a = HnswGetNeighbors(base, element, lc);
a->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false;
}
}
/*
* Free neighbors
* Allocate memory from the allocator
*/
static void
HnswFreeNeighbors(HnswElement element)
void *
HnswAlloc(HnswAllocator * allocator, Size size)
{
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
if (allocator)
return (*(allocator)->alloc) (size, (allocator)->state);
return palloc(size);
}
/*
* Allocate an element
*/
HnswElement
HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel, HnswAllocator * allocator)
{
HnswElement element = palloc(sizeof(HnswElementData));
HnswElement element = HnswAlloc(allocator, sizeof(HnswElementData));
int level = (int) (-log(RandomDouble()) * ml);
@@ -231,42 +266,26 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
if (level > maxLevel)
level = maxLevel;
element->heaptids = NIL;
element->heaptidsLength = 0;
HnswAddHeapTid(element, heaptid);
element->level = level;
element->deleted = 0;
HnswInitNeighbors(element, m);
HnswInitNeighbors(base, element, m, allocator);
element->value = PointerGetDatum(NULL);
HnswPtrStore(base, element->value, (Pointer) NULL);
return element;
}
/*
* Free an element
*/
void
HnswFreeElement(HnswElement element)
{
HnswFreeNeighbors(element);
list_free_deep(element->heaptids);
if (DatumGetPointer(element->value))
pfree(DatumGetPointer(element->value));
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);
element->heaptids[element->heaptidsLength++] = *heaptid;
}
/*
@@ -276,11 +295,12 @@ HnswElement
HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
{
HnswElement element = palloc(sizeof(HnswElementData));
char *base = NULL;
element->blkno = blkno;
element->offno = offno;
element->neighbors = NULL;
element->value = PointerGetDatum(NULL);
HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
HnswPtrStore(base, element->value, (Pointer) NULL);
return element;
}
@@ -379,7 +399,9 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage);
if (!building)
if (building)
MarkBufferDirty(buf);
else
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
}
@@ -388,26 +410,28 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
* Set element tuple, except for neighbor info
*/
void
HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level;
etup->deleted = 0;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
if (i < list_length(element->heaptids))
etup->heaptids[i] = *((ItemPointer) list_nth(element->heaptids, i));
if (i < element->heaptidsLength)
etup->heaptids[i] = element->heaptids[i];
else
ItemPointerSetInvalid(&etup->heaptids[i]);
}
memcpy(&etup->data, DatumGetPointer(element->value), VARSIZE_ANY(DatumGetPointer(element->value)));
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
}
/*
* Set neighbor tuple
*/
void
HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
{
int idx = 0;
@@ -415,7 +439,7 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
for (int lc = e->level; lc >= 0; lc--)
{
HnswNeighborArray *neighbors = &e->neighbors[lc];
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
int lm = HnswGetLayerM(m, lc);
for (int i = 0; i < lm; i++)
@@ -425,8 +449,9 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
if (i < neighbors->length)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement hce = HnswPtrAccess(base, hc->element);
ItemPointerSet(indextid, hc->element->blkno, hc->element->offno);
ItemPointerSet(indextid, hce->blkno, hce->offno);
}
else
ItemPointerSetInvalid(indextid);
@@ -442,12 +467,14 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
static void
LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
{
char *base = NULL;
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
int neighborCount = (element->level + 2) * m;
Assert(HnswIsNeighborTuple(ntup));
HnswInitNeighbors(element, m);
HnswInitNeighbors(base, element, m, NULL);
/* Ensure expected neighbors */
if (ntup->count != neighborCount)
@@ -473,9 +500,9 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
if (level < 0)
level = 0;
neighbors = &element->neighbors[level];
neighbors = HnswGetNeighbors(base, element, level);
hc = &neighbors->items[neighbors->length++];
hc->element = e;
HnswPtrStore(base, hc->element, e);
}
}
@@ -507,7 +534,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
element->deleted = etup->deleted;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptids = NIL;
element->heaptidsLength = 0;
if (loadHeaptids)
{
@@ -522,7 +549,12 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
}
if (loadVec)
element->value = datumCopy(PointerGetDatum(&etup->data), false, -1);
{
char *base = NULL;
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
HnswPtrStore(base, element->value, DatumGetPointer(value));
}
}
/*
@@ -558,24 +590,27 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
* Get the distance for a candidate
*/
static float
GetCandidateDistance(HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
GetCandidateDistance(char *base, HnswCandidate * hc, Datum q, FmgrInfo *procinfo, Oid collation)
{
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, hc->element->value));
HnswElement hce = HnswPtrAccess(base, hc->element);
Datum value = HnswGetValue(base, hce);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, q, value));
}
/*
* Create a candidate for the entry point
*/
HnswCandidate *
HnswEntryCandidate(HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = entryPoint;
HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL)
hc->distance = GetCandidateDistance(hc, q, procinfo, collation);
hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation);
else
HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadVec);
HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec);
return hc;
}
@@ -621,48 +656,97 @@ CreatePairingHeapNode(HnswCandidate * c)
return node;
}
/*
* Init visited
*/
static inline void
InitVisited(char *base, visited_hash * v, Relation index, int ef, int m)
{
if (index != NULL)
v->tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
else if (base != NULL)
v->offsets = offsethash_create(CurrentMemoryContext, ef * m * 2, NULL);
else
v->pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL);
}
/*
* Add to visited
*/
static inline void
AddToVisited(visited_hash v, HnswCandidate * hc, Relation index, bool *found)
AddToVisited(char *base, visited_hash * v, HnswCandidate * hc, Relation index, bool *found)
{
if (index == NULL)
if (index != NULL)
{
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointerData indextid;
ItemPointerSet(&indextid, element->blkno, element->offno);
tidhash_insert(v->tids, indextid, found);
}
else if (base != NULL)
{
#if PG_VERSION_NUM >= 130000
pointerhash_insert_hash(v.pointers, (uintptr_t) hc->element, hc->element->hash, found);
HnswElement element = HnswPtrAccess(base, hc->element);
offsethash_insert_hash(v->offsets, HnswPtrOffset(hc->element), element->hash, found);
#else
pointerhash_insert(v.pointers, (uintptr_t) hc->element, found);
offsethash_insert(v->offsets, HnswPtrOffset(hc->element), found);
#endif
}
else
{
ItemPointerData indextid;
#if PG_VERSION_NUM >= 130000
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointerSet(&indextid, hc->element->blkno, hc->element->offno);
tidhash_insert(v.tids, indextid, found);
pointerhash_insert_hash(v->pointers, (uintptr_t) HnswPtrPointer(hc->element), element->hash, found);
#else
pointerhash_insert(v->pointers, (uintptr_t) HnswPtrPointer(hc->element), found);
#endif
}
}
/*
* Count element towards ef
*/
static inline bool
CountElement(char *base, HnswElement skipElement, HnswCandidate * hc)
{
HnswElement e;
if (skipElement == NULL)
return true;
/* Ensure does not access heaptidsLength during in-memory build */
pg_memory_barrier();
e = HnswPtrAccess(base, hc->element);
return e->heaptidsLength != 0;
}
/*
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
{
ListCell *lc2;
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
visited_hash v;
ListCell *lc2;
HnswNeighborArray *neighborhoodData = NULL;
Size neighborhoodSize;
/* Create hash table */
InitVisited(base, &v, index, ef, m);
/* Create local memory for neighborhood if needed */
if (index == NULL)
v.pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL);
else
v.tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
{
neighborhoodSize = offsetof(HnswNeighborArray, items) + sizeof(HnswCandidate) * HnswGetLayerM(m, lc);
neighborhoodData = palloc(neighborhoodSize);
}
/* Add entry points to v, C, and W */
foreach(lc2, ep)
@@ -670,7 +754,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
bool found;
AddToVisited(v, hc, index, &found);
AddToVisited(base, &v, hc, index, &found);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
@@ -680,7 +764,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* would be ideal to do this for inserts as well, but this could
* affect insert performance.
*/
if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
if (CountElement(base, skipElement, hc))
wlen++;
}
@@ -689,38 +773,51 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswNeighborArray *neighborhood;
HnswCandidate *c = ((HnswPairingHeapNode *) pairingheap_remove_first(C))->inner;
HnswCandidate *f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
HnswElement cElement;
if (c->distance > f->distance)
break;
if (c->element->neighbors == NULL)
HnswLoadNeighbors(c->element, index, m);
cElement = HnswPtrAccess(base, c->element);
if (HnswPtrIsNull(base, cElement->neighbors))
HnswLoadNeighbors(cElement, index, m);
/* Get the neighborhood at layer lc */
neighborhood = &c->element->neighbors[lc];
neighborhood = HnswGetNeighbors(base, cElement, lc);
/* Copy neighborhood to local memory if needed */
if (index == NULL)
{
LWLockAcquire(&cElement->lock, LW_SHARED);
memcpy(neighborhoodData, neighborhood, neighborhoodSize);
LWLockRelease(&cElement->lock);
neighborhood = neighborhoodData;
}
for (int i = 0; i < neighborhood->length; i++)
{
HnswCandidate *e = &neighborhood->items[i];
bool visited;
AddToVisited(v, e, index, &visited);
AddToVisited(base, &v, e, index, &visited);
if (!visited)
{
float eDistance;
HnswElement eElement = HnswPtrAccess(base, e->element);
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (index == NULL)
eDistance = GetCandidateDistance(e, q, procinfo, collation);
eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
else
HnswLoadElement(e->element, &eDistance, &q, index, procinfo, collation, inserting);
HnswLoadElement(eElement, &eDistance, &q, index, procinfo, collation, inserting);
Assert(!e->element->deleted);
Assert(!eElement->deleted);
/* Make robust to issues */
if (e->element->level < lc)
if (eElement->level < lc)
continue;
if (eDistance < f->distance || wlen < ef)
@@ -728,7 +825,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
/* Copy e */
HnswCandidate *ec = palloc(sizeof(HnswCandidate));
ec->element = e->element;
HnswPtrStore(base, ec->element, eElement);
ec->distance = eDistance;
pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
@@ -739,7 +836,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance.
*/
if (skipElement == NULL || list_length(e->element->heaptids) != 0)
if (CountElement(base, skipElement, e))
{
wlen++;
@@ -764,7 +861,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
}
/*
* Compare candidate distances
* Compare candidate distances with pointer tie-breaker
*/
static int
#if PG_VERSION_NUM >= 130000
@@ -782,10 +879,38 @@ CompareCandidateDistances(const void *a, const void *b)
if (hca->distance > hcb->distance)
return -1;
if (hca->element < hcb->element)
if (HnswPtrPointer(hca->element) < HnswPtrPointer(hcb->element))
return 1;
if (hca->element > hcb->element)
if (HnswPtrPointer(hca->element) > HnswPtrPointer(hcb->element))
return -1;
return 0;
}
/*
* Compare candidate distances with offset tie-breaker
*/
static int
#if PG_VERSION_NUM >= 130000
CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
#else
CompareCandidateDistancesOffset(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;
if (HnswPtrOffset(hca->element) < HnswPtrOffset(hcb->element))
return 1;
if (HnswPtrOffset(hca->element) > HnswPtrOffset(hcb->element))
return -1;
return 0;
@@ -795,46 +920,58 @@ CompareCandidateDistances(const void *a, const void *b)
* Calculate the distance between elements
*/
static float
HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid collation)
HnswGetDistance(char *base, HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid collation)
{
Datum aValue;
Datum bValue;
/* Look for cached distance */
if (a->neighbors != NULL)
if (!HnswPtrIsNull(base, a->neighbors))
{
Assert(a->level >= lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(base, a, lc);
for (int i = 0; i < a->neighbors[lc].length; i++)
for (int i = 0; i < neighbors->length; i++)
{
if (a->neighbors[lc].items[i].element == b)
return a->neighbors[lc].items[i].distance;
HnswElement element = HnswPtrAccess(base, neighbors->items[i].element);
if (element == b)
return neighbors->items[i].distance;
}
}
if (b->neighbors != NULL)
if (!HnswPtrIsNull(base, b->neighbors))
{
Assert(b->level >= lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(base, b, lc);
for (int i = 0; i < b->neighbors[lc].length; i++)
for (int i = 0; i < neighbors->length; i++)
{
if (b->neighbors[lc].items[i].element == a)
return b->neighbors[lc].items[i].distance;
HnswElement element = HnswPtrAccess(base, neighbors->items[i].element);
if (element == a)
return neighbors->items[i].distance;
}
}
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, a->value, b->value));
aValue = HnswGetValue(base, a);
bValue = HnswGetValue(base, b);
return DatumGetFloat8(FunctionCall2Coll(procinfo, collation, aValue, bValue));
}
/*
* Check if an element is closer to q than any element from R
*/
static bool
CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid collation)
CheckElementCloser(char *base, HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid collation)
{
HnswElement eElement = HnswPtrAccess(base, e->element);
ListCell *lc2;
foreach(lc2, r)
{
HnswCandidate *ri = lfirst(lc2);
float distance = HnswGetDistance(e->element, ri->element, lc, procinfo, collation);
HnswElement riElement = HnswPtrAccess(base, ri->element);
float distance = HnswGetDistance(base, eElement, riElement, lc, procinfo, collation);
if (distance <= e->distance)
return false;
@@ -847,12 +984,13 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper
*/
static List *
SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
SelectNeighbors(char *base, List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{
List *r = NIL;
List *w = list_copy(c);
pairingheap *wd;
bool mustCalculate = !e2->neighbors[lc].closerSet;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e2, lc);
bool mustCalculate = !neighbors->closerSet;
List *added = NIL;
bool removedAny = false;
@@ -863,7 +1001,12 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
/* Ensure order of candidates is deterministic for closer caching */
if (sortCandidates)
list_sort(w, CompareCandidateDistances);
{
if (base == NULL)
list_sort(w, CompareCandidateDistances);
else
list_sort(w, CompareCandidateDistancesOffset);
}
while (list_length(w) > 0 && list_length(r) < m)
{
@@ -874,7 +1017,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
/* Use previous state of r and wd to skip work when possible */
if (mustCalculate)
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, lc, procinfo, collation);
else if (list_length(added) > 0)
{
/*
@@ -883,7 +1026,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
*/
if (e->closer)
{
e->closer = CheckElementCloser(e, added, lc, procinfo, collation);
e->closer = CheckElementCloser(base, e, added, lc, procinfo, collation);
if (!e->closer)
removedAny = true;
@@ -896,7 +1039,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
*/
if (removedAny)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, lc, procinfo, collation);
if (e->closer)
added = lappend(added, e);
}
@@ -904,7 +1047,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
}
else if (e == newCandidate)
{
e->closer = CheckElementCloser(e, r, lc, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, lc, procinfo, collation);
if (e->closer)
added = lappend(added, e);
}
@@ -916,7 +1059,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
}
/* Cached value can only be used in future if sorted deterministically */
e2->neighbors[lc].closerSet = sortCandidates;
neighbors->closerSet = sortCandidates;
/* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m)
@@ -934,38 +1077,14 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
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 (!datumIsEqual(e->value, neighbor->element->value, false, -1))
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)
AddConnections(char *base, HnswElement element, List *neighbors, int m, int lc)
{
ListCell *lc2;
HnswNeighborArray *a = &element->neighbors[lc];
HnswNeighborArray *a = HnswGetNeighbors(base, element, lc);
foreach(lc2, neighbors)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
@@ -975,13 +1094,13 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
* Update connections
*/
void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
HnswElement hce = HnswPtrAccess(base, hc->element);
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
HnswCandidate hc2;
hc2.element = element;
HnswPtrStore(base, hc2.element, element);
hc2.distance = hc->distance;
if (currentNeighbors->length < m)
@@ -1000,19 +1119,20 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */
if (index != NULL)
{
Datum q = hc->element->value;
Datum q = HnswGetValue(base, hce);
for (int i = 0; i < currentNeighbors->length; i++)
{
HnswCandidate *hc3 = &currentNeighbors->items[i];
HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
if (DatumGetPointer(hc3->element->value) == NULL)
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true);
if (HnswPtrIsNull(base, hc3Element->value))
HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true);
else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
hc3->distance = GetCandidateDistance(base, hc3, q, procinfo, collation);
/* Prune element if being deleted */
if (list_length(hc3->element->heaptids) == 0)
if (hc3Element->heaptidsLength == 0)
{
pruned = &currentNeighbors->items[i];
break;
@@ -1029,7 +1149,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2);
SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true);
SelectNeighbors(base, c, m, lc, procinfo, collation, hce, &hc2, &pruned, true);
/* Should not happen */
if (pruned == NULL)
@@ -1039,7 +1159,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Find and replace the pruned element */
for (int i = 0; i < currentNeighbors->length; i++)
{
if (currentNeighbors->items[i].element == pruned->element)
if (HnswPtrEqual(base, currentNeighbors->items[i].element, pruned->element))
{
currentNeighbors->items[i] = hc2;
@@ -1057,43 +1177,65 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
* Remove elements being deleted or skipped
*/
static List *
RemoveElements(List *w, HnswElement skipElement)
RemoveElements(char *base, List *w, HnswElement skipElement)
{
ListCell *lc2;
List *w2 = NIL;
/* Ensure does not access heaptidsLength during in-memory build */
pg_memory_barrier();
foreach(lc2, w)
{
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
HnswElement hce = HnswPtrAccess(base, hc->element);
/* Skip self for vacuuming update */
if (skipElement != NULL && hc->element->blkno == skipElement->blkno && hc->element->offno == skipElement->offno)
if (skipElement != NULL && hce->blkno == skipElement->blkno && hce->offno == skipElement->offno)
continue;
if (list_length(hc->element->heaptids) != 0)
if (hce->heaptidsLength != 0)
w2 = lappend(w2, hc);
}
return w2;
}
#if PG_VERSION_NUM >= 130000
/*
* Precompute hash
*/
static void
PrecomputeHash(char *base, HnswElement element)
{
HnswElementPtr ptr;
HnswPtrStore(base, ptr, element);
if (base == NULL)
element->hash = hash_pointer((uintptr_t) HnswPtrPointer(ptr));
else
element->hash = hash_offset(HnswPtrOffset(ptr));
}
#endif
/*
* Algorithm 1 from paper
*/
void
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
{
List *ep;
List *w;
int level = element->level;
int entryLevel;
Datum q = element->value;
Datum q = HnswGetValue(base, element);
HnswElement skipElement = existing ? element : NULL;
#if PG_VERSION_NUM >= 130000
/* Precompute hash */
if (index == NULL)
element->hash = hash_pointer((uintptr_t) element);
PrecomputeHash(base, element);
#endif
/* No neighbors if no entry point */
@@ -1101,13 +1243,13 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
return;
/* Get entry point and level */
ep = list_make1(HnswEntryCandidate(entryPoint, q, index, procinfo, collation, true));
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, procinfo, collation, true));
entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement);
ep = w;
}
@@ -1125,12 +1267,12 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *neighbors;
List *lw;
w = HnswSearchLayer(q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement);
/* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */
if (index != NULL)
lw = RemoveElements(w, skipElement);
lw = RemoveElements(base, w, skipElement);
else
lw = w;
@@ -1139,9 +1281,9 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
* sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance.
*/
neighbors = SelectNeighbors(lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
neighbors = SelectNeighbors(base, lw, lm, lc, procinfo, collation, element, NULL, NULL, false);
AddConnections(element, neighbors, lm, lc);
AddConnections(base, element, neighbors, lm, lc);
ep = w;
}

View File

@@ -199,21 +199,22 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
char *base = NULL;
/* 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;
HnswInitNeighbors(base, element, m, NULL);
element->heaptidsLength = 0;
/* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Update neighbor tuple */
/* Do this before getting page to minimize locking */
HnswSetNeighborTuple(ntup, element, m);
HnswSetNeighborTuple(base, ntup, element, m);
/* Get neighbor page */
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
@@ -230,7 +231,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf);
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, true, false);
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false);
}
/*
@@ -301,7 +302,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
{
/* Reset neighbors from previous update */
if (highestPoint != NULL)
highestPoint->neighbors = NULL;
HnswPtrStore((char *) NULL, highestPoint->neighbors, (HnswNeighborArrayPtr *) NULL);
RepairGraphElement(vacuumstate, entryPoint, highestPoint);
}

View File

@@ -543,10 +543,10 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
pfree(list);
}
#ifdef IVFFLAT_KMEANS_DEBUG
/*
* Print k-means metrics
*/
#ifdef IVFFLAT_KMEANS_DEBUG
static void
PrintKmeansMetrics(IvfflatBuildState * buildstate)
{

View File

@@ -6,6 +6,10 @@
#include "ivfflat.h"
#include "miscadmin.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h"
#endif
/*
* Initialize with kmeans++
*
@@ -151,6 +155,23 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
}
}
#ifdef IVFFLAT_MEMORY
/*
* Show memory usage
*/
static void
ShowMemoryUsage(Size estimatedSize)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "total memory: %zu MB",
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
#endif
elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
}
#endif
/*
* Use Elkan for performance. This requires distance function to satisfy triangle inequality.
*
@@ -231,6 +252,10 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
vec->dim = dimensions;
}
#ifdef IVFFLAT_MEMORY
ShowMemoryUsage(totalSize);
#endif
/* Pick initial centers */
InitCenters(index, samples, centers, lowerBound);

View File

@@ -95,11 +95,10 @@ for my $i (0 .. $#operators)
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel
# Build index in parallel in memory
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET client_min_messages = DEBUG;
SET min_parallel_table_scan_size = 1;
SET hnsw.enable_parallel_build = on;
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
is($ret, 0, $stderr);
@@ -109,6 +108,21 @@ for my $i (0 .. $#operators)
test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;");
# Build index in parallel on disk
# Set parallel_workers on table to use workers with low maintenance_work_mem
($ret, $stdout, $stderr) = $node->psql("postgres", qq(
ALTER TABLE tst SET (parallel_workers = 2);
SET client_min_messages = DEBUG;
SET maintenance_work_mem = '4MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
ALTER TABLE tst RESET (parallel_workers);
));
is($ret, 0, $stderr);
like($stderr, qr/using \d+ parallel workers/);
like($stderr, qr/hnsw graph no longer fits into maintenance_work_mem/);
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();