Compare commits

..

2 Commits

Author SHA1 Message Date
Andrew Kane
dad5dffe8e Fixed issue 2024-09-20 22:29:23 -07:00
Andrew Kane
06b641407f Added time budget for HNSW index scans 2024-09-20 22:06:09 -07:00
19 changed files with 183 additions and 800 deletions

View File

@@ -1,8 +1,6 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- Added support for iterative index scans
- Added casts for arrays to `sparsevec` - Added casts for arrays to `sparsevec`
- Improved cost estimation
- Reduced memory usage for HNSW index scans - Reduced memory usage for HNSW index scans
- Dropped support for Postgres 12 - Dropped support for Postgres 12

View File

@@ -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); 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, its possible for rows to be slightly out of order by distance. For strict ordering, use:
```sql
WITH approx_order AS MATERIALIZED (
SELECT *, embedding <-> '[1,2,3]' AS distance FROM items WHERE ... ORDER BY distance LIMIT 5
) SELECT * FROM approx_order ORDER BY distance;
```
## Half-Precision Vectors ## Half-Precision Vectors
*Added in 0.7.0* *Added in 0.7.0*

View File

@@ -159,6 +159,24 @@ CheckStateArray(ArrayType *statearray, const char *caller)
return (float8 *) ARR_DATA_PTR(statearray); 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 * Convert textual representation to internal representation
*/ */

View File

@@ -12,15 +12,13 @@
#include "utils/float.h" #include "utils/float.h"
#include "utils/guc.h" #include "utils/guc.h"
#include "utils/selfuncs.h" #include "utils/selfuncs.h"
#include "utils/spccache.h"
#if PG_VERSION_NUM < 150000 #if PG_VERSION_NUM < 150000
#define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x) #define MarkGUCPrefixReserved(x) EmitWarningsOnPlaceholders(x)
#endif #endif
int hnsw_ef_search; int hnsw_ef_search;
int hnsw_ef_stream; int hnsw_time_budget;
bool hnsw_streaming;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -71,16 +69,9 @@ HnswInit(void)
"Valid range is 1..1000.", &hnsw_ef_search, "Valid range is 1..1000.", &hnsw_ef_search,
HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL); HNSW_DEFAULT_EF_SEARCH, HNSW_MIN_EF_SEARCH, HNSW_MAX_EF_SEARCH, PGC_USERSET, 0, NULL, NULL, NULL);
/* TODO Figure out name */ DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search",
DefineCustomBoolVariable("hnsw.streaming", "Use streaming mode", NULL, &hnsw_time_budget,
NULL, &hnsw_streaming, -1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL);
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);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
} }
@@ -102,33 +93,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 * Estimate the cost of an index scan
*/ */
@@ -140,12 +104,7 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
{ {
GenericCosts costs; GenericCosts costs;
int m; int m;
int ef;
int entryLevel; int entryLevel;
int layer0TuplesMax;
double layer0Selectivity;
double scalingFactor = 0.55;
double spc_seq_page_cost;
Relation index; Relation index;
/* Never use index without order */ /* Never use index without order */
@@ -165,57 +124,15 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
HnswGetMetaPageInfo(index, &m, NULL); HnswGetMetaPageInfo(index, &m, NULL);
index_close(index, NoLock); 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);
/* /* TODO Improve estimate of visited tuples (currently underestimates) */
* HNSW cost estimation follows a formula that accounts for the total /* Account for number of tuples (or entry level), m, and ef_search */
* number of tuples indexed combined with the parameters that most costs.numIndexTuples = (entryLevel + 2) * m;
* 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)));
costs.numIndexTuples = (entryLevel * m) +
(layer0TuplesMax * layer0Selectivity);
genericcostestimate(root, path, loop_count, &costs); 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 */ /* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost; *indexStartupCost = costs.indexTotalCost;
*indexTotalCost = costs.indexTotalCost; *indexTotalCost = costs.indexTotalCost;

View File

@@ -12,10 +12,6 @@
#include "utils/sampling.h" #include "utils/sampling.h"
#include "vector.h" #include "vector.h"
#ifdef HNSW_BENCH
#include "portability/instr_time.h"
#endif
#define HNSW_MAX_DIM 2000 #define HNSW_MAX_DIM 2000
#define HNSW_MAX_NNZ 1000 #define HNSW_MAX_NNZ 1000
@@ -46,10 +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_STREAMING false
#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
@@ -76,21 +68,6 @@
#define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page)) #define HnswPageGetOpaque(page) ((HnswPageOpaque) PageGetSpecialPointer(page))
#define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page)) #define HnswPageGetMeta(page) ((HnswMetaPageData *) PageGetContents(page))
#ifdef HNSW_BENCH
#define HnswBench(name, code) \
do { \
instr_time start; \
instr_time duration; \
INSTR_TIME_SET_CURRENT(start); \
(code); \
INSTR_TIME_SET_CURRENT(duration); \
INSTR_TIME_SUBTRACT(duration, start); \
elog(INFO, "%s: %.3f ms", name, INSTR_TIME_GET_MILLISEC(duration)); \
} while (0)
#else
#define HnswBench(name, code) (code)
#endif
#if PG_VERSION_NUM >= 150000 #if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state) #define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed) #define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
@@ -129,8 +106,7 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_ef_stream; extern int hnsw_time_budget;
extern bool hnsw_streaming;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
@@ -154,7 +130,6 @@ struct HnswElementData
uint8 heaptidsLength; uint8 heaptidsLength;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint8 version;
uint32 hash; uint32 hash;
HnswNeighborsPtr neighbors; HnswNeighborsPtr neighbors;
BlockNumber blkno; BlockNumber blkno;
@@ -181,16 +156,12 @@ struct HnswNeighborArray
HnswCandidate items[FLEXIBLE_ARRAY_MEMBER]; HnswCandidate items[FLEXIBLE_ARRAY_MEMBER];
}; };
typedef struct HnswSearchCandidate typedef struct HnswPairingHeapNode
{ {
HnswCandidate *inner;
pairingheap_node c_node; pairingheap_node c_node;
pairingheap_node w_node; pairingheap_node w_node;
HnswElementPtr element; } HnswPairingHeapNode;
float distance;
} HnswSearchCandidate;
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr)
/* HNSW index options */ /* HNSW index options */
typedef struct HnswOptions typedef struct HnswOptions
@@ -335,10 +306,10 @@ typedef struct HnswElementTupleData
uint8 type; uint8 type;
uint8 level; uint8 level;
uint8 deleted; uint8 deleted;
uint8 version; uint8 unused;
ItemPointerData heaptids[HNSW_HEAPTIDS]; ItemPointerData heaptids[HNSW_HEAPTIDS];
ItemPointerData neighbortid; ItemPointerData neighbortid;
uint16 unused; uint16 unused2;
Vector data; Vector data;
} HnswElementTupleData; } HnswElementTupleData;
@@ -347,32 +318,20 @@ typedef HnswElementTupleData * HnswElementTuple;
typedef struct HnswNeighborTupleData typedef struct HnswNeighborTupleData
{ {
uint8 type; uint8 type;
uint8 version; uint8 unused;
uint16 count; uint16 count;
ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER]; ItemPointerData indextids[FLEXIBLE_ARRAY_MEMBER];
} HnswNeighborTupleData; } HnswNeighborTupleData;
typedef HnswNeighborTupleData * HnswNeighborTuple; typedef HnswNeighborTupleData * HnswNeighborTuple;
typedef union
{
struct pointerhash_hash *pointers;
struct offsethash_hash *offsets;
struct tidhash_hash *tids;
} visited_hash;
typedef struct HnswScanOpaqueData typedef struct HnswScanOpaqueData
{ {
const HnswTypeInfo *typeInfo; const HnswTypeInfo *typeInfo;
bool first; bool first;
List *w; List *w;
visited_hash v;
pairingheap *discarded;
Datum q;
int m;
int64 tuples;
double previousDistance;
MemoryContext tmpCtx; MemoryContext tmpCtx;
instr_time start;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;
@@ -417,14 +376,14 @@ bool HnswCheckNorm(FmgrInfo *procinfo, Oid collation, Datum value);
Buffer HnswNewBuffer(Relation index, ForkNumber forkNum); Buffer HnswNewBuffer(Relation index, ForkNumber forkNum);
void HnswInitPage(Buffer buf, Page page); void HnswInitPage(Buffer buf, Page page);
void HnswInit(void); void HnswInit(void);
List *HnswSearchLayer(char *base, 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, instr_time *start);
HnswElement HnswGetEntryPoint(Relation index); HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint); void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
void *HnswAlloc(HnswAllocator * allocator, Size size); void *HnswAlloc(HnswAllocator * allocator, Size size);
HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc); HnswElement HnswInitElement(char *base, ItemPointer tid, int m, double ml, int maxLevel, HnswAllocator * alloc);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno); HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing); void HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswSearchCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec); HnswCandidate *HnswEntryCandidate(char *base, HnswElement em, Datum q, Relation rel, FmgrInfo *procinfo, Oid collation, bool loadVec);
void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building); void HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, BlockNumber insertPage, ForkNumber forkNum, bool building);
void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m); void HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m);
void HnswAddHeapTid(HnswElement element, ItemPointer heaptid); void HnswAddHeapTid(HnswElement element, ItemPointer heaptid);

View File

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

View File

@@ -1,10 +1,9 @@
#include "postgres.h" #include "postgres.h"
#include <float.h>
#include "access/relscan.h" #include "access/relscan.h"
#include "hnsw.h" #include "hnsw.h"
#include "pgstat.h" #include "pgstat.h"
#include "portability/instr_time.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "storage/lmgr.h" #include "storage/lmgr.h"
#include "utils/memutils.h" #include "utils/memutils.h"
@@ -24,13 +23,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
int m; int m;
HnswElement entryPoint; HnswElement entryPoint;
char *base = NULL; char *base = NULL;
instr_time *start = hnsw_time_budget == -1 ? NULL : &so->start;
/* Get m and entry point */ /* Get m and entry point */
HnswGetMetaPageInfo(index, &m, &entryPoint); HnswGetMetaPageInfo(index, &m, &entryPoint);
so->q = q;
so->m = m;
if (entryPoint == NULL) if (entryPoint == NULL)
return NIL; return NIL;
@@ -38,44 +35,11 @@ GetScanItems(IndexScanDesc scan, Datum q)
for (int lc = entryPoint->level; lc >= 1; lc--) 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, start);
ep = w; 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, start);
}
/*
* Resume scan at ground level with discarded candidates
*/
static List *
ResumeScanItems(IndexScanDesc scan)
{
HnswScanOpaque so = (HnswScanOpaque) scan->opaque;
Relation index = scan->indexRelation;
FmgrInfo *procinfo = so->procinfo;
Oid collation = so->collation;
List *ep = NIL;
char *base = NULL;
int batch_size = hnsw_ef_search;
if (pairingheap_is_empty(so->discarded))
return NIL;
/* Get next batch of candidates */
for (int i = 0; i < batch_size; i++)
{
HnswSearchCandidate *hc;
if (pairingheap_is_empty(so->discarded))
break;
hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded));
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);
} }
/* /*
@@ -119,8 +83,6 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData)); so = (HnswScanOpaque) palloc(sizeof(HnswScanOpaqueData));
so->typeInfo = HnswGetTypeInfo(index); so->typeInfo = HnswGetTypeInfo(index);
so->first = true; so->first = true;
so->v.tids = NULL;
so->discarded = NULL;
so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext, so->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw scan temporary context", "Hnsw scan temporary context",
ALLOCSET_DEFAULT_SIZES); ALLOCSET_DEFAULT_SIZES);
@@ -130,6 +92,9 @@ hnswbeginscan(Relation index, int nkeys, int norderbys)
so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC); so->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
so->collation = index->rd_indcollation[0]; so->collation = index->rd_indcollation[0];
if (hnsw_time_budget != -1)
INSTR_TIME_SET_CURRENT(so->start);
scan->opaque = so; scan->opaque = so;
return scan; return scan;
@@ -143,15 +108,7 @@ 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;
so->tuples = 0;
so->previousDistance = -INFINITY;
MemoryContextReset(so->tmpCtx); MemoryContextReset(so->tmpCtx);
if (keys && scan->numberOfKeys > 0) if (keys && scan->numberOfKeys > 0)
@@ -201,7 +158,7 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
*/ */
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = GetScanItems(scan, value)); so->w = GetScanItems(scan, value);
/* Release shared lock */ /* Release shared lock */
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock); UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
@@ -209,101 +166,26 @@ 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); elog(INFO, "memory: %zu MB", MemoryContextMemAllocated(so->tmpCtx, false) / (1024 * 1024));
#endif #endif
} }
for (;;) while (list_length(so->w) > 0)
{ {
char *base = NULL; char *base = NULL;
HnswSearchCandidate *hc; HnswCandidate *hc = llast(so->w);
HnswElement element; HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid; ItemPointer heaptid;
if (list_length(so->w) == 0)
{
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 (pairingheap_is_empty(so->discarded))
{
ereport(NOTICE,
(errmsg("hnsw index scan exceeded work_mem after " INT64_FORMAT " tuples", so->tuples),
errhint("Increase work_mem to scan more tuples.")));
break;
}
/* Return remaining tuples */
so->w = lappend(so->w, HnswGetSearchCandidate(w_node, pairingheap_remove_first(so->discarded)));
}
else
{
/*
* Locking ensures when neighbors are read, the elements they
* reference will not be deleted (and replaced) during the
* iteration.
*
* Elements loaded into memory on previous iterations may have
* been deleted (and replaced), so when reading neighbors, the
* element version must be checked.
*/
LockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
HnswBench("scan iteration", so->w = ResumeScanItems(scan));
UnlockPage(scan->indexRelation, HNSW_SCAN_LOCK, ShareLock);
#if defined(HNSW_MEMORY)
elog(INFO, "memory: %zu KB", MemoryContextMemAllocated(so->tmpCtx, false) / 1024);
#endif
}
if (list_length(so->w) == 0)
break;
}
hc = llast(so->w);
element = HnswPtrAccess(base, hc->element);
/* Move to next element if no valid heap TIDs */ /* Move to next element if no valid heap TIDs */
if (element->heaptidsLength == 0) if (element->heaptidsLength == 0)
{ {
so->w = list_delete_last(so->w); so->w = list_delete_last(so->w);
/* Mark memory as free for next iteration */
if (hnsw_streaming)
{
pfree(element);
pfree(hc);
}
continue; continue;
} }
heaptid = &element->heaptids[--element->heaptidsLength]; heaptid = &element->heaptids[--element->heaptidsLength];
if (hc->distance < so->previousDistance)
continue;
so->previousDistance = hc->distance;
MemoryContextSwitchTo(oldCtx); MemoryContextSwitchTo(oldCtx);
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;

View File

@@ -9,6 +9,7 @@
#include "fmgr.h" #include "fmgr.h"
#include "hnsw.h" #include "hnsw.h"
#include "lib/pairingheap.h" #include "lib/pairingheap.h"
#include "portability/instr_time.h"
#include "sparsevec.h" #include "sparsevec.h"
#include "storage/bufmgr.h" #include "storage/bufmgr.h"
#include "utils/datum.h" #include "utils/datum.h"
@@ -100,6 +101,13 @@ hash_offset(Size offset)
#define SH_DEFINE #define SH_DEFINE
#include "lib/simplehash.h" #include "lib/simplehash.h"
typedef union
{
pointerhash_hash *pointers;
offsethash_hash *offsets;
tidhash_hash *tids;
} visited_hash;
typedef union typedef union
{ {
HnswElement element; HnswElement element;
@@ -246,8 +254,6 @@ HnswInitElement(char *base, ItemPointer heaptid, int m, double ml, int maxLevel,
element->level = level; element->level = level;
element->deleted = 0; element->deleted = 0;
/* Start at one to make it easier to find issues */
element->version = 1;
HnswInitNeighbors(base, element, m, allocator); HnswInitNeighbors(base, element, m, allocator);
@@ -400,7 +406,6 @@ HnswSetElementTuple(char *base, HnswElementTuple etup, HnswElement element)
etup->type = HNSW_ELEMENT_TUPLE_TYPE; etup->type = HNSW_ELEMENT_TUPLE_TYPE;
etup->level = element->level; etup->level = element->level;
etup->deleted = 0; etup->deleted = 0;
etup->version = element->version;
for (int i = 0; i < HNSW_HEAPTIDS; i++) for (int i = 0; i < HNSW_HEAPTIDS; i++)
{ {
if (i < element->heaptidsLength) if (i < element->heaptidsLength)
@@ -443,7 +448,6 @@ HnswSetNeighborTuple(char *base, HnswNeighborTuple ntup, HnswElement e, int m)
} }
ntup->count = idx; ntup->count = idx;
ntup->version = e->version;
} }
/* /*
@@ -517,7 +521,6 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
{ {
element->level = etup->level; element->level = etup->level;
element->deleted = etup->deleted; element->deleted = etup->deleted;
element->version = etup->version;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid); element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid); element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptidsLength = 0; element->heaptidsLength = 0;
@@ -606,10 +609,10 @@ GetElementDistance(char *base, HnswElement element, Datum q, FmgrInfo *procinfo,
/* /*
* Create a candidate for the entry point * Create a candidate for the entry point
*/ */
HnswSearchCandidate * HnswCandidate *
HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec) HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index, FmgrInfo *procinfo, Oid collation, bool loadVec)
{ {
HnswSearchCandidate *hc = palloc(sizeof(HnswSearchCandidate)); HnswCandidate *hc = palloc(sizeof(HnswCandidate));
HnswPtrStore(base, hc->element, entryPoint); HnswPtrStore(base, hc->element, entryPoint);
if (index == NULL) if (index == NULL)
@@ -619,31 +622,19 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
return hc; return hc;
} }
#define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
#define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
/* /*
* Compare candidate distances * Compare candidate distances
*/ */
static int static int
CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{ {
if (HnswGetSearchCandidateConst(c_node, a)->distance < HnswGetSearchCandidateConst(c_node, b)->distance) if (HnswGetPairingHeapCandidateConst(c_node, a)->distance < HnswGetPairingHeapCandidateConst(c_node, b)->distance)
return 1; return 1;
if (HnswGetSearchCandidateConst(c_node, a)->distance > HnswGetSearchCandidateConst(c_node, b)->distance) if (HnswGetPairingHeapCandidateConst(c_node, a)->distance > HnswGetPairingHeapCandidateConst(c_node, b)->distance)
return -1;
return 0;
}
/*
* Compare discarded candidate distances
*/
static int
CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance)
return 1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance)
return -1; return -1;
return 0; return 0;
@@ -655,15 +646,27 @@ CompareNearestDiscardedCandidates(const pairingheap_node *a, const pairingheap_n
static int static int
CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg) CompareFurthestCandidates(const pairingheap_node *a, const pairingheap_node *b, void *arg)
{ {
if (HnswGetSearchCandidateConst(w_node, a)->distance < HnswGetSearchCandidateConst(w_node, b)->distance) if (HnswGetPairingHeapCandidateConst(w_node, a)->distance < HnswGetPairingHeapCandidateConst(w_node, b)->distance)
return -1; return -1;
if (HnswGetSearchCandidateConst(w_node, a)->distance > HnswGetSearchCandidateConst(w_node, b)->distance) if (HnswGetPairingHeapCandidateConst(w_node, a)->distance > HnswGetPairingHeapCandidateConst(w_node, b)->distance)
return 1; return 1;
return 0; return 0;
} }
/*
* Create a pairing heap node for a candidate
*/
static HnswPairingHeapNode *
CreatePairingHeapNode(HnswCandidate * c)
{
HnswPairingHeapNode *node = palloc(sizeof(HnswPairingHeapNode));
node->inner = c;
return node;
}
/* /*
* Init visited * Init visited
*/ */
@@ -764,30 +767,20 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
int start; int start;
ItemPointerData indextids[HNSW_MAX_M * 2]; ItemPointerData indextids[HNSW_MAX_M * 2];
*unvisitedLength = 0;
buf = ReadBuffer(index, element->neighborPage); buf = ReadBuffer(index, element->neighborPage);
LockBuffer(buf, BUFFER_LOCK_SHARE); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno)); ntup = (HnswNeighborTuple) PageGetItem(page, PageGetItemId(page, element->neighborOffno));
start = (element->level - lc) * m;
/*
* Ensure the neighbor tuple has not been deleted or replaced between
* index scan iterations
*/
if (ntup->version != element->version || ntup->count != (element->level + 2) * m)
{
UnlockReleaseBuffer(buf);
return;
}
/* Copy to minimize lock time */ /* Copy to minimize lock time */
start = (element->level - lc) * m;
memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData)); memcpy(&indextids, ntup->indextids + start, lm * sizeof(ItemPointerData));
UnlockReleaseBuffer(buf); UnlockReleaseBuffer(buf);
*unvisitedLength = 0;
for (int i = 0; i < lm; i++) for (int i = 0; i < lm; i++)
{ {
ItemPointer indextid = &indextids[i]; ItemPointer indextid = &indextids[i];
@@ -807,13 +800,13 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
* Algorithm 2 from paper * Algorithm 2 from paper
*/ */
List * 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, instr_time *start)
{ {
List *w = NIL; List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL); pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL); pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0; int wlen = 0;
visited_hash vh; visited_hash v;
ListCell *lc2; ListCell *lc2;
HnswNeighborArray *localNeighborhood = NULL; HnswNeighborArray *localNeighborhood = NULL;
Size neighborhoodSize = 0; Size neighborhoodSize = 0;
@@ -821,19 +814,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited)); HnswUnvisited *unvisited = palloc(lm * sizeof(HnswUnvisited));
int unvisitedLength; int unvisitedLength;
if (v == NULL) InitVisited(base, &v, index, ef, m);
{
v = &vh;
initVisited = true;
}
if (initVisited)
{
InitVisited(base, v, index, ef, m);
if (discarded != NULL)
*discarded = pairingheap_allocate(CompareNearestDiscardedCandidates, NULL);
}
/* Create local memory for neighborhood if needed */ /* Create local memory for neighborhood if needed */
if (index == NULL) if (index == NULL)
@@ -845,19 +826,15 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Add entry points to v, C, and W */ /* Add entry points to v, C, and W */
foreach(lc2, ep) foreach(lc2, ep)
{ {
HnswSearchCandidate *hc = (HnswSearchCandidate *) lfirst(lc2); HnswCandidate *hc = (HnswCandidate *) lfirst(lc2);
bool found; bool found;
HnswPairingHeapNode *node;
if (initVisited) AddToVisited(base, &v, hc->element, index, &found);
{
AddToVisited(base, v, hc->element, index, &found);
if (tuples != NULL) node = CreatePairingHeapNode(hc);
(*tuples)++; pairingheap_add(C, &node->c_node);
} pairingheap_add(W, &node->w_node);
pairingheap_add(C, &hc->c_node);
pairingheap_add(W, &hc->w_node);
/* /*
* Do not count elements being deleted towards ef when vacuuming. It * Do not count elements being deleted towards ef when vacuuming. It
@@ -870,31 +847,40 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
while (!pairingheap_is_empty(C)) while (!pairingheap_is_empty(C))
{ {
HnswSearchCandidate *c = HnswGetSearchCandidate(c_node, pairingheap_remove_first(C)); HnswCandidate *c = HnswGetPairingHeapCandidate(c_node, pairingheap_remove_first(C));
HnswSearchCandidate *f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); HnswCandidate *f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
HnswElement cElement; HnswElement cElement;
/* Check time budget */
if (start != NULL)
{
instr_time duration;
INSTR_TIME_SET_CURRENT(duration);
INSTR_TIME_SUBTRACT(duration, *start);
if (INSTR_TIME_GET_MILLISEC(duration) >= hnsw_time_budget)
break;
}
if (c->distance > f->distance) if (c->distance > f->distance)
break; break;
cElement = HnswPtrAccess(base, c->element); cElement = HnswPtrAccess(base, c->element);
if (index == NULL) if (index == NULL)
HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, v, lc, localNeighborhood, neighborhoodSize); HnswLoadUnvisitedFromMemory(base, cElement, unvisited, &unvisitedLength, &v, lc, localNeighborhood, neighborhoodSize);
else else
HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, v, index, m, lm, lc); HnswLoadUnvisitedFromDisk(cElement, unvisited, &unvisitedLength, &v, index, m, lm, lc);
if (tuples != NULL)
(*tuples) += unvisitedLength;
for (int i = 0; i < unvisitedLength; i++) for (int i = 0; i < unvisitedLength; i++)
{ {
HnswElement eElement; HnswElement eElement;
HnswSearchCandidate *e; HnswCandidate *e;
HnswPairingHeapNode *node;
float eDistance; float eDistance;
bool alwaysAdd = wlen < ef; bool alwaysAdd = wlen < ef;
f = HnswGetSearchCandidate(w_node, pairingheap_first(W)); f = HnswGetPairingHeapCandidate(w_node, pairingheap_first(W));
if (index == NULL) if (index == NULL)
{ {
@@ -909,33 +895,29 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Avoid any allocations if not adding */ /* Avoid any allocations if not adding */
eElement = NULL; eElement = NULL;
HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd || discarded != NULL ? NULL : &f->distance, &eElement); HnswLoadElementImpl(blkno, offno, &eDistance, &q, index, procinfo, collation, inserting, alwaysAdd ? NULL : &f->distance, &eElement);
if (eElement == NULL)
continue;
} }
if (eElement == NULL || !(eDistance < f->distance || alwaysAdd)) if (!(eDistance < f->distance || alwaysAdd))
{
if (discarded != NULL)
{
/* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate));
HnswPtrStore(base, e->element, eElement);
e->distance = eDistance;
pairingheap_add(*discarded, &e->w_node);
}
continue; continue;
}
Assert(!eElement->deleted);
/* Make robust to issues */ /* Make robust to issues */
if (eElement->level < lc) if (eElement->level < lc)
continue; continue;
/* Create a new candidate */ /* Create a new candidate */
e = palloc(sizeof(HnswSearchCandidate)); e = palloc(sizeof(HnswCandidate));
HnswPtrStore(base, e->element, eElement); HnswPtrStore(base, e->element, eElement);
e->distance = eDistance; e->distance = eDistance;
pairingheap_add(C, &e->c_node);
pairingheap_add(W, &e->w_node); node = CreatePairingHeapNode(e);
pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
/* /*
* Do not count elements being deleted towards ef when vacuuming. * Do not count elements being deleted towards ef when vacuuming.
@@ -948,12 +930,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* No need to decrement wlen */ /* No need to decrement wlen */
if (wlen > ef) if (wlen > ef)
{ pairingheap_remove_first(W);
HnswSearchCandidate *d = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W));
if (discarded != NULL)
pairingheap_add(*discarded, &d->w_node);
}
} }
} }
} }
@@ -961,7 +938,7 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
/* Add each element of W to w */ /* Add each element of W to w */
while (!pairingheap_is_empty(W)) while (!pairingheap_is_empty(W))
{ {
HnswSearchCandidate *hc = HnswGetSearchCandidate(w_node, pairingheap_remove_first(W)); HnswCandidate *hc = HnswGetPairingHeapCandidate(w_node, pairingheap_remove_first(W));
w = lappend(w, hc); w = lappend(w, hc);
} }
@@ -1326,7 +1303,7 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
/* 1st phase: greedy search to insert level */ /* 1st phase: greedy search to insert level */
for (int lc = entryLevel; lc >= level + 1; lc--) for (int lc = entryLevel; lc >= level + 1; lc--)
{ {
w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL, NULL, true, NULL); w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL);
ep = w; ep = w;
} }
@@ -1342,27 +1319,16 @@ HnswFindElementNeighbors(char *base, HnswElement element, HnswElement entryPoint
{ {
int lm = HnswGetLayerM(m, lc); int lm = HnswGetLayerM(m, lc);
List *neighbors; List *neighbors;
List *lw = NIL; List *lw;
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);
/* Convert search candidates to candidates */
foreach(lc2, w)
{
HnswSearchCandidate *sc = lfirst(lc2);
HnswCandidate *hc = palloc(sizeof(HnswCandidate));
hc->element = sc->element;
hc->distance = sc->distance;
lw = lappend(lw, hc);
}
/* Elements being deleted or skipped can help with search */ /* Elements being deleted or skipped can help with search */
/* but should be removed before selecting neighbors */ /* but should be removed before selecting neighbors */
if (index != NULL) if (index != NULL)
lw = RemoveElements(base, lw, skipElement); lw = RemoveElements(base, w, skipElement);
else
lw = w;
/* /*
* Candidates are sorted, but not deterministically. Could set * Candidates are sorted, but not deterministically. Could set

View File

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

View File

@@ -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); 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 */ /* Use total cost since most work happens before first tuple is returned */
*indexStartupCost = costs.indexTotalCost; *indexStartupCost = costs.indexTotalCost;
*indexTotalCost = costs.indexTotalCost; *indexTotalCost = costs.indexTotalCost;

View File

@@ -253,9 +253,8 @@ typedef struct IvfflatScanOpaqueData
/* Sorting */ /* Sorting */
Tuplesortstate *sortstate; Tuplesortstate *sortstate;
TupleDesc tupdesc; TupleDesc tupdesc;
TupleTableSlot *vslot; TupleTableSlot *slot;
TupleTableSlot *mslot; bool isnull;
BufferAccessStrategy bas;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;

View File

@@ -15,19 +15,16 @@
#include "utils/memutils.h" #include "utils/memutils.h"
#endif #endif
#define GetScanList(ptr) pairingheap_container(IvfflatScanList, ph_node, ptr)
#define GetScanListConst(ptr) pairingheap_const_container(IvfflatScanList, ph_node, ptr)
/* /*
* Compare list distances * Compare list distances
*/ */
static int static int
CompareLists(const pairingheap_node *a, const pairingheap_node *b, void *arg) 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; return 1;
if (GetScanListConst(a)->distance < GetScanListConst(b)->distance) if (((const IvfflatScanList *) a)->distance < ((const IvfflatScanList *) b)->distance)
return -1; return -1;
return 0; return 0;
@@ -79,14 +76,14 @@ GetScanLists(IndexScanDesc scan, Datum value)
/* Calculate max distance */ /* Calculate max distance */
if (listCount == so->probes) if (listCount == so->probes)
maxDistance = GetScanList(pairingheap_first(so->listQueue))->distance; maxDistance = ((IvfflatScanList *) pairingheap_first(so->listQueue))->distance;
} }
else if (distance < maxDistance) else if (distance < maxDistance)
{ {
IvfflatScanList *scanlist; IvfflatScanList *scanlist;
/* Remove */ /* Remove */
scanlist = GetScanList(pairingheap_remove_first(so->listQueue)); scanlist = (IvfflatScanList *) pairingheap_remove_first(so->listQueue);
/* Reuse */ /* Reuse */
scanlist->startPage = list->startPage; scanlist->startPage = list->startPage;
@@ -94,7 +91,7 @@ GetScanLists(IndexScanDesc scan, Datum value)
pairingheap_add(so->listQueue, &scanlist->ph_node); pairingheap_add(so->listQueue, &scanlist->ph_node);
/* Update max distance */ /* 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; IvfflatScanOpaque so = (IvfflatScanOpaque) scan->opaque;
TupleDesc tupdesc = RelationGetDescr(scan->indexRelation); TupleDesc tupdesc = RelationGetDescr(scan->indexRelation);
double tuples = 0; 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 */ /* Search closest probes lists */
while (!pairingheap_is_empty(so->listQueue)) 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 */ /* Search all entry pages for list */
while (BlockNumberIsValid(searchPage)) while (BlockNumberIsValid(searchPage))
@@ -127,7 +131,7 @@ GetScanItems(IndexScanDesc scan, Datum value)
Page page; Page page;
OffsetNumber maxoffno; 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); LockBuffer(buf, BUFFER_LOCK_SHARE);
page = BufferGetPage(buf); page = BufferGetPage(buf);
maxoffno = PageGetMaxOffsetNumber(page); maxoffno = PageGetMaxOffsetNumber(page);
@@ -166,6 +170,8 @@ GetScanItems(IndexScanDesc scan, Datum value)
} }
} }
FreeAccessStrategy(bas);
if (tuples < 100) if (tuples < 100)
ereport(DEBUG1, ereport(DEBUG1,
(errmsg("index scan found few tuples"), (errmsg("index scan found few tuples"),
@@ -268,16 +274,7 @@ ivfflatbeginscan(Relation index, int nkeys, int norderbys)
/* Prep sort */ /* Prep sort */
so->sortstate = InitScanSortState(so->tupdesc); so->sortstate = InitScanSortState(so->tupdesc);
/* Need separate slots for puttuple and gettuple */ so->slot = MakeSingleTupleTableSlot(so->tupdesc, &TTSOpsMinimalTuple);
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->listQueue = pairingheap_allocate(CompareLists, scan); so->listQueue = pairingheap_allocate(CompareLists, scan);
@@ -351,10 +348,9 @@ ivfflatgettuple(IndexScanDesc scan, ScanDirection dir)
pfree(DatumGetPointer(value)); 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->slot, 2, &so->isnull));
ItemPointer heaptid = (ItemPointer) DatumGetPointer(slot_getattr(so->mslot, 2, &isnull));
scan->xs_heaptid = *heaptid; scan->xs_heaptid = *heaptid;
scan->xs_recheck = false; scan->xs_recheck = false;
@@ -375,10 +371,6 @@ ivfflatendscan(IndexScanDesc scan)
pairingheap_free(so->listQueue); pairingheap_free(so->listQueue);
tuplesort_end(so->sortstate); tuplesort_end(so->sortstate);
FreeAccessStrategy(so->bas);
FreeTupleDesc(so->tupdesc);
/* TODO Free vslot and mslot without freeing TupleDesc */
pfree(so); pfree(so);
scan->opaque = NULL; scan->opaque = NULL;

View File

@@ -155,6 +155,24 @@ CheckStateArray(ArrayType *statearray, const char *caller)
return (float8 *) ARR_DATA_PTR(statearray); 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 * Convert textual representation to internal representation
*/ */

View File

@@ -94,7 +94,8 @@ like($explain, qr/Seq Scan/);
$explain = $node->safe_psql("postgres", qq( $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query'; 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 # Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);"); $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 = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit; 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(); done_testing();

View File

@@ -18,13 +18,9 @@ $node->start;
# Create table and index # Create table and index
$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), c int4, t text);"); $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", $node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) i;" "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", "CREATE INDEX idx ON tst USING hnsw (v vector_l2_ops);");
$node->safe_psql("postgres", "ANALYZE tst;"); $node->safe_psql("postgres", "ANALYZE tst;");
@@ -100,25 +96,13 @@ $explain = $node->safe_psql("postgres", qq(
)); ));
like($explain, qr/Seq Scan/); 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 # Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);"); $node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq( $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit; EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
)); ));
# Use attribute index # TODO Use attribute index
like($explain, qr/Bitmap Index Scan on attribute_idx/); like($explain, qr/Index Scan using idx/);
# Test partial index # Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);"); $node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);");

View File

@@ -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();

View File

@@ -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();

View File

@@ -1,66 +0,0 @@
use strict;
use warnings FATAL => 'all';
use PostgreSQL::Test::Cluster;
use PostgreSQL::Test::Utils;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = PostgreSQL::Test::Cluster->new('node');
$node->init;
$node->start;
# Create table
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4 PRIMARY KEY, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql] FROM generate_series(1, 100000) i;"
);
$node->safe_psql("postgres", qq(
SET maintenance_work_mem = '128MB';
SET max_parallel_maintenance_workers = 2;
CREATE INDEX ON tst USING hnsw (v vector_l2_ops)
));
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.streaming = on;
SET work_mem = '8MB';
SELECT COUNT(*) FROM (SELECT v FROM tst WHERE i % 10000 = 0 ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 11) t;
));
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/);
done_testing();

View File

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