Compare commits

..

53 Commits

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

View File

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

View File

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

View File

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

View File

@@ -24,5 +24,3 @@ CREATE CAST (double precision[] AS sparsevec)
CREATE CAST (numeric[] AS sparsevec)
WITH FUNCTION array_to_sparsevec(numeric[], integer, boolean) AS ASSIGNMENT;
-- TODO add ivfflat attributes

View File

@@ -916,13 +916,3 @@ CREATE OPERATOR CLASS sparsevec_l1_ops
OPERATOR 1 <+> (sparsevec, sparsevec) FOR ORDER BY float_ops,
FUNCTION 1 l1_distance(sparsevec, sparsevec),
FUNCTION 3 hnsw_sparsevec_support(internal);
-- ivfflat attributes
CREATE OPERATOR CLASS vector_integer_ops
DEFAULT FOR TYPE integer USING ivfflat AS
OPERATOR 2 < ,
OPERATOR 3 <= ,
OPERATOR 4 = ,
OPERATOR 5 >= ,
OPERATOR 6 > ;

View File

@@ -18,7 +18,17 @@
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif
static const struct config_enum_entry hnsw_iterative_scan_options[] = {
{"off", HNSW_ITERATIVE_SCAN_OFF, false},
{"relaxed_order", HNSW_ITERATIVE_SCAN_RELAXED, false},
{"strict_order", HNSW_ITERATIVE_SCAN_STRICT, false},
{NULL, 0, false}
};
int hnsw_ef_search;
int hnsw_iterative_scan;
int hnsw_max_scan_tuples;
double hnsw_scan_mem_multiplier;
int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind;
@@ -67,7 +77,21 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
DefineCustomEnumVariable("hnsw.iterative_scan", "Sets the mode for iterative scans",
NULL, &hnsw_iterative_scan,
HNSW_ITERATIVE_SCAN_OFF, hnsw_iterative_scan_options, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* This is approximate and does not affect the initial scan */
DefineCustomIntVariable("hnsw.max_scan_tuples", "Sets the max number of tuples to visit for iterative scans",
NULL, &hnsw_max_scan_tuples,
20000, 1, INT_MAX, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* Same range as hash_mem_multiplier */
DefineCustomRealVariable("hnsw.scan_mem_multiplier", "Sets the multiple of work_mem to use for iterative scans",
NULL, &hnsw_scan_mem_multiplier,
1, 1, 1000, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw");
}
@@ -113,6 +137,10 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return;
}

View File

@@ -109,8 +109,18 @@
/* Variables */
extern int hnsw_ef_search;
extern int hnsw_iterative_scan;
extern int hnsw_max_scan_tuples;
extern double hnsw_scan_mem_multiplier;
extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeScanMode
{
HNSW_ITERATIVE_SCAN_OFF,
HNSW_ITERATIVE_SCAN_RELAXED,
HNSW_ITERATIVE_SCAN_STRICT
} HnswIterativeScanMode;
typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray;
@@ -132,6 +142,7 @@ struct HnswElementData
uint8 heaptidsLength;
uint8 level;
uint8 deleted;
uint8 version;
uint32 hash;
HnswNeighborsPtr neighbors;
BlockNumber blkno;
@@ -319,10 +330,10 @@ typedef struct HnswElementTupleData
uint8 type;
uint8 level;
uint8 deleted;
uint8 unused;
uint8 version;
ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid;
uint16 unused2;
uint16 unused;
Vector data;
} HnswElementTupleData;
@@ -331,7 +342,7 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData
{
uint8 type;
uint8 unused;
uint8 version;
uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData;
@@ -356,6 +367,13 @@ typedef struct HnswScanOpaqueData
const HnswTypeInfo *typeInfo;
bool first;
List *w;
visited_hash v;
pairingheap *discarded;
HnswQuery q;
int m;
int64 tuples;
double previousDistance;
Size maxMemory;
MemoryContext tmpCtx;
/* Support functions */
@@ -399,7 +417,7 @@ bool HnswCheckNorm(HnswSupport * support, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page);
void HnswInit(void);
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement);
List *HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size);

View File

@@ -36,7 +36,7 @@ GetInsertPage(Relation index)
* Check for a free offset
*/
static bool
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage)
HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size etupSize, Size ntupSize, Buffer *nbuf, Page *npage, OffsetNumber *freeOffno, OffsetNumber *freeNeighborOffno, BlockNumber *newInsertPage, uint8 *tupleVersion)
{
OffsetNumber offno;
OffsetNumber maxoffno = PageGetMaxOffsetNumber(page);
@@ -98,6 +98,7 @@ HnswFreeOffset(Relation index, Buffer buf, Page page, HnswElement element, Size
{
*freeOffno = offno;
*freeNeighborOffno = neighborOffno;
*tupleVersion = etup->version;
return true;
}
else if (*nbuf != buf)
@@ -153,6 +154,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
OffsetNumber freeOffno = InvalidOffsetNumber;
OffsetNumber freeNeighborOffno = InvalidOffsetNumber;
BlockNumber newInsertPage = InvalidBlockNumber;
uint8 tupleVersion;
char *base = NULL;
/* Calculate sizes */
@@ -202,7 +204,7 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
}
/* Next, try space from a deleted element */
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage))
if (HnswFreeOffset(index, buf, page, e, etupSize, ntupSize, &nbuf, &npage, &freeOffno, &freeNeighborOffno, &newInsertPage, &tupleVersion))
{
if (nbuf != buf)
{
@@ -212,6 +214,10 @@ AddElementOnDisk(Relation index, HnswElement e, int m, BlockNumber insertPage, B
npage = GenericXLogRegisterBuffer(state, nbuf, 0);
}
/* Set tuple version */
etup->version = tupleVersion;
ntup->version = tupleVersion;
break;
}

View File

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

View File

@@ -251,6 +251,8 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
element->level = level;
element->deleted = 0;
/* Start at one to make it easier to find issues */
element->version = 1;
HnswInitNeighbors(base, element, m, allocator);
@@ -430,6 +432,7 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level;
etup->deleted = 0;
etup->version = element->version;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
if (i < element->heaptidsLength)
@@ -472,6 +475,7 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
}
ntup->count = idx;
ntup->version = e->version;
}
/*
@@ -482,6 +486,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
{
element->level = etup->level;
element->deleted = etup->deleted;
element->version = etup->version;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptidsLength = 0;
@@ -576,21 +581,34 @@ GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport *
return HnswGetDistance(q->value, value, support);
}
/*
* Allocate a search candidate
*/
static HnswSearchCandidate *
HnswInitSearchCandidate(char *base, HnswElement element, double distance)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, sc->element, element);
sc->distance = distance;
return sc;
}
/*
* Create a candidate for the entry point
*/
HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL;
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, support);
distance = GetElementDistance(base, entryPoint, q, support);
else
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL);
return sc;
HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL);
return HnswInitSearchCandidate(base, entryPoint, distance);
}
/*
@@ -608,6 +626,21 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
return 0;
}
/*
* Compare discarded candidate distances
*/
static int
CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return -1;
return 0;
}
/*
* Compare candidate distances
*/
@@ -728,8 +761,11 @@ HnswLoadNeighborTids(HnswElement element, ItemPointerData *indextids, Relation i
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
/* Ensure expected neighbors */
if (ntup->count != (element->level + 2) * m)
/*
* Ensure the neighbor tuple has not been deleted or replaced between
* index scan iterations
*/
if (ntup->version != element->version || ntup->count != (element->level + 2) * m)
{
UnlockReleaseBuffer(buf);
return false;
@@ -775,13 +811,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper
*/
List *
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement)
HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation index, HnswSupport * support, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
{
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
visited_hash v;
visited_hash vh;
ListCell *lc2;
HnswNeighborArray *localNeighborhood = NULL;
Size neighborhoodSize = 0;
@@ -790,7 +826,19 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
int unvisitedLength;
bool inMemory = index == NULL;
InitVisited(base, &v, inMemory, ef, m);
if (v == NULL)
{
v = &vh;
initVisited = true;
}
if (initVisited)
{
InitVisited(base, v, inMemory, ef, m);
if (discarded != NULL)
*discarded = pairingheap_allocate(CompareNearestDiscardedCandidates, NULL);
}
/* Create local memory for neighborhood if needed */
if (inMemory)
@@ -805,7 +853,14 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
HnswSearchCandidate *sc = (HnswSearchCandidate *) lfirst(lc2);
bool found;
AddToVisited(base, &v, sc->element, inMemory, &found);
if (initVisited)
{
AddToVisited(base, v, sc->element, inMemory, &found);
/* OK to count elements instead of tuples */
if (tuples != NULL)
(*tuples)++;
}
pairingheap_add(C, &sc->c_node);
pairingheap_add(W, &sc->w_node);
@@ -831,9 +886,13 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
cElement = HnswPtrAccess(base, c->element);
if (inMemory)
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize);
else
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc);
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
/* OK to count elements instead of tuples */
if (tuples != NULL)
(*tuples) += unvisitedLength;
for (int i = 0; i < unvisitedLength; i++)
{
@@ -857,25 +916,30 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
/* Avoid any allocations if not adding */
eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
HnswLoadElementImpl(blkno, offno, &eDistance, q, index, support, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement);
if (eElement == NULL)
continue;
}
if (!(eDistance < f->distance || alwaysAdd))
continue;
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd))
{
if (discarded != NULL)
{
/* Create a new candidate */
e = HnswInitSearchCandidate(base, eElement, eDistance);
pairingheap_add(*discarded, &e->w_node);
}
Assert(!eElement->deleted);
continue;
}
/* Make robust to issues */
if (eElement->level < lc)
continue;
/* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
e = HnswInitSearchCandidate(base, eElement, eDistance);
pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node);
@@ -890,7 +954,12 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
/* No need to decrement wlen */
if (wlen > ef)
pairingheap_remove_first(W);
{
HnswSearchCandidate *d = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
if (discarded != NULL)
pairingheap_add(*discarded, &d->w_node);
}
}
}
}
@@ -1225,7 +1294,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
/* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--)
{
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement);
w = HnswSearchLayer(base, &q, ep, 1, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
ep = w;
}
@@ -1244,7 +1313,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
List *lw = NIL;
ListCell *lc2;
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement);
w = HnswSearchLayer(base, &q, ep, efConstruction, lc, index, support, m, true, skipElement, NULL, NULL, true, NULL);
/* Convert search candidates to candidates */
foreach(lc2, w)

View File

@@ -527,6 +527,14 @@ MarkDeleted(HnswVacuumState * vacuumstate)
for (int i = 0; i < ntup->count; i++)
ItemPointerSetInvalid(&ntup->indextids[i]);
/* Increment version */
/* This is used to avoid incorrect reads for iterative scans */
/* Reserve some bits for future use */
etup->version++;
if (etup->version > 15)
etup->version = 1;
ntup->version = etup->version;
/*
* We modified the tuples in place, no need to call
* PageIndexTupleOverwrite

View File

@@ -138,7 +138,7 @@ SampleRows(IvfflatBuildState * buildstate)
* Add tuple to sort
*/
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 minDistance = DBL_MAX;
@@ -184,11 +184,6 @@ AddTupleToSort(Relation index, ItemPointer tid, Datum *values, bool *isnull, Ivf
slot->tts_isnull[1] = false;
slot->tts_values[2] = value;
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);
/*
@@ -220,7 +215,7 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Add tuple to sort */
AddTupleToSort(index, tid, values, isnull, buildstate);
AddTupleToSort(index, tid, values, buildstate);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
@@ -231,20 +226,19 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
* Get index tuple from sort state
*/
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)
{
if (tuplesort_gettupleslot(sortstate, true, false, slot, NULL))
{
bool unused;
Datum value;
bool 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]);
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
value = slot_getattr(slot, 3, &isnull);
/* Form the index tuple */
*itup = index_form_tuple(tupdesc, values, isnull);
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &unused)));
*itup = index_form_tuple(tupdesc, &value, &isnull);
(*itup)->t_tid = *((ItemPointer) DatumGetPointer(slot_getattr(slot, 2, &isnull)));
}
else
*list = -1;
@@ -262,14 +256,12 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = buildstate->tupdesc;
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_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++)
{
@@ -305,7 +297,7 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
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);
@@ -315,9 +307,6 @@ InsertTuples(Relation index, IvfflatBuildState * buildstate, ForkNumber forkNum)
/* Set the start and insert pages */
IvfflatUpdateList(index, buildstate->listInfo[i], insertPage, InvalidBlockNumber, startPage, forkNum);
}
pfree(values);
pfree(isnull);
}
/*
@@ -341,19 +330,6 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
errmsg("type not supported for ivfflat index")));
/* TODO See if needed */
if (IndexRelationGetNumberOfKeyAttributes(index) > 3)
elog(ERROR, "index cannot have more than three columns");
if (!OidIsValid(index_getprocid(index, 1, IVFFLAT_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, IVFFLAT_DISTANCE_PROC)))
elog(ERROR, "column %d cannot be a vector", i + 1);
}
/* Require column to have dimensions to be indexed */
if (buildstate->dimensions < 0)
ereport(ERROR,
@@ -381,11 +357,10 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
errmsg("dimensions must be greater than one for this opclass")));
/* Create tuple description for sorting */
buildstate->sortdesc = CreateTemplateTupleDesc(2 + buildstate->tupdesc->natts);
buildstate->sortdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
for (int i = 0; i < buildstate->tupdesc->natts; i++)
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) (3 + i), NULL, buildstate->tupdesc->attrs[i].atttypid, -1, 0);
TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 3, "vector", buildstate->tupdesc->attrs[0].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);

View File

@@ -17,8 +17,16 @@
#endif
int ivfflat_probes;
int ivfflat_iterative_scan;
int ivfflat_max_probes;
static relopt_kind ivfflat_relopt_kind;
static const struct config_enum_entry ivfflat_iterative_scan_options[] = {
{"off", IVFFLAT_ITERATIVE_SCAN_OFF, false},
{"relaxed_order", IVFFLAT_ITERATIVE_SCAN_RELAXED, false},
{NULL, 0, false}
};
/*
* Initialize index options and variables
*/
@@ -31,7 +39,16 @@ IvfflatInit(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
"Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL);
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
DefineCustomEnumVariable("ivfflat.iterative_scan", "Sets the mode for iterative scans",
NULL, &ivfflat_iterative_scan,
IVFFLAT_ITERATIVE_SCAN_OFF, ivfflat_iterative_scan_options, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
/* If this is less than probes, probes is used */
DefineCustomIntVariable("ivfflat.max_probes", "Sets the max number of probes for iterative scans",
NULL, &ivfflat_max_probes,
IVFFLAT_MAX_LISTS, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
MarkGUCPrefixReserved("ivfflat");
}
@@ -82,6 +99,10 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return;
}
@@ -167,7 +188,7 @@ ivfflathandler(PG_FUNCTION_ARGS)
amroutine->amcanorderbyop = true;
amroutine->amcanbackward = false; /* can change direction mid-scan */
amroutine->amcanunique = false;
amroutine->amcanmulticol = true;
amroutine->amcanmulticol = false;
amroutine->amoptionalkey = true;
amroutine->amsearcharray = false;
amroutine->amsearchnulls = false;

View File

@@ -80,6 +80,14 @@
/* Variables */
extern int ivfflat_probes;
extern int ivfflat_iterative_scan;
extern int ivfflat_max_probes;
typedef enum IvfflatIterativeScanMode
{
IVFFLAT_ITERATIVE_SCAN_OFF,
IVFFLAT_ITERATIVE_SCAN_RELAXED
} IvfflatIterativeScanMode;
typedef struct VectorArrayData
{
@@ -248,8 +256,11 @@ typedef struct IvfflatScanOpaqueData
{
const IvfflatTypeInfo *typeInfo;
int probes;
int maxProbes;
int dimensions;
bool first;
Datum value;
MemoryContext tmpCtx;
/* Sorting */
Tuplesortstate *sortstate;
@@ -266,7 +277,9 @@ typedef struct IvfflatScanOpaqueData
/* Lists */
pairingheap *listQueue;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */
BlockNumber *listPages;
int listIndex;
IvfflatScanList *lists;
} IvfflatScanOpaqueData;
typedef IvfflatScanOpaqueData * IvfflatScanOpaque;

View File

@@ -78,8 +78,6 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
BlockNumber insertPage = InvalidBlockNumber;
ListInfo listInfo;
BlockNumber originalInsertPage;
TupleDesc tupdesc = RelationGetDescr(index);
Datum *newValues = palloc(tupdesc->natts * sizeof(Datum));
/* Detoast once for all calls */
value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -104,12 +102,8 @@ InsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_tid, R
Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage;
newValues[0] = value;
for (int i = 1; i < tupdesc->natts; i++)
newValues[i] = values[i];
/* Form tuple */
itup = index_form_tuple(tupdesc, newValues, isnull);
itup = index_form_tuple(RelationGetDescr(index), &value, isnull);
itup->t_tid = *heap_tid;
/* Get tuple size */

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@@ -1,197 +0,0 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $node;
my @queries = ();
my @where = ();
my @expected;
my $limit = 20;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
my $nc = 100;
my $nc2 = 10;
sub test_recall
{
my ($probes, $min, $operator) = @_;
my $correct = 0;
my $total = 0;
for my $j (0 .. 2)
{
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
EXPLAIN ANALYZE SELECT i FROM tst WHERE $where[$j] ORDER BY v $operator '$queries[$j]' LIMIT $limit;
));
like($explain, qr/Index Cond/);
}
for my $i (0 .. $#queries)
{
my $actual = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = $probes;
SELECT i FROM tst WHERE $where[$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, c2 int4);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, i % $nc2 FROM generate_series(1, 50000) i;"
);
# Generate queries
for my $i (1 .. 100)
{
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
push(@queries, "[" . join(",", @r) . "]");
if ($i % 3 == 0)
{
my $c = int(rand() * $nc);
push(@where, "c = $c");
}
elsif ($i % 3 == 1)
{
my $c2 = int(rand() * $nc2);
push(@where, "c2 = $c2");
}
else
{
# use c2 to ensure results
my $c2 = int(rand() * $nc2);
push(@where, "c = $c2 AND c2 = $c2");
}
}
# Add index
$node->safe_psql("postgres", qq(
CREATE INDEX ON tst USING ivfflat (v vector_l2_ops, c, c2) WITH (lists = 100);
));
# Insert more rows
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, i % $nc2 FROM generate_series(1, 50000) i;"
);
# Get exact results
@expected = ();
for my $i (0 .. $#queries)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst WHERE $where[$i] ORDER BY v <-> '$queries[$i]' LIMIT $limit;
));
push(@expected, $res);
}
# Test recall
test_recall(10, 0.99, '<->');
# Test vacuum
$node->safe_psql("postgres", "DELETE FROM tst WHERE c > 5;");
$node->safe_psql("postgres", "VACUUM tst;");
# Test less than
my $explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 10 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(c < 10\)/);
# Test less than or equal
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c <= 10 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(c <= 10\)/);
# Test greater than or equal
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= 90 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(c >= 90\)/);
# Test greater than
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c > 90 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(c > 90\)/);
# Test multiple attribute columns
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = 1 AND c2 = 1 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(\(c = 1\) AND \(c2 = 1\)\)/);
# Test only last attribute column
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c2 = 1 ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Cond: \(c2 = 1\)/);
# Test only vector column
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$queries[0]' LIMIT $limit;
));
like($explain, qr/Index Scan/);
# Test only attribute columns
$explain = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = 1;
));
like($explain, qr/Seq Scan/);
# Test columns
my ($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (c);");
like($stderr, qr/first column must be a vector/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (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 ivfflat (v vector_cosine_ops, c, c, c);");
like($stderr, qr/index cannot have more than three columns/);
($ret, $stdout, $stderr) = $node->psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_cosine_ops, v vector_cosine_ops);");
like($stderr, qr/column 2 cannot be a vector/);
done_testing();

View File

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

View File

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

View File

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

View File

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