Compare commits

..

43 Commits

Author SHA1 Message Date
Andrew Kane
ea689b92b2 Removed duplicated macros [skip ci] 2024-10-10 09:42:03 -07:00
Andrew Kane
51dad9ff2f Merge branch 'master' into hnsw-streaming 2024-10-10 09:41:35 -07:00
Andrew Kane
d06e75f78c Moved union and macros [skip ci] 2024-10-10 09:39:04 -07:00
Andrew Kane
8e1a2715ea Removed assertions [skip ci] 2024-10-10 09:36:17 -07:00
Andrew Kane
25fae23fc3 Removed bench code [skip ci] 2024-10-10 09:35:46 -07:00
Andrew Kane
46d3164a30 Improved variable name [skip ci] 2024-10-10 02:14:01 -07:00
Andrew Kane
f2b2040306 Fixed assertion 2024-10-10 01:29:34 -07:00
Andrew Kane
c46f078e3c Merge branch 'master' into hnsw-streaming 2024-10-10 01:14:33 -07:00
Andrew Kane
c0f6570c4a Debug updates [skip ci] 2024-10-08 22:58:09 -07:00
Andrew Kane
1d3d0f46ac Use DEBUG1 for exceeding work_mem [skip ci] 2024-10-08 22:52:12 -07:00
Andrew Kane
caabac33b8 Fixed compilation [skip ci] 2024-09-29 21:33:07 -07:00
Andrew Kane
8b9333d468 Merge branch 'master' into hnsw-streaming 2024-09-29 15:13:51 -07:00
Andrew Kane
87f5e40495 Updated name [skip ci] 2024-09-29 13:46:05 -07:00
Andrew Kane
8798c1474c Updated naming [skip ci] 2024-09-29 13:45:25 -07:00
Andrew Kane
1b337ad97d Changed option to enum [skip ci] 2024-09-29 13:39:53 -07:00
Andrew Kane
0047630baf Updated readme [skip ci] 2024-09-29 10:37:07 -07:00
Andrew Kane
8b253359ab Improved code [skip ci] 2024-09-28 16:33:30 -07:00
Andrew Kane
8de5f55b0b Fixed assertion [skip ci] 2024-09-28 16:18:53 -07:00
Andrew Kane
351db562af Improved code [skip ci] 2024-09-28 16:07:31 -07:00
Andrew Kane
e1c2d03dba Improved cost estimation [skip ci] 2024-09-28 16:04:07 -07:00
Andrew Kane
ba8e29600b Added todo [skip ci] 2024-09-28 15:24:21 -07:00
Andrew Kane
49e05fb5ba Updated readme [skip ci] 2024-09-28 13:13:10 -07:00
Andrew Kane
9b42662188 Only adjust cost if scanning less than half of the tuples [skip ci] 2024-09-28 12:21:20 -07:00
Andrew Kane
7265927fd6 Updated readme [skip ci] 2024-09-28 12:07:09 -07:00
Andrew Kane
ab57217f48 Added todo [skip ci] 2024-09-28 10:07:09 -07:00
Andrew Kane
5f6e031ccc Updated changelog [skip ci] 2024-09-28 09:57:05 -07:00
Andrew Kane
1a1221f905 Merge branch 'master' into hnsw-streaming 2024-09-26 08:34:07 -07:00
Andrew Kane
40c3e402c7 Removed todo [skip ci] 2024-09-25 17:29:20 -07:00
Andrew Kane
058248fdcc Improved cost code [skip ci] 2024-09-25 17:23:29 -07:00
Andrew Kane
73c5145b77 Use int for ef [skip ci] 2024-09-25 16:52:41 -07:00
Andrew Kane
ec4a23fe49 Added cost estimation [skip ci] 2024-09-25 16:45:04 -07:00
Andrew Kane
38207f5640 Merge branch 'master' into hnsw-streaming 2024-09-25 16:09:09 -07:00
Andrew Kane
4e35c6abe3 Updated readme [skip ci] 2024-09-24 23:24:48 -07:00
Andrew Kane
11e4d040d9 Fixed test [skip ci] 2024-09-24 19:38:06 -07:00
Andrew Kane
b2fa625255 Fixed crash with empty index [skip ci] 2024-09-23 09:42:31 -07:00
Andrew Kane
a8e699c927 Improved message [skip ci] 2024-09-22 22:31:48 -07:00
Andrew Kane
91541fece6 Fixed example [skip ci] 2024-09-22 22:28:39 -07:00
Andrew Kane
f3de487da2 Started readme updates [skip ci] 2024-09-22 22:26:27 -07:00
Andrew Kane
721d4b7e3f Improved test for ef_stream [skip ci] 2024-09-22 18:51:38 -07:00
Andrew Kane
28066d8fe4 Added test for ef_stream [skip ci] 2024-09-22 18:35:35 -07:00
Andrew Kane
495041e43b Added option to limit tuples [skip ci] 2024-09-22 18:10:19 -07:00
Andrew Kane
52c385c03a Only pass discarded when streaming [skip ci] 2024-09-22 17:47:10 -07:00
Andrew Kane
80cbd32dab Added streaming option for HNSW 2024-09-22 12:02:48 -07:00
24 changed files with 276 additions and 865 deletions

View File

@@ -49,7 +49,7 @@ jobs:
- postgres: 16 - postgres: 16
os: macos-14 os: macos-14
- postgres: 14 - postgres: 14
os: macos-13 os: macos-12
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v4
- uses: ankane/setup-postgres@v1 - uses: ankane/setup-postgres@v1

View File

@@ -1,6 +1,5 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- Added support for inline filtering with HNSW
- Added support for iterative index scans - Added support for iterative index scans
- Added casts for arrays to `sparsevec` - Added casts for arrays to `sparsevec`
- Improved cost estimation - Improved cost estimation

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); 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 Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql ```sql
@@ -457,6 +451,77 @@ 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); CREATE TABLE items (embedding vector(3), category_id int) PARTITION BY LIST(category_id);
``` ```
## Iterative Search [unreleased]
*Added in 0.8.0*
With approximate indexes, you can end up with less results than expected due to filtering conditions in the query.
Starting with 0.8.0, you can enable iterative search. If too few results from the initial index scan match the query filters, it will resume scanning until enough results are found. This can significantly improve recall (at the cost of speed).
```tsql
SET hnsw.streaming = on;
-- or
SET ivfflat.streaming = on;
```
However, there are some important caveats.
### Iterative Caveats
With iterative search, its possible for rows to be slightly out of order by distance. For strict ordering, use:
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order ORDER BY distance;
```
For distance filters, use a CTE and place the filter outside it.
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order WHERE distance < 0.1 ORDER BY distance;
```
### Iterative Options
Since scanning a large portion of the index is expensive, there are options to control when the scan ends.
#### HNSW
Specify the max number of additional tuples visited
```sql
SET hnsw.ef_stream = 10000;
```
The scan will also end if reaches `work_mem`. You can see when this happens by enabling debug messages.
```sql
SET client_min_messages = debug1;
```
```text
DEBUG: hnsw index scan exceeded work_mem after 10000 tuples
HINT: Increase work_mem to scan more tuples.
```
If the server has enough memory, you can adjust this with:
```sql
SET work_mem = '8MB';
```
#### IVFFlat
Specify the max number of probes
```sql
SET ivfflat.max_probes = 100;
```
## Half-Precision Vectors ## Half-Precision Vectors
*Added in 0.7.0* *Added in 0.7.0*
@@ -1195,7 +1260,6 @@ Thanks to:
- [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf) - [k-means++: The Advantage of Careful Seeding](https://theory.stanford.edu/~sergei/papers/kMeansPP-soda.pdf)
- [Concept Decompositions for Large Sparse Text Data using Clustering](https://www.cs.utexas.edu/users/inderjit/public_papers/concept_mlj.pdf) - [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) - [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 ## History

View File

@@ -24,11 +24,3 @@ CREATE CAST (double precision[] AS sparsevec)
CREATE CAST (numeric[] AS sparsevec) CREATE CAST (numeric[] AS sparsevec)
WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT; 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, OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(sparsevec, sparsevec), FUNCTION 1 l1_distance(sparsevec, sparsevec),
FUNCTION 3 hnsw_sparsevec_support(internal); 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

@@ -20,13 +20,15 @@
static const struct config_enum_entry hnsw_iterative_search_options[] = { static const struct config_enum_entry hnsw_iterative_search_options[] = {
{"off", HNSW_ITERATIVE_SEARCH_OFF, false}, {"off", HNSW_ITERATIVE_SEARCH_OFF, false},
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false}, {"strict", HNSW_ITERATIVE_SEARCH_STRICT, false},
{"relaxed", HNSW_ITERATIVE_SEARCH_RELAXED, false},
/* TODO Change to strict before merging */
{"on", HNSW_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false} {NULL, 0, false}
}; };
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_iterative_search_max_tuples; int hnsw_max_iterative_tuples;
int hnsw_iterative_search; int hnsw_iterative_search;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -78,14 +80,16 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomEnumVariable("hnsw.iterative_search", "Sets iterative search", /* TODO Change name */
DefineCustomEnumVariable("hnsw.streaming", "Iterative search mode",
NULL, &hnsw_iterative_search, NULL, &hnsw_iterative_search,
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Change name */
/* TODO Ensure ivfflat.max_probes uses same value for "all" */ /* TODO Ensure ivfflat.max_probes uses same value for "all" */
DefineCustomIntVariable("hnsw.iterative_search_max_tuples", "Sets the max number of candidates to visit for iterative search", DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search",
"-1 means all", &hnsw_iterative_search_max_tuples, "-1 means all", &hnsw_max_iterative_tuples,
-1, -1, INT_MAX, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
} }
@@ -245,13 +249,13 @@ hnswhandler(PG_FUNCTION_ARGS)
IndexAmRoutine *amroutine = makeNode(IndexAmRoutine); IndexAmRoutine *amroutine = makeNode(IndexAmRoutine);
amroutine->amstrategies = 0; amroutine->amstrategies = 0;
amroutine->amsupport = 4; amroutine->amsupport = 3;
amroutine->amoptsprocnum = 0; amroutine->amoptsprocnum = 0;
amroutine->amcanorder = false; amroutine->amcanorder = false;
amroutine->amcanorderbyop = true; amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */ amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false; amroutine->amcanunique = false;
amroutine->amcanmulticol = true; amroutine->amcanmulticol = false;
amroutine->amoptionalkey = true; amroutine->amoptionalkey = true;
amroutine->amsearcharray = false; amroutine->amsearcharray = false;
amroutine->amsearchnulls = false; amroutine->amsearchnulls = false;
@@ -303,17 +307,3 @@ hnswhandler(PG_FUNCTION_ARGS)
PG_RETURN_POINTER(amroutine); 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_DISTANCE_PROC 1
#define HNSW_NORM_PROC 2 #define HNSW_NORM_PROC 2
#define HNSW_TYPE_INFO_PROC 3 #define HNSW_TYPE_INFO_PROC 3
#define HNSW_ATTRIBUTE_DISTANCE_PROC 4
#define HNSW_VERSION 1 #define HNSW_VERSION 1
#define HNSW_MAGIC_NUMBER 0xA953A953 #define HNSW_MAGIC_NUMBER 0xA953A953
@@ -43,6 +42,9 @@
#define HNSW_DEFAULT_EF_SEARCH 40 #define HNSW_DEFAULT_EF_SEARCH 40
#define HNSW_MIN_EF_SEARCH 1 #define HNSW_MIN_EF_SEARCH 1
#define HNSW_MAX_EF_SEARCH 1000 #define HNSW_MAX_EF_SEARCH 1000
#define HNSW_DEFAULT_EF_STREAM -1
#define HNSW_MIN_EF_STREAM -1
#define HNSW_MAX_EF_STREAM INT_MAX
/* Tuple types */ /* Tuple types */
#define HNSW_ELEMENT_TUPLE_TYPE 1 #define HNSW_ELEMENT_TUPLE_TYPE 1
@@ -108,19 +110,17 @@
#define HnswPtrPointer(hp) (hp).ptr #define HnswPtrPointer(hp) (hp).ptr
#define HnswPtrOffset(hp) relptr_offset((hp).relptr) #define HnswPtrOffset(hp) relptr_offset((hp).relptr)
#define HnswUseIndexTuple(index) (IndexRelationGetNumberOfAttributes(index) > 1)
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_max_iterative_tuples;
extern int hnsw_iterative_search; extern int hnsw_iterative_search;
extern int hnsw_iterative_search_max_tuples;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchType typedef enum HnswIterativeSearchType
{ {
HNSW_ITERATIVE_SEARCH_OFF, HNSW_ITERATIVE_SEARCH_OFF,
HNSW_ITERATIVE_SEARCH_RELAXED, HNSW_ITERATIVE_SEARCH_STRICT,
HNSW_ITERATIVE_SEARCH_STRICT HNSW_ITERATIVE_SEARCH_RELAXED
} HnswIterativeSearchType; } HnswIterativeSearchType;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
@@ -136,7 +136,6 @@ HnswPtrDeclare(HnswElementData, HnswElementRelptr, HnswElementPtr);
HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr); HnswPtrDeclare(HnswNeighborArray, HnswNeighborArrayRelptr, HnswNeighborArrayPtr);
HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr); HnswPtrDeclare(HnswNeighborArrayPtr, HnswNeighborsRelptr, HnswNeighborsPtr);
HnswPtrDeclare(char, DatumRelptr, DatumPtr); HnswPtrDeclare(char, DatumRelptr, DatumPtr);
HnswPtrDeclare(IndexTupleData, IndexTupleRelptr, IndexTuplePtr);
struct HnswElementData struct HnswElementData
{ {
@@ -153,7 +152,6 @@ struct HnswElementData
OffsetNumber neighborOffno; OffsetNumber neighborOffno;
BlockNumber neighborPage; BlockNumber neighborPage;
DatumPtr value; DatumPtr value;
IndexTuplePtr itup;
LWLock lock; LWLock lock;
}; };
@@ -179,7 +177,6 @@ typedef struct HnswSearchCandidate
pairingheap_node w_node; pairingheap_node w_node;
HnswElementPtr element; HnswElementPtr element;
double distance; double distance;
bool matches;
} HnswSearchCandidate; } HnswSearchCandidate;
/* HNSW index options */ /* HNSW index options */
@@ -258,16 +255,14 @@ typedef struct HnswTypeInfo
typedef struct HnswSupport typedef struct HnswSupport
{ {
FmgrInfo *procinfo[2]; FmgrInfo *procinfo;
FmgrInfo *normprocinfo; FmgrInfo *normprocinfo;
Oid *collation; Oid collation;
} HnswSupport; } HnswSupport;
typedef struct HnswQuery typedef struct HnswQuery
{ {
Datum value; Datum value;
IndexTuple itup;
ScanKeyData *keyData;
} HnswQuery; } HnswQuery;
typedef struct HnswBuildState typedef struct HnswBuildState
@@ -296,8 +291,6 @@ typedef struct HnswBuildState
HnswGraph *graph; HnswGraph *graph;
double ml; double ml;
int maxLevel; int maxLevel;
bool useIndexTuple;
TupleDesc tupdesc;
/* Memory */ /* Memory */
MemoryContext graphCtx; MemoryContext graphCtx;
@@ -425,32 +418,30 @@ bool HnswCheckNorm(HnswSupport * support, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples); 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); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size); void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, int m, int efConstruction, bool existing, 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, bool inMemory); 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 HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);
HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator); HnswNeighborArray *HnswInitNeighborArray(int lm, HnswAllocator * allocator);
void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc); void HnswInitNeighbors(char *base, HnswElement element, int m, HnswAllocator * alloc);
bool HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc); 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 HnswUpdateNeighborsOnDisk(Relation index, HnswSupport * support, HnswElement e, int m, bool checkExisting, bool building);
void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec, Relation index); void HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHeaptids, bool loadVec);
void HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance); void HnswLoadElement(HnswElement element, double *distance, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, double *maxDistance);
void HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool useIndexTuple); 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); void HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newElement, float distance, int lm, int *updateIdx, Relation index, HnswSupport * support);
bool HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc);
bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc); bool HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation index, int m, int lm, int lc);
void HnswInitLockTranche(void); void HnswInitLockTranche(void);
const HnswTypeInfo *HnswGetTypeInfo(Relation index); const HnswTypeInfo *HnswGetTypeInfo(Relation index);
PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc); PGDLLEXPORT void HnswParallelBuildMain(dsm_segment *seg, shm_toc *toc);
Size HnswGetElementTupleSize(char *base, HnswElement element, bool useIndexTuple);
bool HnswIndexTupleIsEqual(IndexTuple a, IndexTuple b, TupleDesc tupdesc);
/* Index access methods */ /* Index access methods */
IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo); IndexBuildResult *hnswbuild(Relation heap, Relation index, IndexInfo *indexInfo);

View File

@@ -148,7 +148,6 @@ CreateGraphPages(HnswBuildState * buildstate)
Page page; Page page;
HnswElementPtr iter = buildstate->graph->head; HnswElementPtr iter = buildstate->graph->head;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
bool useIndexTuple = buildstate->useIndexTuple;
/* Calculate sizes */ /* Calculate sizes */
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -168,6 +167,7 @@ CreateGraphPages(HnswBuildState * buildstate)
Size etupSize; Size etupSize;
Size ntupSize; Size ntupSize;
Size combinedSize; Size combinedSize;
Pointer valuePtr = HnswPtrAccess(base, element->value);
/* Update iterator */ /* Update iterator */
iter = element->next; iter = element->next;
@@ -176,7 +176,7 @@ CreateGraphPages(HnswBuildState * buildstate)
MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(etup, 0, HNSW_TUPLE_ALLOC_SIZE);
/* Calculate sizes */ /* Calculate sizes */
etupSize = HnswGetElementTupleSize(base, element, useIndexTuple); etupSize = HNSW_ELEMENT_TUPLE_SIZE(VARSIZE_ANY(valuePtr));
ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(element->level, buildstate->m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
@@ -186,7 +186,7 @@ CreateGraphPages(HnswBuildState * buildstate)
(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED), (errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
errmsg("index tuple too large"))); errmsg("index tuple too large")));
HnswSetElementTuple(base, etup, element, useIndexTuple); HnswSetElementTuple(base, etup, element);
/* Keep element and neighbors on the same page if possible */ /* Keep element and neighbors on the same page if possible */
if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize)) if (PageGetFreeSpace(page) < etupSize || (combinedSize <= maxSize && PageGetFreeSpace(page) < combinedSize))
@@ -327,18 +327,19 @@ AddDuplicateInMemory(HnswElement element, HnswElement dup)
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
FindDuplicateInMemory(char *base, HnswElement element, bool useIndexTuple, TupleDesc tupdesc) FindDuplicateInMemory(char *base, HnswElement element)
{ {
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
IndexTuple itup = HnswPtrAccess(base, element->itup); Datum value = HnswGetValue(base, element);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element); HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
/* Exit early since ordered by distance */ /* Exit early since ordered by distance */
if (!HnswIndexTupleIsEqual(itup, HnswPtrAccess(base, neighborElement->itup), tupdesc)) if (!datumIsEqual(value, neighborValue, false, -1))
return false; return false;
/* Check for space */ /* Check for space */
@@ -365,7 +366,7 @@ AddElementInMemory(char *base, HnswGraph * graph, HnswElement element)
* Update neighbors * Update neighbors
*/ */
static void static void
UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswElement e, int m) UpdateNeighborsInMemory(char *base, HnswSupport * support, HnswElement e, int m)
{ {
for (int lc = e->level; lc >= 0; lc--) for (int lc = e->level; lc >= 0; lc--)
{ {
@@ -387,7 +388,7 @@ UpdateNeighborsInMemory(char *base, Relation index, HnswSupport * support, HnswE
Assert(neighborElement); Assert(neighborElement);
LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE); LWLockAcquire(&neighborElement->lock, LW_EXCLUSIVE);
HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, index, support); HnswUpdateConnection(base, HnswGetNeighbors(base, neighborElement, lc), e, hc->distance, lm, NULL, NULL, support);
LWLockRelease(&neighborElement->lock); LWLockRelease(&neighborElement->lock);
} }
} }
@@ -403,14 +404,14 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
/* Look for duplicate */ /* Look for duplicate */
if (FindDuplicateInMemory(base, element, buildstate->useIndexTuple, buildstate->tupdesc)) if (FindDuplicateInMemory(base, element))
return; return;
/* Add element */ /* Add element */
AddElementInMemory(base, graph, element); AddElementInMemory(base, graph, element);
/* Update neighbors */ /* Update neighbors */
UpdateNeighborsInMemory(base, buildstate->index, support, element, m); UpdateNeighborsInMemory(base, support, element, m);
/* Update entry point if needed (already have lock) */ /* Update entry point if needed (already have lock) */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -423,7 +424,6 @@ UpdateGraphInMemory(HnswSupport * support, HnswElement element, int m, int efCon
static void static void
InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element) InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
{ {
Relation index = buildstate->index;
HnswGraph *graph = buildstate->graph; HnswGraph *graph = buildstate->graph;
HnswSupport *support = &buildstate->support; HnswSupport *support = &buildstate->support;
HnswElement entryPoint; HnswElement entryPoint;
@@ -457,7 +457,7 @@ InsertTupleInMemory(HnswBuildState * buildstate, HnswElement element)
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, true); HnswFindElementNeighbors(base, element, entryPoint, NULL, support, m, efConstruction, false);
/* Update graph in memory */ /* Update graph in memory */
UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate); UpdateGraphInMemory(support, element, m, efConstruction, entryPoint, buildstate);
@@ -476,20 +476,18 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
HnswElement element; HnswElement element;
HnswAllocator *allocator = &buildstate->allocator; HnswAllocator *allocator = &buildstate->allocator;
HnswSupport *support = &buildstate->support; HnswSupport *support = &buildstate->support;
Size valueSize;
Pointer valuePtr;
LWLock *flushLock = &graph->flushLock; LWLock *flushLock = &graph->flushLock;
char *base = buildstate->hnswarea; char *base = buildstate->hnswarea;
TupleDesc tupdesc = buildstate->tupdesc; Datum value;
IndexTuple itup;
Size itupSize;
IndexTuple itupShared;
bool unused;
/* Form index tuple */ /* Form index value */
if (!HnswFormIndexTuple(&itup, values, isnull, buildstate->typeInfo, support, tupdesc)) if (!HnswFormIndexValue(&value, values, isnull, buildstate->typeInfo, support))
return false; return false;
/* Get tuple size */ /* Get datum size */
itupSize = IndexTupleSize(itup); valueSize = VARSIZE_ANY(DatumGetPointer(value));
/* Ensure graph not flushed when inserting */ /* Ensure graph not flushed when inserting */
LWLockAcquire(flushLock, LW_SHARED); LWLockAcquire(flushLock, LW_SHARED);
@@ -499,7 +497,7 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
{ {
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc); return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
} }
/* /*
@@ -531,12 +529,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
LWLockRelease(flushLock); LWLockRelease(flushLock);
return HnswInsertTupleOnDisk(index, support, itup, heaptid, true, tupdesc); return HnswInsertTupleOnDisk(index, support, value, heaptid, true);
} }
/* Ok, we can proceed to allocate the element */ /* Ok, we can proceed to allocate the element */
element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator); element = HnswInitElement(base, heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel, allocator);
itupShared = HnswAlloc(allocator, itupSize); valuePtr = HnswAlloc(allocator, valueSize);
/* /*
* We have now allocated the space needed for the element, so we don't * We have now allocated the space needed for the element, so we don't
@@ -545,10 +543,9 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid, Hn
*/ */
LWLockRelease(&graph->allocatorLock); LWLockRelease(&graph->allocatorLock);
/* Copy the tuple */ /* Copy the datum */
memcpy(itupShared, itup, itupSize); memcpy(valuePtr, DatumGetPointer(value), valueSize);
HnswPtrStore(base, element->itup, itupShared); HnswPtrStore(base, element->value, valuePtr);
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itupShared, 1, tupdesc, &unused)));
/* Create a lock for the element */ /* Create a lock for the element */
LWLockInitialize(&element->lock, hnsw_lock_tranche_id); LWLockInitialize(&element->lock, hnsw_lock_tranche_id);
@@ -675,19 +672,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED), (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("type not supported for hnsw index"))); 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 */ /* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0) if (buildstate->dimensions < 0)
ereport(ERROR, ereport(ERROR,
@@ -714,8 +698,6 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->graph = &buildstate->graphData; buildstate->graph = &buildstate->graphData;
buildstate->ml = HnswGetMl(buildstate->m); buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m); buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->useIndexTuple = HnswUseIndexTuple(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext, buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context", "Hnsw build graph context",

View File

@@ -156,10 +156,9 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion; uint8 tupleVersion;
char *base = NULL; char *base = NULL;
bool useIndexTuple = HnswUseIndexTuple(index);
/* Calculate sizes */ /* 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); ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
combinedSize = etupSize + ntupSize + sizeof(ItemIdData); combinedSize = etupSize + ntupSize + sizeof(ItemIdData);
maxSize = HNSW_MAX_SIZE; maxSize = HNSW_MAX_SIZE;
@@ -167,7 +166,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
/* Prepare element tuple */ /* Prepare element tuple */
etup = palloc0(etupSize); etup = palloc0(etupSize);
HnswSetElementTuple(base, etup, e, useIndexTuple); HnswSetElementTuple(base, etup, e);
/* Prepare neighbor tuple */ /* Prepare neighbor tuple */
ntup = palloc0(ntupSize); ntup = palloc0(ntupSize);
@@ -384,9 +383,8 @@ LoadElementsForInsert(HnswNeighborArray * neighbors, HnswQuery * q, int *idx, Re
HnswCandidate *hc = &neighbors->items[i]; HnswCandidate *hc = &neighbors->items[i];
HnswElement element = HnswPtrAccess(base, hc->element); HnswElement element = HnswPtrAccess(base, hc->element);
double distance; double distance;
bool matches;
HnswLoadElement(element, &distance, &matches, q, index, support, true, NULL); HnswLoadElement(element, &distance, q, index, support, true, NULL);
hc->distance = distance; hc->distance = distance;
/* Prune element if being deleted */ /* Prune element if being deleted */
@@ -430,8 +428,6 @@ GetUpdateIndex(HnswElement element, HnswElement newElement, float distance, int
HnswQuery q; HnswQuery q;
q.value = HnswGetValue(base, element); q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
LoadElementsForInsert(neighbors, &q, &idx, index, support); LoadElementsForInsert(neighbors, &q, &idx, index, support);
@@ -637,30 +633,21 @@ AddDuplicateOnDisk(Relation index, HnswElement element, HnswElement dup, bool bu
* Find duplicate element * Find duplicate element
*/ */
static bool static bool
FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDesc tupdesc) FindDuplicateOnDisk(Relation index, HnswElement element, bool building)
{ {
char *base = NULL; char *base = NULL;
HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0); HnswNeighborArray *neighbors = HnswGetNeighbors(base, element, 0);
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
for (int i = 0; i < neighbors->length; i++) for (int i = 0; i < neighbors->length; i++)
{ {
HnswCandidate *neighbor = &neighbors->items[i]; HnswCandidate *neighbor = &neighbors->items[i];
HnswElement neighborElement = HnswPtrAccess(base, neighbor->element); HnswElement neighborElement = HnswPtrAccess(base, neighbor->element);
Datum neighborValue = HnswGetValue(base, neighborElement);
if (HnswUseIndexTuple(index)) /* Exit early since ordered by distance */
{ if (!datumIsEqual(value, neighborValue, false, -1))
/* Exit early since ordered by distance */ return false;
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;
}
if (AddDuplicateOnDisk(index, element, neighborElement, building)) if (AddDuplicateOnDisk(index, element, neighborElement, building))
return true; return true;
@@ -673,12 +660,12 @@ FindDuplicateOnDisk(Relation index, HnswElement element, bool building, TupleDes
* Update graph on disk * Update graph on disk
*/ */
static void static void
UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building, TupleDesc tupdesc) UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
{ {
BlockNumber newInsertPage = InvalidBlockNumber; BlockNumber newInsertPage = InvalidBlockNumber;
/* Look for duplicate */ /* Look for duplicate */
if (FindDuplicateOnDisk(index, element, building, tupdesc)) if (FindDuplicateOnDisk(index, element, building))
return; return;
/* Add element */ /* Add element */
@@ -700,7 +687,7 @@ UpdateGraphOnDisk(Relation index, HnswSupport * support, HnswElement element, in
* Insert a tuple into the index * Insert a tuple into the index
*/ */
bool bool
HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, ItemPointer heaptid, bool building, TupleDesc tupdesc) HnswInsertTupleOnDisk(Relation index, HnswSupport * support, Datum value, ItemPointer heaptid, bool building)
{ {
HnswElement entryPoint; HnswElement entryPoint;
HnswElement element; HnswElement element;
@@ -708,7 +695,6 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
int efConstruction = HnswGetEfConstruction(index); int efConstruction = HnswGetEfConstruction(index);
LOCKMODE lockmode = ShareLock; LOCKMODE lockmode = ShareLock;
char *base = NULL; char *base = NULL;
bool unused;
/* /*
* Get a shared lock. This allows vacuum to ensure no in-flight inserts * Get a shared lock. This allows vacuum to ensure no in-flight inserts
@@ -722,8 +708,7 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
/* Create an element */ /* Create an element */
element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL); element = HnswInitElement(base, heaptid, m, HnswGetMl(m), HnswGetMaxLevel(m), NULL);
HnswPtrStore(base, element->itup, itup); HnswPtrStore(base, element->value, DatumGetPointer(value));
HnswPtrStore(base, element->value, DatumGetPointer(index_getattr(itup, 1, tupdesc, &unused)));
/* Prevent concurrent inserts when likely updating entry point */ /* Prevent concurrent inserts when likely updating entry point */
if (entryPoint == NULL || element->level > entryPoint->level) if (entryPoint == NULL || element->level > entryPoint->level)
@@ -740,10 +725,10 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
} }
/* Find neighbors for element */ /* Find neighbors for element */
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, false);
/* Update graph on disk */ /* Update graph on disk */
UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building, tupdesc); UpdateGraphOnDisk(index, support, element, m, efConstruction, entryPoint, building);
/* Release lock */ /* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode); UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);
@@ -757,18 +742,17 @@ HnswInsertTupleOnDisk(Relation index, HnswSupport * support, IndexTuple itup, It
static void static void
HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid) HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heaptid)
{ {
IndexTuple itup; Datum value;
const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index); const HnswTypeInfo *typeInfo = HnswGetTypeInfo(index);
TupleDesc tupdesc = RelationGetDescr(index);
HnswSupport support; HnswSupport support;
HnswInitSupport(&support, index); HnswInitSupport(&support, index);
/* Form index tuple */ /* Form index value */
if (!HnswFormIndexTuple(&itup, values, isnull, typeInfo, &support, tupdesc)) if (!HnswFormIndexValue(&value, values, isnull, typeInfo, &support))
return; return;
HnswInsertTupleOnDisk(index, &support, itup, heaptid, false, tupdesc); HnswInsertTupleOnDisk(index, &support, value, heaptid, false);
} }
/* /*

View File

@@ -22,30 +22,26 @@ GetScanItems(IndexScanDesc scan, Datum value)
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
bool inMemory = false;
HnswQuery *q = &so->q; HnswQuery *q = &so->q;
q->value = value;
q->itup = NULL;
q->keyData = scan->keyData;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
q->value = value;
so->m = m; so->m = m;
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false, inMemory)); ep = list_make1(HnswEntryCandidate(base, entryPoint, q, index, support, false));
for (int lc = entryPoint->level; lc >= 1; lc--) for (int lc = entryPoint->level; lc >= 1; lc--)
{ {
w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, inMemory, NULL, NULL, true, NULL); w = HnswSearchLayer(base, q, ep, 1, lc, index, support, m, false, NULL, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, inMemory, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples); return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, support, m, false, NULL, &so->v, hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF ? &so->discarded : NULL, true, &so->tuples);
} }
/* /*
@@ -76,7 +72,7 @@ ResumeScanItems(IndexScanDesc scan)
ep = lappend(ep, sc); ep = lappend(ep, sc);
} }
return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, false, &so->v, &so->discarded, false, &so->tuples); return HnswSearchLayer(base, &so->q, ep, batch_size, 0, index, &so->support, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
} }
/* /*
@@ -100,7 +96,7 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */ /* Normalize if needed */
if (so->support.normprocinfo != NULL) if (so->support.normprocinfo != NULL)
value = HnswNormValue(so->typeInfo, so->support.collation[0], value); value = HnswNormValue(so->typeInfo, so->support.collation, value);
} }
return value; return value;
@@ -229,7 +225,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
break; break;
/* Reached max number of additional tuples */ /* Reached max number of additional tuples */
if (hnsw_iterative_search_max_tuples != -1 && so->tuples >= hnsw_iterative_search_max_tuples) if (hnsw_max_iterative_tuples != -1 && so->tuples >= hnsw_ef_search + hnsw_max_iterative_tuples)
{ {
if (pairingheap_is_empty(so->discarded)) if (pairingheap_is_empty(so->discarded))
break; break;
@@ -282,7 +278,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
element = HnswPtrAccess(base, sc->element); element = HnswPtrAccess(base, sc->element);
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (!sc->matches || element->heaptidsLength == 0) if (element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);

View File

@@ -146,39 +146,11 @@ HnswOptionalProcInfo(Relation index, uint16 procnum)
void void
HnswInitSupport(HnswSupport * support, Relation index) HnswInitSupport(HnswSupport * support, Relation index)
{ {
support->procinfo[0] = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC); support->procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
support->collation = index->rd_indcollation[0];
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
support->procinfo[1] = index_getprocinfo(index, 2, HNSW_ATTRIBUTE_DISTANCE_PROC);
support->collation = index->rd_indcollation;
support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); support->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
} }
/*
* 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);
}
/* /*
* Normalize value * Normalize value
*/ */
@@ -194,38 +166,7 @@ HnswNormValue(const HnswTypeInfo * typeInfo, Oid collation, Datum value)
bool bool
HnswCheckNorm(HnswSupport * support, Datum value) HnswCheckNorm(HnswSupport * support, Datum value)
{ {
return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->collation[0], value)) > 0; return DatumGetFloat8(FunctionCall1Coll(support->normprocinfo, support->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;
} }
/* /*
@@ -316,7 +257,6 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
HnswInitNeighbors(base, element, m, allocator); HnswInitNeighbors(base, element, m, allocator);
HnswPtrStore(base, element->value, (Pointer) NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element; return element;
} }
@@ -343,7 +283,6 @@ HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno)
element->offno = offno; element->offno = offno;
HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL); HnswPtrStore(base, element->neighbors, (HnswNeighborArrayPtr *) NULL);
HnswPtrStore(base, element->value, (Pointer) NULL); HnswPtrStore(base, element->value, (Pointer) NULL);
HnswPtrStore(base, element->itup, (IndexTuple) NULL);
return element; return element;
} }
@@ -456,13 +395,11 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
} }
/* /*
* Form index tuple * Form index value
*/ */
bool bool
HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support, TupleDesc tupdesc) HnswFormIndexValue(Datum *out, Datum *values, bool *isnull, const HnswTypeInfo * typeInfo, HnswSupport * support)
{ {
Datum newValues[2];
/* Detoast once for all calls */ /* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -476,14 +413,10 @@ HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeI
if (!HnswCheckNorm(support, value)) if (!HnswCheckNorm(support, value))
return false; return false;
value = HnswNormValue(typeInfo, support->collation[0], value); value = HnswNormValue(typeInfo, support->collation, value);
} }
newValues[0] = value; *out = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
*out = index_form_tuple(tupdesc, newValues, isnull);
return true; return true;
} }
@@ -492,8 +425,10 @@ HnswFormIndexTuple(IndexTuple *out, Datum *values, bool *isnull, const HnswTypeI
* Set element tuple, except for neighbor info * Set element tuple, except for neighbor info
*/ */
void 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->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
etup->deleted = 0; etup->deleted = 0;
@@ -505,19 +440,7 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element, bool
else else
ItemPointerSetInvalid(&etup->heaptids[i]); ItemPointerSetInvalid(&etup->heaptids[i]);
} }
memcpy(&etup->data, valuePtr, VARSIZE_ANY(valuePtr));
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));
}
} }
/* /*
@@ -559,7 +482,7 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
* Load an element from a tuple * Load an element from a tuple
*/ */
void 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->level = etup->level;
element->deleted = etup->deleted; element->deleted = etup->deleted;
@@ -583,128 +506,26 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
if (loadVec) if (loadVec)
{ {
char *base = NULL; char *base = NULL;
Datum value = datumCopy(PointerGetDatum(&etup->data), false, -1);
if (HnswUseIndexTuple(index)) HnswPtrStore(base, element->value, DatumGetPointer(value));
{
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));
}
} }
} }
/*
* 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 * Calculate the distance between values
*/ */
static double static inline double
HnswGetDistance(IndexTuple itup, Datum vec, HnswQuery * q, Relation index, HnswSupport * support, bool *matches) HnswGetDistance(Datum a, Datum b, HnswSupport * support)
{ {
double g; return DatumGetFloat8(FunctionCall2Coll(support->procinfo, support->collation, a, b));
if (DatumGetPointer(q->value) == NULL)
g = 0;
else
g = DatumGetFloat8(FunctionCall2Coll(support->procinfo[0], support->collation[0], q->value, vec));
Assert(PointerIsValid(matches));
*matches = true;
if (IndexRelationGetNumberOfKeyAttributes(index) > 1)
{
double w = 0.25;
double e = 0.0;
TupleDesc tupdesc = RelationGetDescr(index);
if (q->keyData)
{
/* TODO need to pass length of key data */
int keyCount = 1;
for (int i = 0; i < keyCount; i++)
{
ScanKey key = &q->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(support->procinfo[key->sk_attno - 1], support->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 (q->itup)
{
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(q->itup, i + 2, tupdesc, &attnull);
if (isnull || attnull)
{
if (isnull != attnull)
e += 1000;
}
else
e += fabs(DatumGetFloat8(FunctionCall2Coll(support->procinfo[i + 1], support->collation[i + 1], value, value2)));
}
return w * g + AttributeDistance(e);
}
}
return g;
} }
/* /*
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
static void static void
HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, 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; Buffer buf;
Page page; Page page;
@@ -722,23 +543,10 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
/* Calculate distance */ /* Calculate distance */
if (distance != NULL) if (distance != NULL)
{ {
IndexTuple itup = NULL; if (DatumGetPointer(q->value) == NULL)
Datum value; *distance = 0;
if (HnswUseIndexTuple(index))
{
TupleDesc tupdesc = RelationGetDescr(index);
bool unused;
itup = (IndexTuple) &etup->data;
value = index_getattr(itup, 1, tupdesc, &unused);
}
else else
{ *distance = HnswGetDistance(q->value, PointerGetDatum(&etup->data), support);
value = PointerGetDatum(&etup->data);
}
*distance = HnswGetDistance(itup, value, q, index, support, matches);
} }
/* Load element */ /* Load element */
@@ -747,7 +555,7 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
if (*element == NULL) if (*element == NULL)
*element = HnswInitElementFromBlock(blkno, offno); *element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(*element, etup, true, loadVec, index); HnswLoadElementFromTuple(*element, etup, true, loadVec);
} }
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
@@ -757,36 +565,36 @@ HnswLoadElementImpl(BlockNumber blkno, OffsetNumber offno, double *distance, boo
* Load an element and optionally get its distance from q * Load an element and optionally get its distance from q
*/ */
void void
HnswLoadElement(HnswElement element, double *distance, bool *matches, HnswQuery * q, Relation index, HnswSupport * support, 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, index, support, loadVec, maxDistance, &element); HnswLoadElementImpl(element->blkno, element->offno, distance, q, index, support, loadVec, maxDistance, &element);
} }
/* /*
* Get the distance for an element * Get the distance for an element
*/ */
static double static double
GetElementDistance(char *base, HnswElement element, bool *matches, HnswQuery * q, Relation index, HnswSupport * support) GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport * support)
{ {
Datum value = HnswGetValue(base, element); Datum value = HnswGetValue(base, element);
IndexTuple itup = HnswPtrAccess(base, element->itup);
return HnswGetDistance(itup, value, q, index, support, matches); return HnswGetDistance(q->value, value, support);
} }
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec, bool inMemory) HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{ {
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate)); HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL;
HnswPtrStore(base, sc->element, entryPoint); HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, &sc->matches, q, index, support); sc->distance = GetElementDistance(base, entryPoint, q, support);
else else
HnswLoadElement(entryPoint, &sc->distance, &sc->matches, q, index, support, loadVec, NULL); HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL);
return sc; return sc;
} }
@@ -990,7 +798,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * List *
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, bool inMemory, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples) 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; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
@@ -1003,8 +811,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
uint64 additional = 0; bool inMemory = index == NULL;
uint64 maxAdditional = q->keyData && lc == 0 ? 10000 : 0;
if (v == NULL) if (v == NULL)
{ {
@@ -1044,10 +851,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
pairingheap_add(C, &sc->c_node); pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_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 * Do not count elements being deleted towards ef when vacuuming. It
* would be ideal to do this for inserts as well, but this could * would be ideal to do this for inserts as well, but this could
@@ -1081,7 +884,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswElement eElement; HnswElement eElement;
HnswSearchCandidate *e; HnswSearchCandidate *e;
double eDistance; double eDistance;
bool eMatches;
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); f = HnswGetSearchCandidate(w_node, pairingheap_first(W));
@@ -1089,7 +891,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
if (inMemory) if (inMemory)
{ {
eElement = unvisited[i].element; eElement = unvisited[i].element;
eDistance = GetElementDistance(base, eElement, &eMatches, q, index, support); eDistance = GetElementDistance(base, eElement, q, support);
} }
else else
{ {
@@ -1099,7 +901,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
/* Avoid any allocations if not adding */ /* Avoid any allocations if not adding */
eElement = NULL; eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &eMatches, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement); HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
if (eElement == NULL) if (eElement == NULL)
continue; continue;
@@ -1127,7 +929,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
e = palloc(sizeof(HnswSearchCandidate)); e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement); HnswPtrStore(base, e->element, eElement);
e->distance = eDistance; e->distance = eDistance;
e->matches = eMatches;
pairingheap_add(C, &e->c_node); pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); pairingheap_add(W, &e->w_node);
@@ -1138,10 +939,6 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
*/ */
if (CountElement(skipElement, eElement)) if (CountElement(skipElement, eElement))
{ {
/* Do not count elements that do not match filter towards ef */
if (!e->matches && ++additional <= maxAdditional)
continue;
wlen++; wlen++;
/* No need to decrement wlen */ /* No need to decrement wlen */
@@ -1219,24 +1016,18 @@ CompareCandidateDistancesOffset(const ListCell *a, const ListCell *b)
* Check if an element is closer to q than any element from R * Check if an element is closer to q than any element from R
*/ */
static bool static bool
CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswSupport * support) CheckElementCloser(char *base, HnswCandidate * e, List *r, HnswSupport * support)
{ {
HnswElement eElement = HnswPtrAccess(base, e->element); HnswElement eElement = HnswPtrAccess(base, e->element);
HnswQuery q; Datum eValue = HnswGetValue(base, eElement);
ListCell *lc2; ListCell *lc2;
q.value = HnswGetValue(base, eElement);
q.itup = HnswPtrAccess(base, eElement->itup);
q.keyData = NULL;
foreach(lc2, r) foreach(lc2, r)
{ {
HnswCandidate *ri = lfirst(lc2); HnswCandidate *ri = lfirst(lc2);
HnswElement riElement = HnswPtrAccess(base, ri->element); HnswElement riElement = HnswPtrAccess(base, ri->element);
Datum riValue = HnswGetValue(base, riElement); Datum riValue = HnswGetValue(base, riElement);
IndexTuple ritup = HnswPtrAccess(base, riElement->itup); float distance = HnswGetDistance(eValue, riValue, support);
bool matches;
float distance = HnswGetDistance(ritup, riValue, &q, index, support, &matches);
if (distance <= e->distance) if (distance <= e->distance)
return false; return false;
@@ -1249,7 +1040,7 @@ CheckElementCloser(char *base, HnswCandidate * e, List *r, Relation index, HnswS
* Algorithm 4 from paper * Algorithm 4 from paper
*/ */
static List * static List *
SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * support, 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 *r = NIL;
List *w = list_copy(c); List *w = list_copy(c);
@@ -1283,7 +1074,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
/* Use previous state of r and wd to skip work when possible */ /* Use previous state of r and wd to skip work when possible */
if (mustCalculate) if (mustCalculate)
e->closer = CheckElementCloser(base, e, r, index, support); e->closer = CheckElementCloser(base, e, r, support);
else if (list_length(added) > 0) else if (list_length(added) > 0)
{ {
/* Keep Valgrind happy for in-memory, parallel builds */ /* Keep Valgrind happy for in-memory, parallel builds */
@@ -1296,7 +1087,8 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
*/ */
if (e->closer) if (e->closer)
{ {
e->closer = CheckElementCloser(base, e, added, index, support); e->closer = CheckElementCloser(base, e, added, support);
if (!e->closer) if (!e->closer)
removedAny = true; removedAny = true;
} }
@@ -1308,7 +1100,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
*/ */
if (removedAny) if (removedAny)
{ {
e->closer = CheckElementCloser(base, e, r, index, support); e->closer = CheckElementCloser(base, e, r, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1316,7 +1108,7 @@ SelectNeighbors(char *base, List *c, int lm, Relation index, HnswSupport * suppo
} }
else if (e == newCandidate) else if (e == newCandidate)
{ {
e->closer = CheckElementCloser(base, e, r, index, support); e->closer = CheckElementCloser(base, e, r, support);
if (e->closer) if (e->closer)
added = lappend(added, e); added = lappend(added, e);
} }
@@ -1393,7 +1185,7 @@ HnswUpdateConnection(char *base, HnswNeighborArray * neighbors, HnswElement newE
c = lappend(c, &neighbors->items[i]); c = lappend(c, &neighbors->items[i]);
c = lappend(c, &newHc); c = lappend(c, &newHc);
SelectNeighbors(base, c, lm, index, support, &neighbors->closerSet, &newHc, &pruned, true); SelectNeighbors(base, c, lm, support, &neighbors->closerSet, &newHc, &pruned, true);
/* Should not happen */ /* Should not happen */
if (pruned == NULL) if (pruned == NULL)
@@ -1464,19 +1256,17 @@ PrecomputeHash(char *base, HnswElement element)
* Algorithm 1 from paper * Algorithm 1 from paper
*/ */
void void
HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, HnswSupport * support, 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 *ep;
List *w; List *w;
int level = element->level; int level = element->level;
int entryLevel; int entryLevel;
HnswQuery q; HnswQuery q;
HnswElement skipElement = existing ? element : NULL; HnswElement skipElement = existing ? element : NULL;
bool inMemory = index == NULL;
q.value = HnswGetValue(base, element); q.value = HnswGetValue(base, element);
q.itup = HnswPtrAccess(base, element->itup);
q.keyData = NULL;
/* Precompute hash */ /* Precompute hash */
if (inMemory) if (inMemory)
@@ -1487,13 +1277,13 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
return; return;
/* Get entry point and level */ /* Get entry point and level */
ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true, inMemory)); ep = list_make1(HnswEntryCandidate(base, entryPoint, &q, index, support, true));
entryLevel = entryPoint->level; entryLevel = entryPoint->level;
/* 1st phase: greedy search to insert level */ /* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--) for (int lc = entryLevel; lc >= level + 1; lc--)
{ {
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL); w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
ep = w; ep = w;
} }
@@ -1512,7 +1302,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL; List *lw = NIL;
ListCell *lc2; ListCell *lc2;
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, inMemory, NULL, NULL, true, NULL); w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
/* Convert search candidates to candidates */ /* Convert search candidates to candidates */
foreach(lc2, w) foreach(lc2, w)
@@ -1536,7 +1326,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
* sortCandidates to true for in-memory builds to enable closer * sortCandidates to true for in-memory builds to enable closer
* caching, but there does not seem to be a difference in performance. * caching, but there does not seem to be a difference in performance.
*/ */
neighbors = SelectNeighbors(base, lw, lm, index, support, &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); AddConnections(base, element, neighbors, lc);

View File

@@ -204,7 +204,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
element->heaptidsLength = 0; element->heaptidsLength = 0;
/* Find neighbors for element, skipping itself */ /* Find neighbors for element, skipping itself */
HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true, false); HnswFindElementNeighbors(base, element, entryPoint, index, support, m, efConstruction, true);
/* Zero memory for each element */ /* Zero memory for each element */
MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE); MemSet(ntup, 0, HNSW_TUPLE_ALLOC_SIZE);
@@ -256,7 +256,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
LockPage(index, HNSW_UPDATE_LOCK, ShareLock); LockPage(index, HNSW_UPDATE_LOCK, ShareLock);
/* Load element */ /* Load element */
HnswLoadElement(highestPoint, NULL, NULL, NULL, index, support, true, NULL); HnswLoadElement(highestPoint, NULL, NULL, index, support, true, NULL);
/* Repair if needed */ /* Repair if needed */
if (NeedsUpdated(vacuumstate, highestPoint)) if (NeedsUpdated(vacuumstate, highestPoint))
@@ -294,7 +294,7 @@ RepairGraphEntryPoint(HnswVacuumState * vacuumstate)
* is outdated, this can remove connections at higher levels in * is outdated, this can remove connections at higher levels in
* the graph until they are repaired, but this should be fine. * the graph until they are repaired, but this should be fine.
*/ */
HnswLoadElement(entryPoint, NULL, NULL, NULL, index, support, true, NULL); HnswLoadElement(entryPoint, NULL, NULL, index, support, true, NULL);
if (NeedsUpdated(vacuumstate, entryPoint)) if (NeedsUpdated(vacuumstate, entryPoint))
{ {
@@ -370,7 +370,7 @@ RepairGraph(HnswVacuumState * vacuumstate)
/* Create an element */ /* Create an element */
element = HnswInitElementFromBlock(blkno, offno); element = HnswInitElementFromBlock(blkno, offno);
HnswLoadElementFromTuple(element, etup, false, true, index); HnswLoadElementFromTuple(element, etup, false, true);
elements = lappend(elements, element); elements = lappend(elements, element);
} }
@@ -440,7 +440,6 @@ MarkDeleted(HnswVacuumState * vacuumstate)
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
Relation index = vacuumstate->index; Relation index = vacuumstate->index;
BufferAccessStrategy bas = vacuumstate->bas; BufferAccessStrategy bas = vacuumstate->bas;
bool useIndexTuple = HnswUseIndexTuple(index);
/* /*
* Wait for index scans to complete. Scans before this point may contain * Wait for index scans to complete. Scans before this point may contain
@@ -522,14 +521,7 @@ MarkDeleted(HnswVacuumState * vacuumstate)
/* Overwrite element */ /* Overwrite element */
etup->deleted = 1; etup->deleted = 1;
if (useIndexTuple) MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
{
IndexTuple itup = (IndexTuple) &etup->data;
MemSet(itup, 0, IndexTupleSize(itup));
}
else
MemSet(&etup->data, 0, VARSIZE_ANY(&etup->data));
/* Overwrite neighbors */ /* Overwrite neighbors */
for (int i = 0; i < ntup->count; i++) for (int i = 0; i < ntup->count; i++)

View File

@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
* Add tuple to sort * Add tuple to sort
*/ */
static void static void
AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, IvfflatBuildState * buildstate) AddTupleToSort(Relation index, ItemPointer tid, Datum *values, IvfflatBuildState * buildstate)
{ {
double distance; double distance;
double minDistance = DBL_MAX; double minDistance = DBL_MAX;
@@ -184,11 +184,6 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, Ivf
slot->tts_isnull[1] = false; slot->tts_isnull[1] = false;
slot->tts_values[2] = value; slot->tts_values[2] = value;
slot->tts_isnull[2] = false; slot->tts_isnull[2] = false;
for (int i = 1; i < buildstate->tupdesc->natts; i++)
{
slot->tts_values[2 + i] = values[i];
slot->tts_isnull[2 + i] = isnull[i];
}
ExecStoreVirtualTuple(slot); ExecStoreVirtualTuple(slot);
/* /*
@@ -220,7 +215,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx); oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add tuple to sort */ /* Add tuple to sort */
AddTupleToSort(index, tid, values, isnull, buildstate); AddTupleToSort(index, tid, values, buildstate);
/* Reset memory context */ /* Reset memory context */
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
@@ -231,20 +226,19 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
* Get index tuple from sort state * Get index tuple from sort state
*/ */
static inline void static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, Datum *values, bool *isnull, IndexTuple *itup, int *list) GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, IndexTuple *itup, int *list)
{ {
Datum value;
bool isnull;
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL)) if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{ {
bool unused; *list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
value = slot_getattr(slot, 3, &isnull);
*list = DatumGetInt32(slot_getattr(slot, 1, &unused));
for (int i = 0; i < tupdesc->natts; i++)
values[i] = slot_getattr(slot, 3 + i, &isnull[i]);
/* Form the index tuple */ /* Form the index tuple */
*itup = index_form_tuple(tupdesc, values, isnull); *itup = index_form_tuple(tupdesc, &value, &isnull);
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused))); (*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
} }
else else
*list = -1; *list = -1;
@@ -260,16 +254,14 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
IndexTuple itup = NULL; /* silence compiler warning */ IndexTuple itup = NULL; /* silence compiler warning */
int64 inserted = 0; int64 inserted = 0;
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = buildstate->tupdesc; TupleDesc tupdesc = RelationGetDescr(index);
Datum *values = palloc(tupdesc->natts * sizeof(Datum));
bool *isnull = palloc(tupdesc->natts * sizeof(bool));
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_TOTAL, buildstate->indtuples);
GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
for (int i = 0; i < buildstate->centers->length; i++) for (int i = 0; i < buildstate->centers->length; i++)
{ {
@@ -305,7 +297,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted); pgstat_progress_update_param(PROGRESS_CREATEIDX_TUPLES_DONE, ++inserted);
GetNextTuple(buildstate->sortstate, tupdesc, slot, values, isnull, &itup, &list); GetNextTuple(buildstate->sortstate, tupdesc, slot, &itup, &list);
} }
insertPage = BufferGetBlockNumber(buf); insertPage = BufferGetBlockNumber(buf);
@@ -315,9 +307,6 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
/* Set the start and insert pages */ /* Set the start and insert pages */
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum); IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
} }
pfree(values);
pfree(isnull);
} }
/* /*
@@ -330,7 +319,6 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->index = index; buildstate->index = index;
buildstate->indexInfo = indexInfo; buildstate->indexInfo = indexInfo;
buildstate->typeInfo = IvfflatGetTypeInfo(index); buildstate->typeInfo = IvfflatGetTypeInfo(index);
buildstate->tupdesc = RelationGetDescr(index);
buildstate->lists = IvfflatGetLists(index); buildstate->lists = IvfflatGetLists(index);
buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod; buildstate->dimensions = TupleDescAttr(index->rd_att, 0)->atttypmod;
@@ -368,13 +356,12 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
errmsg("dimensions must be greater than one for this opclass"))); errmsg("dimensions must be greater than one for this opclass")));
/* Create tuple description for sorting */ /* Create tuple description for sorting */
buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts); buildstate->tupdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
for (int i = 0; i < buildstate->tupdesc->natts; i++) TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions)); buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions));
buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists); buildstate->listInfo = palloc(sizeof(ListInfo) * buildstate->lists);
@@ -646,7 +633,7 @@ IvfflatParallelScanAndSort(IvfflatSpool * ivfspool, IvfflatShared * ivfshared, S
InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo); InitBuildState(&buildstate, ivfspool->heap, ivfspool->index, indexInfo);
memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen); memcpy(buildstate.centers->items, ivfcenters, buildstate.centers->itemsize * buildstate.centers->maxlen);
buildstate.centers->length = buildstate.centers->maxlen; buildstate.centers->length = buildstate.centers->maxlen;
ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, sortmem, coordinate); ivfspool->sortstate = InitBuildSortState(buildstate.tupdesc, sortmem, coordinate);
buildstate.sortstate = ivfspool->sortstate; buildstate.sortstate = ivfspool->sortstate;
scan = table_beginscan_parallel(ivfspool->heap, scan = table_beginscan_parallel(ivfspool->heap,
ParallelTableScanFromIvfflatShared(ivfshared)); ParallelTableScanFromIvfflatShared(ivfshared));
@@ -963,7 +950,7 @@ AssignTuples(IvfflatBuildState * buildstate)
} }
/* Begin serial/leader tuplesort */ /* Begin serial/leader tuplesort */
buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate); buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate);
/* Add tuples to sort */ /* Add tuples to sort */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)

View File

@@ -17,16 +17,8 @@
#endif #endif
int ivfflat_probes; int ivfflat_probes;
int ivfflat_iterative_search;
int ivfflat_iterative_search_max_probes;
static relopt_kind ivfflat_relopt_kind; static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_search_options[] = {
{"off", IVFFLAT_ITERATIVE_SEARCH_OFF, false},
{"on", IVFFLAT_ITERATIVE_SEARCH_RELAXED, false},
{NULL, 0, false}
};
/* /*
* Initialize index options and variables * Initialize index options and variables
*/ */
@@ -41,14 +33,6 @@ IvfflatInit(void)
"Valid range is 1..lists.", &ivfflat_probes, "Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomEnumVariable("ivfflat.iterative_search", "Sets whether to use iterative search",
NULL, &ivfflat_iterative_search,
IVFFLAT_ITERATIVE_SEARCH_OFF, ivfflat_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
DefineCustomIntVariable("ivfflat.iterative_search_max_probes", "Sets the max number of probes for iterative search",
"Zero sets to the number of lists", &ivfflat_iterative_search_max_probes,
0, 0, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat"); MarkGUCPrefixReserved("ivfflat");
} }

View File

@@ -80,14 +80,6 @@
/* Variables */ /* Variables */
extern int ivfflat_probes; extern int ivfflat_probes;
extern int ivfflat_iterative_search;
extern int ivfflat_iterative_search_max_probes;
typedef enum IvfflatIterativeSearchType
{
IVFFLAT_ITERATIVE_SEARCH_OFF,
IVFFLAT_ITERATIVE_SEARCH_RELAXED
} IvfflatIterativeSearchType;
typedef struct VectorArrayData typedef struct VectorArrayData
{ {
@@ -173,7 +165,6 @@ typedef struct IvfflatBuildState
Relation index; Relation index;
IndexInfo *indexInfo; IndexInfo *indexInfo;
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
TupleDesc tupdesc;
/* Settings */ /* Settings */
int dimensions; int dimensions;
@@ -207,7 +198,7 @@ typedef struct IvfflatBuildState
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc sortdesc; TupleDesc tupdesc;
TupleTableSlot *slot; TupleTableSlot *slot;
/* Memory */ /* Memory */
@@ -256,10 +247,8 @@ typedef struct IvfflatScanOpaqueData
{ {
const IvfflatTypeInfo *typeInfo; const IvfflatTypeInfo *typeInfo;
int probes; int probes;
int maxProbes;
int dimensions; int dimensions;
bool first; bool first;
Datum value;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
@@ -276,8 +265,6 @@ typedef struct IvfflatScanOpaqueData
/* Lists */ /* Lists */
pairingheap *listQueue; pairingheap *listQueue;
BlockNumber *listPages;
int listIndex;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;

View File

@@ -78,8 +78,6 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
ListInfo listInfo; ListInfo listInfo;
BlockNumber originalInsertPage; BlockNumber originalInsertPage;
TupleDesc tupdesc = RelationGetDescr(index);
Datum *newValues = palloc(tupdesc->natts * sizeof(Datum));
/* Detoast once for all calls */ /* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0])); value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -100,16 +98,12 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
IvfflatGetMetaPageInfo(index, NULL, NULL); IvfflatGetMetaPageInfo(index, NULL, NULL);
/* Find the insert page - sets the page and list info */ /* Find the insert page - sets the page and list info */
FindInsertPage(index, &value, &insertPage, &listInfo); FindInsertPage(index, values, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;
newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
/* Form tuple */ /* Form tuple */
itup = index_form_tuple(tupdesc, newValues, isnull); itup = index_form_tuple(RelationGetDescr(index), &value, isnull);
itup->t_tid = *heap_tid; itup->t_tid = *heap_tid;
/* Get tuple size */ /* Get tuple size */

View File

@@ -65,7 +65,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Use procinfo from the index instead of scan key for performance */ /* Use procinfo from the index instead of scan key for performance */
distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value)); distance = DatumGetFloat8(so->distfunc(so->procinfo, so->collation, PointerGetDatum(&list->center), value));
if (listCount < so->maxProbes) if (listCount < so->probes)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
@@ -78,7 +78,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node); pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Calculate max distance */ /* Calculate max distance */
if (listCount == so->maxProbes) if (listCount == so->probes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
@@ -102,11 +102,6 @@ GetScanLists(IndexScanDesc scan, Datum value)
UnlockReleaseBuffer(cbuf); 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));
} }
/* /*
@@ -119,14 +114,11 @@ GetScanItems(IndexScanDesc scan, Datum value)
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; double tuples = 0;
TupleTableSlot *slot = so->vslot; TupleTableSlot *slot = so->vslot;
int batchProbes = 0;
tuplesort_reset(so->sortstate);
/* Search closest probes lists */ /* Search closest probes lists */
while (so->listIndex < so->maxProbes && (++batchProbes) <= so->probes) while (!pairingheap_is_empty(so->listQueue))
{ {
BlockNumber searchPage = so->listPages[so->listIndex++]; BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
/* Search all entry pages for list */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))
@@ -174,17 +166,13 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
if (tuples < 100 && ivfflat_iterative_search == IVFFLAT_ITERATIVE_SEARCH_OFF) if (tuples < 100)
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("index scan found few tuples"), (errmsg("index scan found few tuples"),
errdetail("Index may have been created with little data."), errdetail("Index may have been created with little data."),
errhint("Recreate the index and possibly decrease lists."))); errhint("Recreate the index and possibly decrease lists.")));
tuplesort_performsort(so->sortstate); tuplesort_performsort(so->sortstate);
#if defined(IVFFLAT_MEMORY)
elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#endif
} }
/* /*
@@ -252,7 +240,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
int lists; int lists;
int dimensions; int dimensions;
int probes = ivfflat_probes; int probes = ivfflat_probes;
int maxProbes;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
@@ -262,21 +249,10 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
if (probes > lists) if (probes > lists)
probes = lists; probes = lists;
if (ivfflat_iterative_search != IVFFLAT_ITERATIVE_SEARCH_OFF) so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + probes * sizeof(IvfflatScanList));
{
if (ivfflat_iterative_search_max_probes == 0)
maxProbes = lists;
else
maxProbes = Min(ivfflat_iterative_search_max_probes, lists);
}
else
maxProbes = probes;
so = (IvfflatScanOpaque) palloc(offsetof(IvfflatScanOpaqueData, lists) + maxProbes * sizeof(IvfflatScanList));
so->typeInfo = IvfflatGetTypeInfo(index); so->typeInfo = IvfflatGetTypeInfo(index);
so->first = true; so->first = true;
so->probes = probes; so->probes = probes;
so->maxProbes = maxProbes;
so->dimensions = dimensions; so->dimensions = dimensions;
/* Set support functions */ /* Set support functions */
@@ -304,8 +280,6 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->bas = GetAccessStrategy(BAS_BULKREAD); so->bas = GetAccessStrategy(BAS_BULKREAD);
so->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
so->listPages = palloc(maxProbes * sizeof(BlockNumber));
so->listIndex = 0;
scan->opaque = so; scan->opaque = so;
@@ -320,9 +294,11 @@ ivfflatrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
if (!so->first)
tuplesort_reset(so->sortstate);
so->first = true; so->first = true;
pairingheap_reset(so->listQueue); pairingheap_reset(so->listQueue);
so->listIndex = 0;
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData)); memmove(scan->keyData, keys, scan->numberOfKeys * sizeof(ScanKeyData));
@@ -338,8 +314,6 @@ bool
ivfflatgettuple(IndexScanDesc scan, ScanDirection dir) ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
ItemPointer heaptid;
bool isnull;
/* /*
* Index can be used to scan backward, but Postgres doesn't support * Index can be used to scan backward, but Postgres doesn't support
@@ -367,25 +341,28 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
IvfflatBench("GetScanLists", GetScanLists(scan, value)); IvfflatBench("GetScanLists", GetScanLists(scan, value));
IvfflatBench("GetScanItems", GetScanItems(scan, value)); IvfflatBench("GetScanItems", GetScanItems(scan, value));
so->first = false; so->first = false;
so->value = value;
/* TODO clean up if we allocated a new value */ #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));
} }
while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
if (so->listIndex == so->maxProbes) bool isnull;
return false; ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
IvfflatBench("GetScanItems", GetScanItems(scan, so->value)); scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
} }
heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull)); return false;
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
} }
/* /*
@@ -397,7 +374,6 @@ ivfflatendscan(IndexScanDesc scan)
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
pfree(so->listPages);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas); FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc); FreeTupleDesc(so->tupdesc);

View File

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

View File

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

View File

@@ -26,7 +26,7 @@ $node->safe_psql("postgres", qq(
my $count = $node->safe_psql("postgres", qq( my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = on; SET hnsw.streaming = on;
SET work_mem = '8MB'; SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
)); ));
@@ -34,16 +34,16 @@ is($count, 10);
foreach ((30000, 50000, 70000)) foreach ((30000, 50000, 70000))
{ {
my $max_tuples = $_; my $ef_stream = $_;
my $expected = $max_tuples / 10000; my $expected = $ef_stream / 10000;
my $sum = 0; my $sum = 0;
for my $i (1 .. 20) for my $i (1 .. 20)
{ {
$count = $node->safe_psql("postgres", qq( $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = on; SET hnsw.streaming = on;
SET hnsw.iterative_search_max_tuples = $max_tuples; SET hnsw.ef_stream = $ef_stream;
SET work_mem = '8MB'; SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
)); ));
@@ -57,7 +57,7 @@ foreach ((30000, 50000, 70000))
my ($ret, $stdout, $stderr) = $node->psql("postgres", qq( my ($ret, $stdout, $stderr) = $node->psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.iterative_search = on; SET hnsw.streaming = on;
SET client_min_messages = debug1; SET client_min_messages = debug1;
SET work_mem = '2MB'; SET work_mem = '2MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t; SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;

View File

@@ -1,54 +0,0 @@
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_search = on;
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_search = on;
SET ivfflat.iterative_search_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

@@ -21,7 +21,7 @@ sub test_recall
my $explain = $node->safe_psql("postgres", qq( my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search; SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_search = on; SET hnsw.streaming = on;
EXPLAIN ANALYZE SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[0]' LIMIT $limit; 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/); like($explain, qr/Index Scan using idx on tst/);
@@ -31,7 +31,7 @@ sub test_recall
my $actual = $node->safe_psql("postgres", qq( my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off; SET enable_seqscan = off;
SET hnsw.ef_search = $ef_search; SET hnsw.ef_search = $ef_search;
SET hnsw.iterative_search = on; SET hnsw.streaming = on;
SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit; SELECT i FROM tst WHERE i % $c = 0 ORDER BY v $operator '$queries[$i]' LIMIT $limit;
)); ));
my @actual_ids = split("\n", $actual); my @actual_ids = split("\n", $actual);

View File

@@ -1,125 +0,0 @@
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_search = on;
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_search = on;
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);
}
# 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.58, $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

@@ -1,113 +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 = 1000;
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", qq(
SET maintenance_work_mem = '256MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_cosine_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_cosine_ops);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_cosine_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_cosine_ops, v vector_cosine_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();