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
17 changed files with 184 additions and 294 deletions

View File

@@ -1,7 +1,6 @@
## 0.8.0 (unreleased) ## 0.8.0 (unreleased)
- 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

@@ -1,4 +1,4 @@
ARG PG_MAJOR=17 ARG PG_MAJOR=16
FROM postgres:$PG_MAJOR FROM postgres:$PG_MAJOR
ARG PG_MAJOR ARG PG_MAJOR

View File

@@ -66,7 +66,7 @@ dist:
git archive --format zip --prefix=$(EXTENSION)-$(EXTVERSION)/ --output dist/$(EXTENSION)-$(EXTVERSION).zip master git archive --format zip --prefix=$(EXTENSION)-$(EXTVERSION)/ --output dist/$(EXTENSION)-$(EXTVERSION).zip master
# for Docker # for Docker
PG_MAJOR ?= 17 PG_MAJOR ?= 16
.PHONY: docker .PHONY: docker

View File

@@ -52,8 +52,6 @@ nmake /F Makefile.win
nmake /F Makefile.win install nmake /F Makefile.win install
``` ```
Note: Postgres 17 is not supported yet due to an upstream issue
See the [installation notes](#installation-notes---windows) if you run into issues See the [installation notes](#installation-notes---windows) if you run into issues
You can also install it with [Docker](#docker) or [conda-forge](#conda-forge). You can also install it with [Docker](#docker) or [conda-forge](#conda-forge).
@@ -102,8 +100,6 @@ Or add a vector column to an existing table
ALTER TABLE items ADD COLUMN embedding vector(3); ALTER TABLE items ADD COLUMN embedding vector(3);
``` ```
Also supports [half-precision](#half-precision-vectors), [binary](#binary-vectors), and [sparse](#sparse-vectors) vectors
Insert vectors Insert vectors
```sql ```sql
@@ -149,8 +145,6 @@ Supported distance functions are:
- `<#>` - (negative) inner product - `<#>` - (negative) inner product
- `<=>` - cosine distance - `<=>` - cosine distance
- `<+>` - L1 distance (added in 0.7.0) - `<+>` - L1 distance (added in 0.7.0)
- `<~>` - Hamming distance (binary vectors, added in 0.7.0)
- `<%>` - Jaccard distance (binary vectors, added in 0.7.0)
Get the nearest neighbors to a row Get the nearest neighbors to a row
@@ -989,7 +983,7 @@ l2_normalize(sparsevec) → sparsevec | Normalize with Euclidean norm | 0.7.0
If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with: If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with:
```sh ```sh
export PG_CONFIG=/Library/PostgreSQL/17/bin/pg_config export PG_CONFIG=/Library/PostgreSQL/16/bin/pg_config
``` ```
Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use: Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use:
@@ -1000,11 +994,11 @@ sudo --preserve-env=PG_CONFIG make install
A few common paths on Mac are: A few common paths on Mac are:
- EDB installer - `/Library/PostgreSQL/17/bin/pg_config` - EDB installer - `/Library/PostgreSQL/16/bin/pg_config`
- Homebrew (arm64) - `/opt/homebrew/opt/postgresql@17/bin/pg_config` - Homebrew (arm64) - `/opt/homebrew/opt/postgresql@16/bin/pg_config`
- Homebrew (x86-64) - `/usr/local/opt/postgresql@17/bin/pg_config` - Homebrew (x86-64) - `/usr/local/opt/postgresql@16/bin/pg_config`
Note: Replace `17` with your Postgres server version Note: Replace `16` with your Postgres server version
### Missing Header ### Missing Header
@@ -1013,10 +1007,10 @@ If compilation fails with `fatal error: postgres.h: No such file or directory`,
For Ubuntu and Debian, use: For Ubuntu and Debian, use:
```sh ```sh
sudo apt install postgresql-server-dev-17 sudo apt install postgresql-server-dev-16
``` ```
Note: Replace `17` with your Postgres server version Note: Replace `16` with your Postgres server version
### Missing SDK ### Missing SDK
@@ -1049,17 +1043,17 @@ If installation fails with `Access is denied`, re-run the installation instructi
Get the [Docker image](https://hub.docker.com/r/pgvector/pgvector) with: Get the [Docker image](https://hub.docker.com/r/pgvector/pgvector) with:
```sh ```sh
docker pull pgvector/pgvector:pg17 docker pull pgvector/pgvector:pg16
``` ```
This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (replace `17` with your Postgres server version, and run it the same way). This adds pgvector to the [Postgres image](https://hub.docker.com/_/postgres) (replace `16` with your Postgres server version, and run it the same way).
You can also build the image manually: You can also build the image manually:
```sh ```sh
git clone --branch v0.7.4 https://github.com/pgvector/pgvector.git git clone --branch v0.7.4 https://github.com/pgvector/pgvector.git
cd pgvector cd pgvector
docker build --pull --build-arg PG_MAJOR=17 -t myuser/pgvector . docker build --pull --build-arg PG_MAJOR=16 -t myuser/pgvector .
``` ```
### Homebrew ### Homebrew
@@ -1070,7 +1064,7 @@ With Homebrew Postgres, you can use:
brew install pgvector brew install pgvector
``` ```
Note: This only adds it to the `postgresql@17` and `postgresql@14` formulas Note: This only adds it to the `postgresql@14` formula
### PGXN ### PGXN
@@ -1085,22 +1079,22 @@ pgxn install vector
Debian and Ubuntu packages are available from the [PostgreSQL APT Repository](https://wiki.postgresql.org/wiki/Apt). Follow the [setup instructions](https://wiki.postgresql.org/wiki/Apt#Quickstart) and run: Debian and Ubuntu packages are available from the [PostgreSQL APT Repository](https://wiki.postgresql.org/wiki/Apt). Follow the [setup instructions](https://wiki.postgresql.org/wiki/Apt#Quickstart) and run:
```sh ```sh
sudo apt install postgresql-17-pgvector sudo apt install postgresql-16-pgvector
``` ```
Note: Replace `17` with your Postgres server version Note: Replace `16` with your Postgres server version
### Yum ### Yum
RPM packages are available from the [PostgreSQL Yum Repository](https://yum.postgresql.org/). Follow the [setup instructions](https://www.postgresql.org/download/linux/redhat/) for your distribution and run: RPM packages are available from the [PostgreSQL Yum Repository](https://yum.postgresql.org/). Follow the [setup instructions](https://www.postgresql.org/download/linux/redhat/) for your distribution and run:
```sh ```sh
sudo yum install pgvector_17 sudo yum install pgvector_16
# or # or
sudo dnf install pgvector_17 sudo dnf install pgvector_16
``` ```
Note: Replace `17` with your Postgres server version Note: Replace `16` with your Postgres server version
### pkg ### pkg

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,13 +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_time_budget;
int hnsw_lock_tranche_id; int hnsw_lock_tranche_id;
static relopt_kind hnsw_relopt_kind; static relopt_kind hnsw_relopt_kind;
@@ -69,6 +69,10 @@ 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);
DefineCustomIntVariable("hnsw.time_budget", "Sets the time budget for search",
NULL, &hnsw_time_budget,
-1, -1, INT_MAX, PGC_USERSET, GUC_UNIT_MS, NULL, NULL, NULL);
MarkGUCPrefixReserved("hnsw"); MarkGUCPrefixReserved("hnsw");
} }
@@ -101,10 +105,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
GenericCosts costs; GenericCosts costs;
int m; int m;
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 */
@@ -124,55 +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);
/* /* Approximate entry level */
* HNSW cost estimation follows a formula that accounts for the total entryLevel = (int) -log(1.0 / path->indexinfo->tuples) * HnswGetMl(m);
* 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) * hnsw_ef_search;
layer0Selectivity = (scalingFactor * log(path->indexinfo->tuples + 1)) /
(log(m) * (1 + log(hnsw_ef_search)));
costs.numIndexTuples = (entryLevel * m) + /* TODO Improve estimate of visited tuples (currently underestimates) */
(layer0TuplesMax * layer0Selectivity); /* Account for number of tuples (or entry level), m, and ef_search */
costs.numIndexTuples = (entryLevel + 2) * m;
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

@@ -106,6 +106,7 @@
/* Variables */ /* Variables */
extern int hnsw_ef_search; extern int hnsw_ef_search;
extern int hnsw_time_budget;
extern int hnsw_lock_tranche_id; extern int hnsw_lock_tranche_id;
typedef struct HnswElementData HnswElementData; typedef struct HnswElementData HnswElementData;
@@ -155,13 +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;
/* HNSW index options */ /* HNSW index options */
typedef struct HnswOptions typedef struct HnswOptions
@@ -331,6 +331,7 @@ typedef struct HnswScanOpaqueData
bool first; bool first;
List *w; List *w;
MemoryContext tmpCtx; MemoryContext tmpCtx;
instr_time start;
/* Support functions */ /* Support functions */
FmgrInfo *procinfo; FmgrInfo *procinfo;
@@ -375,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); 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

@@ -3,6 +3,7 @@
#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"
@@ -22,6 +23,7 @@ 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);
@@ -33,11 +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); 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); return HnswSearchLayer(base, q, ep, hnsw_ef_search, 0, index, procinfo, collation, m, false, NULL, start);
} }
/* /*
@@ -90,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;
@@ -161,14 +166,14 @@ 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
} }
while (list_length(so->w) > 0) while (list_length(so->w) > 0)
{ {
char *base = NULL; char *base = NULL;
HnswSearchCandidate *hc = llast(so->w); HnswCandidate *hc = llast(so->w);
HnswElement element = HnswPtrAccess(base, hc->element); HnswElement element = HnswPtrAccess(base, hc->element);
ItemPointer heaptid; ItemPointer 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"
@@ -608,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)
@@ -621,8 +622,8 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
return hc; return hc;
} }
#define HnswGetSearchCandidate(membername, ptr) pairingheap_container(HnswSearchCandidate, membername, ptr) #define HnswGetPairingHeapCandidate(membername, ptr) (pairingheap_container(HnswPairingHeapNode, membername, ptr)->inner)
#define HnswGetSearchCandidateConst(membername, ptr) pairingheap_const_container(HnswSearchCandidate, membername, ptr) #define HnswGetPairingHeapCandidateConst(membername, ptr) (pairingheap_const_container(HnswPairingHeapNode, membername, ptr)->inner)
/* /*
* Compare candidate distances * Compare candidate distances
@@ -630,10 +631,10 @@ HnswEntryCandidate(char *base, HnswElement entryPoint, Datum q, Relation index,
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 -1;
return 0; return 0;
@@ -645,15 +646,27 @@ CompareNearestCandidates(const pairingheap_node *a, const pairingheap_node *b, v
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
*/ */
@@ -759,16 +772,9 @@ HnswLoadUnvisitedFromDisk(HnswElement element, HnswUnvisited * unvisited, int *u
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 expected neighbors */
if (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);
@@ -794,7 +800,7 @@ 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) 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);
@@ -820,13 +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;
AddToVisited(base, &v, hc->element, index, &found); AddToVisited(base, &v, hc->element, index, &found);
pairingheap_add(C, &hc->c_node); node = CreatePairingHeapNode(hc);
pairingheap_add(W, &hc->w_node); pairingheap_add(C, &node->c_node);
pairingheap_add(W, &node->w_node);
/* /*
* Do not count elements being deleted towards ef when vacuuming. It * Do not count elements being deleted towards ef when vacuuming. It
@@ -839,10 +847,21 @@ 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;
@@ -856,11 +875,12 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
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)
{ {
@@ -891,11 +911,13 @@ HnswSearchLayer(char *base, Datum q, List *ep, int ef, int lc, Relation index, F
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.
@@ -916,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);
} }
@@ -1281,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); w = HnswSearchLayer(base, q, ep, 1, lc, index, procinfo, collation, m, true, skipElement, NULL);
ep = w; ep = w;
} }
@@ -1297,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); 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

@@ -85,8 +85,6 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
MemSet(&costs, 0, sizeof(costs)); MemSet(&costs, 0, sizeof(costs));
genericcostestimate(root, path, loop_count, &costs);
index = index_open(path->indexinfo->indexoid, NoLock); index = index_open(path->indexinfo->indexoid, NoLock);
IvfflatGetMetaPageInfo(index, &lists, NULL); IvfflatGetMetaPageInfo(index, &lists, NULL);
index_close(index, NoLock); index_close(index, NoLock);
@@ -96,9 +94,14 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
if (ratio > 1.0) if (ratio > 1.0)
ratio = 1.0; ratio = 1.0;
/* Set startup cost since most work happens before first tuple is returned */ /*
costs.indexStartupCost = costs.indexTotalCost * ratio; * This gives us the subset of tuples to visit. This value is passed into
costs.numIndexPages *= ratio; * the generic cost estimator to determine the number of pages to visit
* during the index scan.
*/
costs.numIndexTuples = path->indexinfo->tuples * ratio;
genericcostestimate(root, path, loop_count, &costs);
get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost); get_tablespace_page_costs(path->indexinfo->reltablespace, NULL, &spc_seq_page_cost);
@@ -106,25 +109,30 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
if (costs.numIndexPages > path->indexinfo->rel->pages && ratio < 0.5) if (costs.numIndexPages > path->indexinfo->rel->pages && ratio < 0.5)
{ {
/* Change all page cost from random to sequential */ /* Change all page cost from random to sequential */
costs.indexStartupCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost); costs.indexTotalCost -= costs.numIndexPages * (costs.spc_random_page_cost - spc_seq_page_cost);
/* Remove cost of extra pages */ /* Remove cost of extra pages */
costs.indexStartupCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost; costs.indexTotalCost -= (costs.numIndexPages - path->indexinfo->rel->pages) * spc_seq_page_cost;
} }
else else
{ {
/* Change some page cost from random to sequential */ /* Change some page cost from random to sequential */
costs.indexStartupCost -= 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);
} }
*indexStartupCost = costs.indexStartupCost; /*
* 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; *indexTotalCost = costs.indexTotalCost;
*indexSelectivity = costs.indexSelectivity; *indexSelectivity = costs.indexSelectivity;
*indexCorrelation = costs.indexCorrelation; *indexCorrelation = costs.indexCorrelation;
*indexPages = costs.numIndexPages; *indexPages = costs.numIndexPages;
Assert(*indexStartupCost > 0);
Assert(*indexTotalCost > 0);
} }
/* /*

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,51 +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 ? 3000 : 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/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 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,50 +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, 6000) 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/);
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
$node->safe_psql("postgres", "DROP TABLE tst;");
}
done_testing();