Compare commits

..

9 Commits

Author SHA1 Message Date
Andrew Kane
6c347b7f3e Removed unused functions 2023-12-19 19:47:03 -05:00
Andrew Kane
3e628986a1 Removed metapage logging [skip ci] 2023-12-19 17:23:35 -05:00
Andrew Kane
858a89575b Removed metapage [skip ci] 2023-12-19 17:20:55 -05:00
Andrew Kane
d81c1c9de0 Improved code [skip ci] 2023-12-19 17:01:11 -05:00
Andrew Kane
4988d04338 Improved code [skip ci] 2023-12-19 17:00:14 -05:00
Andrew Kane
c640def56c Use forkNum [skip ci] 2023-12-19 13:58:58 -05:00
Andrew Kane
87fe23d9ec Added code for Postgres < 11.8 2023-12-19 13:55:56 -05:00
Andrew Kane
041f939bde Fixed warnings 2023-12-19 12:01:09 -05:00
Andrew Kane
5d7bf9509d Reduced WAL generation for HNSW index builds 2023-12-19 11:49:40 -05:00
37 changed files with 413 additions and 924 deletions

View File

@@ -2,9 +2,7 @@
- Improved performance of HNSW
- Added support for on-disk parallel index builds for HNSW
- Reduced memory usage for HNSW index builds
- Reduced WAL generation for HNSW index builds
- Fixed error with logical replication
- Fixed `invalid memory alloc request size` error with HNSW index build
## 0.5.1 (2023-10-10)

221
README.md
View File

@@ -161,12 +161,80 @@ You can add an index to use approximate nearest neighbor search, which trades so
Supported index types are:
- [HNSW](#hnsw) - added in 0.5.0
- [IVFFlat](#ivfflat)
- [HNSW](#hnsw) - added in 0.5.0
## IVFFlat
An IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
Three keys to achieving good recall are:
1. Create the index *after* the table has some data
2. Choose an appropriate number of lists - a good place to start is `rows / 1000` for up to 1M rows and `sqrt(rows)` for over 1M rows
3. When querying, specify an appropriate number of [probes](#query-options) (higher is better for recall, lower is better for speed) - a good place to start is `sqrt(lists)`
Add an index for each distance function you want to use.
L2 distance
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
```
Inner product
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);
```
Cosine distance
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
```
Vectors with up to 2,000 dimensions can be indexed.
### Query Options
Specify the number of probes (1 by default)
```sql
SET ivfflat.probes = 10;
```
A higher value provides better recall at the cost of speed, and it can be set to the number of lists for exact nearest neighbor search (at which point the planner wont use the index)
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL ivfflat.probes = 10;
SELECT ...
COMMIT;
```
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for IVFFlat are:
1. `initializing`
2. `performing k-means`
3. `assigning tuples`
4. `loading tuples`
Note: `%` is only populated during the `loading tuples` phase
## HNSW
An HNSW index creates a multilayer graph. It has better query performance than IVFFlat (in terms of speed-recall tradeoff), but has slower build times and uses more memory. Also, an index can be created without any data in the table since there isnt a training step like IVFFlat.
An HNSW index creates a multilayer graph. It has slower build times and uses more memory than IVFFlat, but has better query performance (in terms of speed-recall tradeoff). Theres no training step like IVFFlat, so the index can be created without any data in the table.
Add an index for each distance function you want to use.
@@ -222,24 +290,6 @@ SELECT ...
COMMIT;
```
### Index Build Time
Indexes build significantly faster when the graph fits into `maintenance_work_mem`
```sql
SET maintenance_work_mem = '8GB';
```
A notice is shown when the graph no longer fits
```text
NOTICE: hnsw graph no longer fits into maintenance_work_mem after 100000 tuples
DETAIL: Building will take significantly more time.
HINT: Increase maintenance_work_mem to speed up builds.
```
Note: Do not set `maintenance_work_mem` so high that it exhausts the memory on the server
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
@@ -253,84 +303,6 @@ The phases for HNSW are:
1. `initializing`
2. `loading tuples`
## IVFFlat
An IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff).
Three keys to achieving good recall are:
1. Create the index *after* the table has some data
2. Choose an appropriate number of lists - a good place to start is `rows / 1000` for up to 1M rows and `sqrt(rows)` for over 1M rows
3. When querying, specify an appropriate number of [probes](#query-options) (higher is better for recall, lower is better for speed) - a good place to start is `sqrt(lists)`
Add an index for each distance function you want to use.
L2 distance
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
```
Inner product
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_ip_ops) WITH (lists = 100);
```
Cosine distance
```sql
CREATE INDEX ON items USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
```
Vectors with up to 2,000 dimensions can be indexed.
### Query Options
Specify the number of probes (1 by default)
```sql
SET ivfflat.probes = 10;
```
A higher value provides better recall at the cost of speed, and it can be set to the number of lists for exact nearest neighbor search (at which point the planner wont use the index)
Use `SET LOCAL` inside a transaction to set it for a single query
```sql
BEGIN;
SET LOCAL ivfflat.probes = 10;
SELECT ...
COMMIT;
```
### Index Build Time
Speed up index creation on large tables by increasing the number of parallel workers (2 by default)
```sql
SET max_parallel_maintenance_workers = 7; -- plus leader
```
For a large number of workers, you may also need to increase `max_parallel_workers` (8 by default)
### Indexing Progress
Check [indexing progress](https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING) with Postgres 12+
```sql
SELECT phase, round(100.0 * tuples_done / nullif(tuples_total, 0), 1) AS "%" FROM pg_stat_progress_create_index;
```
The phases for IVFFlat are:
1. `initializing`
2. `performing k-means`
3. `assigning tuples`
4. `loading tuples`
Note: `%` is only populated during the `loading tuples` phase
## Filtering
There are a few ways to index nearest neighbor queries with a `WHERE` clause
@@ -348,7 +320,8 @@ CREATE INDEX ON items (category_id);
Or a [partial index](https://www.postgresql.org/docs/current/indexes-partial.html) on the vector column for approximate search
```sql
CREATE INDEX ON items USING hnsw (embedding vector_l2_ops) WHERE (category_id = 123);
CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)
WHERE (category_id = 123);
```
Use [partitioning](https://www.postgresql.org/docs/current/ddl-partitioning.html) for approximate search on many different values of the `WHERE` columns
@@ -582,7 +555,7 @@ sum(vector) → vector | sum | 0.5.0
If your machine has multiple Postgres installations, specify the path to [pg_config](https://www.postgresql.org/docs/current/app-pgconfig.html) with:
```sh
export PG_CONFIG=/Library/PostgreSQL/16/bin/pg_config
export PG_CONFIG=/Applications/Postgres.app/Contents/Versions/latest/bin/pg_config
```
Then re-run the installation instructions (run `make clean` before `make` if needed). If `sudo` is needed for `make install`, use:
@@ -591,14 +564,6 @@ Then re-run the installation instructions (run `make clean` before `make` if nee
sudo --preserve-env=PG_CONFIG make install
```
A few common paths on Mac are:
- EDB installer - `/Library/PostgreSQL/16/bin/pg_config`
- Homebrew (arm64) - `/opt/homebrew/opt/postgresql@16/bin/pg_config`
- Homebrew (x86-64) - `/usr/local/opt/postgresql@16/bin/pg_config`
Note: Replace `16` with your Postgres server version
### Missing Header
If compilation fails with `fatal error: postgres.h: No such file or directory`, make sure Postgres development files are installed on the server.
@@ -606,14 +571,10 @@ If compilation fails with `fatal error: postgres.h: No such file or directory`,
For Ubuntu and Debian, use:
```sh
sudo apt install postgresql-server-dev-16
sudo apt install postgresql-server-dev-15
```
Note: Replace `16` with your Postgres server version
### Missing SDK
If compilation fails and the output includes `warning: no such sysroot directory` on Mac, reinstall Xcode Command Line Tools.
Note: Replace `15` with your Postgres server version
### Windows
@@ -628,7 +589,7 @@ Note: The exact path will vary depending on your Visual Studio version and editi
Then use `nmake` to build:
```cmd
set "PGROOT=C:\Program Files\PostgreSQL\16"
set "PGROOT=C:\Program Files\PostgreSQL\15"
git clone --branch v0.5.1 https://github.com/pgvector/pgvector.git
cd pgvector
nmake /F Makefile.win
@@ -678,22 +639,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:
```sh
sudo apt install postgresql-16-pgvector
sudo apt install postgresql-15-pgvector
```
Note: Replace `16` with your Postgres server version
Note: Replace `15` with your Postgres server version
### 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:
```sh
sudo yum install pgvector_16
sudo yum install pgvector_15
# or
sudo dnf install pgvector_16
sudo dnf install pgvector_15
```
Note: Replace `16` with your Postgres server version
Note: Replace `15` with your Postgres server version
### conda-forge
@@ -803,25 +764,7 @@ make prove_installcheck PROVE_TESTS=test/t/001_wal.pl # TAP test
To enable benchmarking:
```sh
make clean && PG_CFLAGS="-DIVFFLAT_BENCH" make && make install
```
To show memory usage:
```sh
make clean && PG_CFLAGS="-DHNSW_MEMORY -DIVFFLAT_MEMORY" make && make install
```
To enable assertions:
```sh
make clean && PG_CFLAGS="-DUSE_ASSERT_CHECKING" make && make install
```
To get k-means metrics:
```sh
make clean && PG_CFLAGS="-DIVFFLAT_KMEANS_DEBUG" make && make install
make clean && PG_CFLAGS=-DIVFFLAT_BENCH make && make install
```
Resources for contributors

View File

@@ -79,14 +79,12 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
int m;
int entryLevel;
Relation index;
double selectivity = 1;
ListCell *lc;
#if PG_VERSION_NUM < 120000
List *qinfos;
#endif
/* Never use index without order or limit */
if (path->indexorderbys == NULL || root->limit_tuples < 0)
/* Never use index without order */
if (path->indexorderbys == NULL)
{
*indexStartupCost = DBL_MAX;
*indexTotalCost = DBL_MAX;
@@ -96,29 +94,6 @@ hnswcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
return;
}
/* Get the selectivity of non-index conditions */
foreach(lc, path->indexinfo->indrestrictinfo)
{
RestrictInfo *rinfo = lfirst(lc);
if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
selectivity *= rinfo->norm_selec;
}
/*
* Do not use index if limit + offset > expected tuples unless
* enable_seqscan = off
*/
if (root->limit_tuples > hnsw_ef_search * selectivity)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
MemSet(&costs, 0, sizeof(costs));
index = index_open(path->indexinfo->indexoid, NoLock);

View File

@@ -6,7 +6,6 @@
#include "access/generic_xlog.h"
#include "access/parallel.h"
#include "access/reloptions.h"
#include "lib/ilist.h"
#include "nodes/execnodes.h"
#include "port.h" /* for random() */
#include "utils/sampling.h"
@@ -73,10 +72,8 @@
#if PG_VERSION_NUM >= 150000
#define RandomDouble() pg_prng_double(&pg_global_prng_state)
#define SeedRandom(seed) pg_prng_seed(&pg_global_prng_state, seed)
#else
#define RandomDouble() (((double) random()) / MAX_RANDOM_VALUE)
#define SeedRandom(seed) srandom(seed)
#endif
#if PG_VERSION_NUM < 130000
@@ -94,23 +91,21 @@
#define HnswGetMl(m) (1 / log(m))
/* Ensure fits on page and in uint8 */
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / (m)) - 2, 255)
#define HnswGetNeighbors(element, lc) (AssertMacro((element)->level >= (lc)), &(element)->neighbors[lc])
#define HnswGetMaxLevel(m) Min(((BLCKSZ - MAXALIGN(SizeOfPageHeaderData) - MAXALIGN(sizeof(HnswPageOpaqueData)) - offsetof(HnswNeighborTupleData, indextids) - sizeof(ItemIdData)) / (sizeof(ItemPointerData)) / m) - 2, 255)
/* Variables */
extern int hnsw_ef_search;
extern bool hnsw_enable_parallel_build;
typedef struct HnswNeighborArray HnswNeighborArray;
typedef struct HnswElementData
{
slist_node next;
ItemPointerData heaptids[HNSW_HEAPTIDS];
uint8 heaptidsLength;
List *heaptids;
uint8 level;
uint8 deleted;
uint32 hash;
struct HnswNeighborArray *neighbors;
HnswNeighborArray *neighbors;
BlockNumber blkno;
OffsetNumber offno;
OffsetNumber neighborOffno;
@@ -148,16 +143,6 @@ typedef struct HnswOptions
int efConstruction; /* size of dynamic candidate list */
} HnswOptions;
typedef struct HnswGraph
{
slist_head elements;
HnswElement entryPoint;
long memoryUsed;
long memoryTotal;
bool flushed;
double indtuples;
} HnswGraph;
typedef struct HnswSpool
{
Relation heap;
@@ -181,7 +166,7 @@ typedef struct HnswShared
/* Mutable state */
int nparticipantsdone;
double reltuples;
HnswGraph graphData;
double indtuples;
#if PG_VERSION_NUM < 120000
ParallelHeapScanDescData heapdesc; /* must come last */
@@ -224,14 +209,15 @@ typedef struct HnswBuildState
Oid collation;
/* Variables */
HnswGraph graphData;
HnswGraph *graph;
List *elements;
HnswElement entryPoint;
double ml;
int maxLevel;
long memoryLeft;
bool flushed;
Vector *normvec;
/* Memory */
MemoryContext graphCtx;
MemoryContext tmpCtx;
/* Parallel builds */
@@ -339,8 +325,10 @@ List *HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, Fmgr
HnswElement HnswGetEntryPoint(Relation index);
void HnswGetMetaPageInfo(Relation index, int *m, HnswElement * entryPoint);
HnswElement HnswInitElement(ItemPointer tid, int m, double ml, int maxLevel);
void HnswFreeElement(HnswElement element);
HnswElement HnswInitElementFromBlock(BlockNumber blkno, OffsetNumber offno);
void HnswInsertElement(HnswElement element, HnswElement entryPoint, Relation index, FmgrInfo *procinfo, Oid collation, int m, int efConstruction, bool existing);
HnswElement HnswFindDuplicate(HnswElement e);
HnswCandidate *HnswEntryCandidate(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 HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m);

View File

@@ -56,9 +56,7 @@
#define PARALLEL_KEY_HNSW_SHARED UINT64CONST(0xA000000000000001)
#define PARALLEL_KEY_QUERY_TEXT UINT64CONST(0xA000000000000002)
#if PG_VERSION_NUM < 130000
#define GENERATIONCHUNK_RAWSIZE (SIZEOF_SIZE_T + SIZEOF_VOID_P * 2)
#endif
#define LIST_MAX_LENGTH ((1 << 26) - 1)
/*
* Create the metapage
@@ -90,7 +88,6 @@ CreateMetaPage(HnswBuildState * buildstate)
((PageHeader) page)->pd_lower =
((char *) metap + sizeof(HnswMetaPageData)) - (char *) page;
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
}
@@ -107,7 +104,6 @@ HnswBuildAppendPage(Relation index, Buffer *buf, Page *page, ForkNumber forkNum)
HnswPageGetOpaque(*page)->nextblkno = BufferGetBlockNumber(newbuf);
/* Commit */
MarkBufferDirty(*buf);
UnlockReleaseBuffer(*buf);
/* Can take a while, so ensure we can interrupt */
@@ -137,7 +133,7 @@ CreateElementPages(HnswBuildState * buildstate)
BlockNumber insertPage;
Buffer buf;
Page page;
slist_iter iter;
ListCell *lc;
/* Calculate sizes */
etupAllocSize = BLCKSZ;
@@ -152,9 +148,9 @@ CreateElementPages(HnswBuildState * buildstate)
page = BufferGetPage(buf);
HnswInitPage(buf, page);
slist_foreach(iter, &buildstate->graph->elements)
foreach(lc, buildstate->elements)
{
HnswElement element = slist_container(HnswElementData, next, iter.cur);
HnswElement element = lfirst(lc);
Size etupSize;
Size ntupSize;
Size combinedSize;
@@ -209,10 +205,9 @@ CreateElementPages(HnswBuildState * buildstate)
insertPage = BufferGetBlockNumber(buf);
/* Commit */
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->graph->entryPoint, insertPage, forkNum, true);
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_ALWAYS, buildstate->entryPoint, insertPage, forkNum, true);
pfree(etup);
pfree(ntup);
@@ -227,15 +222,15 @@ CreateNeighborPages(HnswBuildState * buildstate)
Relation index = buildstate->index;
ForkNumber forkNum = buildstate->forkNum;
int m = buildstate->m;
slist_iter iter;
ListCell *lc;
HnswNeighborTuple ntup;
/* Allocate once */
ntup = palloc0(BLCKSZ);
slist_foreach(iter, &buildstate->graph->elements)
foreach(lc, buildstate->elements)
{
HnswElement e = slist_container(HnswElementData, next, iter.cur);
HnswElement e = lfirst(lc);
Buffer buf;
Page page;
Size ntupSize = HNSW_NEIGHBOR_TUPLE_SIZE(e->level, m);
@@ -254,30 +249,25 @@ CreateNeighborPages(HnswBuildState * buildstate)
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
MarkBufferDirty(buf);
UnlockReleaseBuffer(buf);
}
pfree(ntup);
}
#ifdef HNSW_MEMORY
/*
* Show memory usage
* Free elements
*/
static void
ShowMemoryUsage(HnswBuildState * buildstate)
FreeElements(HnswBuildState * buildstate)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "graph memory: %zu MB, total memory: %zu MB",
MemoryContextMemAllocated(buildstate->graphCtx, false) / (1024 * 1024),
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
elog(INFO, "estimated memory: %zu MB", buildstate->memoryUsed / (1024 * 1024));
#endif
ListCell *lc;
foreach(lc, buildstate->elements)
HnswFreeElement(lfirst(lc));
list_free(buildstate->elements);
}
#endif
/*
* Flush pages
@@ -285,78 +275,26 @@ ShowMemoryUsage(HnswBuildState * buildstate)
static void
FlushPages(HnswBuildState * buildstate)
{
#ifdef HNSW_MEMORY
ShowMemoryUsage(buildstate);
#endif
CreateMetaPage(buildstate);
CreateElementPages(buildstate);
CreateNeighborPages(buildstate);
buildstate->graph->flushed = true;
MemoryContextReset(buildstate->graphCtx);
buildstate->flushed = true;
FreeElements(buildstate);
}
#if PG_VERSION_NUM < 130000
/*
* Get the memory used by an element
*/
static long
HnswElementMemory(HnswElement e, int m)
{
long elementSize = sizeof(HnswElementData);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
/* Each allocation has a chunk header */
elementSize += (e->level + 4) * GENERATIONCHUNK_RAWSIZE;
/* Add an extra 5% for alignment and other overhead */
return elementSize * 1.05;
}
#endif
/*
* Find duplicate element
* Insert tuple
*/
static bool
HnswFindDuplicateInMemory(HnswElement element)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1))
return false;
/* Check for space */
if (neighbor->element->heaptidsLength < HNSW_HEAPTIDS)
{
HnswAddHeapTid(neighbor->element, &element->heaptids[0]);
return true;
}
}
return false;
}
/*
* Insert tuple into in-memory graph
*/
static bool
InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuildState * buildstate)
InsertTuple(Relation index, Datum *values, HnswElement element, HnswBuildState * buildstate, HnswElement * dup, MemoryContext outerCtx)
{
FmgrInfo *procinfo = buildstate->procinfo;
Oid collation = buildstate->collation;
HnswGraph *graph = buildstate->graph;
HnswElement entryPoint = graph->entryPoint;
HnswElement entryPoint = buildstate->entryPoint;
int efConstruction = buildstate->efConstruction;
int m = buildstate->m;
MemoryContext oldCtx;
HnswElement element;
/* Detoast once for all calls */
Datum value = PointerGetDatum(PG_DETOAST_DATUM(values[0]));
@@ -368,67 +306,52 @@ InsertTupleInMemory(Relation index, Datum *values, ItemPointer heaptid, HnswBuil
return false;
}
/* Allocate element in graph memory context */
oldCtx = MemoryContextSwitchTo(buildstate->graphCtx);
element = HnswInitElement(heaptid, buildstate->m, buildstate->ml, buildstate->maxLevel);
/* Copy value to element so accessible outside of memory context */
oldCtx = MemoryContextSwitchTo(outerCtx);
element->value = datumCopy(value, false, -1);
MemoryContextSwitchTo(oldCtx);
/* Update memory usage */
#if PG_VERSION_NUM >= 130000
graph->memoryUsed = MemoryContextMemAllocated(buildstate->graphCtx, false);
#else
graph->memoryUsed += HnswElementMemory(element, buildstate->m);
#endif
/* Insert element in graph */
HnswInsertElement(element, entryPoint, NULL, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
if (HnswFindDuplicateInMemory(element))
*dup = HnswFindDuplicate(element);
/* Update neighbors if needed */
if (*dup == NULL)
{
/* No need to free element since memory unlikely to be reallocated */
return true;
}
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = &element->neighbors[lc];
/* Add element */
slist_push_head(&graph->elements, &element->next);
/* Update neighbors */
for (int lc = element->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(element, lc);
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
for (int i = 0; i < neighbors->length; i++)
HnswUpdateConnection(element, &neighbors->items[i], lm, lc, NULL, NULL, procinfo, collation);
}
}
/* Update entry point if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
graph->entryPoint = element;
if (*dup == NULL && (entryPoint == NULL || element->level > entryPoint->level))
buildstate->entryPoint = element;
return true;
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
return *dup == NULL;
}
/*
* Acquire a lock if needed
* Get the memory used by an element
*/
static inline void
HnswLockAcquire(HnswShared * hnswshared)
static long
HnswElementMemory(HnswElement e, int m)
{
if (hnswshared)
SpinLockAcquire(&hnswshared->mutex);
}
long elementSize = sizeof(HnswElementData);
/*
* Release a lock if needed
*/
static inline void
HnswLockRelease(HnswShared * hnswshared)
{
if (hnswshared)
SpinLockRelease(&hnswshared->mutex);
elementSize += sizeof(HnswNeighborArray) * (e->level + 1);
elementSize += sizeof(HnswCandidate) * (m * (e->level + 2));
elementSize += sizeof(ItemPointerData);
elementSize += VARSIZE_ANY(DatumGetPointer(e->value));
return elementSize;
}
/*
@@ -439,9 +362,9 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
bool *isnull, bool tupleIsAlive, void *state)
{
HnswBuildState *buildstate = (HnswBuildState *) state;
HnswGraph *graph = buildstate->graph;
HnswShared *hnswshared = buildstate->hnswshared;
MemoryContext oldCtx;
HnswElement element;
HnswElement dup = NULL;
bool inserted;
#if PG_VERSION_NUM < 130000
@@ -452,50 +375,70 @@ BuildCallback(Relation index, CALLBACK_ITEM_POINTER, Datum *values,
if (isnull[0])
return;
/* Flush pages if needed */
if (!graph->flushed && graph->memoryUsed >= graph->memoryTotal)
if (buildstate->flushed || buildstate->memoryLeft <= 0 || list_length(buildstate->elements) == LIST_MAX_LENGTH)
{
ereport(NOTICE,
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) graph->indtuples),
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
if (!buildstate->flushed)
{
if (buildstate->memoryLeft <= 0)
ereport(NOTICE,
(errmsg("hnsw graph no longer fits into maintenance_work_mem after " INT64_FORMAT " tuples", (int64) buildstate->indtuples),
errdetail("Building will take significantly more time."),
errhint("Increase maintenance_work_mem to speed up builds.")));
FlushPages(buildstate);
FlushPages(buildstate);
}
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
if (HnswInsertTuple(buildstate->index, values, isnull, tid, buildstate->heap, true))
{
if (buildstate->hnswshared)
{
HnswShared *hnswshared = buildstate->hnswshared;
SpinLockAcquire(&hnswshared->mutex);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++hnswshared->indtuples);
SpinLockRelease(&hnswshared->mutex);
}
else
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++buildstate->indtuples);
}
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
return;
}
/* Allocate necessary memory outside of memory context */
element = HnswInitElement(tid, buildstate->m, buildstate->ml, buildstate->maxLevel);
/* Use memory context since detoast can allocate */
oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
/* Insert tuple */
if (graph->flushed)
inserted = HnswInsertTuple(index, values, isnull, tid, buildstate->heap, true);
else
inserted = InsertTupleInMemory(index, values, tid, buildstate);
/* Update progress */
if (inserted)
{
HnswLockAcquire(hnswshared);
UpdateProgress(PROGRESS_CREATEIDX_TUPLES_DONE, ++graph->indtuples);
HnswLockRelease(hnswshared);
}
inserted = InsertTuple(index, values, element, buildstate, &dup, oldCtx);
/* Reset memory context */
MemoryContextSwitchTo(oldCtx);
MemoryContextReset(buildstate->tmpCtx);
}
/*
* Initialize the graph
*/
static void
InitGraph(HnswGraph * graph)
{
slist_init(&graph->elements);
graph->entryPoint = NULL;
graph->memoryUsed = 0;
graph->memoryTotal = maintenance_work_mem * 1024L;
graph->flushed = false;
graph->indtuples = 0;
/* Add outside memory context */
if (dup != NULL)
{
HnswAddHeapTid(dup, tid);
buildstate->memoryLeft -= sizeof(ItemPointerData);
}
/* Add to buildstate or free */
if (inserted)
{
buildstate->elements = lappend(buildstate->elements, element);
buildstate->memoryLeft -= HnswElementMemory(element, buildstate->m);
}
else
HnswFreeElement(element);
}
/*
@@ -531,20 +474,16 @@ InitBuildState(HnswBuildState * buildstate, Relation heap, Relation index, Index
buildstate->normprocinfo = HnswOptionalProcInfo(index, HNSW_NORM_PROC);
buildstate->collation = index->rd_indcollation[0];
InitGraph(&buildstate->graphData);
buildstate->graph = &buildstate->graphData;
buildstate->elements = NIL;
buildstate->entryPoint = NULL;
buildstate->ml = HnswGetMl(buildstate->m);
buildstate->maxLevel = HnswGetMaxLevel(buildstate->m);
buildstate->memoryLeft = maintenance_work_mem * 1024L;
buildstate->flushed = false;
/* Reuse for each tuple */
buildstate->normvec = InitVector(buildstate->dimensions);
buildstate->graphCtx = GenerationContextCreate(CurrentMemoryContext,
"Hnsw build graph context",
#if PG_VERSION_NUM >= 150000
1024 * 1024, 1024 * 1024,
#endif
1024 * 1024);
buildstate->tmpCtx = AllocSetContextCreate(CurrentMemoryContext,
"Hnsw build temporary context",
ALLOCSET_DEFAULT_SIZES);
@@ -560,7 +499,6 @@ static void
FreeBuildState(HnswBuildState * buildstate)
{
pfree(buildstate->normvec);
MemoryContextDelete(buildstate->graphCtx);
MemoryContextDelete(buildstate->tmpCtx);
}
@@ -580,7 +518,7 @@ ParallelHeapScan(HnswBuildState * buildstate)
SpinLockAcquire(&hnswshared->mutex);
if (hnswshared->nparticipantsdone == nparticipanttuplesorts)
{
buildstate->graph = &hnswshared->graphData;
buildstate->indtuples = hnswshared->indtuples;
reltuples = hnswshared->reltuples;
SpinLockRelease(&hnswshared->mutex);
break;
@@ -615,7 +553,9 @@ HnswParallelScanAndInsert(HnswSpool * hnswspool, HnswShared * hnswshared, bool p
indexInfo = BuildIndexInfo(hnswspool->index);
indexInfo->ii_Concurrent = hnswshared->isconcurrent;
InitBuildState(&buildstate, hnswspool->heap, hnswspool->index, indexInfo, MAIN_FORKNUM);
buildstate.graph = &hnswshared->graphData;
/* TODO Support in-memory builds */
buildstate.memoryLeft = 0;
buildstate.flushed = true;
buildstate.hnswshared = hnswshared;
#if PG_VERSION_NUM >= 120000
scan = table_beginscan_parallel(hnswspool->heap,
@@ -840,10 +780,7 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
/* Initialize mutable state */
hnswshared->nparticipantsdone = 0;
hnswshared->reltuples = 0;
InitGraph(&hnswshared->graphData);
/* TODO Support in-memory builds */
hnswshared->graphData.memoryTotal = 0;
hnswshared->graphData.flushed = true;
hnswshared->indtuples = 0;
#if PG_VERSION_NUM >= 120000
table_parallelscan_initialize(buildstate->heap,
ParallelTableScanFromHnswShared(hnswshared),
@@ -894,27 +831,6 @@ HnswBeginParallel(HnswBuildState * buildstate, bool isconcurrent, int request)
WaitForParallelWorkersToAttach(pcxt);
}
/*
* Compute parallel workers
*/
static int
ComputeParallelWorkers(Relation heap, Relation index)
{
int parallel_workers;
/* Make sure it's safe to use parallel workers */
parallel_workers = plan_create_index_workers(RelationGetRelid(heap), RelationGetRelid(index));
if (parallel_workers == 0)
return 0;
/* Use parallel_workers storage parameter on table if set */
parallel_workers = RelationGetParallelWorkers(heap, -1);
if (parallel_workers != -1)
return Min(parallel_workers, max_parallel_maintenance_workers);
return max_parallel_maintenance_workers;
}
/*
* Build graph
*/
@@ -926,8 +842,8 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
UpdateProgress(PROGRESS_CREATEIDX_SUBPHASE, PROGRESS_HNSW_PHASE_LOAD);
/* Calculate parallel workers */
if (buildstate->heap != NULL && hnsw_enable_parallel_build)
parallel_workers = ComputeParallelWorkers(buildstate->heap, buildstate->index);
if (hnsw_enable_parallel_build)
parallel_workers = plan_create_index_workers(RelationGetRelid(buildstate->heap), RelationGetRelid(buildstate->index));
/* Attempt to launch parallel worker scan when required */
if (parallel_workers > 0)
@@ -937,29 +853,20 @@ BuildGraph(HnswBuildState * buildstate, ForkNumber forkNum)
HnswBeginParallel(buildstate, buildstate->indexInfo->ii_Concurrent, parallel_workers);
}
/* Add tuples to graph */
if (buildstate->heap != NULL)
/* Add tuples to sort */
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
if (buildstate->hnswleader)
buildstate->reltuples = ParallelHeapScan(buildstate);
else
{
#if PG_VERSION_NUM >= 120000
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL);
buildstate->reltuples = table_index_build_scan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, true, BuildCallback, (void *) buildstate, NULL);
#else
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
buildstate->reltuples = IndexBuildHeapScan(buildstate->heap, buildstate->index, buildstate->indexInfo,
true, BuildCallback, (void *) buildstate, NULL);
#endif
}
buildstate->indtuples = buildstate->graph->indtuples;
}
/* Flush pages */
if (!buildstate->graph->flushed)
FlushPages(buildstate);
/* End parallel build */
if (buildstate->hnswleader)
HnswEndParallel(buildstate->hnswleader);
@@ -988,13 +895,13 @@ static void
BuildIndex(Relation heap, Relation index, IndexInfo *indexInfo,
HnswBuildState * buildstate, ForkNumber forkNum)
{
#ifdef HNSW_MEMORY
SeedRandom(42);
#endif
InitBuildState(buildstate, heap, index, indexInfo, forkNum);
BuildGraph(buildstate, forkNum);
if (buildstate->heap != NULL)
BuildGraph(buildstate, forkNum);
if (!buildstate->flushed)
FlushPages(buildstate);
if (RelationNeedsWAL(index))
log_newpage_range(index, forkNum, 0, RelationGetNumberOfBlocks(index), true);

View File

@@ -5,7 +5,6 @@
#include "hnsw.h"
#include "storage/bufmgr.h"
#include "storage/lmgr.h"
#include "utils/datum.h"
#include "utils/memutils.h"
/*
@@ -220,9 +219,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
HnswInsertAppendPage(index, &newbuf, &newpage, state, page, building);
/* Commit */
if (building)
MarkBufferDirty(buf);
else
if (!building)
GenericXLogFinish(state);
/* Unlock previous buffer */
@@ -297,13 +294,7 @@ WriteNewElementPages(Relation index, HnswElement e, int m, BlockNumber insertPag
}
/* Commit */
if (building)
{
MarkBufferDirty(buf);
if (nbuf != buf)
MarkBufferDirty(nbuf);
}
else
if (!building)
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
if (nbuf != buf)
@@ -343,7 +334,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
for (int lc = e->level; lc >= 0; lc--)
{
int lm = HnswGetLayerM(m, lc);
HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc];
for (int i = 0; i < neighbors->length; i++)
{
@@ -430,9 +421,7 @@ HnswUpdateNeighborPages(Relation index, FmgrInfo *procinfo, Oid collation, HnswE
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
if (building)
MarkBufferDirty(buf);
else
if (!building)
GenericXLogFinish(state);
}
else if (!building)
@@ -491,56 +480,34 @@ HnswAddDuplicate(Relation index, HnswElement element, HnswElement dup, bool buil
}
/* Add heap TID */
etup->heaptids[i] = element->heaptids[0];
etup->heaptids[i] = *((ItemPointer) linitial(element->heaptids));
/* Overwrite tuple */
if (!PageIndexTupleOverwrite(page, dup->offno, (Item) etup, etupSize))
elog(ERROR, "failed to add index item to \"%s\"", RelationGetRelationName(index));
/* Commit */
if (building)
MarkBufferDirty(buf);
else
if (!building)
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
return true;
}
/*
* Find duplicate element
*/
static bool
HnswFindDuplicate(Relation index, HnswElement element, bool building)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(element, 0);
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (!datumIsEqual(element->value, neighbor->element->value, false, -1))
return false;
if (HnswAddDuplicate(index, element, neighbor->element, building))
return true;
}
return false;
}
/*
* Write changes to disk
*/
static void
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement entryPoint, bool building)
WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement element, int m, int efConstruction, HnswElement dup, HnswElement entryPoint, bool building)
{
BlockNumber newInsertPage = InvalidBlockNumber;
/* Look for duplicate */
if (HnswFindDuplicate(index, element, building))
return;
/* Try to add to existing page */
if (dup != NULL)
{
if (HnswAddDuplicate(index, element, dup, building))
return;
}
/* Write element and neighbor tuples */
WriteNewElementPages(index, element, m, GetInsertPage(index), &newInsertPage, building);
@@ -552,7 +519,7 @@ WriteElement(Relation index, FmgrInfo *procinfo, Oid collation, HnswElement elem
/* Update neighbors */
HnswUpdateNeighborPages(index, procinfo, collation, element, m, false, building);
/* Update entry point if needed */
/* Update metapage if needed */
if (entryPoint == NULL || element->level > entryPoint->level)
HnswUpdateMetaPage(index, HNSW_UPDATE_ENTRY_GREATER, element, InvalidBlockNumber, MAIN_FORKNUM, building);
}
@@ -571,6 +538,7 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
int efConstruction = HnswGetEfConstruction(index);
FmgrInfo *procinfo = index_getprocinfo(index, 1, HNSW_DISTANCE_PROC);
Oid collation = index->rd_indcollation[0];
HnswElement dup;
LOCKMODE lockmode = ShareLock;
/* Detoast once for all calls */
@@ -615,8 +583,11 @@ HnswInsertTuple(Relation index, Datum *values, bool *isnull, ItemPointer heap_ti
/* Insert element in graph */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, false);
/* Look for duplicate */
dup = HnswFindDuplicate(element);
/* Write to disk */
WriteElement(index, procinfo, collation, element, m, efConstruction, entryPoint, building);
WriteElement(index, procinfo, collation, element, m, efConstruction, dup, entryPoint, building);
/* Release lock */
UnlockPage(index, HNSW_UPDATE_LOCK, lockmode);

View File

@@ -188,13 +188,15 @@ hnswgettuple(IndexScanDesc scan, ScanDirection dir)
ItemPointer heaptid;
/* Move to next element if no valid heap TIDs */
if (hc->element->heaptidsLength == 0)
if (list_length(hc->element->heaptids) == 0)
{
so->w = list_delete_last(so->w);
continue;
}
heaptid = &hc->element->heaptids[--hc->element->heaptidsLength];
heaptid = llast(hc->element->heaptids);
hc->element->heaptids = list_delete_last(hc->element->heaptids);
MemoryContextSwitchTo(oldCtx);

View File

@@ -206,6 +206,17 @@ HnswInitNeighbors(HnswElement element, int m)
}
}
/*
* Free neighbors
*/
static void
HnswFreeNeighbors(HnswElement element)
{
for (int lc = 0; lc <= element->level; lc++)
pfree(element->neighbors[lc].items);
pfree(element->neighbors);
}
/*
* Allocate an element
*/
@@ -220,7 +231,7 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
if (level > maxLevel)
level = maxLevel;
element->heaptidsLength = 0;
element->heaptids = NIL;
HnswAddHeapTid(element, heaptid);
element->level = level;
@@ -233,13 +244,29 @@ HnswInitElement(ItemPointer heaptid, int m, double ml, int maxLevel)
return element;
}
/*
* Free an element
*/
void
HnswFreeElement(HnswElement element)
{
HnswFreeNeighbors(element);
list_free_deep(element->heaptids);
if (DatumGetPointer(element->value))
pfree(DatumGetPointer(element->value));
pfree(element);
}
/*
* Add a heap TID to an element
*/
void
HnswAddHeapTid(HnswElement element, ItemPointer heaptid)
{
element->heaptids[element->heaptidsLength++] = *heaptid;
ItemPointer copy = palloc(sizeof(ItemPointerData));
ItemPointerCopy(heaptid, copy);
element->heaptids = lappend(element->heaptids, copy);
}
/*
@@ -352,9 +379,7 @@ HnswUpdateMetaPage(Relation index, int updateEntry, HnswElement entryPoint, Bloc
HnswUpdateMetaPageInfo(page, updateEntry, entryPoint, insertPage);
if (building)
MarkBufferDirty(buf);
else
if (!building)
GenericXLogFinish(state);
UnlockReleaseBuffer(buf);
}
@@ -370,8 +395,8 @@ HnswSetElementTuple(HnswElementTuple etup, HnswElement element)
etup->deleted = 0;
for (int i = 0; i < HNSW_HEAPTIDS; i++)
{
if (i < element->heaptidsLength)
etup->heaptids[i] = element->heaptids[i];
if (i < list_length(element->heaptids))
etup->heaptids[i] = *((ItemPointer) list_nth(element->heaptids, i));
else
ItemPointerSetInvalid(&etup->heaptids[i]);
}
@@ -390,7 +415,7 @@ HnswSetNeighborTuple(HnswNeighborTuple ntup, HnswElement e, int m)
for (int lc = e->level; lc >= 0; lc--)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(e, lc);
HnswNeighborArray *neighbors = &e->neighbors[lc];
int lm = HnswGetLayerM(m, lc);
for (int i = 0; i < lm; i++)
@@ -448,7 +473,7 @@ LoadNeighborsFromPage(HnswElement element, Relation index, Page page, int m)
if (level < 0)
level = 0;
neighbors = HnswGetNeighbors(element, level);
neighbors = &element->neighbors[level];
hc = &neighbors->items[neighbors->length++];
hc->element = e;
}
@@ -482,7 +507,7 @@ HnswLoadElementFromTuple(HnswElement element, HnswElementTuple etup, bool loadHe
element->deleted = etup->deleted;
element->neighborPage = ItemPointerGetBlockNumber(&etup->neighbortid);
element->neighborOffno = ItemPointerGetOffsetNumber(&etup->neighbortid);
element->heaptidsLength = 0;
element->heaptids = NIL;
if (loadHeaptids)
{
@@ -625,12 +650,13 @@ AddToVisited(visited_hash v, HnswCandidate * hc, Relation index, bool *found)
List *
HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *procinfo, Oid collation, int m, bool inserting, HnswElement skipElement)
{
ListCell *lc2;
List *w = NIL;
pairingheap *C = pairingheap_allocate(CompareNearestCandidates, NULL);
pairingheap *W = pairingheap_allocate(CompareFurthestCandidates, NULL);
int wlen = 0;
visited_hash v;
ListCell *lc2;
/* Create hash table */
if (index == NULL)
@@ -654,7 +680,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* would be ideal to do this for inserts as well, but this could
* affect insert performance.
*/
if (skipElement == NULL || hc->element->heaptidsLength != 0)
if (skipElement == NULL || list_length(hc->element->heaptids) != 0)
wlen++;
}
@@ -671,7 +697,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
HnswLoadNeighbors(c->element, index, m);
/* Get the neighborhood at layer lc */
neighborhood = HnswGetNeighbors(c->element, lc);
neighborhood = &c->element->neighbors[lc];
for (int i = 0; i < neighborhood->length; i++)
{
@@ -713,7 +739,7 @@ HnswSearchLayer(Datum q, List *ep, int ef, int lc, Relation index, FmgrInfo *pro
* vacuuming. It would be ideal to do this for inserts as
* well, but this could affect insert performance.
*/
if (skipElement == NULL || e->element->heaptidsLength != 0)
if (skipElement == NULL || list_length(e->element->heaptids) != 0)
{
wlen++;
@@ -774,23 +800,23 @@ HnswGetDistance(HnswElement a, HnswElement b, int lc, FmgrInfo *procinfo, Oid co
/* Look for cached distance */
if (a->neighbors != NULL)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(a, lc);
Assert(a->level >= lc);
for (int i = 0; i < neighbors->length; i++)
for (int i = 0; i < a->neighbors[lc].length; i++)
{
if (neighbors->items[i].element == b)
return neighbors->items[i].distance;
if (a->neighbors[lc].items[i].element == b)
return a->neighbors[lc].items[i].distance;
}
}
if (b->neighbors != NULL)
{
HnswNeighborArray *neighbors = HnswGetNeighbors(b, lc);
Assert(b->level >= lc);
for (int i = 0; i < neighbors->length; i++)
for (int i = 0; i < b->neighbors[lc].length; i++)
{
if (neighbors->items[i].element == a)
return neighbors->items[i].distance;
if (b->neighbors[lc].items[i].element == a)
return b->neighbors[lc].items[i].distance;
}
}
@@ -826,8 +852,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
List *r = NIL;
List *w = list_copy(c);
pairingheap *wd;
HnswNeighborArray *neighbors = HnswGetNeighbors(e2, lc);
bool mustCalculate = !neighbors->closerSet;
bool mustCalculate = !e2->neighbors[lc].closerSet;
List *added = NIL;
bool removedAny = false;
@@ -891,7 +916,7 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
}
/* Cached value can only be used in future if sorted deterministically */
neighbors->closerSet = sortCandidates;
e2->neighbors[lc].closerSet = sortCandidates;
/* Keep pruned connections */
while (!pairingheap_is_empty(wd) && list_length(r) < m)
@@ -909,6 +934,30 @@ SelectNeighbors(List *c, int m, int lc, FmgrInfo *procinfo, Oid collation, HnswE
return r;
}
/*
* Find duplicate element
*/
HnswElement
HnswFindDuplicate(HnswElement e)
{
HnswNeighborArray *neighbors = &e->neighbors[0];
for (int i = 0; i < neighbors->length; i++)
{
HnswCandidate *neighbor = &neighbors->items[i];
/* Exit early since ordered by distance */
if (!datumIsEqual(e->value, neighbor->element->value, false, -1))
break;
/* Check for space */
if (list_length(neighbor->element->heaptids) < HNSW_HEAPTIDS)
return neighbor->element;
}
return NULL;
}
/*
* Add connections
*/
@@ -916,7 +965,7 @@ static void
AddConnections(HnswElement element, List *neighbors, int m, int lc)
{
ListCell *lc2;
HnswNeighborArray *a = HnswGetNeighbors(element, lc);
HnswNeighborArray *a = &element->neighbors[lc];
foreach(lc2, neighbors)
a->items[a->length++] = *((HnswCandidate *) lfirst(lc2));
@@ -928,7 +977,7 @@ AddConnections(HnswElement element, List *neighbors, int m, int lc)
void
HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int *updateIdx, Relation index, FmgrInfo *procinfo, Oid collation)
{
HnswNeighborArray *currentNeighbors = HnswGetNeighbors(hc->element, lc);
HnswNeighborArray *currentNeighbors = &hc->element->neighbors[lc];
HnswCandidate hc2;
@@ -963,7 +1012,7 @@ HnswUpdateConnection(HnswElement element, HnswCandidate * hc, int m, int lc, int
hc3->distance = GetCandidateDistance(hc3, q, procinfo, collation);
/* Prune element if being deleted */
if (hc3->element->heaptidsLength == 0)
if (list_length(hc3->element->heaptids) == 0)
{
pruned = &currentNeighbors->items[i];
break;
@@ -1021,7 +1070,7 @@ RemoveElements(List *w, HnswElement skipElement)
if (skipElement != NULL && hc->element->blkno == skipElement->blkno && hc->element->offno == skipElement->offno)
continue;
if (hc->element->heaptidsLength != 0)
if (list_length(hc->element->heaptids) != 0)
w2 = lappend(w2, hc);
}

View File

@@ -206,7 +206,7 @@ RepairGraphElement(HnswVacuumState * vacuumstate, HnswElement element, HnswEleme
/* Init fields */
HnswInitNeighbors(element, m);
element->heaptidsLength = 0;
element->heaptids = NIL;
/* Add element to graph, skipping itself */
HnswInsertElement(element, entryPoint, index, procinfo, collation, m, efConstruction, true);

View File

@@ -543,10 +543,10 @@ CreateListPages(Relation index, VectorArray centers, int dimensions,
pfree(list);
}
#ifdef IVFFLAT_KMEANS_DEBUG
/*
* Print k-means metrics
*/
#ifdef IVFFLAT_KMEANS_DEBUG
static void
PrintKmeansMetrics(IvfflatBuildState * buildstate)
{

View File

@@ -72,14 +72,12 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
double ratio;
double spc_seq_page_cost;
Relation index;
double selectivity = 1;
ListCell *lc;
#if PG_VERSION_NUM < 120000
List *qinfos;
#endif
/* Never use index without order or limit */
if (path->indexorderbys == NULL || root->limit_tuples < 0)
/* Never use index without order */
if (path->indexorderbys == NULL)
{
*indexStartupCost = DBL_MAX;
*indexTotalCost = DBL_MAX;
@@ -107,29 +105,6 @@ ivfflatcostestimate(PlannerInfo *root, IndexPath *path, double loop_count,
*/
costs.numIndexTuples = path->indexinfo->tuples * ratio;
/* Get the selectivity of non-index conditions */
foreach(lc, path->indexinfo->indrestrictinfo)
{
RestrictInfo *rinfo = lfirst(lc);
if (rinfo->norm_selec >= 0 && rinfo->norm_selec <= 1 && rinfo->norm_selec != (Selectivity) DEFAULT_INEQ_SEL)
selectivity *= rinfo->norm_selec;
}
/*
* Do not use index if limit + offset > expected tuples unless
* enable_seqscan = off
*/
if (root->limit_tuples > costs.numIndexTuples * selectivity)
{
*indexStartupCost = 1.0e10 - 1;
*indexTotalCost = 1.0e10 - 1;
*indexSelectivity = 0;
*indexCorrelation = 0;
*indexPages = 0;
return;
}
#if PG_VERSION_NUM >= 120000
genericcostestimate(root, path, loop_count, &costs);
#else

View File

@@ -6,10 +6,6 @@
#include "ivfflat.h"
#include "miscadmin.h"
#ifdef IVFFLAT_MEMORY
#include "utils/memutils.h"
#endif
/*
* Initialize with kmeans++
*
@@ -155,23 +151,6 @@ QuickCenters(Relation index, VectorArray samples, VectorArray centers)
}
}
#ifdef IVFFLAT_MEMORY
/*
* Show memory usage
*/
static void
ShowMemoryUsage(Size estimatedSize)
{
#if PG_VERSION_NUM >= 130000
elog(INFO, "total memory: %zu MB",
MemoryContextMemAllocated(CurrentMemoryContext, true) / (1024 * 1024));
#else
MemoryContextStats(CurrentMemoryContext);
#endif
elog(INFO, "estimated memory: %zu MB", estimatedSize / (1024 * 1024));
}
#endif
/*
* Use Elkan for performance. This requires distance function to satisfy triangle inequality.
*
@@ -252,10 +231,6 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
vec->dim = dimensions;
}
#ifdef IVFFLAT_MEMORY
ShowMemoryUsage(totalSize);
#endif
/* Pick initial centers */
InitCenters(index, samples, centers, lowerBound);

View File

@@ -177,15 +177,14 @@ PGDLLEXPORT PG_FUNCTION_INFO_V1(vector_in);
Datum
vector_in(PG_FUNCTION_ARGS)
{
char *lit = PG_GETARG_CSTRING(0);
char *str = PG_GETARG_CSTRING(0);
int32 typmod = PG_GETARG_INT32(2);
float x[VECTOR_MAX_DIM];
int dim = 0;
char *pt;
char *stringEnd;
Vector *result;
char *litcopy = pstrdup(lit);
char *str = litcopy;
char *lit = pstrdup(str);
while (vector_isspace(*str))
str++;
@@ -269,7 +268,7 @@ vector_in(PG_FUNCTION_ARGS)
(errcode(ERRCODE_DATA_EXCEPTION),
errmsg("vector must have at least 1 dimension")));
pfree(litcopy);
pfree(lit);
CheckExpectedDim(typmod, dim);

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
val
---------
[1,1,1]
@@ -11,13 +11,13 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
[1,2,4]
(3 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
count
-------
3
(1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
count
-------
3

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
val
---------
[1,2,4]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2;
count
-------
4

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
[0,0,0]
(4 rows)
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
val
---------
[0,0,0]
@@ -28,7 +28,7 @@ SELECT COUNT(*) FROM t;
(1 row)
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)

View File

@@ -2,7 +2,7 @@ SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_cosine_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
val
---------
[1,1,1]
@@ -11,13 +11,13 @@ SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
[1,2,4]
(3 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
count
-------
3
(1 row)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
count
-------
3

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
val
---------
[1,2,4]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
[0,0,0]
(4 rows)
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2;
count
-------
4

View File

@@ -3,7 +3,7 @@ CREATE TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]
@@ -12,7 +12,7 @@ SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
[0,0,0]
(4 rows)
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
val
---------
[0,0,0]
@@ -31,7 +31,7 @@ TRUNCATE t;
NOTICE: ivfflat index created with little data
DETAIL: This will cause low recall.
HINT: Drop the index until the table has more data.
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
-----
(0 rows)

View File

@@ -2,7 +2,7 @@ SET enable_seqscan = off;
CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
val
---------
[1,2,3]

View File

@@ -6,8 +6,8 @@ CREATE INDEX ON t USING hnsw (val vector_cosine_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
DROP TABLE t;

View File

@@ -6,7 +6,7 @@ CREATE INDEX ON t USING hnsw (val vector_ip_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2;
DROP TABLE t;

View File

@@ -6,11 +6,11 @@ CREATE INDEX ON t USING hnsw (val vector_l2_ops);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
SELECT COUNT(*) FROM t;
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -4,6 +4,6 @@ CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING hnsw (val vector_l2_ops);
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -6,8 +6,8 @@ CREATE INDEX ON t USING ivfflat (val vector_cosine_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <=> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]' LIMIT 5) t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector) LIMIT 5) t2;
SELECT * FROM t ORDER BY val <=> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> '[0,0,0]') t2;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <=> (SELECT NULL::vector)) t2;
DROP TABLE t;

View File

@@ -6,7 +6,7 @@ CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <#> '[3,3,3]' LIMIT 5;
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector) LIMIT 5) t2;
SELECT * FROM t ORDER BY val <#> '[3,3,3]';
SELECT COUNT(*) FROM (SELECT * FROM t ORDER BY val <#> (SELECT NULL::vector)) t2;
DROP TABLE t;

View File

@@ -6,11 +6,11 @@ CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
INSERT INTO t (val) VALUES ('[1,2,4]');
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector) LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
SELECT * FROM t ORDER BY val <-> (SELECT NULL::vector);
SELECT COUNT(*) FROM t;
TRUNCATE t;
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -4,6 +4,6 @@ CREATE UNLOGGED TABLE t (val vector(3));
INSERT INTO t (val) VALUES ('[0,0,0]'), ('[1,2,3]'), ('[1,1,1]'), (NULL);
CREATE INDEX ON t USING ivfflat (val vector_l2_ops) WITH (lists = 1);
SELECT * FROM t ORDER BY val <-> '[3,3,3]' LIMIT 5;
SELECT * FROM t ORDER BY val <-> '[3,3,3]';
DROP TABLE t;

View File

@@ -49,7 +49,7 @@ is(idx_scan(), 0);
$count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 100;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 20000) t;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
is($count, $expected);
is(idx_scan(), 1);

View File

@@ -42,7 +42,7 @@ for my $i (1 .. 20)
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 20) t;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
is($count, 10);
@@ -63,7 +63,7 @@ $node->pgbench(
my $count = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1000;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1) LIMIT 1000) t;
SELECT COUNT(*) FROM (SELECT v FROM tst ORDER BY v <-> (SELECT v FROM tst LIMIT 1)) t;
));
# Elements may lose all incoming connections with the HNSW algorithm
# Vacuuming can fix this if one of the elements neighbors is deleted

View File

@@ -26,7 +26,7 @@ sub test_duplicates
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET hnsw.ef_search = 1;
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]' LIMIT 20) t;
SELECT COUNT(*) FROM (SELECT * FROM tst ORDER BY v <-> '[1,1,1]') t;
));
is($res, 10);
}

View File

@@ -1,111 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my $nc = 50;
my $limit = 20;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4, t text);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc, 'test ' || i FROM generate_series(1, 10000) 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 $c = int(rand() * $nc);
# Test attribute filtering
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c != $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%test%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed like
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE t LIKE '%%other%%' ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test distance filtering
$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/);
# Test distance filtering greater than distance
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' > 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test distance filtering without order
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1;
));
like($explain, qr/Seq Scan/);
# Test distance filtering without limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
));
like($explain, qr/Seq Scan/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Bitmap Index Scan on attribute_idx/);
# Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING hnsw (v vector_l2_ops) WHERE (c = $c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using partial_idx/);
done_testing();

View File

@@ -0,0 +1,43 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i serial, v vector($dim));");
$node->safe_psql("postgres",
"INSERT INTO tst (v) SELECT ARRAY[$array_sql] FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops);");
# Delete data
$node->safe_psql("postgres", "DELETE FROM tst WHERE i % 100 != 0;");
my $exp = $node->safe_psql("postgres", qq(
SET enable_indexscan = off;
SELECT i FROM tst ORDER BY v <-> '[0,0,0]';
));
# Run twice to make sure correct tuples marked as dead
for (1 .. 2)
{
my $res = $node->safe_psql("postgres", qq(
SET enable_seqscan = off;
SET ivfflat.probes = 100;
SELECT i FROM tst ORDER BY v <-> '[0,0,0]';
));
is($res, $exp);
}
done_testing();

View File

@@ -1,99 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
my $dim = 3;
my $nc = 50;
my $limit = 20;
my $array_sql = join(",", ('random()') x $dim);
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (i int4, v vector($dim), c int4);");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT i, ARRAY[$array_sql], i % $nc FROM generate_series(1, 10000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 100);");
$node->safe_psql("postgres", "ANALYZE tst;");
# Generate query
my @r = ();
for (1 .. $dim)
{
push(@r, rand());
}
my $query = "[" . join(",", @r) . "]";
my $c = int(rand() * $nc);
# Test attribute filtering
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test attribute filtering with few rows removed
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c != $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with few rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c >= 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using idx/);
# Test attribute filtering with many rows removed comparison
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c < 1 ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Seq Scan/);
# Test distance filtering
$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/);
# Test distance filtering greater than distance
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' > 1 ORDER BY v <-> '$query' LIMIT $limit;
));
# TODO Do not use index
like($explain, qr/Index Scan using idx/);
# Test distance filtering without order
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1;
));
like($explain, qr/Seq Scan/);
# Test distance filtering without limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE v <-> '$query' < 1 ORDER BY v <-> '$query';
));
like($explain, qr/Seq Scan/);
# Test attribute index
$node->safe_psql("postgres", "CREATE INDEX attribute_idx ON tst (c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Bitmap Index Scan on attribute_idx/);
# Test partial index
$node->safe_psql("postgres", "CREATE INDEX partial_idx ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 5) WHERE (c = $c);");
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT i FROM tst WHERE c = $c ORDER BY v <-> '$query' LIMIT $limit;
));
like($explain, qr/Index Scan using partial_idx/);
done_testing();

View File

@@ -1,64 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 1000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING ivfflat (v vector_l2_ops) WITH (lists = 10);");
# Test limit
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 100;
));
like($explain, qr/Index Scan/);
# Test limit with probes
$explain = $node->safe_psql("postgres", qq(
SET ivfflat.probes = 2;
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 200;
));
like($explain, qr/Index Scan/);
# Test limit + offset
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 90 OFFSET 10;
));
like($explain, qr/Index Scan/);
# Test limit > expected tuples
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 101;
));
like($explain, qr/Seq Scan/);
# Test limit > expected tuples with probes
$explain = $node->safe_psql("postgres", qq(
SET ivfflat.probes = 2;
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 201;
));
like($explain, qr/Seq Scan/);
# Test limit + offset > expected tuples
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 91 OFFSET 10;
));
like($explain, qr/Seq Scan/);
# Test no limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]';
));
like($explain, qr/Seq Scan/);
done_testing();

View File

@@ -1,62 +0,0 @@
use strict;
use warnings;
use PostgresNode;
use TestLib;
use Test::More;
# Initialize node
my $node = get_new_node('node');
$node->init;
$node->start;
# Create table and index
$node->safe_psql("postgres", "CREATE EXTENSION vector;");
$node->safe_psql("postgres", "CREATE TABLE tst (v vector(3));");
$node->safe_psql("postgres",
"INSERT INTO tst SELECT ARRAY[random(), random(), random()] FROM generate_series(1, 1000) i;"
);
$node->safe_psql("postgres", "CREATE INDEX ON tst USING hnsw (v vector_l2_ops);");
# Test limit
my $explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 40;
));
like($explain, qr/Index Scan/);
# Test limit with CTE
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE WITH cte AS (SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 40) SELECT * FROM cte;
));
like($explain, qr/Index Scan/);
# Test limit + offset
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 30 OFFSET 10;
));
like($explain, qr/Index Scan/);
# Test limit > ef_search
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 41;
));
like($explain, qr/Seq Scan/);
# Test limit > ef_search with CTE
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE WITH cte AS (SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 41) SELECT * FROM cte;
));
like($explain, qr/Seq Scan/);
# Test limit + offset > ef_search
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]' LIMIT 31 OFFSET 10;
));
like($explain, qr/Seq Scan/);
# Test no limit
$explain = $node->safe_psql("postgres", qq(
EXPLAIN ANALYZE SELECT * FROM tst ORDER BY v <-> '[1,2,3]';
));
like($explain, qr/Seq Scan/);
done_testing();