Compare commits

..

61 Commits

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

View File

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

View File

@@ -451,67 +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] ## Iterative Index Scans
*Added in 0.8.0* *Unreleased*
With approximate indexes, you can end up with less results than expected due to filtering conditions in the query. 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.
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). There are two modes for iterative scans: strict and relaxed.
```tsql Strict ensures results are in the exact order by distance
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 ```sql
WITH approx_order AS MATERIALIZED ( SET hnsw.iterative_scan = strict_order;
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. Relaxed allows results to be slightly out of order by distance, but provides better recall
```sql ```sql
WITH approx_order AS MATERIALIZED ( SET hnsw.iterative_scan = relaxed_order;
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5 # or
) SELECT * FROM approx_order WHERE distance < 0.1 ORDER BY distance; SET ivfflat.iterative_scan = relaxed_order;
``` ```
### Iterative Options 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
Since scanning a large portion of the index is expensive, there are options to control when the scan ends. ```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 #### HNSW
Specify the max number of additional tuples visited Specify the max number of tuples to visit (20,000 by default)
```sql ```sql
SET hnsw.ef_stream = 10000; SET hnsw.max_scan_tuples = 20000;
``` ```
The scan will also end if reaches `work_mem`. You can see when this happens by enabling debug messages. 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 ```sql
SET client_min_messages = debug1; SET client_min_messages = debug1;
``` ```
which will show when a scan reaches the memory limit
```text ```text
DEBUG: hnsw index scan exceeded work_mem after 10000 tuples DEBUG: hnsw index scan reached memory limit after 20000 tuples
HINT: Increase work_mem to scan more tuples. HINT: Increase hnsw.scan_mem_multiplier to scan more tuples.
```
If the server has enough memory, you can adjust this with:
```sql
SET work_mem = '8MB';
``` ```
#### IVFFlat #### IVFFlat
@@ -522,6 +532,8 @@ Specify the max number of probes
SET ivfflat.max_probes = 100; SET ivfflat.max_probes = 100;
``` ```
Note: If this is lower than `ivfflat.probes`, `ivfflat.probes` will be used
## Half-Precision Vectors ## Half-Precision Vectors
*Added in 0.7.0* *Added in 0.7.0*

View File

@@ -18,18 +18,17 @@
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x) #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
static const struct config_enum_entry hnsw_iterative_search_options[] = { static const struct config_enum_entry hnsw_iterative_scan_options[] = {
{"off", HNSW_ITERATIVE_SEARCH_OFF, false}, {"off", HNSW_ITERATIVE_SCAN_OFF, false},
{"strict", HNSW_ITERATIVE_SEARCH_STRICT, false}, {"relaxed_order", HNSW_ITERATIVE_SCAN_RELAXED, false},
{"relaxed", HNSW_ITERATIVE_SEARCH_RELAXED, false}, {"strict_order", HNSW_ITERATIVE_SCAN_STRICT, 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_max_iterative_tuples; int hnsw_iterative_scan;
int hnsw_iterative_search; int hnsw_max_scan_tuples;
double hnsw_scan_mem_multiplier;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -78,18 +77,21 @@ HnswInit(void)
DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search", DefineCustomIntVariable("hnsw.ef_search", "Sets the size of the dynamic candidate list for search",
"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, GUC_EXPLAIN, NULL, NULL, NULL);
/* TODO Change name */ DefineCustomEnumVariable("hnsw.iterative_scan", "Sets the mode for iterative scans",
DefineCustomEnumVariable("hnsw.streaming", "Iterative search mode", NULL, &hnsw_iterative_scan,
NULL, &hnsw_iterative_search, HNSW_ITERATIVE_SCAN_OFF, hnsw_iterative_scan_options, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
HNSW_ITERATIVE_SEARCH_OFF, hnsw_iterative_search_options, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Change name */ /* This is approximate and does not affect the initial scan */
/* TODO Ensure ivfflat.max_probes uses same value for "all" */ DefineCustomIntVariable("hnsw.max_scan_tuples", "Sets the max number of tuples to visit for iterative scans",
DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search", NULL, &hnsw_max_scan_tuples,
"-1 means all", &hnsw_max_iterative_tuples, 20000, 1, INT_MAX, PGC_USERSET, GUC_EXPLAIN, NULL, NULL, NULL);
HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, 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"); MarkGUCPrefixReserved("hnsw");
} }
@@ -135,6 +137,10 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0; *indexSelectivity = 0;
*indexCorrelation = 0; *indexCorrelation = 0;
*indexPages = 0; *indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return; return;
} }

View File

@@ -42,9 +42,6 @@
#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
@@ -112,16 +109,17 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_max_iterative_tuples; extern int hnsw_iterative_scan;
extern int hnsw_iterative_search; extern int hnsw_max_scan_tuples;
extern double hnsw_scan_mem_multiplier;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef enum HnswIterativeSearchType typedef enum HnswIterativeScanMode
{ {
HNSW_ITERATIVE_SEARCH_OFF, HNSW_ITERATIVE_SCAN_OFF,
HNSW_ITERATIVE_SEARCH_STRICT, HNSW_ITERATIVE_SCAN_RELAXED,
HNSW_ITERATIVE_SEARCH_RELAXED HNSW_ITERATIVE_SCAN_STRICT
} HnswIterativeSearchType; } HnswIterativeScanMode;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
typedef struct HnswNeighborArray HnswNeighborArray; typedef struct HnswNeighborArray HnswNeighborArray;
@@ -375,6 +373,7 @@ typedef struct HnswScanOpaqueData
int m; int m;
int64 tuples; int64 tuples;
double previousDistance; double previousDistance;
Size maxMemory;
MemoryContext tmpCtx; MemoryContext tmpCtx;
/* Support functions */ /* Support functions */

View File

@@ -41,7 +41,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
ep = w; ep = w;
} }
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); 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);
} }
/* /*
@@ -102,6 +102,17 @@ GetScanValue(IndexScanDesc scan)
return value; return value;
} }
#if defined(HNSW_MEMORY)
/*
* Show memory usage
*/
static void
ShowMemoryUsage(HnswScanOpaque so)
{
elog(INFO, "memory: %zu KB, tuples: " INT64_FORMAT, MemoryContextMemAllocated(so->tmpCtx, false) / 1024, so->tuples);
}
#endif
/* /*
* Prepare for an index scan * Prepare for an index scan
*/ */
@@ -110,21 +121,29 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
{ {
IndexScanDesc scan; IndexScanDesc scan;
HnswScanOpaque so; HnswScanOpaque so;
double maxMemory;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index); so->typeInfo = HnswGetTypeInfo(index);
so->first = true;
so->v.tids = NULL;
so->discarded = NULL;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES);
/* Set support functions */ /* Set support functions */
HnswInitSupport(&so->support, index); 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; scan->opaque = so;
return scan; return scan;
@@ -138,13 +157,10 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
{ {
HnswScanOpaque so = (HnswScanOpaque) scan->opaque; HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
if (so->v.tids != NULL)
tidhash_reset(so->v.tids);
if (so->discarded != NULL)
pairingheap_reset(so->discarded);
so->first = true; so->first = true;
/* v and discarded are allocated in tmpCtx */
so->v.tids = NULL;
so->discarded = NULL;
so->tuples = 0; so->tuples = 0;
so->previousDistance = -get_float8_infinity(); so->previousDistance = -get_float8_infinity();
MemoryContextReset(so->tmpCtx); MemoryContextReset(so->tmpCtx);
@@ -204,7 +220,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->first = false; so->first = false;
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); ShowMemoryUsage(so);
#endif #endif
} }
@@ -217,15 +233,15 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
if (list_length(so->w) == 0) if (list_length(so->w) == 0)
{ {
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_OFF) if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_OFF)
break; break;
/* Empty index */ /* Empty index */
if (so->discarded == NULL) if (so->discarded == NULL)
break; break;
/* Reached max number of additional tuples */ /* Reached max number of tuples */
if (hnsw_max_iterative_tuples != -1 && so->tuples >= hnsw_ef_search + hnsw_max_iterative_tuples) if (so->tuples >= hnsw_max_scan_tuples)
{ {
if (pairingheap_is_empty(so->discarded)) if (pairingheap_is_empty(so->discarded))
break; break;
@@ -234,13 +250,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded))); so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
} }
/* Prevent scans from consuming too much memory */ /* Prevent scans from consuming too much memory */
else if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L) else if (MemoryContextMemAllocated(so->tmpCtx, false) > so->maxMemory)
{ {
if (pairingheap_is_empty(so->discarded)) if (pairingheap_is_empty(so->discarded))
{ {
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples), (errmsg("hnsw index scan reached memory limit after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples."))); errhint("Increase hnsw.scan_mem_multiplier to scan more tuples.")));
break; break;
} }
@@ -266,7 +282,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY) #if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024); ShowMemoryUsage(so);
#endif #endif
} }
@@ -283,7 +299,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */ /* Mark memory as free for next iteration */
if (hnsw_iterative_search != HNSW_ITERATIVE_SEARCH_OFF) if (hnsw_iterative_scan != HNSW_ITERATIVE_SCAN_OFF)
{ {
pfree(element); pfree(element);
pfree(sc); pfree(sc);
@@ -294,7 +310,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
heaptid = &element->heaptids[--element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
if (hnsw_iterative_search == HNSW_ITERATIVE_SEARCH_STRICT) if (hnsw_iterative_scan == HNSW_ITERATIVE_SCAN_STRICT)
{ {
if (sc->distance < so->previousDistance) if (sc->distance < so->previousDistance)
continue; continue;

View File

@@ -581,21 +581,34 @@ GetElementDistance(char *base, HnswElement element, HnswQuery * q, HnswSupport *
return HnswGetDistance(q->value, value, support); 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 * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswSearchCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec) HnswEntryCandidate(char *base, HnswElement entryPoint, HnswQuery * q, Relation index, HnswSupport * support, bool loadVec)
{ {
HnswSearchCandidate *sc = palloc(sizeof(HnswSearchCandidate));
bool inMemory = index == NULL; bool inMemory = index == NULL;
double distance;
HnswPtrStore(base, sc->element, entryPoint);
if (inMemory) if (inMemory)
sc->distance = GetElementDistance(base, entryPoint, q, support); distance = GetElementDistance(base, entryPoint, q, support);
else else
HnswLoadElement(entryPoint, &sc->distance, q, index, support, loadVec, NULL); HnswLoadElement(entryPoint, &distance, q, index, support, loadVec, NULL);
return sc;
return HnswInitSearchCandidate(base, entryPoint, distance);
} }
/* /*
@@ -844,6 +857,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
{ {
AddToVisited(base, v, sc->element, inMemory, &found); AddToVisited(base, v, sc->element, inMemory, &found);
/* OK to count elements instead of tuples */
if (tuples != NULL) if (tuples != NULL)
(*tuples)++; (*tuples)++;
} }
@@ -876,6 +890,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
else 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) if (tuples != NULL)
(*tuples) += unvisitedLength; (*tuples) += unvisitedLength;
@@ -912,9 +927,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
if (discarded != NULL) if (discarded != NULL)
{ {
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node); pairingheap_add(*discarded, &e->w_node);
} }
@@ -926,9 +939,7 @@ HnswSearchLayer(char *base, HnswQuery * q, List *ep, int ef, int lc, Relation in
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = HnswInitSearchCandidate(base, eElement, eDistance);
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(C, &e->c_node); pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); pairingheap_add(W, &e->w_node);

View File

@@ -228,11 +228,11 @@ BuildCallback(Relation index, ItemPointer tid, Datum *values,
static inline void static inline void
GetNextTuple(Tuplesortstate *sortstate, TupleDesc tupdesc, TupleTableSlot *slot, 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))
{ {
Datum value;
bool isnull;
*list = DatumGetInt32(slot_getattr(slot, 1, &isnull)); *list = DatumGetInt32(slot_getattr(slot, 1, &isnull));
value = slot_getattr(slot, 3, &isnull); value = slot_getattr(slot, 3, &isnull);
@@ -254,8 +254,8 @@ 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->tupdesc, &TTSOpsMinimalTuple); TupleTableSlot *slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsMinimalTuple);
TupleDesc tupdesc = RelationGetDescr(index); TupleDesc tupdesc = buildstate->tupdesc;
pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD); pgstat_progress_update_param(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_IVFFLAT_PHASE_LOAD);
@@ -319,6 +319,7 @@ InitBuildState(IvfflatBuildState * buildstate, Relation heap, Relation index, In
buildstate->index = index; buildstate->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;
@@ -356,12 +357,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->tupdesc = CreateTemplateTupleDesc(3); buildstate->sortdesc = CreateTemplateTupleDesc(3);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 1, "list", INT4OID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 1, "list", INT4OID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 2, "tid", TIDOID, -1, 0);
TupleDescInitEntry(buildstate->tupdesc, (AttrNumber) 3, "vector", RelationGetDescr(index)->attrs[0].atttypid, -1, 0); TupleDescInitEntry(buildstate->sortdesc, (AttrNumber) 3, "vector", buildstate->tupdesc->attrs[0].atttypid, -1, 0);
buildstate->slot = MakeSingleTupleTableSlot(buildstate->tupdesc, &TTSOpsVirtual); buildstate->slot = MakeSingleTupleTableSlot(buildstate->sortdesc, &TTSOpsVirtual);
buildstate->centers = VectorArrayInit(buildstate->lists, buildstate->dimensions, buildstate->typeInfo->itemSize(buildstate->dimensions)); buildstate->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);
@@ -633,7 +634,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.tupdesc, sortmem, coordinate); ivfspool->sortstate = InitBuildSortState(buildstate.sortdesc, 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));
@@ -950,7 +951,7 @@ AssignTuples(IvfflatBuildState * buildstate)
} }
/* Begin serial/leader tuplesort */ /* Begin serial/leader tuplesort */
buildstate->sortstate = InitBuildSortState(buildstate->tupdesc, maintenance_work_mem, coordinate); buildstate->sortstate = InitBuildSortState(buildstate->sortdesc, maintenance_work_mem, coordinate);
/* Add tuples to sort */ /* Add tuples to sort */
if (buildstate->heap != NULL) if (buildstate->heap != NULL)

View File

@@ -17,8 +17,16 @@
#endif #endif
int ivfflat_probes; int ivfflat_probes;
int ivfflat_iterative_scan;
int ivfflat_max_probes;
static relopt_kind ivfflat_relopt_kind; 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 * Initialize index options and variables
*/ */
@@ -31,7 +39,16 @@ IvfflatInit(void)
DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes", DefineCustomIntVariable("ivfflat.probes", "Sets the number of probes",
"Valid range is 1..lists.", &ivfflat_probes, "Valid range is 1..lists.", &ivfflat_probes,
IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 0, NULL, NULL, NULL); IVFFLAT_DEFAULT_PROBES, IVFFLAT_MIN_LISTS, IVFFLAT_MAX_LISTS, PGC_USERSET, 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"); MarkGUCPrefixReserved("ivfflat");
} }
@@ -82,6 +99,10 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*indexSelectivity = 0; *indexSelectivity = 0;
*indexCorrelation = 0; *indexCorrelation = 0;
*indexPages = 0; *indexPages = 0;
#if PG_VERSION_NUM >= 180000
/* See "On disable_cost" thread on pgsql-hackers */
path->path.disabled_nodes = 2;
#endif
return; return;
} }

View File

@@ -80,6 +80,14 @@
/* Variables */ /* Variables */
extern int ivfflat_probes; 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 typedef struct VectorArrayData
{ {
@@ -165,6 +173,7 @@ 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;
@@ -198,7 +207,7 @@ typedef struct IvfflatBuildState
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc sortdesc;
TupleTableSlot *slot; TupleTableSlot *slot;
/* Memory */ /* Memory */
@@ -247,8 +256,11 @@ 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;
MemoryContext tmpCtx;
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
@@ -265,7 +277,9 @@ typedef struct IvfflatScanOpaqueData
/* Lists */ /* Lists */
pairingheap *listQueue; pairingheap *listQueue;
IvfflatScanList lists[FLEXIBLE_ARRAY_MEMBER]; /* must come last */ BlockNumber *listPages;
int listIndex;
IvfflatScanList *lists;
} IvfflatScanOpaqueData; } IvfflatScanOpaqueData;
typedef IvfflatScanOpaqueData * IvfflatScanOpaque; typedef IvfflatScanOpaqueData * IvfflatScanOpaque;

View File

@@ -98,7 +98,7 @@ 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, values, &insertPage, &listInfo); FindInsertPage(index, &value, &insertPage, &listInfo);
Assert(BlockNumberIsValid(insertPage)); Assert(BlockNumberIsValid(insertPage));
originalInsertPage = insertPage; originalInsertPage = insertPage;

View File

@@ -10,10 +10,7 @@
#include "miscadmin.h" #include "miscadmin.h"
#include "pgstat.h" #include "pgstat.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h" #include "utils/memutils.h"
#endif
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr) #define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
#define GetScanListConst(ptr) pairingheap_const_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 */ /* 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->probes) if (listCount < so->maxProbes)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
@@ -78,7 +75,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->probes) if (listCount == so->maxProbes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance; maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
@@ -102,6 +99,11 @@ 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));
} }
/* /*
@@ -114,11 +116,14 @@ 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 (!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 */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))
@@ -166,13 +171,17 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
if (tuples < 100) if (tuples < 100 && ivfflat_iterative_scan == IVFFLAT_ITERATIVE_SCAN_OFF)
ereport(DEBUG1, 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
} }
/* /*
@@ -209,7 +218,13 @@ GetScanValue(IndexScanDesc scan)
/* Normalize if needed */ /* Normalize if needed */
if (so->normprocinfo != NULL) if (so->normprocinfo != NULL)
{
MemoryContext oldCtx = MemoryContextSwitchTo(so->tmpCtx);
value = IvfflatNormValue(so->typeInfo, so->collation, value); value = IvfflatNormValue(so->typeInfo, so->collation, value);
MemoryContextSwitchTo(oldCtx);
}
} }
return value; return value;
@@ -240,19 +255,30 @@ 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;
MemoryContext oldCtx;
scan = RelationGetIndexScan(index, nkeys, norderbys); scan = RelationGetIndexScan(index, nkeys, norderbys);
/* Get lists and dimensions from metapage */ /* Get lists and dimensions from metapage */
IvfflatGetMetaPageInfo(index, &lists, &dimensions); IvfflatGetMetaPageInfo(index, &lists, &dimensions);
if (ivfflat_iterative_scan != IVFFLAT_ITERATIVE_SCAN_OFF)
maxProbes = Max(ivfflat_max_probes, probes);
else
maxProbes = probes;
if (probes > lists) if (probes > lists)
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->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 */
@@ -260,6 +286,12 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC); so->normprocinfo = IvfflatOptionalProcInfo(index, IVFFLAT_NORM_PROC);
so->collation = index->rd_indcollation[0]; 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 */ /* Create tuple description for sorting */
so->tupdesc = CreateTemplateTupleDesc(2); so->tupdesc = CreateTemplateTupleDesc(2);
TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0); TupleDescInitEntry(so->tupdesc, (AttrNumber) 1, "distance", FLOAT8OID, -1, 0);
@@ -280,6 +312,11 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
so->bas = GetAccessStrategy(BAS_BULKREAD); so->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;
so->lists = palloc(maxProbes * sizeof(IvfflatScanList));
MemoryContextSwitchTo(oldCtx);
scan->opaque = so; scan->opaque = so;
@@ -294,11 +331,9 @@ 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));
@@ -314,6 +349,8 @@ 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
@@ -341,28 +378,23 @@ 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;
#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));
} }
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL)) while (!tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
{ {
bool isnull; if (so->listIndex == so->maxProbes)
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull)); return false;
scan->xs_heaptid = *heaptid; IvfflatBench("GetScanItems", GetScanItems(scan, so->value));
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
} }
return false; heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_heaptid = *heaptid;
scan->xs_recheck = false;
scan->xs_recheckorderby = false;
return true;
} }
/* /*
@@ -373,12 +405,10 @@ ivfflatendscan(IndexScanDesc scan)
{ {
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
pairingheap_free(so->listQueue); /* Free any temporary files */
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
/* TODO Free vslot and mslot without freeing TupleDesc */ MemoryContextDelete(so->tmpCtx);
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;

View File

@@ -26,7 +26,7 @@ ivfflatbulkdelete(IndexVacuumInfo *info, IndexBulkDeleteResult *stats,
Page cpage; Page cpage;
OffsetNumber coffno; OffsetNumber coffno;
OffsetNumber cmaxoffno; OffsetNumber cmaxoffno;
BlockNumber startPages[MaxOffsetNumber]; BlockNumber listPages[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));
startPages[coffno - FirstOffsetNumber] = list->startPage; listPages[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 = startPages[coffno - FirstOffsetNumber]; BlockNumber searchPage = listPages[coffno - FirstOffsetNumber];
BlockNumber insertPage = InvalidBlockNumber; BlockNumber insertPage = InvalidBlockNumber;
/* Iterate over entry pages */ /* Iterate over entry pages */

View File

@@ -99,6 +99,32 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <+> (SELECT NULL::vector)) t2
4 4
(1 row) (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; DROP TABLE t;
-- unlogged -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); 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) ERROR: 0 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000)
SET hnsw.ef_search = 1001; SET hnsw.ef_search = 1001;
ERROR: 1001 is outside the valid range for parameter "hnsw.ef_search" (1 .. 1000) 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; DROP TABLE t;

View File

@@ -81,6 +81,37 @@ SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2
3 3
(1 row) (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; DROP TABLE t;
-- unlogged -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); CREATE UNLOGGED TABLE t (val vector(3));
@@ -109,4 +140,27 @@ SHOW ivfflat.probes;
1 1
(1 row) (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; 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; 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 -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); CREATE UNLOGGED TABLE t (val vector(3));
@@ -81,4 +98,16 @@ SHOW hnsw.ef_search;
SET hnsw.ef_search = 0; SET hnsw.ef_search = 0;
SET hnsw.ef_search = 1001; 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; 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; 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 -- unlogged
CREATE UNLOGGED TABLE t (val vector(3)); 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; 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; DROP TABLE t;

View File

@@ -6,13 +6,7 @@ use Test::More;
my $dim = 3; my $dim = 3;
my @r = (); my $array_sql = join(",", ('random()') x $dim);
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');
@@ -23,19 +17,20 @@ $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 % 10, ARRAY[$array_sql] FROM generate_series(1, 100000) i;" "INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
); );
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);"); $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", $node->safe_psql("postgres", "INSERT INTO tst SELECT * FROM tmp;");
"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

@@ -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

@@ -26,25 +26,26 @@ $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.streaming = on; SET hnsw.iterative_scan = relaxed_order;
SET work_mem = '8MB'; 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; 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); is($count, 10);
foreach ((30000, 50000, 70000)) foreach ((30000, 50000, 70000))
{ {
my $ef_stream = $_; my $max_tuples = $_;
my $expected = $ef_stream / 10000; my $expected = $max_tuples / 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.streaming = on; SET hnsw.iterative_scan = relaxed_order;
SET hnsw.ef_stream = $ef_stream; SET hnsw.max_scan_tuples = $max_tuples;
SET work_mem = '8MB'; 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; 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; $sum += $count;
@@ -57,11 +58,11 @@ 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.streaming = on; SET hnsw.iterative_scan = relaxed_order;
SET client_min_messages = debug1; SET client_min_messages = debug1;
SET work_mem = '2MB'; 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; 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 exceeded work_mem after \d+ tuples/); like($stderr, qr/hnsw index scan reached memory limit after \d+ tuples/);
done_testing(); done_testing();

View File

@@ -10,18 +10,18 @@ my @expected;
my $limit = 20; my $limit = 20;
my $dim = 3; my $dim = 3;
my $array_sql = join(",", ('random()') x $dim); my $array_sql = join(",", ('random()') x $dim);
my @cs = (100, 1000); my @cs = (50, 500);
sub test_recall sub test_recall
{ {
my ($c, $ef_search, $min, $operator) = @_; my ($c, $ef_search, $min, $operator, $mode) = @_;
my $correct = 0; my $correct = 0;
my $total = 0; my $total = 0;
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.streaming = on; SET hnsw.iterative_scan = $mode;
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.streaming = on; SET hnsw.iterative_scan = $mode;
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);
@@ -50,7 +50,7 @@ sub test_recall
$total += $limit; $total += $limit;
} }
cmp_ok($correct / $total, ">=", $min, $operator); cmp_ok($correct / $total, ">=", $min, "$operator $mode $c");
} }
# Initialize node # Initialize node
@@ -62,7 +62,7 @@ $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, ARRAY[$array_sql] FROM generate_series(1, 50000) i;"
); );
# Generate queries # Generate queries
@@ -108,21 +108,8 @@ for my $i (0 .. $#operators)
push(@expected, $res); push(@expected, $res);
} }
if ($c == 100) test_recall($c, 40, 0.99, $operator, "strict_order");
{ test_recall($c, 40, 0.99, $operator, "relaxed_order");
test_recall($c, 40, 0.99, $operator);
}
else
{
if ($operator eq "<->")
{
test_recall($c, 40, 0.99, $operator);
}
else
{
test_recall($c, 40, 0.99, $operator);
}
}
} }
$node->safe_psql("postgres", "DROP INDEX idx;"); $node->safe_psql("postgres", "DROP INDEX idx;");