Compare commits

..

69 Commits

Author SHA1 Message Date
Andrew Kane
fc0d3e7fdb Added search parameters to EXPLAIN [skip ci] 2024-10-28 01:20:59 -07:00
Andrew Kane
c530a3c490 Updated comment [skip ci] 2024-10-28 00:56:10 -07:00
Andrew Kane
d8b9e8ef73 Updated readme [skip ci] 2024-10-28 00:55:25 -07:00
Andrew Kane
00894efed5 Updated readme [skip ci] 2024-10-28 00:45:26 -07:00
Andrew Kane
0aa0f6619b Updated readme [skip ci] 2024-10-28 00:43:39 -07:00
Andrew Kane
38d053001e Updated readme [skip ci] 2024-10-28 00:32:12 -07:00
Andrew Kane
ccb95407e7 Updated readme [skip ci] 2024-10-28 00:09:26 -07:00
Andrew Kane
04d5e934a1 Scan 2000 more tuples with default work_mem 2024-10-27 22:23:28 -07:00
Andrew Kane
b163b5b196 Moved code [skip ci] 2024-10-27 22:21:43 -07:00
Andrew Kane
6a30c1e824 Fixed test [skip ci] 2024-10-27 21:10:51 -07:00
Andrew Kane
2db1b19644 Use greater than [skip ci] 2024-10-27 21:07:22 -07:00
Andrew Kane
305d62146e Updated comment [skip ci] 2024-10-27 21:05:32 -07:00
Andrew Kane
f9d627c9a9 Updated default value of hnsw.scan_mem_multiplier [skip ci] 2024-10-27 21:05:04 -07:00
Andrew Kane
38f42820be Added test for hnsw.scan_mem_multiplier [skip ci] 2024-10-27 20:05:58 -07:00
Andrew Kane
15c8245b42 Updated readme [skip ci] 2024-10-27 19:30:40 -07:00
Andrew Kane
572a9ab404 Updated readme [skip ci] 2024-10-27 18:58:50 -07:00
Andrew Kane
00492d7e57 Ensure max memory fits into Size for HNSW index scans 2024-10-27 14:21:15 -07:00
Andrew Kane
857d716d9e Renamed iterative_search to iterative_scan 2024-10-27 14:02:22 -07:00
Andrew Kane
c5dd2af750 Added comments [skip ci] 2024-10-25 21:39:03 -07:00
Andrew Kane
78b877bdaf Revert "Renamed iterative_search to iterative_scan"
This reverts commit 7043cce893.
2024-10-24 20:32:07 -07:00
Andrew Kane
7043cce893 Renamed iterative_search to iterative_scan 2024-10-24 20:31:43 -07:00
Andrew Kane
62039d74f6 Added iterative search section to readme [skip ci] 2024-10-24 18:05:29 -07:00
Andrew Kane
ac6576e53a Added hnsw.search_mem_multiplier option 2024-10-24 18:02:20 -07:00
Andrew Kane
67eff41c44 Updated changelog [skip ci] 2024-10-23 15:12:40 -07:00
Andrew Kane
1291b12090 Added Postgres 18 to CI [skip ci] 2024-10-22 00:38:19 -07:00
Andrew Kane
24522700b8 Improved hnswrescan 2024-10-21 23:41:32 -07:00
Andrew Kane
bfb3a45b31 Use consistent order [skip ci] 2024-10-21 21:47:03 -07:00
Andrew Kane
e718eb8da4 Updated range and defaults for iterative search parameters 2024-10-21 20:38:50 -07:00
Andrew Kane
049972a4a3 Improved test output [skip ci] 2024-10-13 17:22:49 -07:00
Andrew Kane
61027645e9 Improved test output [skip ci] 2024-10-13 17:21:38 -07:00
Andrew Kane
a41b327b33 Speed up test [skip ci] 2024-10-13 17:12:12 -07:00
Andrew Kane
7f735ebd9b Added test for strict order [skip ci] 2024-10-13 17:04:03 -07:00
Andrew Kane
02b01e1ca9 Show tuples with memory usage [skip ci] 2024-10-12 22:05:13 -07:00
Andrew Kane
388e42f6e6 Fixed flaky test [skip ci] 2024-10-11 15:48:19 -07:00
Andrew Kane
bf379eed86 Use a memory context for IVFFlat index scans 2024-10-11 15:46:38 -07:00
Andrew Kane
e1bc929429 Simplify lists for IvfflatScanOpaque [skip ci] 2024-10-11 15:29:23 -07:00
Andrew Kane
38285aacc7 Revert "Updated IVFFlat to support multiple attributes (not enabled yet)"
This reverts commit 772ab69de6.
2024-10-11 15:01:54 -07:00
Andrew Kane
a2408e60fa Revert "Added IndexTuple to HNSW elements (first step to support multiple attributes)"
This reverts commit 53a8734bac.
2024-10-11 14:57:57 -07:00
Andrew Kane
53a8734bac Added IndexTuple to HNSW elements (first step to support multiple attributes) 2024-10-11 14:12:01 -07:00
Andrew Kane
7484625227 Added comments [skip ci] 2024-10-11 11:59:36 -07:00
Andrew Kane
d1ebb8db73 Use -1 for no limit for ivfflat.max_probes [skip ci] 2024-10-11 11:43:32 -07:00
Andrew Kane
42af8aa1d1 Updated GUC descriptions [skip ci] 2024-10-11 11:26:27 -07:00
Andrew Kane
9d15a76b60 Improved enum naming [skip ci] 2024-10-11 11:20:36 -07:00
Andrew Kane
a3a20f9816 Simplified GUC names [skip ci] 2024-10-11 11:18:01 -07:00
Andrew Kane
b26a21b848 Added regression tests for iterative search [skip ci] 2024-10-11 11:07:11 -07:00
Andrew Kane
2dc392ed6c Updated GUC names [skip ci] 2024-10-10 23:50:11 -07:00
Andrew Kane
960d2848cb Updated comment [skip ci] 2024-10-10 21:02:33 -07:00
Andrew Kane
8e88b481a6 Use a lower max allocation size than default to allow scanning more tuples for iterative search before exceeding work_mem [skip ci] 2024-10-10 20:57:57 -07:00
Andrew Kane
124018b8dd Added HnswInitSearchCandidate function 2024-10-10 19:30:47 -07:00
Andrew Kane
35b252a3e3 Switched to macos-13 on CI [skip ci] 2024-10-10 18:33:35 -07:00
Andrew Kane
2832e746f0 Use get_float8_infinity instead of INFINITY
Co-authored-by: "Jonathan S. Katz" <jkatz@amazon.com>
2024-10-10 18:16:39 -07:00
Andrew Kane
961cb17d80 Added iterative search for HNSW [skip ci] 2024-10-10 18:14:39 -07:00
Andrew Kane
c91ed7b2c3 Added iterative search for IVFFlat [skip ci] 2024-10-10 18:12:27 -07:00
Andrew Kane
48fe70c219 Fixed sortdesc for multiple attributes [skip ci] 2024-10-10 15:22:23 -07:00
Andrew Kane
29908405ab List assignment does not change with scaling [skip ci] 2024-10-10 15:16:08 -07:00
Andrew Kane
08d0340655 Improved IVFFlat vacuum test [skip ci] 2024-10-10 14:24:26 -07:00
Andrew Kane
7d2eb49c2a Fixed list assignment for cosine distance and non-normalized vectors for IVFFlat 2024-10-10 12:54:13 -07:00
Andrew Kane
772ab69de6 Updated IVFFlat to support multiple attributes (not enabled yet) 2024-10-10 12:34:03 -07:00
Andrew Kane
e13e9a9614 Improved scoping [skip ci] 2024-10-10 12:19:08 -07:00
Andrew Kane
e2fab306ac Added tupdesc to buildstate [skip ci] 2024-10-10 12:18:41 -07:00
Andrew Kane
edc2126a4a Renamed tupdesc to sortdesc [skip ci] 2024-10-10 12:16:57 -07:00
Andrew Kane
2cbd08b6c0 Moved unions and macros [skip ci] 2024-10-10 09:41:26 -07:00
Andrew Kane
fa6782985a Added HnswQuery struct for query data 2024-10-09 23:45:47 -07:00
Andrew Kane
32ab27d72a Added HnswSupport struct for support functions 2024-10-09 23:10:26 -07:00
Andrew Kane
064db12de7 Moved procinfo initialization for inserts [skip ci] 2024-10-09 21:59:21 -07:00
Andrew Kane
45a6eef9e0 Improved variable name [skip ci] 2024-10-09 21:52:10 -07:00
Andrew Kane
17266ed409 Use inMemory for conditionals 2024-10-09 21:49:32 -07:00
Andrew Kane
a98534e5ab DRY HNSW procinfo 2024-10-09 21:03:18 -07:00
Andrew Kane
57c05c59a2 DRY code for forming index value 2024-10-09 20:50:17 -07:00
28 changed files with 1240 additions and 736 deletions

View File

@@ -8,8 +8,8 @@ jobs:
fail-fast: false
matrix:
include:
# - postgres: 18
# os: ubuntu-24.04
- postgres: 18
os: ubuntu-24.04
- postgres: 17
os: ubuntu-24.04
- postgres: 16
@@ -49,7 +49,7 @@ jobs:
- postgres: 16
os: macos-14
- postgres: 14
os: macos-12
os: macos-13
steps:
- uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1

View File

@@ -1,6 +1,6 @@
## 0.8.0 (unreleased)
- Added support for inline filtering with HNSW
- Added support for iterative index scans
- Added casts for arrays to `sparsevec`
- Improved cost estimation
- Improved performance of HNSW inserts and on-disk index builds

View File

@@ -439,12 +439,6 @@ Create an index on one [or more](https://www.postgresql.org/docs/current/indexes
CREATE INDEX ON items (category_id);
```
Or a composite HNSW index for approximate search (added in 0.8.0)
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops, category_id);
```
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql
@@ -457,6 +451,89 @@ Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html
CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id);
```
## Iterative Index Scans
*Unreleased*
With approximate indexes, queries with filtering can return less results (due to post-filtering). Starting with 0.8.0, you can enable iterative index scans. If too few results from the initial scan match the filters, the scan will resume until enough results are found (or it reaches `hnsw.max_scan_tuples` or `ivfflat.max_probes`). This can significantly improve recall.
There are two modes for iterative scans: strict and relaxed.
Strict ensures results are in the exact order by distance
```sql
SET hnsw.iterative_scan = strict_order;
```
Relaxed allows results to be slightly out of order by distance, but provides better recall
```sql
SET hnsw.iterative_scan = relaxed_order;
# or
SET ivfflat.iterative_scan = relaxed_order;
```
With relaxed ordering, you can use a [materialized CTE](https://www.postgresql.org/docs/current/queries-with.html#QUERIES-WITH-CTE-MATERIALIZATION) to get strict ordering
```sql
WITH relaxed_results AS MATERIALIZED (
SELECT id, embedding <-> '[1,2,3]' AS distance FROM items WHERE category_id = 123 ORDER BY distance LIMIT 5
) SELECT * FROM relaxed_results ORDER BY distance;
```
For queries that filter by distance, use a materialized CTE and place the distance filter outside of it for best performance (due to the [current behavior](https://www.postgresql.org/message-id/flat/CAOdR5yGUoMQ6j7M5hNUXrySzaqZVGf_Ne%2B8fwZMRKTFxU1nbJg%40mail.gmail.com) of the Postgres executor)
```sql
WITH nearest_results AS MATERIALIZED (
SELECT id, embedding <-> '[1,2,3]' AS distance FROM items ORDER BY distance LIMIT 5
) SELECT * FROM nearest_results WHERE distance < 5 ORDER BY distance;
```
Note: Place any other filters inside the CTE
### Iterative Scan Options
Since scanning a large portion of an approximate index is expensive, there are options to control when a scan ends
#### HNSW
Specify the max number of tuples to visit (20,000 by default)
```sql
SET hnsw.max_scan_tuples = 20000;
```
Note: This is approximate and does not affect the initial scan
Specify the max amount of memory to use, as a multiple of `work_mem` (1 by default)
```sql
SET hnsw.scan_mem_multiplier = 2;
```
You can see when increasing this is needed by enabling debug messages
```sql
SET client_min_messages = debug1;
```
which will show when a scan reaches the memory limit
```text
DEBUG: hnsw index scan reached memory limit after 20000 tuples
HINT: Increase hnsw.scan_mem_multiplier to scan more tuples.
```
#### IVFFlat
Specify the max number of probes
```sql
SET ivfflat.max_probes = 100;
```
Note: If this is lower than `ivfflat.probes`, `ivfflat.probes` will be used
## Half-Precision Vectors
*Added in 0.7.0*
@@ -1195,7 +1272,6 @@ Thanks to:
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf)
- [Efficient and Robust Approximate Nearest Neighbor Search using Hierarchical Navigable Small World Graphs](https://arxiv.org/ftp/arxiv/papers/1603/1603.09320.pdf)
- [HQANN: Efficient and Robust Similarity Search for Hybrid Queries with Structured and Unstructured Constraints](https://arxiv.org/pdf/2207.07940.pdf)
## History

View File

@@ -24,11 +24,3 @@ CREATE CAST (double precision[] AS sparsevec)
CREATE CAST (numeric[] AS sparsevec)
WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT;
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -916,13 +916,3 @@ CREATE OPERATOR CLASS sparsevec_l1_ops
OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(sparsevec, sparsevec),
FUNCTION 3 hnsw_sparsevec_support(internal);
-- hnsw attributes
CREATE FUNCTION hnsw_attribute_distance(integer, integer) RETURNS float8
AS 'MODULE_PATHNAME', 'hnsw_int4_attribute_distance' LANGUAGE C IMMUTABLE STRICT PARALLEL SAFE;
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING hnsw AS
OPERATOR 2 = (integer, integer),
FUNCTION 4 hnsw_attribute_distance(integer, integer);

View File

@@ -18,7 +18,17 @@
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif
static const struct config_enum_entry hnsw_iterative_scan_options[] = {
{"off", HNSW_ITERATIVE_SCAN_OFF, false},
{"relaxed_order", HNSW_ITERATIVE_SCAN_RELAXED, false},
{"strict_order", HNSW_ITERATIVE_SCAN_STRICT, false},
{NULL, 0, false}
};
int hnsw_ef_search;
int hnsw_iterative_scan;
int hnsw_max_scan_tuples;
double hnsw_scan_mem_multiplier;
int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind;
@@ -67,7 +77,21 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
DefineCustomEnumVariable("hnsw.iterative_scan", "Sets the mode for iterative scans",
NULL, &hnsw_iterative_scan,
HNSW_ITERATIVE_SCAN_OFF, hnsw_iterative_scan_options, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* This is approximate and does not affect the initial scan */
DefineCustomIntVariable("hnsw.max_scan_tuples", "Sets the max number of tuples to visit for iterative scans",
NULL, &hnsw_max_scan_tuples,
20000, 1, INT_MAX, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* Same range as hash_mem_multiplier */
DefineCustomRealVariable("hnsw.scan_mem_multiplier", "Sets the multiple of work_mem to use for iterative scans",
NULL, &hnsw_scan_mem_multiplier,
1, 1, 1000, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw");
}
@@ -113,6 +137,10 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return;
}
@@ -227,13 +255,13 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0;
amroutine->amsupport = 4;
amroutine->amsupport = 3;
amroutine->amoptsprocnum = 0;
amroutine->amcanorder = false;
amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false;
amroutine->amcanmulticol = true;
amroutine->amcanmulticol = false;
amroutine->amoptionalkey = true;
amroutine->amsearcharray = false;
amroutine->amsearchnulls = false;
@@ -285,17 +313,3 @@ hnswhandler(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(amroutine);
}
/*
* Get the distance between two int4 attributes
*/
PGDLLEXPORT PG_FUNCTION_INFO_V1(hnsw_int4_attribute_distance);
Datum
hnsw_int4_attribute_distance(PG_FUNCTION_ARGS)
{
int32 a = PG_GETARG_INT32(0);
int32 b = PG_GETARG_INT32(1);
double distance = ((double) a) - ((double) b);
PG_RETURN_FLOAT8(distance);
}

View File

@@ -19,7 +19,6 @@
#define HNSW_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2
#define HNSW_TYPE_INFO_PROC 3
#define HNSW_ATTRIBUTE_DISTANCE_PROC 4
#define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953
@@ -89,6 +88,9 @@
/* 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 HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
#define HnswGetValue(base, element) PointerGetDatum(HnswPtrAccess(base, (element)->value))
#if PG_VERSION_NUM < 140005
@@ -105,12 +107,20 @@
#define HnswPtrPointer(hp) (hp).ptr
#define HnswPtrOffset(hp) relptr_offset((hp).relptr)
#define HnswUseIndexTuple(index) (IndexRelationGetNumberOfAttributes(index) > 1)
/* Variables */
extern int hnsw_ef_search;
extern int hnsw_iterative_scan;
extern int hnsw_max_scan_tuples;
extern double hnsw_scan_mem_multiplier;
extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeScanMode
{
HNSW_ITERATIVE_SCAN_OFF,
HNSW_ITERATIVE_SCAN_RELAXED,
HNSW_ITERATIVE_SCAN_STRICT
} HnswIterativeScanMode;
typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray;
@@ -124,7 +134,6 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr);
HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);
struct HnswElementData
{
@@ -133,6 +142,7 @@ struct HnswElementData
uint8 heaptidsLength;
uint8 level;
uint8 deleted;
uint8 version;
uint32 hash;
HnswNeighborsPtr neighbors;
BlockNumber blkno;
@@ -140,7 +150,6 @@ struct HnswElementData
OffsetNumber neighborOffno;
BlockNumber neighborPage;
DatumPtr value;
IndexTuplePtr itup;
LWLock lock;
};
@@ -166,7 +175,6 @@ typedef struct HnswSearchCandidate
pairingheap_node w_node;
HnswElementPtr element;
double distance;
bool matches;
} HnswSearchCandidate;
/* HNSW index options */
@@ -243,6 +251,18 @@ typedef struct HnswTypeInfo
void (*checkValue) (Pointer v);
} HnswTypeInfo;
typedef struct HnswSupport
{
FmgrInfo *procinfo;
FmgrInfo *normprocinfo;
Oid collation;
} HnswSupport;
typedef struct HnswQuery
{
Datum value;
} HnswQuery;
typedef struct HnswBuildState
{
/* Info */
@@ -262,17 +282,13 @@ typedef struct HnswBuildState
double reltuples;
/* Support functions */
FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo;
Oid *collation;
HnswSupport support;
/* Variables */
HnswGraph graphData;
HnswGraph *graph;
double ml;
int maxLevel;
bool useIndexTuple;
TupleDesc tupdesc;
/* Memory */
MemoryContext graphCtx;
@@ -314,10 +330,10 @@ typedef struct HnswElementTupleData
uint8 type;
uint8 level;
uint8 deleted;
uint8 unused;
uint8 version;
ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid;
uint16 unused2;
uint16 unused;
Vector data;
} HnswElementTupleData;
@@ -326,24 +342,42 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData
{
uint8 type;
uint8 unused;
uint8 version;
uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData;
typedef HnswNeighborTupleData * HnswNeighborTuple;
typedef union
{
struct pointerhash_hash *pointers;
struct offsethash_hash *offsets;
struct tidhash_hash *tids;
} visited_hash;
typedef union
{
HnswElement element;
ItemPointerData indextid;
} HnswUnvisited;
typedef struct HnswScanOpaqueData
{
const HnswTypeInfo *typeInfo;
bool first;
List *w;
visited_hash v;
pairingheap *discarded;
HnswQuery q;
int m;
int64 tuples;
double previousDistance;
Size maxMemory;
MemoryContext tmpCtx;
/* Support functions */
FmgrInfo *procinfo[2];
FmgrInfo *normprocinfo;
Oid *collation;
HnswSupport support;
} HnswScanOpaqueData;
typedef HnswScanOpaqueData * HnswScanOpaque;
@@ -361,8 +395,7 @@ typedef struct HnswVacuumState
int efConstruction;
/* Support functions */
FmgrInfo *procinfo[2];
Oid *collation;
HnswSupport support;
/* Variables */
struct tidhash_hash *deleted;
@@ -378,37 +411,36 @@ typedef struct HnswVacuumState
int HnswGetM(Relation index);
int HnswGetEfConstruction(Relation index);
FmgrInfo *HnswOptionalProcInfo(Relation index, uint16 procnum);
void HnswInitSupport(HnswSupport * support, Relation index);
Datum HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value);
bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
bool HnswCheckNorm(HnswSupport * support, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);
List *HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, int m, bool inserting, HnswElement skipElement, bool inMemory);
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfo, Oid *collation, int m, int efConstruction, bool existing, bool inMemory);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation rel, FmgrInfo **procinfo, Oid *collation, bool loadVec, bool inMemory);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, HnswQuery * q, Relation rel, HnswSupport * support, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, 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, Relation index);
void HnswLoadElement(HnswElement element, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo **procinfo, Oid *collation);
bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building);
void HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
bool HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element);
void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
void HnswInitLockTranche(void);
const HnswTypeInfo *HnswGetTypeInfo(Relation index);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
void HnswInitProcinfo(FmgrInfo **procinfo, Oid **collation, Relation index);
Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);
/* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -148,7 +148,6 @@ CreateGraphPages(HnswBuildState * buildstate)
Page page;
HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
/* Calculate sizes */
maxSize = HNSW_MAX_SIZE;
@@ -168,6 +167,7 @@ CreateGraphPages(HnswBuildState * buildstate)
Size etupSize;
Size ntupSize;
Size combinedSize;
Pointer valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */
iter = element->next;
@@ -176,7 +176,7 @@ CreateGraphPages(HnswBuildState * buildstate)
MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE);
/* Calculate sizes */
etupSize = HnswGetElementTupleSize(base, element, useIndexTuple);
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -186,7 +186,7 @@ CreateGraphPages(HnswBuildState * buildstate)
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("index tuple too large")));
HnswSetElementTuple(base, etup, element, useIndexTuple);
HnswSetElementTuple(base, etup, element);
/* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -327,29 +327,20 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
* Find duplicate element
*/
static bool
FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc)
FindDuplicateInMemory(char *base, HnswElement element)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
if (useIndexTuple)
{
/* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
/* Check for space */
if (AddDuplicateInMemory(element, neighborElement))
@@ -375,7 +366,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors
*/
static void
UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m)
UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
{
for (int lc = e->level; lc >= 0; lc--)
{
@@ -397,7 +388,7 @@ UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *co
Assert(neighborElement);
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, procinfo, collation);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, NULL, support);
LWLockRelease(&neighborElement->lock);
}
}
@@ -407,20 +398,20 @@ UpdateNeighborsInMemory(char *base, Relation index, FmgrInfo **procinfo, Oid *co
* Update graph in memory
*/
static void
UpdateGraphInMemory(FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, HnswBuildState * buildstate)
UpdateGraphInMemory(HnswSupport * support, 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, buildstate->useIndexTuple, buildstate->tupdesc))
if (FindDuplicateInMemory(base, element))
return;
/* Add element */
AddElementInMemory(base, graph, element);
/* Update neighbors */
UpdateNeighborsInMemory(base, buildstate->index, procinfo, collation, element, m);
UpdateNeighborsInMemory(base, support, element, m);
/* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -433,10 +424,8 @@ UpdateGraphInMemory(FmgrInfo **procinfo, Oid *collation, HnswElement element, in
static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{
Relation index = buildstate->index;
FmgrInfo **procinfo = buildstate->procinfo;
Oid *collation = buildstate->collation;
HnswGraph *graph = buildstate->graph;
HnswSupport *support = &buildstate->support;
HnswElement entryPoint;
LWLock *entryLock = &graph->entryLock;
LWLock *entryWaitLock = &graph->entryWaitLock;
@@ -468,10 +457,10 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, true);
HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false);
/* Update graph in memory */
UpdateGraphInMemory(procinfo, collation, element, m, efConstruction, entryPoint, buildstate);
UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
/* Release entry lock */
LWLockRelease(entryLock);
@@ -483,35 +472,19 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
static bool
InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, HnswBuildState * buildstate)
{
const HnswTypeInfo *typeInfo = buildstate->typeInfo;
HnswGraph *graph = buildstate->graph;
HnswElement element;
HnswAllocator *allocator = &buildstate->allocator;
HnswSupport *support = &buildstate->support;
Size valueSize;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
TupleDesc tupdesc = buildstate->tupdesc;
IndexTuple itup;
Size itupSize;
IndexTuple itupPtr;
Datum value;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (buildstate->normprocinfo != NULL)
{
if (!HnswCheckNorm(buildstate->normprocinfo, buildstate->collation[0], value))
return false;
value = HnswNormValue(typeInfo, buildstate->collation[0], value);
}
/* Form index value */
if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support))
return false;
/* Get datum size */
valueSize = VARSIZE_ANY(DatumGetPointer(value));
@@ -524,7 +497,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
{
LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true);
return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
}
/*
@@ -556,22 +529,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, value, values, isnull, heaptid, true);
return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
}
/* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
if (useIndexTuple)
{
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
itupSize = IndexTupleSize(itup);
itupPtr = HnswAlloc(allocator, itupSize);
}
else
valuePtr = HnswAlloc(allocator, valueSize);
valuePtr = HnswAlloc(allocator, valueSize);
/*
* We have now allocated the space needed for the element, so we don't
@@ -581,19 +544,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(&graph->allocatorLock);
/* Copy the datum */
if (useIndexTuple)
{
bool unused;
memcpy(itupPtr, itup, itupSize);
HnswPtrStore(base, element->itup, itupPtr);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupPtr, 1, tupdesc, &unused)));
}
else
{
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
}
memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->value, valuePtr);
/* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -720,19 +672,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("type not supported for hnsw index")));
/* TODO See if needed */
if (IndexRelationGetNumberOfKeyAttributes(index) > 2)
elog(ERROR, "index cannot have more than two columns");
if (!OidIsValid(index_getprocid(index, 1, HNSW_DISTANCE_PROC)))
elog(ERROR, "first column must be a vector");
for (int i = 1; i < IndexRelationGetNumberOfKeyAttributes(index); i++)
{
if (!OidIsValid(index_getprocid(index, i + 1, HNSW_ATTRIBUTE_DISTANCE_PROC)))
elog(ERROR, "column %d cannot be a vector", i + 1);
}
/* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0)
ereport(ERROR,
@@ -753,15 +692,12 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->indtuples = 0;
/* Get support functions */
HnswInitProcinfo(buildstate->procinfo, &buildstate->collation, index);
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
HnswInitSupport(&buildstate->support, index);
InitGraph(&buildstate->graphData, NULL, (Size) maintenance_work_mem * 1024L);
buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->useIndexTuple = HnswUseIndexTuple(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",

View File

@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
* Check for a free offset
*/
static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion)
{
OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,6 +98,7 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
{
*freeOffno = offno;
*freeNeighborOffno = neighborOffno;
*tupleVersion = etup->version;
return true;
}
else if (*nbuf != buf)
@@ -153,11 +154,11 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion;
char *base = NULL;
bool useIndexTuple = HnswUseIndexTuple(index);
/* Calculate sizes */
etupSize = HnswGetElementTupleSize(base, e, useIndexTuple);
etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(HnswPtrAccess(base, e->value)));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE;
@@ -165,7 +166,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
/* Prepare element tuple */
etup = palloc0(etupSize);
HnswSetElementTuple(base, etup, e, useIndexTuple);
HnswSetElementTuple(base, etup, e);
/* Prepare neighbor tuple */
ntup = palloc0(ntupSize);
@@ -203,7 +204,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
}
/* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
{
if (nbuf != buf)
{
@@ -213,6 +214,10 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
/* Set tuple version */
etup->version = tupleVersion;
ntup->version = tupleVersion;
break;
}
@@ -369,7 +374,7 @@ HnswLoadNeighbors(HnswElement element, Relation index, int m, int lm, int lc)
* Load elements for insert
*/
static void
LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, int *idx, Relation index, FmgrInfo **procinfo, Oid *collation)
LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Relation index, HnswSupport * support)
{
char *base = NULL;
@@ -378,9 +383,8 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, i
HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element);
double distance;
bool matches;
HnswLoadElement(element, &distance, &matches, &q, qtup, NULL, index, procinfo, collation, true, NULL);
HnswLoadElement(element, &distance, q, index, support, true, NULL);
hc->distance = distance;
/* Prune element if being deleted */
@@ -396,7 +400,7 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, Datum q, IndexTuple qtup, i
* Get update index
*/
static int
GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, MemoryContext updateCtx)
GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int m, int lm, int lc, Relation index, HnswSupport * support, MemoryContext updateCtx)
{
char *base = NULL;
int idx = -1;
@@ -421,13 +425,14 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
idx = -2;
else
{
Datum q = HnswGetValue(base, element);
IndexTuple qtup = HnswPtrAccess(base, element->itup);;
HnswQuery q;
LoadElementsForInsert(neighbors, q, qtup, &idx, index, procinfo, collation);
q.value = HnswGetValue(base, element);
LoadElementsForInsert(neighbors, &q, &idx, index, support);
if (idx == -1)
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, procinfo, collation);
HnswUpdateConnection(base, neighbors, newElement, distance, lm, &idx, index, support);
}
MemoryContextSwitchTo(oldCtx);
@@ -532,7 +537,7 @@ UpdateNeighborOnDisk(HnswElement element, HnswElement newElement, int idx, int m
* Update neighbors
*/
void
HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement e, int m, bool checkExisting, bool building)
HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building)
{
char *base = NULL;
@@ -555,7 +560,7 @@ HnswUpdateNeighborsOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, H
HnswElement neighborElement = HnswPtrAccess(base, hc->element);
int idx;
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, procinfo, collation, updateCtx);
idx = GetUpdateIndex(neighborElement, e, hc->distance, m, lm, lc, index, support, updateCtx);
/* New element was not selected as a neighbor */
if (idx == -1)
@@ -633,26 +638,16 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
TupleDesc tupdesc = RelationGetDescr(index);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
if (HnswUseIndexTuple(index))
{
/* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc))
return false;
}
else
{
/* Exit early since ordered by distance */
if (!datumIsEqual(value, HnswGetValue(base, neighborElement), false, -1))
return false;
}
/* Exit early since ordered by distance */
if (!datumIsEqual(value, neighborValue, false, -1))
return false;
if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true;
@@ -665,7 +660,7 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
* Update graph on disk
*/
static void
UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{
BlockNumber newInsertPage = InvalidBlockNumber;
@@ -681,7 +676,7 @@ UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswEleme
HnswUpdateMetaPage(index, 0, NULL, newInsertPage, MAIN_FORKNUM, building);
/* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, false, building);
HnswUpdateNeighborsOnDisk(index, support, element, m, false, building);
/* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -692,19 +687,15 @@ UpdateGraphOnDisk(Relation index, FmgrInfo **procinfo, Oid *collation, HnswEleme
* Insert a tuple into the index
*/
bool
HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull, ItemPointer heap_tid, bool building)
HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building)
{
HnswElement entryPoint;
HnswElement element;
int m;
int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo[2];
Oid *collation;
LOCKMODE lockmode = ShareLock;
char *base = NULL;
HnswInitProcinfo(procinfo, &collation, index);
/*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts
* before repairing graph. Use a page lock so it does not interfere with
@@ -716,24 +707,8 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
HnswGetMetaPageInfo(index, &m, &entryPoint);
/* Create an element */
element = HnswInitElement(base, heap_tid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
if (HnswUseIndexTuple(index))
{
/* TODO no toast */
TupleDesc tupdesc = RelationGetDescr(index);
IndexTuple itup;
bool unused;
/* TODO fix */
values[0] = value;
itup = index_form_tuple(tupdesc, values, isnull);
HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
}
else
HnswPtrStore(base, element->value, DatumGetPointer(value));
element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->value, DatumGetPointer(value));
/* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level)
@@ -750,10 +725,10 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
}
/* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, false, false);
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false);
/* Update graph on disk */
UpdateGraphOnDisk(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -765,31 +740,19 @@ HnswInsertTupleOnDisk(Relation index, Datum value, Datum *values, bool *isnull,
* Insert a tuple into the index
*/
static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid)
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
{
Datum value;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
FmgrInfo *normprocinfo;
Oid *collation = index->rd_indcollation;
HnswSupport support;
/* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
HnswInitSupport(&support, index);
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Form index value */
if (!HnswFormIndexValue(&value, values, isnull, typeInfo, &support))
return;
/* Normalize if needed */
normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
if (normprocinfo != NULL)
{
if (!HnswCheckNorm(normprocinfo, collation[0], value))
return;
value = HnswNormValue(typeInfo, collation[0], value);
}
HnswInsertTupleOnDisk(index, value, values, isnull, heap_tid, false);
HnswInsertTupleOnDisk(index, &support, value, heaptid, false);
}
/*

View File

@@ -5,41 +5,74 @@
#include "pgstat.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/float.h"
#include "utils/memutils.h"
/*
* Algorithm 5 from paper
*/
static List *
GetScanItems(IndexScanDesc scan, Datum q)
GetScanItems(IndexScanDesc scan, Datum value)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo **procinfo = so->procinfo;
Oid *collation = so->collation;
HnswSupport *support = &so->support;
List *ep;
List *w;
int m;
HnswElement entryPoint;
char *base = NULL;
bool inMemory = false;
ScanKeyData *keyData = scan->keyData;
HnswQuery *q = &so->q;
/* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint);
q->value = value;
so->m = m;
if (entryPoint == NULL)
return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, NULL, keyData, index, procinfo, collation, false, inMemory));
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false));
for (int lc = entryPoint->level; lc >= 1; lc--)
{
w = HnswSearchLayer(base, q, NULL, keyData, ep, 1, lc, index, procinfo, collation, m, false, NULL, inMemory);
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL);
ep = w;
}
return HnswSearchLayer(base, q, NULL, keyData, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, inMemory);
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF ? &so->discarded : NULL, true, &so->tuples);
}
/*
* Resume scan at ground level with discarded candidates
*/
static List *
ResumeScanItems(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
List *ep = NIL;
char *base = NULL;
int batch_size = hnsw_ef_search;
if (pairingheap_is_empty(so->discarded))
return NIL;
/* Get next batch of candidates */
for (int i = 0; i < batch_size; i++)
{
HnswSearchCandidate *sc;
if (pairingheap_is_empty(so->discarded))
break;
sc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
ep = lappend(ep, sc);
}
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
}
/*
@@ -62,13 +95,24 @@ GetScanValue(IndexScanDesc scan)
Assert(!VARATT_IS_EXTENDED(DatumGetPointer(value)));
/* Normalize if needed */
if (so->normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->collation[0], value);
if (so->support.normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->support.collation, value);
}
return value;
}
#if defined(HNSW_MEMORY)
/*
* Show memory usage
*/
static void
ShowMemoryUsage(HnswScanOpaque so)
{
elog(INFO, "memory: %zu KB, tuples: " INT64_FORMAT, MemoryContextMemAllocated(so->tmpCtx, false) / 1024, so->tuples);
}
#endif
/*
* Prepare for an index scan
*/
@@ -77,19 +121,28 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
{
IndexScanDesc scan;
HnswScanOpaque so;
double maxMemory;
scan = RelationGetIndexScan(index, nkeys, norderbys);
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index);
so->first = true;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */
HnswInitProcinfo(so->procinfo, &so->collation, index);
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
HnswInitSupport(&so->support, index);
/*
* Use a lower max allocation size than default to allow scanning more
* tuples for iterative search before exceeding work_mem
*/
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
0, 8 * 1024, 256 * 1024);
/* Calculate max memory */
/* Add 256 extra bytes to fill last block when close */
maxMemory = (double) work_mem * hnsw_scan_mem_multiplier * 1024.0 + 256;
so->maxMemory = Min(maxMemory, (double) SIZE_MAX);
scan->opaque = so;
@@ -105,6 +158,11 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
so->first = true;
/* v and discarded are allocated in tmpCtx */
so->v.tids = NULL;
so->discarded = NULL;
so->tuples = 0;
so->previousDistance = -get_float8_infinity();
MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0)
@@ -162,26 +220,104 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false;
#if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
ShowMemoryUsage(so);
#endif
}
while (list_length(so->w) > 0)
for (;;)
{
char *base = NULL;
HnswSearchCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element);
HnswSearchCandidate *sc;
HnswElement element;
ItemPointer heaptid;
if (list_length(so->w) == 0)
{
if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_OFF)
break;
/* Empty index */
if (so->discarded == NULL)
break;
/* Reached max number of tuples */
if (so->tuples >= hnsw_max_scan_tuples)
{
if (pairingheap_is_empty(so->discarded))
break;
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
/* Prevent scans from consuming too much memory */
else if (MemoryContextMemAllocated(so->tmpCtx, false) > so->maxMemory)
{
if (pairingheap_is_empty(so->discarded))
{
ereport(DEBUG1,
(errmsg("hnsw index scan reached memory limit after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase hnsw.scan_mem_multiplier to scan more tuples.")));
break;
}
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
else
{
/*
* Locking ensures when neighbors are read, the elements they
* reference will not be deleted (and replaced) during the
* iteration.
*
* Elements loaded into memory on previous iterations may have
* been deleted (and replaced), so when reading neighbors, the
* element version must be checked.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
so->w = ResumeScanItems(scan);
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY)
ShowMemoryUsage(so);
#endif
}
if (list_length(so->w) == 0)
break;
}
sc = llast(so->w);
element = HnswPtrAccess(base, sc->element);
/* Move to next element if no valid heap TIDs */
if (!hc->matches || element->heaptidsLength == 0)
if (element->heaptidsLength == 0)
{
so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */
if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF)
{
pfree(element);
pfree(sc);
}
continue;
}
heaptid = &element->heaptids[--element->heaptidsLength];
if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_STRICT)
{
if (sc->distance < so->previousDistance)
continue;
so->previousDistance = sc->distance;
}
MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid;

View File

@@ -100,19 +100,6 @@ hash_offset(Size offset)
#define SH_DEFINE
#include "lib/simplehash.h"
typedef union
{
pointerhash_hash *pointers;
offsethash_hash *offsets;
tidhash_hash *tids;
} visited_hash;
typedef union
{
HnswElement element;
ItemPointerData indextid;
} HnswUnvisited;
/*
* Get the max number of connections in an upper layer for each element in the index
*/
@@ -154,41 +141,14 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
}
/*
* Init procinfo
* Init support functions
*/
void
HnswInitProcinfo(FmgrInfo **procinfo, Oid **collation, Relation index)
HnswInitSupport(HnswSupport * support, Relation index)
{
procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
*collation = index->rd_indcollation;
}
/*
* Get element tuple size
*/
Size
HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple)
{
Size size;
if (useIndexTuple)
{
IndexTuple itup = HnswPtrAccess(base, element->itup);
size = IndexTupleSize(itup);
}
else
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
size = VARSIZE_ANY(valuePtr);
}
return HNSW_ELEMENT_TUPLE_SIZE(size);
support->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
support->collation = index->rd_indcollation[0];
support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
}
/*
@@ -204,40 +164,9 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
* Check if non-zero norm
*/
bool
HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value)
HnswCheckNorm(HnswSupport * support, Datum value)
{
return DatumGetFloat8(FunctionCall1Coll(procinfo, collation, value)) > 0;
}
/*
* Check if index tuples are equal
*/
bool
HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc)
{
for (int i = 0; i < tupdesc->natts; i++)
{
bool nullA;
bool nullB;
Datum datumA = index_getattr(a, i + 1, tupdesc, &nullA);
Datum datumB = index_getattr(b, i + 1, tupdesc, &nullB);
if (nullA || nullB)
{
if (nullA != nullB)
return false;
}
else
{
Form_pg_attribute att = TupleDescAttr(tupdesc, i);
if (!datumIsEqual(datumA, datumB, att->attbyval, att->attlen))
return false;
}
}
return true;
return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation, value)) > 0;
}
/*
@@ -322,11 +251,12 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
element->level = level;
element->deleted = 0;
/* Start at one to make it easier to find issues */
element->version = 1;
HnswInitNeighbors(base, element, m, allocator);
HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element;
}
@@ -353,7 +283,6 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->offno = offno;
HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element;
}
@@ -465,15 +394,45 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
UnlockReleaseBuffer(buf);
}
/*
* Form index value
*/
bool
HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support)
{
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
/* Check value */
if (typeInfo->checkValue != NULL)
typeInfo->checkValue(DatumGetPointer(value));
/* Normalize if needed */
if (support->normprocinfo != NULL)
{
if (!HnswCheckNorm(support, value))
return false;
value = HnswNormValue(typeInfo, support->collation, value);
}
*out = value;
return true;
}
/*
* Set element tuple, except for neighbor info
*/
void
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple)
HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level;
etup->deleted = 0;
etup->version = element->version;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
if (i < element->heaptidsLength)
@@ -481,19 +440,7 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool
else
ItemPointerSetInvalid(&etup->heaptids[i]);
}
if (useIndexTuple)
{
IndexTuple itup = HnswPtrAccess(base, element->itup);
memcpy(&etup->data, itup, IndexTupleSize(itup));
}
else
{
Pointer valuePtr = HnswPtrAccess(base, element->value);
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
}
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
}
/*
@@ -528,16 +475,18 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
}
ntup->count = idx;
ntup->version = e->version;
}
/*
* Load an element from a tuple
*/
void
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index)
HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec)
{
element->level = etup->level;
element->deleted = etup->deleted;
element->version = etup->version;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptidsLength = 0;
@@ -557,128 +506,26 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
if (loadVec)
{
char *base = NULL;
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
if (HnswUseIndexTuple(index))
{
IndexTuple itup = CopyIndexTuple((IndexTuple) &etup->data);
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
HnswPtrStore(base, element->itup, itup);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
}
else
{
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
HnswPtrStore(base, element->value, DatumGetPointer(value));
}
HnswPtrStore(base, element->value, DatumGetPointer(value));
}
}
/*
* Get the attribute distance
*/
static inline double
AttributeDistance(double e)
{
/* TODO Better bias */
/* must be >> max(w * g) + 1 / log10(2) */
double bias = 4.32;
return e > 0 ? bias - 1.0 / log10(e + 1) : 0;
}
/*
* Calculate the distance between values
*/
static double
HnswGetDistance(IndexTuple itup, Datum vec, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool *matches)
static inline double
HnswGetDistance(Datum a, Datum b, HnswSupport * support)
{
double g;
if (DatumGetPointer(q) == NULL)
g = 0;
else
g = DatumGetFloat8(FunctionCall2Coll(procinfo[0], collation[0], q, vec));
Assert(PointerIsValid(matches));
*matches = true;
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
{
double w = 0.25;
double e = 0.0;
TupleDesc tupdesc = RelationGetDescr(index);
if (keyData)
{
/* TODO need to pass length of key data */
int keyCount = 1;
for (int i = 0; i < keyCount; i++)
{
ScanKey key = &keyData[i];
bool isnull;
Datum value = index_getattr(itup, key->sk_attno, tupdesc, &isnull);
bool attnull = key->sk_flags & SK_ISNULL;
if (isnull || attnull)
{
if (isnull != attnull)
{
e += 1000;
*matches = false;
}
}
else if (!DatumGetBool(FunctionCall2Coll(&key->sk_func, key->sk_collation, value, key->sk_argument)))
{
double ei = fabs(DatumGetFloat8(FunctionCall2Coll(procinfo[key->sk_attno - 1], collation[key->sk_attno - 1], value, key->sk_argument)));
if (ei > 0)
e += ei;
else
/* Distance is zero for inequality */
e += 1000;
*matches = false;
}
}
return w * g + AttributeDistance(e);
}
else if (qtup)
{
int keyCount = IndexRelationGetNumberOfKeyAttributes(index) - 1;
for (int i = 0; i < keyCount; i++)
{
bool isnull;
bool attnull;
Datum value = index_getattr(itup, i + 2, tupdesc, &isnull);
Datum value2 = index_getattr(qtup, i + 2, tupdesc, &attnull);
if (isnull || attnull)
{
if (isnull != attnull)
e += 1000;
}
else
e += fabs(DatumGetFloat8(FunctionCall2Coll(procinfo[i + 1], collation[i + 1], value, value2)));
}
return w * g + AttributeDistance(e);
}
}
return g;
return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, a, b));
}
/*
* Load an element and optionally get its distance from q
*/
static void
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance, HnswElement * element)
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance, HnswElement * element)
{
Buffer buf;
Page page;
@@ -696,23 +543,10 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
/* Calculate distance */
if (distance != NULL)
{
IndexTuple itup = NULL;
Datum value;
if (HnswUseIndexTuple(index))
{
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
itup = (IndexTuple) &etup->data;
value = index_getattr(itup, 1, tupdesc, &unused);
}
if (DatumGetPointer(q->value) == NULL)
*distance = 0;
else
{
value = PointerGetDatum(&etup->data);
}
*distance = HnswGetDistance(itup, value, *q, qtup, keyData, index, procinfo, collation, matches);
*distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
}
/* Load element */
@@ -721,7 +555,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
if (*element == NULL)
*element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(*element, etup, true, loadVec, index);
HnswLoadElementFromTuple(*element, etup, true, loadVec);
}
UnlockReleaseBuffer(buf);
@@ -731,41 +565,51 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
* Load an element and optionally get its distance from q
*/
void
HnswLoadElement(HnswElement element, double *distance, bool *matches, Datum *q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, double *maxDistance)
HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance)
{
HnswLoadElementImpl(element->blkno, element->offno, distance, matches, q, qtup, keyData, index, procinfo, collation, loadVec, maxDistance, &element);
HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element);
}
/*
* Get the distance for an element
*/
static double
GetElementDistance(char *base, HnswElement element, bool *matches, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation)
GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support)
{
Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
return HnswGetDistance(itup, value, q, qtup, keyData, index, procinfo, collation, matches);
return HnswGetDistance(q->value, value, support);
}
/*
* Allocate a search candidate
*/
static HnswSearchCandidate *
HnswInitSearchCandidate(char *base, HnswElement element, double distance)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, sc->element, element);
sc->distance = distance;
return sc;
}
/*
* Create a candidate for the entry point
*/
HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, IndexTuple qtup, ScanKeyData *keyData, Relation index, FmgrInfo **procinfo, Oid *collation, bool loadVec, bool inMemory)
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL;
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, qtup, keyData, index, procinfo, collation);
distance = GetElementDistance(base, entryPoint, q, support);
else
HnswLoadElement(entryPoint, &sc->distance, &sc->matches, &q, qtup, keyData, index, procinfo, collation, loadVec, NULL);
return sc;
}
HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL);
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
return HnswInitSearchCandidate(base, entryPoint, distance);
}
/*
* Compare candidate distances
@@ -782,6 +626,21 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
return 0;
}
/*
* Compare discarded candidate distances
*/
static int
CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return -1;
return 0;
}
/*
* Compare candidate distances
*/
@@ -902,8 +761,11 @@ HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation i
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
/* Ensure expected neighbors */
if (ntup->count != (element->level + 2) * m)
/*
* Ensure the neighbor tuple has not been deleted or replaced between
* index scan iterations
*/
if (ntup->version != element->version || ntup->count != (element->level + 2) * m)
{
UnlockReleaseBuffer(buf);
return false;
@@ -949,23 +811,34 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List *ep, int ef, int lc, Relation index, FmgrInfo **procinfo, Oid *collation, int m, bool inserting, HnswElement skipElement, bool inMemory)
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
{
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
visited_hash v;
visited_hash vh;
ListCell *lc2;
HnswNeighborArray *localNeighborhood = NULL;
Size neighborhoodSize = 0;
int lm = HnswGetLayerM(m, lc);
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength;
uint64 additional = 0;
uint64 maxAdditional = keyData && lc == 0 ? 10000 : 0;
bool inMemory = index == NULL;
InitVisited(base, &v, inMemory, ef, m);
if (v == NULL)
{
v = &vh;
initVisited = true;
}
if (initVisited)
{
InitVisited(base, v, inMemory, ef, m);
if (discarded != NULL)
*discarded = pairingheap_allocate(CompareNearestDiscardedCandidates, NULL);
}
/* Create local memory for neighborhood if needed */
if (inMemory)
@@ -980,15 +853,18 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2);
bool found;
AddToVisited(base, &v, sc->element, inMemory, &found);
if (initVisited)
{
AddToVisited(base, v, sc->element, inMemory, &found);
/* OK to count elements instead of tuples */
if (tuples != NULL)
(*tuples)++;
}
pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_node);
/* Do not count elements that do not match filter towards ef */
if (!sc->matches && ++additional <= maxAdditional)
continue;
/*
* Do not count elements being deleted towards ef when vacuuming. It
* would be ideal to do this for inserts as well, but this could
@@ -1010,16 +886,19 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
cElement = HnswPtrAccess(base, c->element);
if (inMemory)
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize);
else
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc);
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
/* OK to count elements instead of tuples */
if (tuples != NULL)
(*tuples) += unvisitedLength;
for (int i = 0; i < unvisitedLength; i++)
{
HnswElement eElement;
HnswSearchCandidate *e;
double eDistance;
bool eMatches;
bool alwaysAdd = wlen < ef;
f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
@@ -1027,7 +906,7 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
if (inMemory)
{
eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, &eMatches, q, qtup, keyData, index, procinfo, collation);
eDistance = GetElementDistance(base, eElement, q, support);
}
else
{
@@ -1037,26 +916,30 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, &q, qtup, keyData, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
if (eElement == NULL)
continue;
}
if (!(eDistance < f->distance || alwaysAdd))
continue;
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
{
if (discarded != NULL)
{
/* Create a new candidate */
e = HnswInitSearchCandidate(base, eElement, eDistance);
pairingheap_add(*discarded, &e->w_node);
}
Assert(!eElement->deleted);
continue;
}
/* Make robust to issues */
if (eElement->level < lc)
continue;
/* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
e->matches = eMatches;
e = HnswInitSearchCandidate(base, eElement, eDistance);
pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node);
@@ -1067,15 +950,16 @@ HnswSearchLayer(char *base, Datum q, IndexTuple qtup, ScanKeyData *keyData, List
*/
if (CountElement(skipElement, eElement))
{
/* Do not count elements that do not match filter towards ef */
if (!e->matches && ++additional <= maxAdditional)
continue;
wlen++;
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
{
HnswSearchCandidate *d = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
if (discarded != NULL)
pairingheap_add(*discarded, &d->w_node);
}
}
}
}
@@ -1143,11 +1027,10 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
* Check if an element is closer to q than any element from R
*/
static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrInfo **procinfo, Oid *collation)
CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support)
{
HnswElement eElement = HnswPtrAccess(base, e->element);
Datum eValue = HnswGetValue(base, eElement);
IndexTuple etup = HnswPtrAccess(base, eElement->itup);
ListCell *lc2;
foreach(lc2, r)
@@ -1155,9 +1038,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrI
HnswCandidate *ri = lfirst(lc2);
HnswElement riElement = HnswPtrAccess(base, ri->element);
Datum riValue = HnswGetValue(base, riElement);
IndexTuple ritup = HnswPtrAccess(base, riElement->itup);
bool matches;
float distance = HnswGetDistance(etup, eValue, riValue, ritup, NULL, index, procinfo, collation, &matches);
float distance = HnswGetDistance(eValue, riValue, support);
if (distance <= e->distance)
return false;
@@ -1170,7 +1051,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, FmgrI
* Algorithm 4 from paper
*/
static List *
SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo, Oid *collation, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
SelectNeighbors(char *base, List *c, int lm, HnswSupport * support, bool *closerSet, HnswCandidate * newCandidate, HnswCandidate * *pruned, bool sortCandidates)
{
List *r = NIL;
List *w = list_copy(c);
@@ -1204,7 +1085,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
/* Use previous state of r and wd to skip work when possible */
if (mustCalculate)
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, support);
else if (list_length(added) > 0)
{
/* Keep Valgrind happy for in-memory, parallel builds */
@@ -1217,7 +1098,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
*/
if (e->closer)
{
e->closer = CheckElementCloser(base, e, added, index, procinfo, collation);
e->closer = CheckElementCloser(base, e, added, support);
if (!e->closer)
removedAny = true;
@@ -1230,7 +1111,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
*/
if (removedAny)
{
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, support);
if (e->closer)
added = lappend(added, e);
}
@@ -1238,7 +1119,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, FmgrInfo **procinfo
}
else if (e == newCandidate)
{
e->closer = CheckElementCloser(base, e, r, index, procinfo, collation);
e->closer = CheckElementCloser(base, e, r, support);
if (e->closer)
added = lappend(added, e);
}
@@ -1289,7 +1170,7 @@ AddConnections(char *base, HnswElement element, List *neighbors, int lc)
* Update connections
*/
void
HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, FmgrInfo **procinfo, Oid *collation)
HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support)
{
HnswCandidate newHc;
@@ -1315,7 +1196,7 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, index, procinfo, collation, &neighbors->closerSet, &newHc, &pruned, true);
SelectNeighbors(base, c, lm, support, &neighbors->closerSet, &newHc, &pruned, true);
/* Should not happen */
if (pruned == NULL)
@@ -1386,16 +1267,17 @@ PrecomputeHash(char *base, HnswElement element)
* Algorithm 1 from paper
*/
void
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo **procinfo, Oid *collation, int m, int efConstruction, bool existing, bool inMemory)
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing)
{
List *ep;
List *w;
int level = element->level;
int entryLevel;
Datum q = HnswGetValue(base, element);
IndexTuple qtup = HnswPtrAccess(base, element->itup);
ScanKeyData *keyData = NULL;
HnswQuery q;
HnswElement skipElement = existing ? element : NULL;
bool inMemory = index == NULL;
q.value = HnswGetValue(base, element);
/* Precompute hash */
if (inMemory)
@@ -1406,13 +1288,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
return;
/* Get entry point and level */
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, qtup, keyData, index, procinfo, collation, true, inMemory));
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true));
entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(base, q, qtup, keyData, ep, 1, lc, index, procinfo, collation, m, true, skipElement, inMemory);
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
ep = w;
}
@@ -1431,7 +1313,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL;
ListCell *lc2;
w = HnswSearchLayer(base, q, qtup, keyData, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, inMemory);
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
/* Convert search candidates to candidates */
foreach(lc2, w)
@@ -1455,7 +1337,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
* sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance.
*/
neighbors = SelectNeighbors(base, lw, lm, index, procinfo, collation, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
neighbors = SelectNeighbors(base, lw, lm, support, &HnswGetNeighbors(base, element, lc)->closerSet, NULL, NULL, false);
AddConnections(base, element, neighbors, lc);

View File

@@ -184,13 +184,12 @@ static void
RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswElement entryPoint)
{
Relation index = vacuumstate->index;
HnswSupport *support = &vacuumstate->support;
Buffer buf;
Page page;
GenericXLogState *state;
int m = vacuumstate->m;
int efConstruction = vacuumstate->efConstruction;
FmgrInfo **procinfo = vacuumstate->procinfo;
Oid *collation = vacuumstate->collation;
BufferAccessStrategy bas = vacuumstate->bas;
HnswNeighborTuple ntup = vacuumstate->ntup;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, m);
@@ -205,7 +204,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, procinfo, collation, m, efConstruction, true, false);
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true);
/* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -229,7 +228,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
UnlockReleaseBuffer(buf);
/* Update neighbors */
HnswUpdateNeighborsOnDisk(index, procinfo, collation, element, m, true, false);
HnswUpdateNeighborsOnDisk(index, support, element, m, true, false);
}
/*
@@ -239,6 +238,7 @@ static void
RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
{
Relation index = vacuumstate->index;
HnswSupport *support = &vacuumstate->support;
HnswElement highestPoint = &vacuumstate->highestPoint;
HnswElement entryPoint;
MemoryContext oldCtx = MemoryContextSwitchTo(vacuumstate->tmpCtx);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */
HnswLoadElement(highestPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
HnswLoadElement(highestPoint, NULL, NULL, index, support, true, NULL);
/* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine.
*/
HnswLoadElement(entryPoint, NULL, NULL, NULL, NULL, NULL, index, vacuumstate->procinfo, vacuumstate->collation, true, NULL);
HnswLoadElement(entryPoint, NULL, NULL, index, support, true, NULL);
if (NeedsUpdated(vacuumstate, entryPoint))
{
@@ -370,7 +370,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Create an element */
element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true, index);
HnswLoadElementFromTuple(element, etup, false, true);
elements = lappend(elements, element);
}
@@ -440,7 +440,6 @@ MarkDeleted(HnswVacuumState * vacuumstate)
BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas;
bool useIndexTuple = HnswUseIndexTuple(index);
/*
* Wait for index scans to complete. Scans before this point may contain
@@ -522,19 +521,20 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */
etup->deleted = 1;
if (useIndexTuple)
{
IndexTuple itup = (IndexTuple) &etup->data;
MemSet(itup, 0, IndexTupleSize(itup));
}
else
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
/* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++)
ItemPointerSetInvalid(&ntup->indextids[i]);
/* Increment version */
/* This is used to avoid incorrect reads for iterative scans */
/* Reserve some bits for future use */
etup->version++;
if (etup->version > 15)
etup->version = 1;
ntup->version = etup->version;
/*
* We modified the tuples in place, no need to call
* PageIndexTupleOverwrite
@@ -581,12 +581,13 @@ InitVacuumState(HnswVacuumState * vacuumstate, IndexVacuumInfo *info, IndexBulkD
vacuumstate->callback_state = callback_state;
vacuumstate->efConstruction = HnswGetEfConstruction(index);
vacuumstate->bas = GetAccessStrategy(BAS_BULKREAD);
HnswInitProcinfo(vacuumstate->procinfo, &vacuumstate->collation, index);
vacuumstate->ntup = palloc0(HNSW_TUPLE_ALLOC_SIZE);
vacuumstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw vacuum temporary context",
ALLOCSET_DEFAULT_SIZES);
HnswInitSupport(&vacuumstate->support, index);
/* Get m from metapage */
HnswGetMetaPageInfo(index, &vacuumstate->m, NULL);

View File

@@ -228,11 +228,11 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list)
{
Datum value;
bool isnull;
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{
Datum value;
bool isnull;
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
value = slot_getattr(slot, 3, &isnull);
@@ -254,8 +254,8 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IndexTuple itup = NULL; /* silence compiler warning */
int64 inserted = 0;
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = RelationGetDescr(index);
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = buildstate->tupdesc;
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
@@ -319,6 +319,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->index = index;
buildstate->indexInfo = indexInfo;
buildstate->typeInfo = IvfflatGetTypeInfo(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
@@ -356,12 +357,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
errmsg("dimensions must be greater than one for this opclass")));
/* Create tuple description for sorting */
buildstate->tupdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
buildstate->sortdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 3, "vector", buildstate->tupdesc->attrs[0].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -633,7 +634,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate);
ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate);
buildstate.sortstate = ivfspool->sortstate;
scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared));
@@ -950,7 +951,7 @@ AssignTuples(IvfflatBuildState * buildstate)
}
/* Begin serial/leader tuplesort */
buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate);
buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);
/* Add tuples to sort */
if (buildstate->heap != NULL)

View File

@@ -17,8 +17,16 @@
#endif
int ivfflat_probes;
int ivfflat_iterative_scan;
int ivfflat_max_probes;
static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_scan_options[] = {
{"off", IVFFLAT_ITERATIVE_SCAN_OFF, false},
{"relaxed_order", IVFFLAT_ITERATIVE_SCAN_RELAXED, false},
{NULL, 0, false}
};
/*
* Initialize index options and variables
*/
@@ -31,7 +39,16 @@ IvfflatInit(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of 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, GUC_EXPLAIN, NULL, NULL, NULL);
DefineCustomEnumVariable("ivfflat.iterative_scan", "Sets the mode for iterative scans",
NULL, &ivfflat_iterative_scan,
IVFFLAT_ITERATIVE_SCAN_OFF, ivfflat_iterative_scan_options, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* If this is less than probes, probes is used */
DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative scans",
NULL, &ivfflat_max_probes,
IVFFLAT_MAX_LISTS, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat");
}
@@ -82,6 +99,10 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return;
}

View File

@@ -80,6 +80,14 @@
/* Variables */
extern int ivfflat_probes;
extern int ivfflat_iterative_scan;
extern int ivfflat_max_probes;
typedef enum IvfflatIterativeScanMode
{
IVFFLAT_ITERATIVE_SCAN_OFF,
IVFFLAT_ITERATIVE_SCAN_RELAXED
} IvfflatIterativeScanMode;
typedef struct VectorArrayData
{
@@ -165,6 +173,7 @@ typedef struct IvfflatBuildState
Relation index;
IndexInfo *indexInfo;
const IvfflatTypeInfo *typeInfo;
TupleDesc tupdesc;
/* Settings */
int dimensions;
@@ -198,7 +207,7 @@ typedef struct IvfflatBuildState
/* Sorting */
Tuplesortstate *sortstate;
TupleDesc tupdesc;
TupleDesc sortdesc;
TupleTableSlot *slot;
/* Memory */
@@ -247,8 +256,11 @@ typedef struct IvfflatScanOpaqueData
{
const IvfflatTypeInfo *typeInfo;
int probes;
int maxProbes;
int dimensions;
bool first;
Datum value;
MemoryContext tmpCtx;
/* Sorting */
Tuplesortstate *sortstate;
@@ -265,7 +277,9 @@ typedef struct IvfflatScanOpaqueData
/* Lists */
pairingheap *listQueue;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
BlockNumber *listPages;
int listIndex;
IvfflatScanList *lists;
} IvfflatScanOpaqueData;
typedef IvfflatScanOpaqueData * IvfflatScanOpaque;

View File

@@ -98,7 +98,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
IvfflatGetMetaPageInfo(index, NULL, NULL);
/* Find the insert page - sets the page and list info */
FindInsertPage(index, values, &insertPage, &listInfo);
FindInsertPage(index, &value, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage;

View File

@@ -10,10 +10,7 @@
#include "miscadmin.h"
#include "pgstat.h"
#include "storage/bufmgr.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h"
#endif
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
@@ -65,7 +62,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Use procinfo from the index instead of scan key for performance */
distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
if (listCount < so->probes)
if (listCount < so->maxProbes)
{
IvfflatScanList *scanlist;
@@ -78,7 +75,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Calculate max distance */
if (listCount == so->probes)
if (listCount == so->maxProbes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
}
else if (distance < maxDistance)
@@ -102,6 +99,11 @@ GetScanLists(IndexScanDesc scan, Datum value)
UnlockReleaseBuffer(cbuf);
}
for (int i = listCount - 1; i >= 0; i--)
so->listPages[i] = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
Assert(pairingheap_is_empty(so->listQueue));
}
/*
@@ -114,11 +116,14 @@ GetScanItems(IndexScanDesc scan, Datum value)
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0;
TupleTableSlot *slot = so->vslot;
int batchProbes = 0;
tuplesort_reset(so->sortstate);
/* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue))
while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes)
{
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
BlockNumber searchPage = so->listPages[so->listIndex++];
/* Search all entry pages for list */
while (BlockNumberIsValid(searchPage))
@@ -166,13 +171,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
}
}
if (tuples < 100)
if (tuples < 100 && ivfflat_iterative_scan == IVFFLAT_ITERATIVE_SCAN_OFF)
ereport(DEBUG1,
(errmsg("index scan found few tuples"),
errdetail("Index may have been created with little data."),
errhint("Recreate the index and possibly decrease lists.")));
tuplesort_performsort(so->sortstate);
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
}
/*
@@ -209,7 +218,13 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */
if (so->normprocinfo != NULL)
{
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
value = IvfflatNormValue(so->typeInfo, so->collation, value);
MemoryContextSwitchTo(oldCtx);
}
}
return value;
@@ -240,19 +255,30 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
int lists;
int dimensions;
int probes = ivfflat_probes;
int maxProbes;
MemoryContext oldCtx;
scan = RelationGetIndexScan(index, nkeys, norderbys);
/* Get lists and dimensions from metapage */
IvfflatGetMetaPageInfo(index, &lists, &dimensions);
if (ivfflat_iterative_scan != IVFFLAT_ITERATIVE_SCAN_OFF)
maxProbes = Max(ivfflat_max_probes, probes);
else
maxProbes = probes;
if (probes > lists)
probes = lists;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
if (maxProbes > lists)
maxProbes = lists;
so = (IvfflatScanOpaque) palloc(sizeof(IvfflatScanOpaqueData));
so->typeInfo = IvfflatGetTypeInfo(index);
so->first = true;
so->probes = probes;
so->maxProbes = maxProbes;
so->dimensions = dimensions;
/* Set support functions */
@@ -260,6 +286,12 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
so->collation = index->rd_indcollation[0];
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Ivfflat scan temporary context",
ALLOCSET_DEFAULT_SIZES);
oldCtx = MemoryContextSwitchTo(so->tmpCtx);
/* Create tuple description for sorting */
so->tupdesc = CreateTemplateTupleDesc(2);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
@@ -280,6 +312,11 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->bas = GetAccessStrategy(BAS_BULKREAD);
so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listPages = palloc(maxProbes * sizeof(BlockNumber));
so->listIndex = 0;
so->lists = palloc(maxProbes * sizeof(IvfflatScanList));
MemoryContextSwitchTo(oldCtx);
scan->opaque = so;
@@ -294,11 +331,9 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
if (!so->first)
tuplesort_reset(so->sortstate);
so->first = true;
pairingheap_reset(so->listQueue);
so->listIndex = 0;
if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -314,6 +349,8 @@ bool
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
ItemPointer heaptid;
bool isnull;
/*
* Index can be used to scan backward, but Postgres doesn't support
@@ -341,28 +378,23 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanLists", GetScanLists(scan, value));
IvfflatBench("GetScanItems", GetScanItems(scan, value));
so->first = false;
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
/* Clean up if we allocated a new value */
if (value != scan->orderByData->sk_argument)
pfree(DatumGetPointer(value));
so->value = value;
}
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{
bool isnull;
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
if (so->listIndex == so->maxProbes)
return false;
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
}
return false;
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
}
/*
@@ -373,12 +405,10 @@ ivfflatendscan(IndexScanDesc scan)
{
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue);
/* Free any temporary files */
tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
/* TODO Free vslot and mslot without freeing TupleDesc */
MemoryContextDelete(so->tmpCtx);
pfree(so);
scan->opaque = NULL;

View File

@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
Page cpage;
OffsetNumber coffno;
OffsetNumber cmaxoffno;
BlockNumber startPages[MaxOffsetNumber];
BlockNumber listPages[MaxOffsetNumber];
ListInfo listInfo;
cbuf = ReadBuffer(index, blkno);
@@ -40,7 +40,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
{
IvfflatList list = (IvfflatList) PageGetItem(cpage, PageGetItemId(cpage, coffno));
startPages[coffno - FirstOffsetNumber] = list->startPage;
listPages[coffno - FirstOffsetNumber] = list->startPage;
}
listInfo.blkno = blkno;
@@ -50,7 +50,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
for (coffno = FirstOffsetNumber; coffno <= cmaxoffno; coffno = OffsetNumberNext(coffno))
{
BlockNumber searchPage = startPages[coffno - FirstOffsetNumber];
BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
BlockNumber insertPage = InvalidBlockNumber;
/* Iterate over entry pages */

View File

@@ -99,6 +99,32 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
4
(1 row)
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SET hnsw.iterative_scan = strict_order;
SET hnsw.ef_search = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
SET hnsw.iterative_scan = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
RESET hnsw.iterative_scan;
RESET hnsw.ef_search;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -139,4 +165,29 @@ SET hnsw.ef_search = 0;
ERROR: 0 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SET hnsw.ef_search = 1001;
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SHOW hnsw.iterative_scan;
hnsw.iterative_scan
---------------------
off
(1 row)
SET hnsw.iterative_scan = on;
ERROR: invalid value for parameter "hnsw.iterative_scan": "on"
HINT: Available values: off, relaxed_order, strict_order.
SHOW hnsw.max_scan_tuples;
hnsw.max_scan_tuples
----------------------
20000
(1 row)
SET hnsw.max_scan_tuples = 0;
ERROR: 0 is outside the valid range for parameter "hnsw.max_scan_tuples" (1 .. 2147483647)
SHOW hnsw.scan_mem_multiplier;
hnsw.scan_mem_multiplier
--------------------------
1
(1 row)
SET hnsw.scan_mem_multiplier = 0;
ERROR: 0 is outside the valid range for parameter "hnsw.scan_mem_multiplier" (1 .. 1000)
DROP TABLE t;

View File

@@ -81,6 +81,37 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
3
(1 row)
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 3);
SET ivfflat.iterative_scan = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
[0,0,0]
(3 rows)
SET ivfflat.max_probes = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
(1 row)
SET ivfflat.max_probes = 2;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
[1,1,1]
(2 rows)
RESET ivfflat.iterative_scan;
RESET ivfflat.max_probes;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -109,4 +140,27 @@ SHOW ivfflat.probes;
1
(1 row)
SET ivfflat.probes = 0;
ERROR: 0 is outside the valid range for parameter "ivfflat.probes" (1 .. 32768)
SET ivfflat.probes = 32769;
ERROR: 32769 is outside the valid range for parameter "ivfflat.probes" (1 .. 32768)
SHOW ivfflat.iterative_scan;
ivfflat.iterative_scan
------------------------
off
(1 row)
SET ivfflat.iterative_scan = on;
ERROR: invalid value for parameter "ivfflat.iterative_scan": "on"
HINT: Available values: off, relaxed_order.
SHOW ivfflat.max_probes;
ivfflat.max_probes
--------------------
32768
(1 row)
SET ivfflat.max_probes = 0;
ERROR: 0 is outside the valid range for parameter "ivfflat.max_probes" (1 .. 32768)
SET ivfflat.max_probes = 32769;
ERROR: 32769 is outside the valid range for parameter "ivfflat.max_probes" (1 .. 32768)
DROP TABLE t;

View File

@@ -57,6 +57,23 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SET hnsw.iterative_scan = strict_order;
SET hnsw.ef_search = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET hnsw.iterative_scan = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
RESET hnsw.iterative_scan;
RESET hnsw.ef_search;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -81,4 +98,16 @@ SHOW hnsw.ef_search;
SET hnsw.ef_search = 0;
SET hnsw.ef_search = 1001;
SHOW hnsw.iterative_scan;
SET hnsw.iterative_scan = on;
SHOW hnsw.max_scan_tuples;
SET hnsw.max_scan_tuples = 0;
SHOW hnsw.scan_mem_multiplier;
SET hnsw.scan_mem_multiplier = 0;
DROP TABLE t;

View File

@@ -44,6 +44,25 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
DROP TABLE t;
-- iterative
CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 3);
SET ivfflat.iterative_scan = relaxed_order;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET ivfflat.max_probes = 1;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SET ivfflat.max_probes = 2;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
RESET ivfflat.iterative_scan;
RESET ivfflat.max_probes;
DROP TABLE t;
-- unlogged
CREATE UNLOGGED TABLE t (val vector(3));
@@ -62,4 +81,16 @@ CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 32769);
SHOW ivfflat.probes;
SET ivfflat.probes = 0;
SET ivfflat.probes = 32769;
SHOW ivfflat.iterative_scan;
SET ivfflat.iterative_scan = on;
SHOW ivfflat.max_probes;
SET ivfflat.max_probes = 0;
SET ivfflat.max_probes = 32769;
DROP TABLE t;

View File

@@ -6,13 +6,7 @@ use Test::More;
my $dim = 3;
my @r = ();
for (1 .. $dim)
{
my $v = int(rand(1000)) + 1;
push(@r, "i % $v");
}
my $array_sql = join(", ", @r);
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
@@ -23,19 +17,20 @@ $node->start;
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
# Get size
my $size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");
# Store values
$node->safe_psql("postgres", "CREATE TABLE tmp AS SELECT * FROM tst;");
# Delete all, vacuum, and insert same data
$node->safe_psql("postgres", "DELETE FROM tst;");
$node->safe_psql("postgres", "VACUUM tst;");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");
# Check size
my $new_size = $node->safe_psql("postgres", "SELECT pg_total_relation_size('tst_v_idx');");

View File

@@ -1,109 +0,0 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @cs = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my $nc = 50;
sub test_recall
{
my ($min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $cs[0] ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my %actual_set = map { $_ => 1 } @actual_ids;
is(scalar(@actual_ids), $limit);
my @expected_ids = split("\n", $expected[$i]);
foreach (@expected_ids)
{
if (exists($actual_set{$_}))
{
$correct++;
}
$total++;
}
}
cmp_ok($correct / $total, ">=", $min, $operator);
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 20000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
push(@cs, int(rand() * $nc));
}
# Get exact results
@expected = ();
for my $i (0 .. $#queries)
{
my $res = $node->safe_psql("postgres", "SELECT i FROM tst WHERE c = $cs[$i] ORDER BY v <-> '$queries[$i]' LIMIT $limit;");
push(@expected, $res);
}
# Add index
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c);");
# Test recall
test_recall(0.99, '<->');
# Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
$node->safe_psql("postgres", "VACUUM tst;");
# Test columns
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (c, v vector_l2_ops);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, c, c);");
like($stderr, qr/index cannot have more than two columns/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops, v vector_l2_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();

View File

@@ -0,0 +1,54 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_scan = relaxed_order;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
foreach ((30, 50, 70))
{
my $max_probes = $_;
my $expected = $max_probes / 10;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 10;
SET ivfflat.iterative_scan = relaxed_order;
SET ivfflat.max_probes = $max_probes;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
done_testing();

View File

@@ -0,0 +1,125 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my @cs = (100, 1000);
sub test_recall
{
my ($c, $probes, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SET ivfflat.iterative_scan = relaxed_order;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SET ivfflat.iterative_scan = relaxed_order;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, "$operator $c");
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[random(), random(), random()] FROM generate_series(1, 100000) i;"
);
# Generate queries
for (1 .. 20)
{
my $r1 = rand();
my $r2 = rand();
my $r3 = rand();
push(@queries, "[$r1,$r2,$r3]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v $opclass);");
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
if ($c == 100)
{
test_recall($c, 1, 0.57, $operator);
test_recall($c, 10, 0.98, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 1, 0.80, $operator);
}
else
{
test_recall($c, 1, 0.88, $operator);
}
}
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();

View File

@@ -0,0 +1,68 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_scan = relaxed_order;
SET hnsw.max_scan_tuples = 100000;
SET hnsw.scan_mem_multiplier = 2;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
is($count, 10);
foreach ((30000, 50000, 70000))
{
my $max_tuples = $_;
my $expected = $max_tuples / 10000;
my $sum = 0;
for my $i (1 .. 20)
{
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_scan = relaxed_order;
SET hnsw.max_scan_tuples = $max_tuples;
SET hnsw.scan_mem_multiplier = 2;
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
));
$sum += $count;
}
my $avg = $sum / 20;
cmp_ok($avg, '>', $expected - 2);
cmp_ok($avg, '<', $expected + 2);
}
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.iterative_scan = relaxed_order;
SET client_min_messages = debug1;
SET work_mem = '1MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
like($stderr, qr/hnsw index scan reached memory limit after \d+ tuples/);
done_testing();

View File

@@ -0,0 +1,118 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my @cs = (50, 500);
sub test_recall
{
my ($c, $ef_search, $min, $operator, $mode) = @_;
my $correct = 0;
my $total = 0;
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_scan = $mode;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan using idx on tst/);
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_scan = $mode;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
));
my @actual_ids = split("\n", $actual);
my @expected_ids = split("\n", $expected[$i]);
my %expected_set = map { $_ => 1 } @expected_ids;
foreach (@actual_ids)
{
if (exists($expected_set{$_}))
{
$correct++;
}
}
$total += $limit;
}
cmp_ok($correct / $total, ">=", $min, "$operator $mode $c");
}
# Initialize node
$node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 50000) i;"
);
# Generate queries
for (1 .. 20)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
}
# Check each index type
my @operators = ("<->", "<=>");
my @opclasses = ("vector_l2_ops", "vector_cosine_ops");
for my $i (0 .. $#operators)
{
my $operator = $operators[$i];
my $opclass = $opclasses[$i];
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
CREATE INDEX idx ON tst USING hnsw (v $opclass);
));
foreach (@cs)
{
my $c = $_;
# Get exact results
@expected = ();
foreach (@queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
WITH top AS (
SELECT v $operator '$_' AS distance FROM tst WHERE i % $c = 0 ORDER BY distance LIMIT $limit
)
SELECT i FROM tst WHERE (v $operator '$_') <= (SELECT MAX(distance) FROM top)
));
push(@expected, $res);
}
test_recall($c, 40, 0.99, $operator, "strict_order");
test_recall($c, 40, 0.99, $operator, "relaxed_order");
}
$node->safe_psql("postgres", "DROP INDEX idx;");
}
done_testing();