Compare commits

..

29 Commits

Author SHA1 Message Date
Andrew Kane
76418fc093 Improved test [skip ci] 2024-01-23 15:13:14 -08:00
Andrew Kane
34d5a8cf3f Improved types 2024-01-23 15:08:50 -08:00
Andrew Kane
f9ae736c57 Updated comment [skip ci] 2024-01-23 15:01:48 -08:00
Andrew Kane
76c6dbb0a0 Added hamming_distance function 2024-01-23 15:01:11 -08:00
Heikki Linnakangas
c8be3a369b Include generic_xlog.h directly in the .c files where it's needed
There are no references to anything that's in generic_xlog.h in the
header files.
2024-01-23 13:04:03 +02:00
Heikki Linnakangas
e5d1a6bdbb Include reloptions.h directly in the .c files where it's needed
There are no references to anything that's in reloptions.h in the
header files. They need to include genam.h instead, which defines
IndexScanDesc.
2024-01-23 13:02:24 +02:00
Heikki Linnakangas
f31d708c2b Add direct include to pairingheap.h in headers
ivfflat.h and hnsw.h have references to pairingheap_node, so they need
to include lib/pairingheap.h. It happened to work, because
lib/pairingheap.h was being included indirectly through
nodes/execnodes.h, but let's be explicit.

Remove the include from hnswbuild.c, because there are no calls to
pairingheap functions in that file. Instead, add the includes to
hnswutils.c and ivfscan.c, which do have such calls. They are not
strictly necessary again because of the indirect include from hnsw.h
and ivfflat.h, but let's be explicit while we're messing with this.
2024-01-23 12:53:22 +02:00
Heikki Linnakangas
a1b1c99ff7 Remove unused #include
pg_list.h has not been used in hnswbuild.c since commit cb4c770df2.
2024-01-23 12:53:13 +02:00
Andrew Kane
3ace98add6 Changed storage for vector from extended to external 2024-01-23 00:00:12 -08:00
Andrew Kane
083008c21e Added validation for GUC parameters 2024-01-22 23:55:30 -08:00
Andrew Kane
a1e526ef82 Dropped support for Postgres 11 2024-01-22 23:52:54 -08:00
Andrew Kane
8ffb3718a4 Leave more space for other shared memory 2024-01-22 23:31:55 -08:00
Andrew Kane
2d0f162bd7 Added support for in-memory parallel index builds for HNSW
Co-authored-by: Heikki Linnakangas <heikki.linnakangas@iki.fi>
2024-01-22 23:19:10 -08:00
Heikki Linnakangas
4c6928bd3c Remove HnswSpool
It was just used to pass heap/index relations to
HnswParallelScanAndInsert. I think it was copied from nbtsort.c, which
is more complicated. I don't think we need a struct like this.

(That said, I actually think that we should have a state object that
would hold fields like 'heap', 'index', 'procinfo', 'collation'
etc. Passing that object around would simplify the signatures of many
functions. But that's a different story).
2024-01-22 23:11:25 -08:00
Heikki Linnakangas
6fd05dd6f6 Remove unused 'scantuplesortstates' field 2024-01-22 23:08:20 -08:00
Andrew Kane
70106f5413 Use assert checking for scan-build [skip ci] 2024-01-22 23:05:49 -08:00
Andrew Kane
44b90be452 Made variable name consistent across functions [skip ci] 2024-01-22 19:02:33 -08:00
Andrew Kane
31572a7b28 Removed unused parameter [skip ci] 2024-01-22 19:00:45 -08:00
Andrew Kane
2427290ea9 Pass hash by reference 2024-01-22 18:34:40 -08:00
Andrew Kane
ca71ef7a51 Added analyze to filtering tests 2024-01-21 18:08:16 -08:00
Andrew Kane
8bd01ff006 Added filtering tests for like [skip ci] 2024-01-21 18:01:42 -08:00
Andrew Kane
cdb1c9a6d3 Fixed test logic [skip ci] 2024-01-21 16:07:44 -08:00
Andrew Kane
bf34ceef7c Added more filtering tests 2024-01-21 16:03:22 -08:00
Andrew Kane
61b1566ea2 Updated CI to Debian 12 [skip ci] 2024-01-21 12:40:55 -08:00
Andrew Kane
91acc3c178 Added test for filtering with IVFFlat 2024-01-20 21:38:40 -08:00
Andrew Kane
885dd5b665 Fixed test warning [skip ci] 2024-01-20 21:10:24 -08:00
Andrew Kane
2a7b38bf1f Added test for filtering with few rows removed 2024-01-20 21:09:14 -08:00
Andrew Kane
4bd4a0996b Improved filtering test [skip ci] 2024-01-20 21:01:57 -08:00
Andrew Kane
490522b883 Added test for filtering with HNSW 2024-01-20 18:06:54 -08:00
28 changed files with 1219 additions and 789 deletions

View File

@@ -20,8 +20,6 @@ jobs:
os: ubuntu-20.04 os: ubuntu-20.04
- postgres: 12 - postgres: 12
os: ubuntu-20.04 os: ubuntu-20.04
- postgres: 11
os: ubuntu-20.04
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1
@@ -62,7 +60,7 @@ jobs:
wget -q https://github.com/postgres/postgres/archive/refs/tags/REL_14_5.tar.gz wget -q https://github.com/postgres/postgres/archive/refs/tags/REL_14_5.tar.gz
tar xf 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 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: windows:
runs-on: windows-latest runs-on: windows-latest
if: ${{ !startsWith(github.ref_name, 'mac') }} if: ${{ !startsWith(github.ref_name, 'mac') }}
@@ -83,10 +81,10 @@ jobs:
if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }} if: ${{ !startsWith(github.ref_name, 'mac') && !startsWith(github.ref_name, 'windows') }}
runs-on: ubuntu-latest runs-on: ubuntu-latest
container: container:
image: debian:11 image: debian:12
options: --platform linux/386 options: --platform linux/386
steps: steps:
- run: apt-get update && apt-get install -y build-essential git libipc-run-perl postgresql-13 postgresql-server-dev-13 sudo - run: apt-get update && apt-get install -y build-essential git libipc-run-perl postgresql-15 postgresql-server-dev-15 sudo
- run: service postgresql start - run: service postgresql start
- run: | - run: |
git clone https://github.com/${{ github.repository }}.git pgvector git clone https://github.com/${{ github.repository }}.git pgvector

View File

@@ -1,11 +1,17 @@
## 0.5.2 (unreleased) ## 0.6.0 (unreleased)
If upgrading with Postgres < 13, see [this note](https://github.com/pgvector/pgvector#060).
- Changed storage for vector from `extended` to `external`
- Improved performance of HNSW - Improved performance of HNSW
- Added support for on-disk parallel index builds for HNSW - Added support for parallel index builds for HNSW
- Added `hamming_distance` function
- Added validation for GUC parameters
- Reduced memory usage for HNSW index builds - Reduced memory usage for HNSW index builds
- Reduced WAL generation for HNSW index builds - Reduced WAL generation for HNSW index builds
- Fixed error with logical replication - Fixed error with logical replication
- Fixed `invalid memory alloc request size` error with HNSW index build - Fixed `invalid memory alloc request size` error with HNSW index build
- Dropped support for Postgres 11
## 0.5.1 (2023-10-10) ## 0.5.1 (2023-10-10)

View File

@@ -12,7 +12,7 @@
"prereqs": { "prereqs": {
"runtime": { "runtime": {
"requires": { "requires": {
"PostgreSQL": "11.0.0" "PostgreSQL": "12.0.0"
} }
} }
}, },

View File

@@ -729,6 +729,16 @@ SELECT extversion FROM pg_extension WHERE extname = 'vector';
## Upgrade Notes ## Upgrade Notes
### 0.6.0 [unreleased]
If upgrading with Postgres < 13, remove this line from `sql/vector--0.5.1--0.6.0.sql`:
```sql
ALTER TYPE vector SET (STORAGE = external);
```
Then run `make install` and `ALTER EXTENSION vector UPDATE;`.
### 0.4.0 ### 0.4.0
If upgrading with Postgres < 13, remove this line from `sql/vector--0.3.2--0.4.0.sql`: If upgrading with Postgres < 13, remove this line from `sql/vector--0.3.2--0.4.0.sql`:

View File

@@ -0,0 +1,8 @@
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "ALTER EXTENSION vector UPDATE TO '0.6.0'" to load this file. \quit
-- remove this single line for Postgres < 13
ALTER TYPE vector SET (STORAGE = external);
CREATE FUNCTION hamming_distance(bytea, bytea) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

View File

@@ -26,7 +26,7 @@ CREATE TYPE vector (
TYPMOD_IN = vector_typmod_in, TYPMOD_IN = vector_typmod_in,
RECEIVE = vector_recv, RECEIVE = vector_recv,
SEND = vector_send, SEND = vector_send,
STORAGE = extended STORAGE = external
); );
-- functions -- functions
@@ -290,3 +290,8 @@ CREATE OPERATOR CLASS vector_cosine_ops
OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops, OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
FUNCTION 1 vector_negative_inner_product(vector, vector), FUNCTION 1 vector_negative_inner_product(vector, vector),
FUNCTION 2 vector_norm(vector); FUNCTION 2 vector_norm(vector);
-- bytea functions
CREATE FUNCTION hamming_distance(bytea, bytea) RETURNS float8
AS 'MODULE_PATHNAME' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;

View File

@@ -4,25 +4,57 @@
#include <math.h> #include <math.h>
#include "access/amapi.h" #include "access/amapi.h"
#include "access/reloptions.h"
#include "commands/progress.h"
#include "commands/vacuum.h" #include "commands/vacuum.h"
#include "hnsw.h" #include "hnsw.h"
#include "utils/guc.h" #include "utils/guc.h"
#include "utils/selfuncs.h" #include "utils/selfuncs.h"
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM < 150000
#include "commands/progress.h" #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
bool hnsw_enable_parallel_build; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; 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 * Initialize index options and variables
*/ */
void void
HnswInit(void) HnswInit(void)
{ {
HnswInitLockTranche();
hnsw_relopt_kind = add_reloption_kind(); hnsw_relopt_kind = add_reloption_kind();
add_int_reloption(hnsw_relopt_kind, "m", "Max number of connections", add_int_reloption(hnsw_relopt_kind, "m", "Max number of connections",
HNSW_DEFAULT_M, HNSW_MIN_M, HNSW_MAX_M HNSW_DEFAULT_M, HNSW_MIN_M, HNSW_MAX_M
@@ -41,16 +73,12 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* Behind a variable for now since can be slower than building in memory */ MarkGUCPrefixReserved("hnsw");
DefineCustomBoolVariable("hnsw.enable_parallel_build", "Enables or disables building indexes in parallel",
NULL, &hnsw_enable_parallel_build,
false, PGC_USERSET, 0, NULL, NULL, NULL);
} }
/* /*
* Get the name of index build phase * Get the name of index build phase
*/ */
#if PG_VERSION_NUM >= 120000
static char * static char *
hnswbuildphasename(int64 phasenum) hnswbuildphasename(int64 phasenum)
{ {
@@ -64,7 +92,6 @@ hnswbuildphasename(int64 phasenum)
return NULL; return NULL;
} }
} }
#endif
/* /*
* Estimate the cost of an index scan * Estimate the cost of an index scan
@@ -79,9 +106,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
int m; int m;
int entryLevel; int entryLevel;
Relation index; Relation index;
#if PG_VERSION_NUM < 120000
List *qinfos;
#endif
/* Never use index without order */ /* Never use index without order */
if (path->indexorderbys == NULL) if (path->indexorderbys == NULL)
@@ -94,20 +118,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
return; return;
} }
/*
* Do not use index if no limit or limit + offset > ef_search unless
* enable_seqscan = off
*/
if (root->limit_tuples < 0 || root->limit_tuples > hnsw_ef_search)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
MemSet(&costs, 0, sizeof(costs)); MemSet(&costs, 0, sizeof(costs));
index = index_open(path->indexinfo->indexoid, NoLock); index = index_open(path->indexinfo->indexoid, NoLock);
@@ -121,12 +131,7 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
/* Account for number of tuples (or entry level), m, and ef_search */ /* Account for number of tuples (or entry level), m, and ef_search */
costs.numIndexTuples = (entryLevel + 2) * m; costs.numIndexTuples = (entryLevel + 2) * m;
#if PG_VERSION_NUM >= 120000
genericcostestimate(root, path, loop_count, &costs); genericcostestimate(root, path, loop_count, &costs);
#else
qinfos = deconstruct_indexquals(path);
genericcostestimate(root, path, loop_count, qinfos, &costs);
#endif
/* Use total cost since most work happens before first tuple is returned */ /* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost; *indexStartupCost = costs.indexTotalCost;
@@ -220,9 +225,7 @@ hnswhandler(PG_FUNCTION_ARGS)
amroutine->amcostestimate = hnswcostestimate; amroutine->amcostestimate = hnswcostestimate;
amroutine->amoptions = hnswoptions; amroutine->amoptions = hnswoptions;
amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */ amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */
#if PG_VERSION_NUM >= 120000
amroutine->ambuildphasename = hnswbuildphasename; amroutine->ambuildphasename = hnswbuildphasename;
#endif
amroutine->amvalidate = hnswvalidate; amroutine->amvalidate = hnswvalidate;
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
amroutine->amadjustmembers = NULL; amroutine->amadjustmembers = NULL;

View File

@@ -3,21 +3,17 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h" #include "access/genam.h"
#include "access/parallel.h" #include "access/parallel.h"
#include "access/reloptions.h" #include "lib/pairingheap.h"
#include "lib/ilist.h"
#include "nodes/execnodes.h" #include "nodes/execnodes.h"
#include "port.h" /* for random() */ #include "port.h" /* for random() */
#include "utils/relptr.h"
#include "utils/sampling.h" #include "utils/sampling.h"
#include "vector.h" #include "vector.h"
#if PG_VERSION_NUM < 110000
#error "Requires PostgreSQL 11+"
#endif
#if PG_VERSION_NUM < 120000 #if PG_VERSION_NUM < 120000
#include "access/relscan.h" #error "Requires PostgreSQL 12+"
#endif #endif
#define HNSW_MAX_DIM 2000 #define HNSW_MAX_DIM 2000
@@ -96,33 +92,62 @@
/* Ensure fits on page and in uint8 */ /* Ensure fits on page and in uint8 */
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255) #define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
#define HnswGetNeighbors(element, lc) (AssertMacro((element)->level >= (lc)), &(element)->neighbors[lc]) #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 */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build; extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray;
#define HnswPtrDeclare(type, relptrtype, ptrtype) \
relptr_declare(type, relptrtype); \
typedef union { type *ptr; relptrtype relptr; } ptrtype;
/* Pointers that can be absolute or relative */
/* Use char for DatumPtr so works with Pointer */
HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr);
typedef struct HnswElementData typedef struct HnswElementData
{ {
slist_node next; HnswElementPtr next;
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
uint8 heaptidsLength; uint8 heaptidsLength;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint32 hash; uint32 hash;
struct HnswNeighborArray *neighbors; HnswNeighborsPtr neighbors;
BlockNumber blkno; BlockNumber blkno;
OffsetNumber offno; OffsetNumber offno;
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
Datum value; DatumPtr value;
LWLock lock;
} HnswElementData; } HnswElementData;
typedef HnswElementData * HnswElement; typedef HnswElementData * HnswElement;
typedef struct HnswCandidate typedef struct HnswCandidate
{ {
HnswElement element; HnswElementPtr element;
float distance; float distance;
bool closer; bool closer;
} HnswCandidate; } HnswCandidate;
@@ -131,7 +156,7 @@ typedef struct HnswNeighborArray
{ {
int length; int length;
bool closerSet; bool closerSet;
HnswCandidate *items; HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborArray; } HnswNeighborArray;
typedef struct HnswPairingHeapNode typedef struct HnswPairingHeapNode
@@ -150,19 +175,24 @@ typedef struct HnswOptions
typedef struct HnswGraph typedef struct HnswGraph
{ {
slist_head elements; /* Graph state */
HnswElement entryPoint; slock_t lock;
HnswElementPtr head;
double indtuples;
/* Entry state */
LWLock entryLock;
HnswElementPtr entryPoint;
/* Allocations state */
LWLock allocatorLock;
long memoryUsed; long memoryUsed;
long memoryTotal; long memoryTotal;
bool flushed;
double indtuples;
} HnswGraph;
typedef struct HnswSpool /* Flushed state */
{ LWLock flushLock;
Relation heap; bool flushed;
Relation index; } HnswGraph;
} HnswSpool;
typedef struct HnswShared typedef struct HnswShared
{ {
@@ -170,7 +200,6 @@ typedef struct HnswShared
Oid heaprelid; Oid heaprelid;
Oid indexrelid; Oid indexrelid;
bool isconcurrent; bool isconcurrent;
int scantuplesortstates;
/* Worker progress */ /* Worker progress */
ConditionVariable workersdonecv; ConditionVariable workersdonecv;
@@ -182,16 +211,10 @@ typedef struct HnswShared
int nparticipantsdone; int nparticipantsdone;
double reltuples; double reltuples;
HnswGraph graphData; HnswGraph graphData;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} HnswShared; } HnswShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromHnswShared(shared) \ #define ParallelTableScanFromHnswShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared))) (ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(HnswShared)))
#endif
typedef struct HnswLeader typedef struct HnswLeader
{ {
@@ -199,8 +222,15 @@ typedef struct HnswLeader
int nparticipanttuplesorts; int nparticipanttuplesorts;
HnswShared *hnswshared; HnswShared *hnswshared;
Snapshot snapshot; Snapshot snapshot;
char *hnswarea;
} HnswLeader; } HnswLeader;
typedef struct HnswAllocator
{
void *(*alloc) (Size size, void *state);
void *state;
} HnswAllocator;
typedef struct HnswBuildState typedef struct HnswBuildState
{ {
/* Info */ /* Info */
@@ -233,10 +263,12 @@ typedef struct HnswBuildState
/* Memory */ /* Memory */
MemoryContext graphCtx; MemoryContext graphCtx;
MemoryContext tmpCtx; MemoryContext tmpCtx;
HnswAllocator allocator;
/* Parallel builds */ /* Parallel builds */
HnswLeader *hnswleader; HnswLeader *hnswleader;
HnswShared *hnswshared; HnswShared *hnswshared;
char *hnswarea;
} HnswBuildState; } HnswBuildState;
typedef struct HnswMetaPageData typedef struct HnswMetaPageData
@@ -335,23 +367,24 @@ bool HnswNormValue(FmgrInfo *procinfo, Oid collation, Datum *value, Vector * re
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement); List *HnswSearchLayer(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); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel); 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); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswCandidate *HnswEntryCandidate(HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); 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 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 HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
void HnswInitNeighbors(HnswElement element, int m); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, bool building); 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 HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec); void HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswSetElementTuple(HnswElementTuple etup, HnswElement element); void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation); void HnswUpdateConnection(char *base, HnswElement element, HnswCandidate * hc, int lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation);
void HnswLoadNeighbors(HnswElement element, Relation index, int m); void HnswLoadNeighbors(HnswElement element, Relation index, int m);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
@@ -371,6 +404,16 @@ void hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys,
bool hnswgettuple(IndexScanDesc scan, ScanDirection dir); bool hnswgettuple(IndexScanDesc scan, ScanDirection dir);
void hnswendscan(IndexScanDesc scan); void hnswendscan(IndexScanDesc scan);
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 */ /* Hash tables */
typedef struct TidHashEntry typedef struct TidHashEntry
{ {
@@ -398,4 +441,17 @@ typedef struct PointerHashEntry
#define SH_DECLARE #define SH_DECLARE
#include "lib/simplehash.h" #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 #endif

View File

@@ -1,14 +1,53 @@
/*
* The HNSW build happens in two phases:
*
* 1. In-memory phase
*
* In this first phase, the graph is held completely in memory. When the graph
* is fully built, or we run out of memory reserved for the build (determined
* by maintenance_work_mem), we materialize the graph to disk (see
* FlushPages()), and switch to the on-disk phase.
*
* In a parallel build, a large contiguous chunk of shared memory is allocated
* to hold the graph. Each worker process has its own HnswBuildState struct in
* private memory, which contains information that doesn't change throughout
* the build, and pointers to the shared structs in shared memory. The shared
* memory area is mapped to a different address in each worker process, and
* 'HnswBuildState.hnswarea' points to the beginning of the shared area in the
* worker process's address space. All pointers used in the graph are
* "relative pointers", stored as an offset from 'hnswarea'.
*
* Each element is protected by an LWLock. It must be held when reading or
* modifying the element's neighbors or 'heaptids'.
*
* In a non-parallel build, the graph is held in backend-private memory. All
* the elements are allocated in a dedicated memory context, 'graphCtx', and
* the pointers used in the graph are regular pointers.
*
* 2. On-disk phase
*
* In the on-disk phase, the index is built by inserting each vector to the
* index one by one, just like on INSERT. The only difference is that we don't
* WAL-log the individual inserts. If the graph fit completely in memory and
* was fully built in the in-memory phase, the on-disk phase is skipped.
*
* After we have finished building the graph, we perform one more scan through
* the index and write all the pages to the WAL.
*/
#include "postgres.h" #include "postgres.h"
#include <math.h> #include <math.h>
#include "access/parallel.h" #include "access/parallel.h"
#include "access/table.h"
#include "access/tableam.h"
#include "access/xact.h" #include "access/xact.h"
#include "access/xloginsert.h"
#include "catalog/index.h" #include "catalog/index.h"
#include "commands/progress.h"
#include "hnsw.h" #include "hnsw.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "lib/pairingheap.h" #include "optimizer/optimizer.h"
#include "nodes/pg_list.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h" #include "tcop/tcopprot.h"
#include "utils/datum.h" #include "utils/datum.h"
@@ -16,15 +55,8 @@
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
#include "utils/backend_progress.h" #include "utils/backend_progress.h"
#elif PG_VERSION_NUM >= 120000
#include "pgstat.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/tableam.h"
#include "commands/progress.h"
#else #else
#define PROGRESS_CREATEIDX_TUPLES_DONE 0 #include "pgstat.h"
#endif #endif
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
@@ -33,28 +65,16 @@
#define CALLBACK_ITEM_POINTER HeapTuple hup #define CALLBACK_ITEM_POINTER HeapTuple hup
#endif #endif
#if PG_VERSION_NUM >= 120000
#define UpdateProgress(index, val) pgstat_progress_update_param(index, val) #define UpdateProgress(index, val) pgstat_progress_update_param(index, val)
#else
#define UpdateProgress(index, val) ((void)val)
#endif
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
#include "utils/backend_status.h" #include "utils/backend_status.h"
#include "utils/wait_event.h" #include "utils/wait_event.h"
#endif #endif
#if PG_VERSION_NUM >= 120000
#include "access/table.h"
#include "optimizer/optimizer.h"
#else
#include "access/heapam.h"
#include "optimizer/planner.h"
#include "pgstat.h"
#endif
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001) #define PARALLEL_KEY_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)
#if PG_VERSION_NUM < 130000 #if PG_VERSION_NUM < 130000
#define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2) #define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2)
@@ -135,9 +155,11 @@ CreateElementPages(HnswBuildState * buildstate)
HnswElementTuple etup; HnswElementTuple etup;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
BlockNumber insertPage; BlockNumber insertPage;
HnswElement entryPoint;
Buffer buf; Buffer buf;
Page page; Page page;
slist_iter iter; HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
/* Calculate sizes */ /* Calculate sizes */
etupAllocSize = BLCKSZ; etupAllocSize = BLCKSZ;
@@ -152,18 +174,22 @@ CreateElementPages(HnswBuildState * buildstate)
page = BufferGetPage(buf); page = BufferGetPage(buf);
HnswInitPage(buf, page); HnswInitPage(buf, page);
slist_foreach(iter, &buildstate->graph->elements) while (!HnswPtrIsNull(base, iter))
{ {
HnswElement element = slist_container(HnswElementData, next, iter.cur); HnswElement element = HnswPtrAccess(base, iter);
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
void *valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */
iter = element->next;
/* Zero memory for each element */ /* Zero memory for each element */
MemSet(etup, 0, etupAllocSize); MemSet(etup, 0, etupAllocSize);
/* Calculate sizes */ /* 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); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -171,7 +197,7 @@ CreateElementPages(HnswBuildState * buildstate)
if (etupSize > etupAllocSize) if (etupSize > etupAllocSize)
elog(ERROR, "index tuple too large"); elog(ERROR, "index tuple too large");
HnswSetElementTuple(etup, element); HnswSetElementTuple(base, etup, element);
/* Keep element and neighbors on the same page if possible */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize)) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -212,7 +238,8 @@ CreateElementPages(HnswBuildState * buildstate)
MarkBufferDirty(buf); MarkBufferDirty(buf);
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->graph->entryPoint, insertPage, forkNum, true); entryPoint = HnswPtrAccess(base, buildstate->graph->entryPoint);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, entryPoint, insertPage, forkNum, true);
pfree(etup); pfree(etup);
pfree(ntup); pfree(ntup);
@@ -227,19 +254,23 @@ CreateNeighborPages(HnswBuildState * buildstate)
Relation index = buildstate->index; Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum; ForkNumber forkNum = buildstate->forkNum;
int m = buildstate->m; int m = buildstate->m;
slist_iter iter; HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
HnswNeighborTuple ntup; HnswNeighborTuple ntup;
/* Allocate once */ /* Allocate once */
ntup = palloc0(BLCKSZ); ntup = palloc0(BLCKSZ);
slist_foreach(iter, &buildstate->graph->elements) while (!HnswPtrIsNull(base, iter))
{ {
HnswElement e = slist_container(HnswElementData, next, iter.cur); HnswElement e = HnswPtrAccess(base, iter);
Buffer buf; Buffer buf;
Page page; Page page;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
/* Update iterator */
iter = e->next;
/* Can take a while, so ensure we can interrupt */ /* Can take a while, so ensure we can interrupt */
/* Needs to be called when no buffer locks are held */ /* Needs to be called when no buffer locks are held */
CHECK_FOR_INTERRUPTS(); CHECK_FOR_INTERRUPTS();
@@ -248,7 +279,7 @@ CreateNeighborPages(HnswBuildState * buildstate)
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
HnswSetNeighborTuple(ntup, e, m); HnswSetNeighborTuple(base, ntup, e, m);
if (!PageIndexTupleOverwrite(page, e->neighborOffno, (Item) ntup, ntupSize)) if (!PageIndexTupleOverwrite(page, e->neighborOffno, (Item) ntup, ntupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index)); elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
@@ -261,24 +292,6 @@ CreateNeighborPages(HnswBuildState * buildstate)
pfree(ntup); pfree(ntup);
} }
#ifdef HNSW_MEMORY
/*
* Show memory usage
*/
static void
ShowMemoryUsage(HnswBuildState * buildstate)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "graph memory: %zu MB, total memory: %zu MB",
MemoryContextMemAllocated(buildstate->graphCtx, false) / (1024 * 1024),
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
elog(INFO, "estimated memory: %zu MB", buildstate->memoryUsed / (1024 * 1024));
#endif
}
#endif
/* /*
* Flush pages * Flush pages
*/ */
@@ -286,7 +299,7 @@ static void
FlushPages(HnswBuildState * buildstate) FlushPages(HnswBuildState * buildstate)
{ {
#ifdef HNSW_MEMORY #ifdef HNSW_MEMORY
ShowMemoryUsage(buildstate); elog(INFO, "memory: %zu MB", buildstate->graph->memoryUsed / (1024 * 1024));
#endif #endif
CreateMetaPage(buildstate); CreateMetaPage(buildstate);
@@ -297,66 +310,172 @@ FlushPages(HnswBuildState * buildstate)
MemoryContextReset(buildstate->graphCtx); MemoryContextReset(buildstate->graphCtx);
} }
#if PG_VERSION_NUM < 130000
/* /*
* Get the memory used by an element * Add a heap TID to an existing element
*/ */
static long static bool
HnswElementMemory(HnswElement e, int m) AddDuplicateInMemory(HnswElement element, HnswElement dup)
{ {
long elementSize = sizeof(HnswElementData); LWLockAcquire(&dup->lock, LW_EXCLUSIVE);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1); if (dup->heaptidsLength == HNSW_HEAPTIDS)
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2)); {
elementSize += VARSIZE_ANY(DatumGetPointer(e->value)); LWLockRelease(&dup->lock);
/* Each allocation has a chunk header */ return false;
elementSize += (e->level + 4) * GENERATIONCHUNK_RAWSIZE; }
/* Add an extra 5% for alignment and other overhead */
return elementSize * 1.05; HnswAddHeapTid(dup, &element->heaptids[0]);
LWLockRelease(&dup->lock);
return true;
} }
#endif
/* /*
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
HnswFindDuplicateInMemory(HnswElement element) FindDuplicateInMemory(char *base, HnswElement element)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1)) if (!datumIsEqual(value, neighborValue, false, -1))
return false; return false;
/* Check for space */ /* Check for space */
if (neighbor->element->heaptidsLength < HNSW_HEAPTIDS) if (AddDuplicateInMemory(element, neighborElement))
{
HnswAddHeapTid(neighbor->element, &element->heaptids[0]);
return true; return true;
}
} }
return false; return false;
} }
/* /*
* Insert tuple into in-memory graph * Add to element list
*/ */
static bool static void
InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuildState * buildstate) AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
{
SpinLockAcquire(&graph->lock);
element->next = graph->head;
HnswPtrStore(base, graph->head, element);
SpinLockRelease(&graph->lock);
}
/*
* Update neighbors
*/
static void
UpdateNeighborsInMemory(char *base, FmgrInfo *procinfo, Oid collation, HnswElement e, int m)
{
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *hc = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, hc->element);
/* Keep scan-build happy on Mac x86-64 */
Assert(neighborElement);
/* Use element for lock instead of hc since hc can be replaced */
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, e, hc, lm, lc, NULL, NULL, procinfo, collation);
LWLockRelease(&neighborElement->lock);
}
}
}
/*
* Update graph in memory
*/
static void
UpdateGraphInMemory(FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
{
HnswGraph *graph = buildstate->graph;
char *base = buildstate->hnswarea;
/* Look for duplicate */
if (FindDuplicateInMemory(base, element))
return;
/* Add element */
AddElementInMemory(base, graph, element);
/* Update neighbors */
UpdateNeighborsInMemory(base, procinfo, collation, element, m);
/* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswPtrStore(base, graph->entryPoint, element);
}
/*
* Insert tuple in memory
*/
static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{ {
FmgrInfo *procinfo = buildstate->procinfo; FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation; Oid collation = buildstate->collation;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswElement entryPoint = graph->entryPoint; HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock;
int efConstruction = buildstate->efConstruction; int efConstruction = buildstate->efConstruction;
int m = buildstate->m; int m = buildstate->m;
MemoryContext oldCtx; char *base = buildstate->hnswarea;
/* Get entry point */
LWLockAcquire(entryLock, LW_SHARED);
entryPoint = HnswPtrAccess(base, graph->entryPoint);
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
{
/* Release shared lock */
LWLockRelease(entryLock);
/* Get exclusive lock */
LWLockAcquire(entryLock, LW_EXCLUSIVE);
/* Get latest entry point after lock is acquired */
entryPoint = HnswPtrAccess(base, graph->entryPoint);
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
/* Release entry lock */
LWLockRelease(entryLock);
}
/*
* Insert tuple
*/
static bool
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
{
HnswGraph *graph = buildstate->graph;
HnswElement element; HnswElement element;
HnswAllocator *allocator = &buildstate->allocator;
Size valueSize;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea;
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -364,73 +483,83 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
/* Normalize if needed */ /* Normalize if needed */
if (buildstate->normprocinfo != NULL) if (buildstate->normprocinfo != NULL)
{ {
if (!HnswNormValue(buildstate->normprocinfo, collation, &value, buildstate->normvec)) if (!HnswNormValue(buildstate->normprocinfo, buildstate->collation, &value, buildstate->normvec))
return false; return false;
} }
/* Allocate element in graph memory context */ /* Get datum size */
oldCtx = MemoryContextSwitchTo(buildstate->graphCtx); valueSize = VARSIZE_ANY(DatumGetPointer(value));
element = HnswInitElement(heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel);
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Update memory usage */ /* Ensure graph not flushed when inserting */
#if PG_VERSION_NUM >= 130000 LWLockAcquire(flushLock, LW_SHARED);
graph->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
#else
graph->memoryUsed += HnswElementMemory(element, buildstate->m);
#endif
/* Insert element in graph */ /* Are we in the on-disk phase? */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false); if (graph->flushed)
/* Look for duplicate */
if (HnswFindDuplicateInMemory(element))
{ {
/* No need to free element since memory unlikely to be reallocated */ LWLockRelease(flushLock);
return true;
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
} }
/* Add element */ /*
slist_push_head(&graph->elements, &element->next); * In a parallel build, the HnswElement is allocated from the shared
* memory area, so we need to coordinate with other processes.
*/
LWLockAcquire(&graph->allocatorLock, LW_EXCLUSIVE);
/* Update neighbors */ /*
for (int lc = element->level; lc >= 0; lc--) * Check that we have enough memory available for the new element now that
* we have the allocator lock, and flush pages if needed.
*/
if (graph->memoryUsed >= graph->memoryTotal)
{ {
int lm = HnswGetLayerM(m, lc); LWLockRelease(&graph->allocatorLock);
HnswNeighborArray *neighbors = HnswGetNeighbors(element, lc);
for (int i = 0; i < neighbors->length; i++) LWLockRelease(flushLock);
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation); LWLockAcquire(flushLock, LW_EXCLUSIVE);
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.")));
FlushPages(buildstate);
}
LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, buildstate->heap, true);
} }
/* Update entry point if needed */ /* Ok, we can proceed to allocate the element */
if (entryPoint == NULL || element->level > entryPoint->level) element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
graph->entryPoint = element; valuePtr = HnswAlloc(allocator, valueSize);
/*
* We have now allocated the space needed for the element, so we don't
* need the allocator lock anymore. Release it and initialize the rest of
* the element.
*/
LWLockRelease(&graph->allocatorLock);
/* Copy the datum */
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
/* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
/* Insert tuple */
InsertTupleInMemory(buildstate, element);
/* Release flush lock */
LWLockRelease(flushLock);
return true; return true;
} }
/*
* Acquire a lock if needed
*/
static inline void
HnswLockAcquire(HnswShared * hnswshared)
{
if (hnswshared)
SpinLockAcquire(&hnswshared->mutex);
}
/*
* Release a lock if needed
*/
static inline void
HnswLockRelease(HnswShared * hnswshared)
{
if (hnswshared)
SpinLockRelease(&hnswshared->mutex);
}
/* /*
* Callback for table_index_build_scan * Callback for table_index_build_scan
*/ */
@@ -440,9 +569,7 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
{ {
HnswBuildState *buildstate = (HnswBuildState *) state; HnswBuildState *buildstate = (HnswBuildState *) state;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswShared *hnswshared = buildstate->hnswshared;
MemoryContext oldCtx; MemoryContext oldCtx;
bool inserted;
#if PG_VERSION_NUM < 130000 #if PG_VERSION_NUM < 130000
ItemPointer tid = &hup->t_self; ItemPointer tid = &hup->t_self;
@@ -452,31 +579,16 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0]) if (isnull[0])
return; return;
/* Flush pages if needed */ /* Use memory context */
if (!graph->flushed && graph->memoryUsed >= graph->memoryTotal)
{
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.")));
FlushPages(buildstate);
}
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */ /* Insert tuple */
if (graph->flushed) if (InsertTuple(index, values, isnull, tid, buildstate))
inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true);
else
inserted = InsertTupleInMemory(index, values, tid, buildstate);
/* Update progress */
if (inserted)
{ {
HnswLockAcquire(hnswshared); /* Update progress */
SpinLockAcquire(&graph->lock);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples); UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples);
HnswLockRelease(hnswshared); SpinLockRelease(&graph->lock);
} }
/* Reset memory context */ /* Reset memory context */
@@ -488,14 +600,59 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
* Initialize the graph * Initialize the graph
*/ */
static void static void
InitGraph(HnswGraph * graph) InitGraph(HnswGraph * graph, char *base, long memoryTotal)
{ {
slist_init(&graph->elements); HnswPtrStore(base, graph->head, (HnswElement) NULL);
graph->entryPoint = NULL; HnswPtrStore(base, graph->entryPoint, (HnswElement) NULL);
graph->memoryUsed = 0; graph->memoryUsed = 0;
graph->memoryTotal = maintenance_work_mem * 1024L; graph->memoryTotal = memoryTotal;
graph->flushed = false; graph->flushed = false;
graph->indtuples = 0; 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);
}
/*
* 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;
} }
/* /*
@@ -531,7 +688,7 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0]; buildstate->collation = index->rd_indcollation[0];
InitGraph(&buildstate->graphData); InitGraph(&buildstate->graphData, NULL, maintenance_work_mem * 1024L);
buildstate->graph = &buildstate->graphData; buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
@@ -549,8 +706,11 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
"Hnsw build temporary context", "Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
InitAllocator(&buildstate->allocator, &HnswMemoryContextAlloc, buildstate);
buildstate->hnswleader = NULL; buildstate->hnswleader = NULL;
buildstate->hnswshared = NULL; buildstate->hnswshared = NULL;
buildstate->hnswarea = NULL;
} }
/* /*
@@ -581,6 +741,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
if (hnswshared->nparticipantsdone == nparticipanttuplesorts) if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{ {
buildstate->graph = &hnswshared->graphData; buildstate->graph = &hnswshared->graphData;
buildstate->hnswarea = buildstate->hnswleader->hnswarea;
reltuples = hnswshared->reltuples; reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex); SpinLockRelease(&hnswshared->mutex);
break; break;
@@ -600,35 +761,25 @@ ParallelHeapScan(HnswBuildState * buildstate)
* Perform a worker's portion of a parallel insert * Perform a worker's portion of a parallel insert
*/ */
static void static void
HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool progress) HnswParallelScanAndInsert(Relation heapRel, Relation indexRel, HnswShared * hnswshared, char *hnswarea, bool progress)
{ {
HnswBuildState buildstate; HnswBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan; TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples; double reltuples;
IndexInfo *indexInfo; IndexInfo *indexInfo;
/* Join parallel scan */ /* Join parallel scan */
indexInfo = BuildIndexInfo(hnswspool->index); indexInfo = BuildIndexInfo(indexRel);
indexInfo->ii_Concurrent = hnswshared->isconcurrent; indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM); InitBuildState(&buildstate, heapRel, indexRel, indexInfo, MAIN_FORKNUM);
buildstate.graph = &hnswshared->graphData; buildstate.graph = &hnswshared->graphData;
buildstate.hnswshared = hnswshared; buildstate.hnswarea = hnswarea;
#if PG_VERSION_NUM >= 120000 InitAllocator(&buildstate.allocator, &HnswSharedMemoryAlloc, &buildstate);
scan = table_beginscan_parallel(hnswspool->heap, scan = table_beginscan_parallel(heapRel,
ParallelTableScanFromHnswShared(hnswshared)); ParallelTableScanFromHnswShared(hnswshared));
reltuples = table_index_build_scan(hnswspool->heap, hnswspool->index, indexInfo, reltuples = table_index_build_scan(heapRel, indexRel, indexInfo,
true, progress, BuildCallback, true, progress, BuildCallback,
(void *) &buildstate, scan); (void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(hnswspool->heap, &hnswshared->heapdesc);
reltuples = IndexBuildHeapScan(hnswspool->heap, hnswspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Record statistics */ /* Record statistics */
SpinLockAcquire(&hnswshared->mutex); SpinLockAcquire(&hnswshared->mutex);
@@ -655,8 +806,8 @@ void
HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc) HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
{ {
char *sharedquery; char *sharedquery;
HnswSpool *hnswspool;
HnswShared *hnswshared; HnswShared *hnswshared;
char *hnswarea;
Relation heapRel; Relation heapRel;
Relation indexRel; Relation indexRel;
LOCKMODE heapLockmode; LOCKMODE heapLockmode;
@@ -685,28 +836,17 @@ HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc)
} }
/* Open relations within worker */ /* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(hnswshared->heaprelid, heapLockmode); heapRel = table_open(hnswshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(hnswshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(hnswshared->indexrelid, indexLockmode); indexRel = index_open(hnswshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */ hnswarea = shm_toc_lookup(toc, PARALLEL_KEY_HNSW_AREA, false);
hnswspool = (HnswSpool *) palloc0(sizeof(HnswSpool));
hnswspool->heap = heapRel;
hnswspool->index = indexRel;
/* Perform inserts */ /* Perform inserts */
HnswParallelScanAndInsert(hnswspool, hnswshared, false); HnswParallelScanAndInsert(heapRel, indexRel, hnswshared, hnswarea, false);
/* Close relations within worker */ /* Close relations within worker */
index_close(indexRel, indexLockmode); index_close(indexRel, indexLockmode);
#if PG_VERSION_NUM >= 120000
table_close(heapRel, heapLockmode); table_close(heapRel, heapLockmode);
#else
heap_close(heapRel, heapLockmode);
#endif
} }
/* /*
@@ -731,19 +871,7 @@ HnswEndParallel(HnswLeader * hnswleader)
static Size static Size
ParallelEstimateShared(Relation heap, Snapshot snapshot) ParallelEstimateShared(Relation heap, Snapshot snapshot)
{ {
#if PG_VERSION_NUM >= 120000
return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot)); return add_size(BUFFERALIGN(sizeof(HnswShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(HnswShared);
}
return add_size(offsetof(HnswShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
} }
/* /*
@@ -753,15 +881,9 @@ static void
HnswLeaderParticipateAsWorker(HnswBuildState * buildstate) HnswLeaderParticipateAsWorker(HnswBuildState * buildstate)
{ {
HnswLeader *hnswleader = buildstate->hnswleader; HnswLeader *hnswleader = buildstate->hnswleader;
HnswSpool *leaderworker;
/* Allocate memory and initialize private spool */
leaderworker = (HnswSpool *) palloc0(sizeof(HnswSpool));
leaderworker->heap = buildstate->heap;
leaderworker->index = buildstate->index;
/* Perform work common to all participants */ /* Perform work common to all participants */
HnswParallelScanAndInsert(leaderworker, hnswleader->hnswshared, true); HnswParallelScanAndInsert(buildstate->heap, buildstate->index, hnswleader->hnswshared, hnswleader->hnswarea, true);
} }
/* /*
@@ -771,10 +893,12 @@ static void
HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request) HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
{ {
ParallelContext *pcxt; ParallelContext *pcxt;
int scantuplesortstates;
Snapshot snapshot; Snapshot snapshot;
Size esthnswshared; Size esthnswshared;
Size esthnswarea;
Size estother;
HnswShared *hnswshared; HnswShared *hnswshared;
char *hnswarea;
HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader)); HnswLeader *hnswleader = (HnswLeader *) palloc0(sizeof(HnswLeader));
bool leaderparticipates = true; bool leaderparticipates = true;
int querylen; int querylen;
@@ -786,13 +910,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
/* Enter parallel mode and create context */ /* Enter parallel mode and create context */
EnterParallelMode(); EnterParallelMode();
Assert(request > 0); Assert(request > 0);
#if PG_VERSION_NUM >= 120000
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request); pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "HnswParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request;
/* Get snapshot for table scan */ /* Get snapshot for table scan */
if (!isconcurrent) if (!isconcurrent)
@@ -803,7 +921,17 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
/* Estimate size of workspaces */ /* Estimate size of workspaces */
esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot); esthnswshared = ParallelEstimateShared(buildstate->heap, snapshot);
shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared); shm_toc_estimate_chunk(&pcxt->estimator, esthnswshared);
shm_toc_estimate_keys(&pcxt->estimator, 1);
/* Leave space for other objects in shared memory */
/* Docker has a default limit of 64 MB for shm_size */
/* which happens to be the default value of maintenance_work_mem */
esthnswarea = maintenance_work_mem * 1024L;
estother = 3 * 1024 * 1024;
if (esthnswarea > estother)
esthnswarea -= estother;
shm_toc_estimate_chunk(&pcxt->estimator, esthnswarea);
shm_toc_estimate_keys(&pcxt->estimator, 2);
/* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */ /* Finally, estimate PARALLEL_KEY_QUERY_TEXT space */
if (debug_query_string) if (debug_query_string)
@@ -834,25 +962,21 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
hnswshared->heaprelid = RelationGetRelid(buildstate->heap); hnswshared->heaprelid = RelationGetRelid(buildstate->heap);
hnswshared->indexrelid = RelationGetRelid(buildstate->index); hnswshared->indexrelid = RelationGetRelid(buildstate->index);
hnswshared->isconcurrent = isconcurrent; hnswshared->isconcurrent = isconcurrent;
hnswshared->scantuplesortstates = scantuplesortstates;
ConditionVariableInit(&hnswshared->workersdonecv); ConditionVariableInit(&hnswshared->workersdonecv);
SpinLockInit(&hnswshared->mutex); SpinLockInit(&hnswshared->mutex);
/* Initialize mutable state */ /* Initialize mutable state */
hnswshared->nparticipantsdone = 0; hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0; hnswshared->reltuples = 0;
InitGraph(&hnswshared->graphData);
/* TODO Support in-memory builds */
hnswshared->graphData.memoryTotal = 0;
hnswshared->graphData.flushed = true;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap, table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared), ParallelTableScanFromHnswShared(hnswshared),
snapshot); snapshot);
#else
heap_parallelscan_initialize(&hnswshared->heapdesc, buildstate->heap, snapshot); hnswarea = (char *) shm_toc_allocate(pcxt->toc, esthnswarea);
#endif /* 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_SHARED, hnswshared);
shm_toc_insert(pcxt->toc, PARALLEL_KEY_HNSW_AREA, hnswarea);
/* Store query string for workers */ /* Store query string for workers */
if (debug_query_string) if (debug_query_string)
@@ -872,6 +996,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
hnswleader->nparticipanttuplesorts++; hnswleader->nparticipanttuplesorts++;
hnswleader->hnswshared = hnswshared; hnswleader->hnswshared = hnswshared;
hnswleader->snapshot = snapshot; hnswleader->snapshot = snapshot;
hnswleader->hnswarea = hnswarea;
/* If no workers were successfully launched, back out (do serial build) */ /* If no workers were successfully launched, back out (do serial build) */
if (pcxt->nworkers_launched == 0) if (pcxt->nworkers_launched == 0)
@@ -926,16 +1051,12 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD); UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */ /* Calculate parallel workers */
if (buildstate->heap != NULL && hnsw_enable_parallel_build) if (buildstate->heap != NULL)
parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index); parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);
/* Attempt to launch parallel worker scan when required */ /* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0) if (parallel_workers > 0)
{
/* TODO Support in-memory builds */
FlushPages(buildstate);
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers); HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to graph */ /* Add tuples to graph */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)
@@ -943,15 +1064,8 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
if (buildstate->hnswleader) if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate); buildstate->reltuples = ParallelHeapScan(buildstate);
else else
{
#if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL); true, true, BuildCallback, (void *) buildstate, NULL);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
buildstate->indtuples = buildstate->graph->indtuples; buildstate->indtuples = buildstate->graph->indtuples;
} }
@@ -965,22 +1079,6 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
HnswEndParallel(buildstate->hnswleader); HnswEndParallel(buildstate->hnswleader);
} }
#if PG_VERSION_NUM < 110008
void
log_newpage_range(Relation rel, ForkNumber forkNum, BlockNumber startblk, BlockNumber endblk, bool page_std)
{
for (BlockNumber blkno = startblk; blkno < endblk; blkno++)
{
Buffer buf = ReadBufferExtended(rel, forkNum, blkno, RBM_NORMAL, NULL);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
MarkBufferDirty(buf);
log_newpage_buffer(buf, page_std);
UnlockReleaseBuffer(buf);
}
}
#endif
/* /*
* Build the index * Build the index
*/ */

View File

@@ -2,6 +2,7 @@
#include <math.h> #include <math.h>
#include "access/generic_xlog.h"
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"
@@ -116,7 +117,7 @@ HnswInsertAppendPage(Relation index, Buffer *nbuf, Page *npage, GenericXLogState
* Add to element and neighbor pages * Add to element and neighbor pages
*/ */
static void static void
WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building) AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, BlockNumber *updatedInsertPage, bool building)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -134,9 +135,10 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
OffsetNumber freeOffno = InvalidOffsetNumber; OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber; OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
char *base = NULL;
/* Calculate sizes */ /* 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); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -144,11 +146,11 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
/* Prepare element tuple */ /* Prepare element tuple */
etup = palloc0(etupSize); etup = palloc0(etupSize);
HnswSetElementTuple(etup, e); HnswSetElementTuple(base, etup, e);
/* Prepare neighbor tuple */ /* Prepare neighbor tuple */
ntup = palloc0(ntupSize); ntup = palloc0(ntupSize);
HnswSetNeighborTuple(ntup, e, m); HnswSetNeighborTuple(base, ntup, e, m);
/* Find a page (or two if needed) to insert the tuples */ /* Find a page (or two if needed) to insert the tuples */
for (;;) for (;;)
@@ -338,12 +340,14 @@ ConnectionExists(HnswElement e, HnswNeighborTuple ntup, int startIdx, int lm)
* Update neighbors * Update neighbors
*/ */
void void
HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building) HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement e, int m, bool checkExisting, bool building)
{ {
char *base = NULL;
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc); HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
@@ -356,11 +360,12 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
Size ntupSize; Size ntupSize;
int idx = -1; int idx = -1;
int startIdx; 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 */ /* Get latest neighbors since they may have changed */
/* Do not lock yet since selecting neighbors can take time */ /* 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 * Could improve performance for vacuuming by checking neighbors
@@ -370,14 +375,14 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
*/ */
/* Select neighbors */ /* 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 */ /* New element was not selected as a neighbor */
if (idx == -1) if (idx == -1)
continue; continue;
/* Register page */ /* Register page */
buf = ReadBuffer(index, hc->element->neighborPage); buf = ReadBuffer(index, neighborElement->neighborPage);
LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE); LockBuffer(buf, BUFFER_LOCK_EXCLUSIVE);
if (building) if (building)
{ {
@@ -396,7 +401,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
ntupSize = ItemIdGetLength(itemid); ntupSize = ItemIdGetLength(itemid);
/* Calculate index for update */ /* Calculate index for update */
startIdx = (hc->element->level - lc) * m; startIdx = (neighborElement->level - lc) * m;
/* Check for existing connection */ /* Check for existing connection */
if (checkExisting && ConnectionExists(e, ntup, startIdx, lm)) if (checkExisting && ConnectionExists(e, ntup, startIdx, lm))
@@ -447,7 +452,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
* Add a heap TID to an existing element * Add a heap TID to an existing element
*/ */
static bool static bool
HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool building) AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool building)
{ {
Buffer buf; Buffer buf;
Page page; Page page;
@@ -511,19 +516,23 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool buil
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
HnswFindDuplicate(Relation index, HnswElement element, bool building) FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0); char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1)) if (!datumIsEqual(value, neighborValue, false, -1))
return false; return false;
if (HnswAddDuplicate(index, element, neighbor->element, building)) if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true; return true;
} }
@@ -531,26 +540,26 @@ HnswFindDuplicate(Relation index, HnswElement element, bool building)
} }
/* /*
* Write changes to disk * Update graph on disk
*/ */
static void static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building) UpdateGraphOnDisk(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Look for duplicate */ /* Look for duplicate */
if (HnswFindDuplicate(index, element, building)) if (FindDuplicateOnDisk(index, element, building))
return; return;
/* Write element and neighbor tuples */ /* Add element */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage, building); AddElementOnDisk(index, element, m, GetInsertPage(index), &newInsertPage, building);
/* Update insert page if needed */ /* Update insert page if needed */
if (BlockNumberIsValid(newInsertPage)) if (BlockNumberIsValid(newInsertPage))
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building); HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false, building); HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building);
/* Update entry point if needed */ /* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -561,10 +570,8 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, Relation heapRel, 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 entryPoint;
HnswElement element; HnswElement element;
int m; int m;
@@ -572,17 +579,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0]; Oid collation = index->rd_indcollation[0];
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
char *base = NULL;
/* 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;
}
/* /*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -595,8 +592,8 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
/* Create an element */ /* Create an element */
element = HnswInitElement(heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m)); element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
element->value = value; HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -612,11 +609,11 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
entryPoint = HnswGetEntryPoint(index); entryPoint = HnswGetEntryPoint(index);
} }
/* Insert element in graph */ /* Find neighbors for element */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false); HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Write to disk */ /* Update graph on disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, entryPoint, building); UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -624,6 +621,30 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
return true; 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 * Insert a tuple into the index
*/ */
@@ -650,7 +671,7 @@ hnswinsert(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid,
oldCtx = MemoryContextSwitchTo(insertCtx); oldCtx = MemoryContextSwitchTo(insertCtx);
/* Insert tuple */ /* Insert tuple */
HnswInsertTuple(index, values, isnull, heap_tid, heap, false); HnswInsertTuple(index, values, isnull, heap_tid, heap);
/* Delete memory context */ /* Delete memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);

View File

@@ -21,6 +21,7 @@ GetScanItems(IndexScanDesc scan, Datum q)
List *w; List *w;
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
@@ -28,15 +29,15 @@ GetScanItems(IndexScanDesc scan, Datum q)
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; 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--) 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; 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,26 +185,23 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
while (list_length(so->w) > 0) while (list_length(so->w) > 0)
{ {
char *base = NULL;
HnswCandidate *hc = llast(so->w); HnswCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid; ItemPointer heaptid;
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (hc->element->heaptidsLength == 0) if (element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);
continue; continue;
} }
heaptid = &hc->element->heaptids[--hc->element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
#if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
#else
scan->xs_ctup.t_self = *heaptid;
#endif
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
} }

View File

@@ -2,9 +2,12 @@
#include <math.h> #include <math.h>
#include "access/generic_xlog.h"
#include "hnsw.h" #include "hnsw.h"
#include "lib/pairingheap.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h" #include "utils/datum.h"
#include "utils/rel.h"
#include "vector.h" #include "vector.h"
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
@@ -56,13 +59,6 @@ hash_tid(ItemPointerData tid)
#define SH_DEFINE #define SH_DEFINE
#include "lib/simplehash.h" #include "lib/simplehash.h"
/* Needed to include simplehash.h twice */
#if PG_VERSION_NUM < 120000
#undef SH_EQUAL
#define sh_log2 pointerhash_sh_log2
#define sh_pow2 pointerhash_sh_pow2
#endif
/* Pointer hash table */ /* Pointer hash table */
static uint32 static uint32
hash_pointer(uintptr_t ptr) hash_pointer(uintptr_t ptr)
@@ -84,9 +80,31 @@ hash_pointer(uintptr_t ptr)
#define SH_DEFINE #define SH_DEFINE
#include "lib/simplehash.h" #include "lib/simplehash.h"
/* 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 typedef union
{ {
pointerhash_hash *pointers; pointerhash_hash *pointers;
offsethash_hash *offsets;
tidhash_hash *tids; tidhash_hash *tids;
} visited_hash; } visited_hash;
@@ -188,31 +206,46 @@ HnswInitPage(Buffer buf, Page page)
* Allocate neighbors * Allocate neighbors
*/ */
void void
HnswInitNeighbors(HnswElement element, int m) HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * allocator)
{ {
int level = element->level; 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++) for (int lc = 0; lc <= level; lc++)
{ {
HnswNeighborArray *a; HnswNeighborArray *a;
int lm = HnswGetLayerM(m, lc); 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->length = 0;
a->items = palloc(sizeof(HnswCandidate) * lm);
a->closerSet = false; a->closerSet = false;
} }
} }
/*
* Allocate memory from the allocator
*/
void *
HnswAlloc(HnswAllocator * allocator, Size size)
{
if (allocator)
return (*(allocator)->alloc) (size, (allocator)->state);
return palloc(size);
}
/* /*
* Allocate an element * Allocate an element
*/ */
HnswElement 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); int level = (int) (-log(RandomDouble()) * ml);
@@ -226,9 +259,9 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
element->level = level; element->level = level;
element->deleted = 0; element->deleted = 0;
HnswInitNeighbors(element, m); HnswInitNeighbors(base, element, m, allocator);
element->value = PointerGetDatum(NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
return element; return element;
} }
@@ -249,11 +282,12 @@ HnswElement
HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno) HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
{ {
HnswElement element = palloc(sizeof(HnswElementData)); HnswElement element = palloc(sizeof(HnswElementData));
char *base = NULL;
element->blkno = blkno; element->blkno = blkno;
element->offno = offno; element->offno = offno;
element->neighbors = NULL; HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
element->value = PointerGetDatum(NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
return element; return element;
} }
@@ -363,8 +397,10 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */
void 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->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
etup->deleted = 0; etup->deleted = 0;
@@ -375,14 +411,14 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
else else
ItemPointerSetInvalid(&etup->heaptids[i]); 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 * Set neighbor tuple
*/ */
void void
HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m) HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
{ {
int idx = 0; int idx = 0;
@@ -390,7 +426,7 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc); HnswNeighborArray *neighbors = HnswGetNeighbors(base, e, lc);
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
for (int i = 0; i < lm; i++) for (int i = 0; i < lm; i++)
@@ -400,8 +436,9 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
if (i < neighbors->length) if (i < neighbors->length)
{ {
HnswCandidate *hc = &neighbors->items[i]; 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 else
ItemPointerSetInvalid(indextid); ItemPointerSetInvalid(indextid);
@@ -417,12 +454,14 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
static void static void
LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m) LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
{ {
char *base = NULL;
HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno)); HnswNeighborTuple ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
int neighborCount = (element->level + 2) * m; int neighborCount = (element->level + 2) * m;
Assert(HnswIsNeighborTuple(ntup)); Assert(HnswIsNeighborTuple(ntup));
HnswInitNeighbors(element, m); HnswInitNeighbors(base, element, m, NULL);
/* Ensure expected neighbors */ /* Ensure expected neighbors */
if (ntup->count != neighborCount) if (ntup->count != neighborCount)
@@ -448,9 +487,9 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
if (level < 0) if (level < 0)
level = 0; level = 0;
neighbors = HnswGetNeighbors(element, level); neighbors = HnswGetNeighbors(base, element, level);
hc = &neighbors->items[neighbors->length++]; hc = &neighbors->items[neighbors->length++];
hc->element = e; HnswPtrStore(base, hc->element, e);
} }
} }
@@ -497,7 +536,12 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
} }
if (loadVec) 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));
}
} }
/* /*
@@ -533,24 +577,27 @@ HnswLoadElement(HnswElement element, float *distance, Datum *q, Relation index,
* Get the distance for a candidate * Get the distance for a candidate
*/ */
static float 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 * Create a candidate for the entry point
*/ */
HnswCandidate * 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)); HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = entryPoint; HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL) if (index == NULL)
hc->distance = GetCandidateDistance(hc, q, procinfo, collation); hc->distance = GetCandidateDistance(base, hc, q, procinfo, collation);
else else
HnswLoadElement(hc->element, &hc->distance, &q, index, procinfo, collation, loadVec); HnswLoadElement(entryPoint, &hc->distance, &q, index, procinfo, collation, loadVec);
return hc; return hc;
} }
@@ -596,34 +643,79 @@ CreatePairingHeapNode(HnswCandidate * c)
return node; return node;
} }
/*
* Init visited
*/
static inline void
InitVisited(char *base, visited_hash * v, Relation index, int ef, int m)
{
if (index != NULL)
v->tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL);
else if (base != NULL)
v->offsets = offsethash_create(CurrentMemoryContext, ef * m * 2, NULL);
else
v->pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL);
}
/* /*
* Add to visited * Add to visited
*/ */
static inline void 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 #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 #else
pointerhash_insert(v.pointers, (uintptr_t) hc->element, found); offsethash_insert(v->offsets, HnswPtrOffset(hc->element), found);
#endif #endif
} }
else else
{ {
ItemPointerData indextid; #if PG_VERSION_NUM >= 130000
HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointerSet(&indextid, hc->element->blkno, hc->element->offno); pointerhash_insert_hash(v->pointers, (uintptr_t) HnswPtrPointer(hc->element), element->hash, found);
tidhash_insert(v.tids, indextid, 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 * Algorithm 2 from paper
*/ */
List * 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)
{ {
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -631,12 +723,17 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
int wlen = 0; int wlen = 0;
visited_hash v; visited_hash v;
ListCell *lc2; ListCell *lc2;
HnswNeighborArray *neighborhoodData = NULL;
Size neighborhoodSize;
/* Create hash table */ InitVisited(base, &v, index, ef, m);
/* Create local memory for neighborhood if needed */
if (index == NULL) if (index == NULL)
v.pointers = pointerhash_create(CurrentMemoryContext, ef * m * 2, NULL); {
else neighborhoodSize = offsetof(HnswNeighborArray, items) + sizeof(HnswCandidate) * HnswGetLayerM(m, lc);
v.tids = tidhash_create(CurrentMemoryContext, ef * m * 2, NULL); neighborhoodData = palloc(neighborhoodSize);
}
/* Add entry points to v, C, and W */ /* Add entry points to v, C, and W */
foreach(lc2, ep) foreach(lc2, ep)
@@ -644,7 +741,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
bool found; bool found;
AddToVisited(v, hc, index, &found); AddToVisited(base, &v, hc, index, &found);
pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(C, &(CreatePairingHeapNode(hc)->ph_node));
pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node)); pairingheap_add(W, &(CreatePairingHeapNode(hc)->ph_node));
@@ -654,7 +751,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* would be ideal to do this for inserts as well, but this could * would be ideal to do this for inserts as well, but this could
* affect insert performance. * affect insert performance.
*/ */
if (skipElement == NULL || hc->element->heaptidsLength != 0) if (CountElement(base, skipElement, hc))
wlen++; wlen++;
} }
@@ -663,38 +760,51 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswNeighborArray *neighborhood; HnswNeighborArray *neighborhood;
HnswCandidate *c = ((HnswPairingHeapNode *) pairingheap_remove_first(C))->inner; HnswCandidate *c = ((HnswPairingHeapNode *) pairingheap_remove_first(C))->inner;
HnswCandidate *f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; HnswCandidate *f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
HnswElement cElement;
if (c->distance > f->distance) if (c->distance > f->distance)
break; break;
if (c->element->neighbors == NULL) cElement = HnswPtrAccess(base, c->element);
HnswLoadNeighbors(c->element, index, m);
if (HnswPtrIsNull(base, cElement->neighbors))
HnswLoadNeighbors(cElement, index, m);
/* Get the neighborhood at layer lc */ /* Get the neighborhood at layer lc */
neighborhood = HnswGetNeighbors(c->element, lc); neighborhood = HnswGetNeighbors(base, cElement, lc);
/* Copy neighborhood to local memory if needed */
if (index == NULL)
{
LWLockAcquire(&cElement->lock, LW_SHARED);
memcpy(neighborhoodData, neighborhood, neighborhoodSize);
LWLockRelease(&cElement->lock);
neighborhood = neighborhoodData;
}
for (int i = 0; i < neighborhood->length; i++) for (int i = 0; i < neighborhood->length; i++)
{ {
HnswCandidate *e = &neighborhood->items[i]; HnswCandidate *e = &neighborhood->items[i];
bool visited; bool visited;
AddToVisited(v, e, index, &visited); AddToVisited(base, &v, e, index, &visited);
if (!visited) if (!visited)
{ {
float eDistance; float eDistance;
HnswElement eElement = HnswPtrAccess(base, e->element);
f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner; f = ((HnswPairingHeapNode *) pairingheap_first(W))->inner;
if (index == NULL) if (index == NULL)
eDistance = GetCandidateDistance(e, q, procinfo, collation); eDistance = GetCandidateDistance(base, e, q, procinfo, collation);
else 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 */ /* Make robust to issues */
if (e->element->level < lc) if (eElement->level < lc)
continue; continue;
if (eDistance < f->distance || wlen < ef) if (eDistance < f->distance || wlen < ef)
@@ -702,7 +812,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
/* Copy e */ /* Copy e */
HnswCandidate *ec = palloc(sizeof(HnswCandidate)); HnswCandidate *ec = palloc(sizeof(HnswCandidate));
ec->element = e->element; HnswPtrStore(base, ec->element, eElement);
ec->distance = eDistance; ec->distance = eDistance;
pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node)); pairingheap_add(C, &(CreatePairingHeapNode(ec)->ph_node));
@@ -713,7 +823,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* vacuuming. It would be ideal to do this for inserts as * vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance. * well, but this could affect insert performance.
*/ */
if (skipElement == NULL || e->element->heaptidsLength != 0) if (CountElement(base, skipElement, e))
{ {
wlen++; wlen++;
@@ -738,7 +848,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 static int
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
@@ -756,10 +866,38 @@ CompareCandidateDistances(const void *a, const void *b)
if (hca->distance > hcb->distance) if (hca->distance > hcb->distance)
return -1; return -1;
if (hca->element < hcb->element) if (HnswPtrPointer(hca->element) < HnswPtrPointer(hcb->element))
return 1; 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 -1;
return 0; return 0;
@@ -769,46 +907,58 @@ CompareCandidateDistances(const void *a, const void *b)
* Calculate the distance between elements * Calculate the distance between elements
*/ */
static float 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 */ /* Look for cached distance */
if (a->neighbors != NULL) if (!HnswPtrIsNull(base, a->neighbors))
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(a, lc); HnswNeighborArray *neighbors = HnswGetNeighbors(base, a, lc);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
if (neighbors->items[i].element == b) HnswElement element = HnswPtrAccess(base, neighbors->items[i].element);
if (element == b)
return neighbors->items[i].distance; return neighbors->items[i].distance;
} }
} }
if (b->neighbors != NULL) if (!HnswPtrIsNull(base, b->neighbors))
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(b, lc); HnswNeighborArray *neighbors = HnswGetNeighbors(base, b, lc);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
if (neighbors->items[i].element == a) HnswElement element = HnswPtrAccess(base, neighbors->items[i].element);
if (element == a)
return neighbors->items[i].distance; 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 * Check if an element is closer to q than any element from R
*/ */
static bool 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; ListCell *lc2;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); 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) if (distance <= e->distance)
return false; return false;
@@ -821,26 +971,31 @@ CheckElementCloser(HnswCandidate * e, List *r, int lc, FmgrInfo *procinfo, Oid c
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * 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 lm, int lc, FmgrInfo *procinfo, Oid collation, HnswElement e2, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{ {
List *r = NIL; List *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
pairingheap *wd; pairingheap *wd;
HnswNeighborArray *neighbors = HnswGetNeighbors(e2, lc); HnswNeighborArray *neighbors = HnswGetNeighbors(base, e2, lc);
bool mustCalculate = !neighbors->closerSet; bool mustCalculate = !neighbors->closerSet;
List *added = NIL; List *added = NIL;
bool removedAny = false; bool removedAny = false;
if (list_length(w) <= m) if (list_length(w) <= lm)
return w; return w;
wd = pairingheap_allocate(CompareNearestCandidates, NULL); wd = pairingheap_allocate(CompareNearestCandidates, NULL);
/* Ensure order of candidates is deterministic for closer caching */ /* Ensure order of candidates is deterministic for closer caching */
if (sortCandidates) 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) while (list_length(w) > 0 && list_length(r) < lm)
{ {
/* Assumes w is already ordered desc */ /* Assumes w is already ordered desc */
HnswCandidate *e = llast(w); HnswCandidate *e = llast(w);
@@ -849,7 +1004,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 */ /* Use previous state of r and wd to skip work when possible */
if (mustCalculate) 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) else if (list_length(added) > 0)
{ {
/* /*
@@ -858,7 +1013,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
*/ */
if (e->closer) if (e->closer)
{ {
e->closer = CheckElementCloser(e, added, lc, procinfo, collation); e->closer = CheckElementCloser(base, e, added, lc, procinfo, collation);
if (!e->closer) if (!e->closer)
removedAny = true; removedAny = true;
@@ -871,7 +1026,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
*/ */
if (removedAny) if (removedAny)
{ {
e->closer = CheckElementCloser(e, r, lc, procinfo, collation); e->closer = CheckElementCloser(base, e, r, lc, procinfo, collation);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -879,7 +1034,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
} }
else if (e == newCandidate) else if (e == newCandidate)
{ {
e->closer = CheckElementCloser(e, r, lc, procinfo, collation); e->closer = CheckElementCloser(base, e, r, lc, procinfo, collation);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -894,7 +1049,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
neighbors->closerSet = sortCandidates; neighbors->closerSet = sortCandidates;
/* Keep pruned connections */ /* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m) while (!pairingheap_is_empty(wd) && list_length(r) < lm)
r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner); r = lappend(r, ((HnswPairingHeapNode *) pairingheap_remove_first(wd))->inner);
/* Return pruned for update connections */ /* Return pruned for update connections */
@@ -913,10 +1068,10 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
* Add connections * Add connections
*/ */
static void static void
AddConnections(HnswElement element, List *neighbors, int m, int lc) AddConnections(char *base, HnswElement element, List *neighbors, int lc)
{ {
ListCell *lc2; ListCell *lc2;
HnswNeighborArray *a = HnswGetNeighbors(element, lc); HnswNeighborArray *a = HnswGetNeighbors(base, element, lc);
foreach(lc2, neighbors) foreach(lc2, neighbors)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2)); a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
@@ -926,16 +1081,16 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
* Update connections * Update connections
*/ */
void 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 lm, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{ {
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(hc->element, lc); HnswElement hce = HnswPtrAccess(base, hc->element);
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(base, hce, lc);
HnswCandidate hc2; HnswCandidate hc2;
hc2.element = element; HnswPtrStore(base, hc2.element, element);
hc2.distance = hc->distance; hc2.distance = hc->distance;
if (currentNeighbors->length < m) if (currentNeighbors->length < lm)
{ {
currentNeighbors->items[currentNeighbors->length++] = hc2; currentNeighbors->items[currentNeighbors->length++] = hc2;
@@ -951,19 +1106,20 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Load elements on insert */ /* Load elements on insert */
if (index != NULL) if (index != NULL)
{ {
Datum q = hc->element->value; Datum q = HnswGetValue(base, hce);
for (int i = 0; i < currentNeighbors->length; i++) for (int i = 0; i < currentNeighbors->length; i++)
{ {
HnswCandidate *hc3 = &currentNeighbors->items[i]; HnswCandidate *hc3 = &currentNeighbors->items[i];
HnswElement hc3Element = HnswPtrAccess(base, hc3->element);
if (DatumGetPointer(hc3->element->value) == NULL) if (HnswPtrIsNull(base, hc3Element->value))
HnswLoadElement(hc3->element, &hc3->distance, &q, index, procinfo, collation, true); HnswLoadElement(hc3Element, &hc3->distance, &q, index, procinfo, collation, true);
else else
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation); hc3->distance = GetCandidateDistance(base, hc3, q, procinfo, collation);
/* Prune element if being deleted */ /* Prune element if being deleted */
if (hc3->element->heaptidsLength == 0) if (hc3Element->heaptidsLength == 0)
{ {
pruned = &currentNeighbors->items[i]; pruned = &currentNeighbors->items[i];
break; break;
@@ -980,7 +1136,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
c = lappend(c, &currentNeighbors->items[i]); c = lappend(c, &currentNeighbors->items[i]);
c = lappend(c, &hc2); c = lappend(c, &hc2);
SelectNeighbors(c, m, lc, procinfo, collation, hc->element, &hc2, &pruned, true); SelectNeighbors(base, c, lm, lc, procinfo, collation, hce, &hc2, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -990,7 +1146,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
/* Find and replace the pruned element */ /* Find and replace the pruned element */
for (int i = 0; i < currentNeighbors->length; i++) 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; currentNeighbors->items[i] = hc2;
@@ -1008,43 +1164,65 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
* Remove elements being deleted or skipped * Remove elements being deleted or skipped
*/ */
static List * static List *
RemoveElements(List *w, HnswElement skipElement) RemoveElements(char *base, List *w, HnswElement skipElement)
{ {
ListCell *lc2; ListCell *lc2;
List *w2 = NIL; List *w2 = NIL;
/* Ensure does not access heaptidsLength during in-memory build */
pg_memory_barrier();
foreach(lc2, w) foreach(lc2, w)
{ {
HnswCandidate *hc = (HnswCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
HnswElement hce = HnswPtrAccess(base, hc->element);
/* Skip self for vacuuming update */ /* 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; continue;
if (hc->element->heaptidsLength != 0) if (hce->heaptidsLength != 0)
w2 = lappend(w2, hc); w2 = lappend(w2, hc);
} }
return w2; 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 * Algorithm 1 from paper
*/ */
void void
HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing) HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing)
{ {
List *ep; List *ep;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
Datum q = element->value; Datum q = HnswGetValue(base, element);
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
/* Precompute hash */ /* Precompute hash */
if (index == NULL) if (index == NULL)
element->hash = hash_pointer((uintptr_t) element); PrecomputeHash(base, element);
#endif #endif
/* No neighbors if no entry point */ /* No neighbors if no entry point */
@@ -1052,13 +1230,13 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
return; return;
/* Get entry point and level */ /* 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; entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */ /* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--) 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; ep = w;
} }
@@ -1076,12 +1254,12 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
List *neighbors; List *neighbors;
List *lw; 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 */ /* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */ /* but should be removed before selecting neighbors */
if (index != NULL) if (index != NULL)
lw = RemoveElements(w, skipElement); lw = RemoveElements(base, w, skipElement);
else else
lw = w; lw = w;
@@ -1090,9 +1268,9 @@ HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, F
* sortCandidates to true for in-memory builds to enable closer * sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance. * 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, lc);
ep = w; ep = w;
} }

View File

@@ -2,6 +2,7 @@
#include <math.h> #include <math.h>
#include "access/generic_xlog.h"
#include "commands/vacuum.h" #include "commands/vacuum.h"
#include "hnsw.h" #include "hnsw.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
@@ -199,21 +200,22 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup; HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m); Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
char *base = NULL;
/* Skip if element is entry point */ /* Skip if element is entry point */
if (entryPoint != NULL && element->blkno == entryPoint->blkno && element->offno == entryPoint->offno) if (entryPoint != NULL && element->blkno == entryPoint->blkno && element->offno == entryPoint->offno)
return; return;
/* Init fields */ /* Init fields */
HnswInitNeighbors(element, m); HnswInitNeighbors(base, element, m, NULL);
element->heaptidsLength = 0; element->heaptidsLength = 0;
/* Add element to graph, skipping itself */ /* Find neighbors for element, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true); HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true);
/* Update neighbor tuple */ /* Update neighbor tuple */
/* Do this before getting page to minimize locking */ /* Do this before getting page to minimize locking */
HnswSetNeighborTuple(ntup, element, m); HnswSetNeighborTuple(base, ntup, element, m);
/* Get neighbor page */ /* Get neighbor page */
buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas); buf = ReadBufferExtended(index, MAIN_FORKNUM, element->neighborPage, RBM_NORMAL, bas);
@@ -230,7 +232,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
/* Update neighbors */ /* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, true, false); HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false);
} }
/* /*
@@ -301,7 +303,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
{ {
/* Reset neighbors from previous update */ /* Reset neighbors from previous update */
if (highestPoint != NULL) if (highestPoint != NULL)
highestPoint->neighbors = NULL; HnswPtrStore((char *) NULL, highestPoint->neighbors, (HnswNeighborArrayPtr *) NULL);
RepairGraphElement(vacuumstate, entryPoint, highestPoint); RepairGraphElement(vacuumstate, entryPoint, highestPoint);
} }

View File

@@ -2,30 +2,25 @@
#include <float.h> #include <float.h>
#include "access/table.h"
#include "access/tableam.h"
#include "access/parallel.h" #include "access/parallel.h"
#include "access/xact.h" #include "access/xact.h"
#include "catalog/index.h" #include "catalog/index.h"
#include "catalog/pg_operator_d.h" #include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h" #include "catalog/pg_type_d.h"
#include "commands/progress.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "optimizer/optimizer.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "tcop/tcopprot.h" #include "tcop/tcopprot.h"
#include "utils/memutils.h" #include "utils/memutils.h"
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
#include "utils/backend_progress.h" #include "utils/backend_progress.h"
#elif PG_VERSION_NUM >= 120000
#include "pgstat.h"
#endif
#if PG_VERSION_NUM >= 120000
#include "access/tableam.h"
#include "commands/progress.h"
#else #else
#define PROGRESS_CREATEIDX_SUBPHASE 0 #include "pgstat.h"
#define PROGRESS_CREATEIDX_TUPLES_TOTAL 0
#define PROGRESS_CREATEIDX_TUPLES_DONE 0
#endif #endif
#if PG_VERSION_NUM >= 130000 #if PG_VERSION_NUM >= 130000
@@ -34,26 +29,13 @@
#define CALLBACK_ITEM_POINTER HeapTuple hup #define CALLBACK_ITEM_POINTER HeapTuple hup
#endif #endif
#if PG_VERSION_NUM >= 120000
#define UpdateProgress(index, val) pgstat_progress_update_param(index, val) #define UpdateProgress(index, val) pgstat_progress_update_param(index, val)
#else
#define UpdateProgress(index, val) ((void)val)
#endif
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
#include "utils/backend_status.h" #include "utils/backend_status.h"
#include "utils/wait_event.h" #include "utils/wait_event.h"
#endif #endif
#if PG_VERSION_NUM >= 120000
#include "access/table.h"
#include "optimizer/optimizer.h"
#else
#include "access/heapam.h"
#include "optimizer/planner.h"
#include "pgstat.h"
#endif
#define PARALLEL_KEY_IVFFLAT_SHARED UINT64CONST(0xA000000000000001) #define PARALLEL_KEY_IVFFLAT_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_TUPLESORT UINT64CONST(0xA000000000000002) #define PARALLEL_KEY_TUPLESORT UINT64CONST(0xA000000000000002)
#define PARALLEL_KEY_IVFFLAT_CENTERS UINT64CONST(0xA000000000000003) #define PARALLEL_KEY_IVFFLAT_CENTERS UINT64CONST(0xA000000000000003)
@@ -150,13 +132,8 @@ SampleRows(IvfflatBuildState * buildstate)
{ {
BlockNumber targblock = BlockSampler_Next(&buildstate->bs); BlockNumber targblock = BlockSampler_Next(&buildstate->bs);
#if PG_VERSION_NUM >= 120000
table_index_build_range_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, table_index_build_range_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
false, true, false, targblock, 1, SampleCallback, (void *) buildstate, NULL); false, true, false, targblock, 1, SampleCallback, (void *) buildstate, NULL);
#else
IndexBuildHeapRangeScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
false, true, targblock, 1, SampleCallback, (void *) buildstate, NULL);
#endif
} }
} }
@@ -282,11 +259,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IndexTuple itup = NULL; /* silence compiler warning */ IndexTuple itup = NULL; /* silence compiler warning */
int64 inserted = 0; int64 inserted = 0;
#if PG_VERSION_NUM >= 120000
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple);
#else
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc);
#endif
TupleDesc tupdesc = RelationGetDescr(index); TupleDesc tupdesc = RelationGetDescr(index);
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
@@ -375,20 +348,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
elog(ERROR, "dimensions must be greater than one for this opclass"); elog(ERROR, "dimensions must be greater than one for this opclass");
/* Create tuple description for sorting */ /* Create tuple description for sorting */
#if PG_VERSION_NUM >= 120000
buildstate->tupdesc = CreateTemplateTupleDesc(3); buildstate->tupdesc = CreateTemplateTupleDesc(3);
#else
buildstate->tupdesc = CreateTemplateTupleDesc(3, false);
#endif
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
#if PG_VERSION_NUM >= 120000
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
#else
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc);
#endif
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions); buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions);
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists); buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -631,11 +596,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
{ {
SortCoordinate coordinate; SortCoordinate coordinate;
IvfflatBuildState buildstate; IvfflatBuildState buildstate;
#if PG_VERSION_NUM >= 120000
TableScanDesc scan; TableScanDesc scan;
#else
HeapScanDesc scan;
#endif
double reltuples; double reltuples;
IndexInfo *indexInfo; IndexInfo *indexInfo;
@@ -659,18 +620,11 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
buildstate.centers->length = buildstate.centers->maxlen; buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = tuplesort_begin_heap(buildstate.tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, sortmem, coordinate, false); ivfspool->sortstate = tuplesort_begin_heap(buildstate.tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, sortmem, coordinate, false);
buildstate.sortstate = ivfspool->sortstate; buildstate.sortstate = ivfspool->sortstate;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(ivfspool->heap, scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared)); ParallelTableScanFromIvfflatShared(ivfshared));
reltuples = table_index_build_scan(ivfspool->heap, ivfspool->index, indexInfo, reltuples = table_index_build_scan(ivfspool->heap, ivfspool->index, indexInfo,
true, progress, BuildCallback, true, progress, BuildCallback,
(void *) &buildstate, scan); (void *) &buildstate, scan);
#else
scan = heap_beginscan_parallel(ivfspool->heap, &ivfshared->heapdesc);
reltuples = IndexBuildHeapScan(ivfspool->heap, ivfspool->index, indexInfo,
true, BuildCallback,
(void *) &buildstate, scan);
#endif
/* Execute this worker's part of the sort */ /* Execute this worker's part of the sort */
tuplesort_performsort(ivfspool->sortstate); tuplesort_performsort(ivfspool->sortstate);
@@ -740,11 +694,7 @@ IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc)
} }
/* Open relations within worker */ /* Open relations within worker */
#if PG_VERSION_NUM >= 120000
heapRel = table_open(ivfshared->heaprelid, heapLockmode); heapRel = table_open(ivfshared->heaprelid, heapLockmode);
#else
heapRel = heap_open(ivfshared->heaprelid, heapLockmode);
#endif
indexRel = index_open(ivfshared->indexrelid, indexLockmode); indexRel = index_open(ivfshared->indexrelid, indexLockmode);
/* Initialize worker's own spool */ /* Initialize worker's own spool */
@@ -764,11 +714,7 @@ IvfflatParallelBuildMain(dsm_segment *seg, shm_toc *toc)
/* Close relations within worker */ /* Close relations within worker */
index_close(indexRel, indexLockmode); index_close(indexRel, indexLockmode);
#if PG_VERSION_NUM >= 120000
table_close(heapRel, heapLockmode); table_close(heapRel, heapLockmode);
#else
heap_close(heapRel, heapLockmode);
#endif
} }
/* /*
@@ -793,19 +739,7 @@ IvfflatEndParallel(IvfflatLeader * ivfleader)
static Size static Size
ParallelEstimateShared(Relation heap, Snapshot snapshot) ParallelEstimateShared(Relation heap, Snapshot snapshot)
{ {
#if PG_VERSION_NUM >= 120000
return add_size(BUFFERALIGN(sizeof(IvfflatShared)), table_parallelscan_estimate(heap, snapshot)); return add_size(BUFFERALIGN(sizeof(IvfflatShared)), table_parallelscan_estimate(heap, snapshot));
#else
if (!IsMVCCSnapshot(snapshot))
{
Assert(snapshot == SnapshotAny);
return sizeof(IvfflatShared);
}
return add_size(offsetof(IvfflatShared, heapdesc) +
offsetof(ParallelHeapScanDescData, phs_snapshot_data),
EstimateSnapshotSpace(snapshot));
#endif
} }
/* /*
@@ -856,11 +790,7 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
/* Enter parallel mode and create context */ /* Enter parallel mode and create context */
EnterParallelMode(); EnterParallelMode();
Assert(request > 0); Assert(request > 0);
#if PG_VERSION_NUM >= 120000
pcxt = CreateParallelContext("vector", "IvfflatParallelBuildMain", request); pcxt = CreateParallelContext("vector", "IvfflatParallelBuildMain", request);
#else
pcxt = CreateParallelContext("vector", "IvfflatParallelBuildMain", request, true);
#endif
scantuplesortstates = leaderparticipates ? request + 1 : request; scantuplesortstates = leaderparticipates ? request + 1 : request;
@@ -918,13 +848,9 @@ IvfflatBeginParallel(IvfflatBuildState * buildstate, bool isconcurrent, int requ
#ifdef IVFFLAT_KMEANS_DEBUG #ifdef IVFFLAT_KMEANS_DEBUG
ivfshared->inertia = 0; ivfshared->inertia = 0;
#endif #endif
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap, table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromIvfflatShared(ivfshared), ParallelTableScanFromIvfflatShared(ivfshared),
snapshot); snapshot);
#else
heap_parallelscan_initialize(&ivfshared->heapdesc, buildstate->heap, snapshot);
#endif
/* Store shared tuplesort-private state, for which we reserved space */ /* Store shared tuplesort-private state, for which we reserved space */
sharedsort = (Sharedsort *) shm_toc_allocate(pcxt->toc, estsort); sharedsort = (Sharedsort *) shm_toc_allocate(pcxt->toc, estsort);
@@ -1023,15 +949,8 @@ AssignTuples(IvfflatBuildState * buildstate)
if (buildstate->ivfleader) if (buildstate->ivfleader)
buildstate->reltuples = ParallelHeapScan(buildstate); buildstate->reltuples = ParallelHeapScan(buildstate);
else else
{
#if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo, buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL); true, true, BuildCallback, (void *) buildstate, NULL);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
#ifdef IVFFLAT_KMEANS_DEBUG #ifdef IVFFLAT_KMEANS_DEBUG
PrintKmeansMetrics(buildstate); PrintKmeansMetrics(buildstate);

View File

@@ -3,14 +3,16 @@
#include <float.h> #include <float.h>
#include "access/amapi.h" #include "access/amapi.h"
#include "access/reloptions.h"
#include "commands/progress.h"
#include "commands/vacuum.h" #include "commands/vacuum.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "utils/guc.h" #include "utils/guc.h"
#include "utils/selfuncs.h" #include "utils/selfuncs.h"
#include "utils/spccache.h" #include "utils/spccache.h"
#if PG_VERSION_NUM >= 120000 #if PG_VERSION_NUM < 150000
#include "commands/progress.h" #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
int ivfflat_probes; int ivfflat_probes;
@@ -33,12 +35,13 @@ IvfflatInit(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes", DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
"Valid range is 1..lists.", &ivfflat_probes, "Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat");
} }
/* /*
* Get the name of index build phase * Get the name of index build phase
*/ */
#if PG_VERSION_NUM >= 120000
static char * static char *
ivfflatbuildphasename(int64 phasenum) ivfflatbuildphasename(int64 phasenum)
{ {
@@ -56,7 +59,6 @@ ivfflatbuildphasename(int64 phasenum)
return NULL; return NULL;
} }
} }
#endif
/* /*
* Estimate the cost of an index scan * Estimate the cost of an index scan
@@ -72,9 +74,6 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
double ratio; double ratio;
double spc_seq_page_cost; double spc_seq_page_cost;
Relation index; Relation index;
#if PG_VERSION_NUM < 120000
List *qinfos;
#endif
/* Never use index without order */ /* Never use index without order */
if (path->indexorderbys == NULL) if (path->indexorderbys == NULL)
@@ -105,26 +104,7 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*/ */
costs.numIndexTuples = path->indexinfo->tuples * ratio; costs.numIndexTuples = path->indexinfo->tuples * ratio;
/*
* Do not use index if no limit or limit + offset > expected tuples unless
* enable_seqscan = off
*/
if (root->limit_tuples < 0 || root->limit_tuples > costs.numIndexTuples)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
#if PG_VERSION_NUM >= 120000
genericcostestimate(root, path, loop_count, &costs); genericcostestimate(root, path, loop_count, &costs);
#else
qinfos = deconstruct_indexquals(path);
genericcostestimate(root, path, loop_count, qinfos, &costs);
#endif
get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost); get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
@@ -241,9 +221,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
amroutine->amcostestimate = ivfflatcostestimate; amroutine->amcostestimate = ivfflatcostestimate;
amroutine->amoptions = ivfflatoptions; amroutine->amoptions = ivfflatoptions;
amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */ amroutine->amproperty = NULL; /* TODO AMPROP_DISTANCE_ORDERABLE */
#if PG_VERSION_NUM >= 120000
amroutine->ambuildphasename = ivfflatbuildphasename; amroutine->ambuildphasename = ivfflatbuildphasename;
#endif
amroutine->amvalidate = ivfflatvalidate; amroutine->amvalidate = ivfflatvalidate;
#if PG_VERSION_NUM >= 140000 #if PG_VERSION_NUM >= 140000
amroutine->amadjustmembers = NULL; amroutine->amadjustmembers = NULL;

View File

@@ -3,9 +3,10 @@
#include "postgres.h" #include "postgres.h"
#include "access/genam.h"
#include "access/generic_xlog.h" #include "access/generic_xlog.h"
#include "access/parallel.h" #include "access/parallel.h"
#include "access/reloptions.h" #include "lib/pairingheap.h"
#include "nodes/execnodes.h" #include "nodes/execnodes.h"
#include "port.h" /* for random() */ #include "port.h" /* for random() */
#include "utils/sampling.h" #include "utils/sampling.h"
@@ -16,10 +17,6 @@
#include "common/pg_prng.h" #include "common/pg_prng.h"
#endif #endif
#if PG_VERSION_NUM < 120000
#include "access/relscan.h"
#endif
#ifdef IVFFLAT_BENCH #ifdef IVFFLAT_BENCH
#include "portability/instr_time.h" #include "portability/instr_time.h"
#endif #endif
@@ -135,16 +132,10 @@ typedef struct IvfflatShared
#ifdef IVFFLAT_KMEANS_DEBUG #ifdef IVFFLAT_KMEANS_DEBUG
double inertia; double inertia;
#endif #endif
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
#endif
} IvfflatShared; } IvfflatShared;
#if PG_VERSION_NUM >= 120000
#define ParallelTableScanFromIvfflatShared(shared) \ #define ParallelTableScanFromIvfflatShared(shared) \
(ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(IvfflatShared))) (ParallelTableScanDesc) ((char *) (shared) + BUFFERALIGN(sizeof(IvfflatShared)))
#endif
typedef struct IvfflatLeader typedef struct IvfflatLeader
{ {

View File

@@ -2,6 +2,7 @@
#include <float.h> #include <float.h>
#include "access/generic_xlog.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"

View File

@@ -5,6 +5,7 @@
#include "access/relscan.h" #include "access/relscan.h"
#include "catalog/pg_operator_d.h" #include "catalog/pg_operator_d.h"
#include "catalog/pg_type_d.h" #include "catalog/pg_type_d.h"
#include "lib/pairingheap.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "miscadmin.h" #include "miscadmin.h"
#include "pgstat.h" #include "pgstat.h"
@@ -105,12 +106,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; double tuples = 0;
#if PG_VERSION_NUM >= 120000
TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual); TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
#else
TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc);
#endif
/* /*
* Reuse same set of shared buffers for scan * Reuse same set of shared buffers for scan
@@ -216,22 +212,14 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->collation = index->rd_indcollation[0]; so->collation = index->rd_indcollation[0];
/* Create tuple description for sorting */ /* Create tuple description for sorting */
#if PG_VERSION_NUM >= 120000
so->tupdesc = CreateTemplateTupleDesc(2); so->tupdesc = CreateTemplateTupleDesc(2);
#else
so->tupdesc = CreateTemplateTupleDesc(2, false);
#endif
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 2, "heaptid", TIDOID, -1, 0);
/* Prep sort */ /* Prep sort */
so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false); so->sortstate = tuplesort_begin_heap(so->tupdesc, 1, attNums, sortOperators, sortCollations, nullsFirstFlags, work_mem, NULL, false);
#if PG_VERSION_NUM >= 120000
so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple); so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
#else
so->slot = MakeSingleTupleTableSlot(so->tupdesc);
#endif
so->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
@@ -321,12 +309,7 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull)); ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
#if PG_VERSION_NUM >= 120000
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
#else
scan->xs_ctup.t_self = *heaptid;
#endif
scan->xs_recheckorderby = false; scan->xs_recheckorderby = false;
return true; return true;
} }

View File

@@ -1,5 +1,6 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "vector.h" #include "vector.h"

View File

@@ -1,5 +1,6 @@
#include "postgres.h" #include "postgres.h"
#include "access/generic_xlog.h"
#include "commands/vacuum.h" #include "commands/vacuum.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"

View File

@@ -3,14 +3,17 @@
#include <math.h> #include <math.h>
#include "catalog/pg_type.h" #include "catalog/pg_type.h"
#include "common/shortest_dec.h"
#include "fmgr.h" #include "fmgr.h"
#include "hnsw.h" #include "hnsw.h"
#include "ivfflat.h" #include "ivfflat.h"
#include "lib/stringinfo.h" #include "lib/stringinfo.h"
#include "libpq/pqformat.h" #include "libpq/pqformat.h"
#include "port.h" /* for strtof() */ #include "port.h" /* for strtof() */
#include "port/pg_bitutils.h"
#include "utils/array.h" #include "utils/array.h"
#include "utils/builtins.h" #include "utils/builtins.h"
#include "utils/float.h"
#include "utils/lsyscache.h" #include "utils/lsyscache.h"
#include "utils/numeric.h" #include "utils/numeric.h"
#include "vector.h" #include "vector.h"
@@ -19,13 +22,6 @@
#include "varatt.h" #include "varatt.h"
#endif #endif
#if PG_VERSION_NUM >= 120000
#include "common/shortest_dec.h"
#include "utils/float.h"
#else
#include <float.h>
#endif
#if PG_VERSION_NUM < 130000 #if PG_VERSION_NUM < 130000
#define TYPALIGN_DOUBLE 'd' #define TYPALIGN_DOUBLE 'd'
#define TYPALIGN_INT 'i' #define TYPALIGN_INT 'i'
@@ -293,15 +289,6 @@ vector_out(PG_FUNCTION_ARGS)
char *ptr; char *ptr;
int n; int n;
#if PG_VERSION_NUM < 120000
int ndig = FLT_DIG + extra_float_digits;
if (ndig < 1)
ndig = 1;
#define FLOAT_SHORTEST_DECIMAL_LEN (ndig + 10)
#endif
/* /*
* Need: * Need:
* *
@@ -325,11 +312,7 @@ vector_out(PG_FUNCTION_ARGS)
ptr++; ptr++;
} }
#if PG_VERSION_NUM >= 120000
n = float_to_shortest_decimal_bufn(vector->x[i], ptr); n = float_to_shortest_decimal_bufn(vector->x[i], ptr);
#else
n = sprintf(ptr, "%.*g", ndig, vector->x[i]);
#endif
ptr += n; ptr += n;
} }
*ptr = ']'; *ptr = ']';
@@ -1152,3 +1135,43 @@ vector_avg(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(result); PG_RETURN_POINTER(result);
} }
/*
* Ensure same number of bytes
*/
static inline void
CheckByteLengths(uint32 aLen, uint32 bLen)
{
if (aLen != bLen)
ereport(ERROR,
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("different byte lengths %u and %u", aLen, bLen)));
}
/*
* Get the Hamming distance between two binary strings
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hamming_distance);
Datum
hamming_distance(PG_FUNCTION_ARGS)
{
bytea *a = PG_GETARG_BYTEA_PP(0);
bytea *b = PG_GETARG_BYTEA_PP(1);
char *ax = VARDATA_ANY(a);
char *bx = VARDATA_ANY(b);
uint32 aLen = VARSIZE_ANY_EXHDR(a);
uint32 bLen = VARSIZE_ANY_EXHDR(b);
uint64 distance = 0;
CheckByteLengths(aLen, bLen);
for (uint32 i = 0; i < aLen; i++)
{
unsigned char diff = (unsigned char) (ax[i] ^ bx[i]);
distance += pg_number_of_ones[diff];
}
/* TODO Decide on return type */
PG_RETURN_FLOAT8((double) distance);
}

View File

@@ -158,6 +158,32 @@ SELECT l1_distance('[3e38]', '[-3e38]');
Infinity Infinity
(1 row) (1 row)
SELECT hamming_distance('\xFFFF', '\xFFFF');
hamming_distance
------------------
0
(1 row)
SELECT hamming_distance('\xFFFF', '\xFFFE');
hamming_distance
------------------
1
(1 row)
SELECT hamming_distance('\xFFFF', '\xFFFC');
hamming_distance
------------------
2
(1 row)
SELECT hamming_distance('\xFFFF', '\x0000');
hamming_distance
------------------
16
(1 row)
SELECT hamming_distance('\xFFFF', '\x00');
ERROR: different byte lengths 2 and 1
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v; SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
avg avg
----------- -----------

View File

@@ -37,6 +37,12 @@ SELECT l1_distance('[0,0]', '[0,1]');
SELECT l1_distance('[1,2]', '[3]'); SELECT l1_distance('[1,2]', '[3]');
SELECT l1_distance('[3e38]', '[-3e38]'); SELECT l1_distance('[3e38]', '[-3e38]');
SELECT hamming_distance('\xFFFF', '\xFFFF');
SELECT hamming_distance('\xFFFF', '\xFFFE');
SELECT hamming_distance('\xFFFF', '\xFFFC');
SELECT hamming_distance('\xFFFF', '\x0000');
SELECT hamming_distance('\xFFFF', '\x00');
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v; SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v;
SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v; SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]', NULL]) v;
SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v; SELECT avg(v) FROM unnest(ARRAY[]::vector[]) v;

View File

@@ -95,11 +95,10 @@ for my $i (0 .. $#operators)
$node->safe_psql("postgres", "DROP INDEX idx;"); $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( my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET client_min_messages = DEBUG; SET client_min_messages = DEBUG;
SET min_parallel_table_scan_size = 1; SET min_parallel_table_scan_size = 1;
SET hnsw.enable_parallel_build = on;
CREATE INDEX idx ON tst USING hnsw (v $opclass); CREATE INDEX idx ON tst USING hnsw (v $opclass);
)); ));
is($ret, 0, $stderr); is($ret, 0, $stderr);
@@ -109,6 +108,21 @@ for my $i (0 .. $#operators)
test_recall($min, $operator); test_recall($min, $operator);
$node->safe_psql("postgres", "DROP INDEX idx;"); $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(); done_testing();

View File

@@ -0,0 +1,114 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my $nc = 50;
my $limit = 20;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4, t text);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;");
# Generate query
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $query = "[" . join(",", @r) . "]";
my $c = int(rand() * $nc);
# Test attribute filtering
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c != $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%test%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%other%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test distance filtering
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test distance filtering greater than distance
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' > 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test distance filtering without order
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1;
));
like($explain, qr/Seq Scan/);
# Test distance filtering without limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
));
like($explain, qr/Seq Scan/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Use attribute index
like($explain, qr/Index Scan using idx/);
# Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using partial_idx/);
done_testing();

View File

@@ -0,0 +1,116 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my $nc = 50;
my $limit = 20;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4, t text);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 100);");
$node->safe_psql("postgres", "ANALYZE tst;");
# Generate query
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $query = "[" . join(",", @r) . "]";
my $c = int(rand() * $nc);
# Test attribute filtering
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c != $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%test%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%other%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test distance filtering
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test distance filtering greater than distance
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' > 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test distance filtering without order
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1;
));
like($explain, qr/Seq Scan/);
# Test distance filtering without limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Use attribute index
like($explain, qr/Index Scan using idx/);
# Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5) WHERE (c = $c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Use partial index
like($explain, qr/Index Scan using idx/);
done_testing();

View File

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

View File

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