Compare commits

..

62 Commits

Author SHA1 Message Date
Andrew Kane
a7b6be73e1 Use DSA 2024-01-16 21:59:11 -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 1076 additions and 487 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,66 @@
#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;
/* This is initialized when the module is loaded */
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 +162,7 @@ typedef struct HnswNeighborArray
{
int length;
bool closerSet;
HnswCandidate *items;
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborArray;
typedef struct HnswPairingHeapNode
@@ -143,6 +179,27 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */
} HnswOptions;
typedef struct HnswGraph
{
/* 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 HnswSpool
{
Relation heap;
@@ -156,6 +213,7 @@ typedef struct HnswShared
Oid indexrelid;
bool isconcurrent;
int scantuplesortstates;
dsa_pointer dsaptr;
/* Worker progress */
ConditionVariable workersdonecv;
@@ -166,7 +224,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 +242,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 +274,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 +387,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 HnswInsertElement(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 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 HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void 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 +424,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 +461,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

@@ -54,9 +54,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 +91,7 @@ CreateMetaPage(HnswBuildState * buildstate)
((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
}
@@ -104,6 +108,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 +136,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 +155,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 +178,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 +216,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 +235,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 +260,166 @@ 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
HnswAddDuplicateInMemory(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
HnswFindDuplicateInMemory(char *base, HnswElement element)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum value = HnswGetValue(base, element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
/* Check for space */
if (HnswAddDuplicateInMemory(element, neighborElement))
return true;
}
return false;
}
/*
* Add to element and neighbor pages
*/
static void
WriteNewElementPagesInMemory(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
HnswUpdateNeighborPagesInMemory(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);
}
}
}
/*
* Write changes in memory
*/
static void
WriteElementInMemory(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate, HnswGraph * graph, bool updateEntryPoint)
{
char *base = buildstate->hnswarea;
/* Try to add to existing page */
if (HnswFindDuplicateInMemory(base, element))
return;
/* Write element and neighbor tuples */
WriteNewElementPagesInMemory(base, graph, element);
/* Update neighbors */
HnswUpdateNeighborPagesInMemory(base, procinfo, collation, element, m);
/* Update entry point if needed (already have lock) */
if (updateEntryPoint)
HnswPtrStore(base, graph->entryPoint, element);
}
/*
* 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;
HnswGraph *graph = buildstate->graph;
HnswElement entryPoint;
int efConstruction = buildstate->efConstruction;
int m = buildstate->m;
MemoryContext oldCtx;
HnswElement element;
HnswAllocator *allocator = &buildstate->allocator;
Size valueSize;
Pointer valuePtr;
bool updateEntryPoint;
LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -306,52 +431,82 @@ InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState *
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)
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;
/* Get lock for allocator */
LWLockAcquire(&graph->allocatorLock, LW_EXCLUSIVE);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
/* 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);
}
/* Create an element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
valuePtr = HnswAlloc(allocator, valueSize);
/* Release allocator lock */
LWLockRelease(&graph->allocatorLock);
/* Copy datum */
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
/* Create element lock */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
/* Get entry point */
LWLockAcquire(&graph->entryLock, LW_EXCLUSIVE);
entryPoint = HnswPtrAccess(base, graph->entryPoint);
updateEntryPoint = entryPoint == NULL || element->level > entryPoint->level;
/* Release lock if not updating entry point */
if (!updateEntryPoint)
LWLockRelease(&graph->entryLock);
/* Insert element in graph */
HnswInsertElement(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Write to memory */
WriteElementInMemory(index, procinfo, collation, element, m, efConstruction, entryPoint, buildstate, graph, updateEntryPoint);
/* Release lock if needed */
if (updateEntryPoint)
LWLockRelease(&graph->entryLock);
/* Release flush lock */
LWLockRelease(flushLock);
return true;
}
/*
@@ -362,10 +517,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 +528,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 +637,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 +669,7 @@ static void
FreeBuildState(HnswBuildState * buildstate)
{
pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx);
}
@@ -518,7 +689,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 +710,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
* Perform a worker's portion of a parallel insert
*/
static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress)
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, char *hnswarea, bool progress)
{
HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
@@ -553,10 +725,9 @@ HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool p
indexInfo = BuildIndexInfo(hnswspool->index);
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;
buildstate.graph = &hnswshared->graphData;
buildstate.hnswarea = hnswarea;
InitAllocator(&buildstate.allocator, &HnswSharedMemoryAlloc, &buildstate);
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
ParallelTableScanFromHnswShared(hnswshared));
@@ -597,6 +768,9 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared;
dsa_area *dsa;
void *dsaspace;
char *hnswarea;
Relation heapRel;
Relation indexRel;
LOCKMODE heapLockmode;
@@ -637,8 +811,12 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
dsaspace = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false);
dsa = dsa_attach_in_place(dsaspace, seg);
hnswarea = dsa_get_address(dsa, hnswshared->dsaptr);
/* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false);
HnswParallelScanAndInsert(hnswspool, hnswshared, hnswarea, false);
/* Close relations within worker */
index_close(indexRel, indexLockmode);
@@ -701,7 +879,7 @@ HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
leaderworker->index = buildstate->index;
/* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true);
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, hnswleader->hnswarea, true);
}
/*
@@ -714,7 +892,12 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
int scantuplesortstates;
Snapshot snapshot;
Size esthnswshared;
Size esthnswarea;
Size estdsaspace;
HnswShared *hnswshared;
void *dsaspace;
dsa_area *dsa;
char *hnswarea;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true;
int querylen;
@@ -743,7 +926,11 @@ 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);
/* Start with a minimal DSA so InitializeParallelDSM does not fail */
estdsaspace = dsa_minimum_size();
shm_toc_estimate_chunk(&pcxt->estimator, estdsaspace);
shm_toc_estimate_keys(&pcxt->estimator, 2);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string)
@@ -780,7 +967,6 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
/* 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 +975,28 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot);
#endif
dsaspace = shm_toc_allocate(pcxt->toc, estdsaspace);
dsa = dsa_create_in_place(dsaspace, estdsaspace, hnsw_lock_tranche_id, pcxt->seg);
esthnswarea = maintenance_work_mem * 1024L;
hnswshared->dsaptr = dsa_allocate_extended(dsa, esthnswarea, DSA_ALLOC_HUGE | DSA_ALLOC_NO_OOM);
/* If not enough shared memory, back out (do serial build) */
if (!DsaPointerIsValid(hnswshared->dsaptr))
{
if (IsMVCCSnapshot(snapshot))
UnregisterSnapshot(snapshot);
DestroyParallelContext(pcxt);
ExitParallelMode();
return;
}
hnswarea = dsa_get_address(dsa, hnswshared->dsaptr);
/* 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, dsaspace);
/* Store query string for workers */
if (debug_query_string)
@@ -809,6 +1016,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 +1039,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 +1071,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 +1129,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"
/*
@@ -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)
@@ -331,10 +341,12 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
void
HnswUpdateNeighborPages(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)
@@ -480,34 +495,60 @@ 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);
return true;
}
/*
* Find duplicate element
*/
static bool
HnswFindDuplicate(Relation index, HnswElement element, bool building)
{
char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum value = HnswGetValue(base, element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
if (HnswAddDuplicate(index, element, neighborElement, building))
return true;
}
return false;
}
/*
* Write changes to disk
*/
static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint, bool building)
WriteElement(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 (HnswFindDuplicate(index, element, building))
return;
/* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage, building);
@@ -519,7 +560,7 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
/* Update neighbors */
HnswUpdateNeighborPages(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)
@@ -581,13 +609,10 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
}
/* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
dup = HnswFindDuplicate(element);
HnswInsertElement(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Write to disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint, building);
WriteElement(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;
}
@@ -625,44 +660,85 @@ CreatePairingHeapNode(HnswCandidate * c)
* 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 */
if (index == NULL)
v.pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL);
else
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);
/* Create local memory for neighborhood if needed */
if (index == 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 +746,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 +756,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 +765,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_EXCLUSIVE);
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 +817,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 +828,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 +853,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 +871,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 +912,59 @@ 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 +977,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 +994,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 +1010,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 +1019,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 +1032,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 +1040,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 +1052,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 +1070,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 +1087,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 +1112,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 +1142,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 +1152,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 +1170,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)
HnswInsertElement(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 +1236,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 +1260,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 +1274,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

@@ -205,15 +205,15 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
return;
/* Init fields */
HnswInitNeighbors(element, m);
element->heaptids = NIL;
HnswInitNeighbors(NULL, element, m, NULL);
element->heaptidsLength = 0;
/* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);
HnswInsertElement(NULL, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Update neighbor tuple */
/* Do this before getting page to minimize locking */
HnswSetNeighborTuple(ntup, element, m);
HnswSetNeighborTuple(NULL, ntup, element, m);
/* Get neighbor page */
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
@@ -301,7 +301,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();