mirror of
https://github.com/pgvector/pgvector.git
synced 2026-07-22 12:07:34 +08:00
Compare commits
37 Commits
hnsw-strea
...
hnsw-strea
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dd127da62a | ||
|
|
9c59d09a97 | ||
|
|
c5b6769bed | ||
|
|
d7152448f9 | ||
|
|
08420b6404 | ||
|
|
3fba403f8f | ||
|
|
f53e65fb6a | ||
|
|
12015dbe53 | ||
|
|
794ce3a54c | ||
|
|
be51204a45 | ||
|
|
e78ce02dc9 | ||
|
|
7412ee6cee | ||
|
|
ba0196ba10 | ||
|
|
64b0e6359b | ||
|
|
0e83213e88 | ||
|
|
f1c7494c35 | ||
|
|
e07d17fe38 | ||
|
|
96ff1b992f | ||
|
|
569bf93355 | ||
|
|
72aa53ef54 | ||
|
|
4773931b11 | ||
|
|
b7564dfc19 | ||
|
|
7dc1b36adc | ||
|
|
5266c208a6 | ||
|
|
a0fddf2d1d | ||
|
|
63d5d121a3 | ||
|
|
39d2ef624a | ||
|
|
aa0b7ddf70 | ||
|
|
d499ead3c6 | ||
|
|
f20f5e28b8 | ||
|
|
3dde18a883 | ||
|
|
0af1c7dd55 | ||
|
|
f6ebc5d708 | ||
|
|
c6493415b2 | ||
|
|
aaff3de409 | ||
|
|
4d1c6ff955 | ||
|
|
af1727775d |
@@ -1,8 +1,6 @@
|
||||
## 0.8.0 (unreleased)
|
||||
|
||||
- Added support for iterative index scans
|
||||
- Added casts for arrays to `sparsevec`
|
||||
- Improved cost estimation
|
||||
- Reduced memory usage for HNSW index scans
|
||||
- Dropped support for Postgres 12
|
||||
|
||||
|
||||
57
README.md
57
README.md
@@ -445,63 +445,6 @@ 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);
|
||||
```
|
||||
|
||||
## Streaming Queries [unreleased]
|
||||
|
||||
*Added in 0.8.0*
|
||||
|
||||
With approximate indexes, you can end up with less results than expected due to filtering conditions in the query.
|
||||
|
||||
Starting with 0.8.0, you can enable streaming queries. If too few results from the initial index scan match the query filters, it will resume scanning until enough results are found. This can significantly improve recall (at the cost of speed).
|
||||
|
||||
```tsql
|
||||
SET hnsw.streaming = on;
|
||||
-- or
|
||||
SET ivfflat.streaming = on;
|
||||
```
|
||||
|
||||
### Streaming Options
|
||||
|
||||
Since scanning a large portion of the index is expensive, there are options to control when the scan ends.
|
||||
|
||||
#### HNSW
|
||||
|
||||
Specify the max number of additional tuples visited
|
||||
|
||||
```sql
|
||||
SET hnsw.ef_stream = 10000;
|
||||
```
|
||||
|
||||
The scan will also end if reaches `work_mem`, at which point a notice is shown
|
||||
|
||||
```text
|
||||
NOTICE: hnsw index scan exceeded work_mem after 50000 tuples
|
||||
HINT: Increase work_mem to scan more tuples.
|
||||
```
|
||||
|
||||
Adjust this with:
|
||||
|
||||
```sql
|
||||
SET work_mem = '8MB';
|
||||
```
|
||||
|
||||
#### IVFFlat
|
||||
|
||||
Specify the max number of probes
|
||||
|
||||
```sql
|
||||
SET ivfflat.max_probes = 100;
|
||||
```
|
||||
|
||||
### Streaming Order
|
||||
|
||||
With streaming queries, it’s possible for rows to be slightly out of order by distance. For strict ordering, use:
|
||||
|
||||
```sql
|
||||
WITH approx_order AS MATERIALIZED (
|
||||
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
|
||||
) SELECT * FROM approx_order ORDER BY distance;
|
||||
```
|
||||
|
||||
## Half-Precision Vectors
|
||||
|
||||
*Added in 0.7.0*
|
||||
|
||||
@@ -159,6 +159,24 @@ CheckStateArray(ArrayType *statearray, const char *caller)
|
||||
return (float8 *) ARR_DATA_PTR(statearray);
|
||||
}
|
||||
|
||||
#if PG_VERSION_NUM < 120003
|
||||
static pg_noinline void
|
||||
float_overflow_error(void)
|
||||
{
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("value out of range: overflow")));
|
||||
}
|
||||
|
||||
static pg_noinline void
|
||||
float_underflow_error(void)
|
||||
{
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("value out of range: underflow")));
|
||||
}
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Convert textual representation to internal representation
|
||||
*/
|
||||
|
||||
92
src/hnsw.c
92
src/hnsw.c
@@ -12,14 +12,12 @@
|
||||
#include "utils/float.h"
|
||||
#include "utils/guc.h"
|
||||
#include "utils/selfuncs.h"
|
||||
#include "utils/spccache.h"
|
||||
|
||||
#if PG_VERSION_NUM < 150000
|
||||
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
|
||||
#endif
|
||||
|
||||
int hnsw_ef_search;
|
||||
int hnsw_ef_stream;
|
||||
bool hnsw_streaming;
|
||||
int hnsw_lock_tranche_id;
|
||||
static relopt_kind hnsw_relopt_kind;
|
||||
@@ -76,11 +74,7 @@ HnswInit(void)
|
||||
NULL, &hnsw_streaming,
|
||||
HNSW_DEFAULT_STREAMING, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||
|
||||
/* TODO Figure out name */
|
||||
/* TODO Use same value as ivfflat.max_probes for "all" */
|
||||
DefineCustomIntVariable("hnsw.ef_stream", "Sets the max number of additional candidates to visit for streaming search",
|
||||
"-1 means all", &hnsw_ef_stream,
|
||||
HNSW_DEFAULT_EF_STREAM, HNSW_MIN_EF_STREAM, HNSW_MAX_EF_STREAM, PGC_USERSET, 0, NULL, NULL, NULL);
|
||||
/* TODO Add option for limiting iterative search */
|
||||
|
||||
MarkGUCPrefixReserved("hnsw");
|
||||
}
|
||||
@@ -102,33 +96,6 @@ hnswbuildphasename(int64 phasenum)
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
* Estimate ef needed for iterative scans
|
||||
*/
|
||||
static int
|
||||
EstimateEf(PlannerInfo *root, IndexPath *path)
|
||||
{
|
||||
double selectivity = 1;
|
||||
ListCell *lc;
|
||||
|
||||
/* Cannot estimate without limit */
|
||||
/* limit_tuples includes offset */
|
||||
if (root->limit_tuples < 0)
|
||||
return 0;
|
||||
|
||||
/* Get the selectivity of non-index conditions */
|
||||
foreach(lc, path->indexinfo->indrestrictinfo)
|
||||
{
|
||||
RestrictInfo *rinfo = lfirst(lc);
|
||||
|
||||
/* Skip DEFAULT_INEQ_SEL since it may be a distance filter */
|
||||
if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
|
||||
selectivity *= rinfo->norm_selec;
|
||||
}
|
||||
|
||||
return root->limit_tuples / Max(selectivity, 0.00001);
|
||||
}
|
||||
|
||||
/*
|
||||
* Estimate the cost of an index scan
|
||||
*/
|
||||
@@ -140,12 +107,7 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
||||
{
|
||||
GenericCosts costs;
|
||||
int m;
|
||||
int ef;
|
||||
int entryLevel;
|
||||
int layer0TuplesMax;
|
||||
double layer0Selectivity;
|
||||
double scalingFactor = 0.55;
|
||||
double spc_seq_page_cost;
|
||||
Relation index;
|
||||
|
||||
/* Never use index without order */
|
||||
@@ -165,57 +127,17 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
||||
HnswGetMetaPageInfo(index, &m, NULL);
|
||||
index_close(index, NoLock);
|
||||
|
||||
ef = hnsw_streaming ? Max(hnsw_ef_search, EstimateEf(root, path)) : hnsw_ef_search;
|
||||
/* Approximate entry level */
|
||||
entryLevel = (int) -log(1.0 / path->indexinfo->tuples) * HnswGetMl(m);
|
||||
|
||||
/*
|
||||
* HNSW cost estimation follows a formula that accounts for the total
|
||||
* number of tuples indexed combined with the parameters that most
|
||||
* influence the duration of the index scan, namely: m - the number of
|
||||
* tuples that are scanned in each step of the HNSW graph traversal
|
||||
* ef_search - which influences the total number of steps taken at layer 0
|
||||
*
|
||||
* The source of the vector data can impact how many steps it takes to
|
||||
* converge on the set of vectors to return to the executor. Currently, we
|
||||
* use a hardcoded scaling factor (HNSWScanScalingFactor) to help
|
||||
* influence that, but this could later become a configurable parameter
|
||||
* based on the cost estimations.
|
||||
*
|
||||
* The tuple estimator formula is below:
|
||||
*
|
||||
* numIndexTuples = entryLevel * m + layer0TuplesMax * layer0Selectivity
|
||||
*
|
||||
* "entryLevel * m" represents the floor of tuples we need to scan to get
|
||||
* to layer 0 (L0).
|
||||
*
|
||||
* "layer0TuplesMax" is the estimated total number of tuples we'd scan at
|
||||
* L0 if we weren't discarding already visited tuples as part of the scan.
|
||||
*
|
||||
* "layer0Selectivity" estimates the percentage of tuples that are scanned
|
||||
* at L0, accounting for previously visited tuples, multiplied by the
|
||||
* "scalingFactor" (currently hardcoded).
|
||||
*/
|
||||
entryLevel = (int) (log(path->indexinfo->tuples + 1) * HnswGetMl(m));
|
||||
layer0TuplesMax = HnswGetLayerM(m, 0) * ef;
|
||||
layer0Selectivity = (scalingFactor * log(path->indexinfo->tuples + 1)) /
|
||||
(log(m) * (1 + log(ef)));
|
||||
/* TODO Improve estimate of visited tuples (currently underestimates) */
|
||||
/* Account for number of tuples (or entry level), m, and ef_search */
|
||||
costs.numIndexTuples = (entryLevel + 2) * m;
|
||||
|
||||
costs.numIndexTuples = (entryLevel * m) +
|
||||
(layer0TuplesMax * layer0Selectivity);
|
||||
/* TODO Adjust for selectivity for iterative scans */
|
||||
|
||||
genericcostestimate(root, path, loop_count, &costs);
|
||||
|
||||
get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
|
||||
|
||||
/* Adjust cost if needed since TOAST not included in seq scan cost */
|
||||
if (costs.numIndexPages > path->indexinfo->rel->pages)
|
||||
{
|
||||
/* Change all page cost from random to sequential */
|
||||
costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
|
||||
/* Remove cost of extra pages */
|
||||
costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
|
||||
}
|
||||
|
||||
/* Use total cost since most work happens before first tuple is returned */
|
||||
*indexStartupCost = costs.indexTotalCost;
|
||||
*indexTotalCost = costs.indexTotalCost;
|
||||
|
||||
@@ -47,9 +47,6 @@
|
||||
#define HNSW_MIN_EF_SEARCH 1
|
||||
#define HNSW_MAX_EF_SEARCH 1000
|
||||
#define HNSW_DEFAULT_STREAMING false
|
||||
#define HNSW_DEFAULT_EF_STREAM -1
|
||||
#define HNSW_MIN_EF_STREAM -1
|
||||
#define HNSW_MAX_EF_STREAM INT_MAX
|
||||
|
||||
/* Tuple types */
|
||||
#define HNSW_ELEMENT_TUPLE_TYPE 1
|
||||
@@ -129,7 +126,6 @@
|
||||
|
||||
/* Variables */
|
||||
extern int hnsw_ef_search;
|
||||
extern int hnsw_ef_stream;
|
||||
extern bool hnsw_streaming;
|
||||
extern int hnsw_lock_tranche_id;
|
||||
|
||||
@@ -371,7 +367,6 @@ typedef struct HnswScanOpaqueData
|
||||
Datum q;
|
||||
int m;
|
||||
int64 tuples;
|
||||
double previousDistance;
|
||||
MemoryContext tmpCtx;
|
||||
|
||||
/* Support functions */
|
||||
@@ -417,7 +412,7 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
|
||||
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
|
||||
void HnswInitPage(Buffer buf, Page page);
|
||||
void HnswInit(void);
|
||||
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples);
|
||||
List *HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited);
|
||||
HnswElement HnswGetEntryPoint(Relation index);
|
||||
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
|
||||
void *HnswAlloc(HnswAllocator * allocator, Size size);
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
#include "postgres.h"
|
||||
|
||||
#include <float.h>
|
||||
|
||||
#include "access/relscan.h"
|
||||
#include "hnsw.h"
|
||||
#include "pgstat.h"
|
||||
@@ -38,11 +36,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
|
||||
|
||||
for (int lc = entryPoint->level; lc >= 1; lc--)
|
||||
{
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, false, NULL, NULL, NULL, true);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, hnsw_streaming ? &so->discarded : NULL, true, &so->tuples);
|
||||
return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, &so->v, &so->discarded, true);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -75,7 +73,7 @@ ResumeScanItems(IndexScanDesc scan)
|
||||
ep = lappend(ep, hc);
|
||||
}
|
||||
|
||||
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false, &so->tuples);
|
||||
return HnswSearchLayer(base, so->q, ep, batch_size, 0, index, procinfo, collation, so->m, false, NULL, &so->v, &so->discarded, false);
|
||||
}
|
||||
|
||||
/*
|
||||
@@ -119,8 +117,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
|
||||
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
|
||||
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);
|
||||
@@ -143,15 +139,13 @@ hnswrescan(IndexScanDesc scan, ScanKey keys, int nkeys, ScanKey orderbys, int no
|
||||
{
|
||||
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
|
||||
|
||||
if (so->v.tids != NULL)
|
||||
tidhash_reset(so->v.tids);
|
||||
|
||||
if (so->discarded != NULL)
|
||||
if (!so->first)
|
||||
{
|
||||
pairingheap_reset(so->discarded);
|
||||
|
||||
tidhash_reset(so->v.tids);
|
||||
}
|
||||
so->first = true;
|
||||
so->tuples = 0;
|
||||
so->previousDistance = -INFINITY;
|
||||
MemoryContextReset(so->tmpCtx);
|
||||
|
||||
if (keys && scan->numberOfKeys > 0)
|
||||
@@ -225,26 +219,13 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
if (!hnsw_streaming)
|
||||
break;
|
||||
|
||||
/* Empty index */
|
||||
if (so->discarded == NULL)
|
||||
break;
|
||||
|
||||
/* Reached max number of additional tuples */
|
||||
if (hnsw_ef_stream != -1 && so->tuples >= hnsw_ef_search + hnsw_ef_stream)
|
||||
{
|
||||
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) > (Size) work_mem * 1024L)
|
||||
if (MemoryContextMemAllocated(so->tmpCtx, false) > (Size) work_mem * 1024L)
|
||||
{
|
||||
if (pairingheap_is_empty(so->discarded))
|
||||
{
|
||||
ereport(NOTICE,
|
||||
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
|
||||
(errmsg("hnsw iterative search exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
|
||||
errhint("Increase work_mem to scan more tuples.")));
|
||||
|
||||
break;
|
||||
@@ -297,13 +278,10 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
continue;
|
||||
}
|
||||
|
||||
so->tuples++;
|
||||
|
||||
heaptid = &element->heaptids[--element->heaptidsLength];
|
||||
|
||||
if (hc->distance < so->previousDistance)
|
||||
continue;
|
||||
|
||||
so->previousDistance = hc->distance;
|
||||
|
||||
MemoryContextSwitchTo(oldCtx);
|
||||
|
||||
scan->xs_heaptid = *heaptid;
|
||||
|
||||
@@ -776,7 +776,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
|
||||
* 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)
|
||||
if (ntup->version != element->version)
|
||||
{
|
||||
UnlockReleaseBuffer(buf);
|
||||
return;
|
||||
@@ -807,7 +807,7 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
|
||||
* Algorithm 2 from paper
|
||||
*/
|
||||
List *
|
||||
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited, int64 *tuples)
|
||||
HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement, visited_hash * v, pairingheap **discarded, bool initVisited)
|
||||
{
|
||||
List *w = NIL;
|
||||
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
|
||||
@@ -849,13 +849,8 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
||||
bool found;
|
||||
|
||||
if (initVisited)
|
||||
{
|
||||
AddToVisited(base, v, hc->element, index, &found);
|
||||
|
||||
if (tuples != NULL)
|
||||
(*tuples)++;
|
||||
}
|
||||
|
||||
pairingheap_add(C, &hc->c_node);
|
||||
pairingheap_add(W, &hc->w_node);
|
||||
|
||||
@@ -884,9 +879,6 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
|
||||
else
|
||||
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc);
|
||||
|
||||
if (tuples != NULL)
|
||||
(*tuples) += unvisitedLength;
|
||||
|
||||
for (int i = 0; i < unvisitedLength; i++)
|
||||
{
|
||||
HnswElement eElement;
|
||||
@@ -1326,7 +1318,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, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
|
||||
ep = w;
|
||||
}
|
||||
|
||||
@@ -1345,7 +1337,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
|
||||
List *lw = NIL;
|
||||
ListCell *lc2;
|
||||
|
||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL);
|
||||
w = HnswSearchLayer(base, q, ep, efConstruction, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true);
|
||||
|
||||
/* Convert search candidates to candidates */
|
||||
foreach(lc2, w)
|
||||
|
||||
@@ -120,6 +120,13 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
|
||||
costs.indexTotalCost -= 0.5 * costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
|
||||
}
|
||||
|
||||
/*
|
||||
* If the list selectivity is lower than what is returned from the generic
|
||||
* cost estimator, use that.
|
||||
*/
|
||||
if (ratio < costs.indexSelectivity)
|
||||
costs.indexSelectivity = ratio;
|
||||
|
||||
/* Use total cost since most work happens before first tuple is returned */
|
||||
*indexStartupCost = costs.indexTotalCost;
|
||||
*indexTotalCost = costs.indexTotalCost;
|
||||
|
||||
@@ -253,9 +253,8 @@ typedef struct IvfflatScanOpaqueData
|
||||
/* Sorting */
|
||||
Tuplesortstate *sortstate;
|
||||
TupleDesc tupdesc;
|
||||
TupleTableSlot *vslot;
|
||||
TupleTableSlot *mslot;
|
||||
BufferAccessStrategy bas;
|
||||
TupleTableSlot *slot;
|
||||
bool isnull;
|
||||
|
||||
/* Support functions */
|
||||
FmgrInfo *procinfo;
|
||||
|
||||
@@ -15,19 +15,16 @@
|
||||
#include "utils/memutils.h"
|
||||
#endif
|
||||
|
||||
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
|
||||
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
|
||||
|
||||
/*
|
||||
* Compare list distances
|
||||
*/
|
||||
static int
|
||||
CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg)
|
||||
{
|
||||
if (GetScanListConst(a)->distance > GetScanListConst(b)->distance)
|
||||
if (((const IvfflatScanList *) a)->distance > ((const IvfflatScanList *) b)->distance)
|
||||
return 1;
|
||||
|
||||
if (GetScanListConst(a)->distance < GetScanListConst(b)->distance)
|
||||
if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
|
||||
return -1;
|
||||
|
||||
return 0;
|
||||
@@ -79,14 +76,14 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
||||
|
||||
/* Calculate max distance */
|
||||
if (listCount == so->probes)
|
||||
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
|
||||
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
|
||||
}
|
||||
else if (distance < maxDistance)
|
||||
{
|
||||
IvfflatScanList *scanlist;
|
||||
|
||||
/* Remove */
|
||||
scanlist = GetScanList(pairingheap_remove_first(so->listQueue));
|
||||
scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
|
||||
|
||||
/* Reuse */
|
||||
scanlist->startPage = list->startPage;
|
||||
@@ -94,7 +91,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
|
||||
pairingheap_add(so->listQueue, &scanlist->ph_node);
|
||||
|
||||
/* Update max distance */
|
||||
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance;
|
||||
maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -113,12 +110,19 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
|
||||
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
|
||||
double tuples = 0;
|
||||
TupleTableSlot *slot = so->vslot;
|
||||
TupleTableSlot *slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
|
||||
|
||||
/*
|
||||
* Reuse same set of shared buffers for scan
|
||||
*
|
||||
* See postgres/src/backend/storage/buffer/README for description
|
||||
*/
|
||||
BufferAccessStrategy bas = GetAccessStrategy(BAS_BULKREAD);
|
||||
|
||||
/* Search closest probes lists */
|
||||
while (!pairingheap_is_empty(so->listQueue))
|
||||
{
|
||||
BlockNumber searchPage = GetScanList(pairingheap_remove_first(so->listQueue))->startPage;
|
||||
BlockNumber searchPage = ((IvfflatScanList *) pairingheap_remove_first(so->listQueue))->startPage;
|
||||
|
||||
/* Search all entry pages for list */
|
||||
while (BlockNumberIsValid(searchPage))
|
||||
@@ -127,7 +131,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
Page page;
|
||||
OffsetNumber maxoffno;
|
||||
|
||||
buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, so->bas);
|
||||
buf = ReadBufferExtended(scan->indexRelation, MAIN_FORKNUM, searchPage, RBM_NORMAL, bas);
|
||||
LockBuffer(buf, BUFFER_LOCK_SHARE);
|
||||
page = BufferGetPage(buf);
|
||||
maxoffno = PageGetMaxOffsetNumber(page);
|
||||
@@ -166,6 +170,8 @@ GetScanItems(IndexScanDesc scan, Datum value)
|
||||
}
|
||||
}
|
||||
|
||||
FreeAccessStrategy(bas);
|
||||
|
||||
if (tuples < 100)
|
||||
ereport(DEBUG1,
|
||||
(errmsg("index scan found few tuples"),
|
||||
@@ -268,16 +274,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
|
||||
/* Prep sort */
|
||||
so->sortstate = InitScanSortState(so->tupdesc);
|
||||
|
||||
/* Need separate slots for puttuple and gettuple */
|
||||
so->vslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsVirtual);
|
||||
so->mslot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
|
||||
|
||||
/*
|
||||
* Reuse same set of shared buffers for scan
|
||||
*
|
||||
* See postgres/src/backend/storage/buffer/README for description
|
||||
*/
|
||||
so->bas = GetAccessStrategy(BAS_BULKREAD);
|
||||
so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
|
||||
|
||||
so->listQueue = pairingheap_allocate(CompareLists, scan);
|
||||
|
||||
@@ -351,10 +348,9 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
|
||||
pfree(DatumGetPointer(value));
|
||||
}
|
||||
|
||||
if (tuplesort_gettupleslot(so->sortstate, true, false, so->mslot, NULL))
|
||||
if (tuplesort_gettupleslot(so->sortstate, true, false, so->slot, NULL))
|
||||
{
|
||||
bool isnull;
|
||||
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
|
||||
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->slot, 2, &so->isnull));
|
||||
|
||||
scan->xs_heaptid = *heaptid;
|
||||
scan->xs_recheck = false;
|
||||
@@ -375,10 +371,6 @@ ivfflatendscan(IndexScanDesc scan)
|
||||
|
||||
pairingheap_free(so->listQueue);
|
||||
tuplesort_end(so->sortstate);
|
||||
FreeAccessStrategy(so->bas);
|
||||
FreeTupleDesc(so->tupdesc);
|
||||
|
||||
/* TODO Free vslot and mslot without freeing TupleDesc */
|
||||
|
||||
pfree(so);
|
||||
scan->opaque = NULL;
|
||||
|
||||
18
src/vector.c
18
src/vector.c
@@ -155,6 +155,24 @@ CheckStateArray(ArrayType *statearray, const char *caller)
|
||||
return (float8 *) ARR_DATA_PTR(statearray);
|
||||
}
|
||||
|
||||
#if PG_VERSION_NUM < 120003
|
||||
static pg_noinline void
|
||||
float_overflow_error(void)
|
||||
{
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("value out of range: overflow")));
|
||||
}
|
||||
|
||||
static pg_noinline void
|
||||
float_underflow_error(void)
|
||||
{
|
||||
ereport(ERROR,
|
||||
(errcode(ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE),
|
||||
errmsg("value out of range: underflow")));
|
||||
}
|
||||
#endif
|
||||
|
||||
/*
|
||||
* Convert textual representation to internal representation
|
||||
*/
|
||||
|
||||
@@ -94,7 +94,8 @@ like($explain, qr/Seq Scan/);
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
|
||||
));
|
||||
like($explain, qr/Seq Scan/);
|
||||
# TODO Do not use index
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test attribute index
|
||||
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
|
||||
@@ -109,6 +110,7 @@ $node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v v
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using partial_idx/);
|
||||
# TODO Use partial index
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
done_testing();
|
||||
|
||||
@@ -18,13 +18,9 @@ $node->start;
|
||||
# Create table and index
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4, t text);");
|
||||
$node->safe_psql("postgres", "CREATE TABLE cat (i int4 PRIMARY KEY, t text, b boolean);");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) i;"
|
||||
);
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO cat SELECT i, 'cat ' || i, i % 5 = 0 FROM generate_series(1, $nc) i;"
|
||||
);
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
|
||||
$node->safe_psql("postgres", "ANALYZE tst;");
|
||||
|
||||
@@ -100,25 +96,13 @@ $explain = $node->safe_psql("postgres", qq(
|
||||
));
|
||||
like($explain, qr/Seq Scan/);
|
||||
|
||||
# Test join
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT cat.t FROM cat INNER JOIN tst ON cat.i = tst.c ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test join with attribute filtering
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT cat.t FROM cat INNER JOIN tst ON cat.i = tst.c WHERE cat.b = 't' ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test attribute index
|
||||
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
|
||||
$explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
# Use attribute index
|
||||
like($explain, qr/Bitmap Index Scan on attribute_idx/);
|
||||
# TODO Use attribute index
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
# Test partial index
|
||||
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);");
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
use strict;
|
||||
use warnings FATAL => 'all';
|
||||
use PostgreSQL::Test::Cluster;
|
||||
use PostgreSQL::Test::Utils;
|
||||
use Test::More;
|
||||
|
||||
my @dims = (384, 1536);
|
||||
my $limit = 10;
|
||||
|
||||
# Initialize node
|
||||
my $node = PostgreSQL::Test::Cluster->new('node');
|
||||
$node->init;
|
||||
$node->start;
|
||||
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
|
||||
for my $dim (@dims)
|
||||
{
|
||||
my $array_sql = join(",", ('random()') x $dim);
|
||||
my $n = $dim == 384 ? 2000 : 1000;
|
||||
|
||||
# Create table and index
|
||||
$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, $n) i;"
|
||||
);
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
|
||||
$node->safe_psql("postgres", "ANALYZE tst;");
|
||||
|
||||
# Generate query
|
||||
my @r = ();
|
||||
for (1 .. $dim)
|
||||
{
|
||||
push(@r, rand());
|
||||
}
|
||||
my $query = "[" . join(",", @r) . "]";
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
$node->safe_psql("postgres", "DROP TABLE tst;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
@@ -14,7 +14,7 @@ $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", "CREATE TABLE tst (i int4, v vector($dim));");
|
||||
$node->safe_psql("postgres",
|
||||
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
|
||||
);
|
||||
@@ -32,35 +32,12 @@ my $count = $node->safe_psql("postgres", qq(
|
||||
));
|
||||
is($count, 10);
|
||||
|
||||
foreach ((30000, 50000, 70000))
|
||||
{
|
||||
my $ef_stream = $_;
|
||||
my $expected = $ef_stream / 10000;
|
||||
my $sum = 0;
|
||||
|
||||
for my $i (1 .. 20)
|
||||
{
|
||||
$count = $node->safe_psql("postgres", qq(
|
||||
SET enable_seqscan = off;
|
||||
SET hnsw.streaming = on;
|
||||
SET hnsw.ef_stream = $ef_stream;
|
||||
SET work_mem = '8MB';
|
||||
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst WHERE i = $i) LIMIT 11) t;
|
||||
));
|
||||
$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.streaming = on;
|
||||
SET work_mem = '2MB';
|
||||
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
|
||||
));
|
||||
like($stderr, qr/hnsw index scan exceeded work_mem after \d+ tuples/);
|
||||
like($stderr, qr/iterative search exceeded work_mem after \d+ tuples/);
|
||||
|
||||
done_testing();
|
||||
@@ -1,45 +0,0 @@
|
||||
use strict;
|
||||
use warnings FATAL => 'all';
|
||||
use PostgreSQL::Test::Cluster;
|
||||
use PostgreSQL::Test::Utils;
|
||||
use Test::More;
|
||||
|
||||
my @dims = (384, 1536);
|
||||
my $limit = 10;
|
||||
|
||||
# Initialize node
|
||||
my $node = PostgreSQL::Test::Cluster->new('node');
|
||||
$node->init;
|
||||
$node->start;
|
||||
|
||||
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
|
||||
|
||||
for my $dim (@dims)
|
||||
{
|
||||
my $array_sql = join(",", ('random()') x $dim);
|
||||
|
||||
# Create table and index
|
||||
$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, 5000) i;"
|
||||
);
|
||||
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5);");
|
||||
$node->safe_psql("postgres", "ANALYZE tst;");
|
||||
|
||||
# Generate query
|
||||
my @r = ();
|
||||
for (1 .. $dim)
|
||||
{
|
||||
push(@r, rand());
|
||||
}
|
||||
my $query = "[" . join(",", @r) . "]";
|
||||
|
||||
my $explain = $node->safe_psql("postgres", qq(
|
||||
EXPLAIN ANALYZE SELECT i FROM tst ORDER BY v <-> '$query' LIMIT $limit;
|
||||
));
|
||||
like($explain, qr/Index Scan using idx/);
|
||||
|
||||
$node->safe_psql("postgres", "DROP TABLE tst;");
|
||||
}
|
||||
|
||||
done_testing();
|
||||
Reference in New Issue
Block a user