pgvecto.rs Documentation

• 1: Getting Started
  • 1.1: Overview
  • 1.2: Installation
• 2: Usage
  • 2.1: Indexing
  • 2.2: Search
  • 2.3: Monitoring
  • 2.4: Quantization
  • 2.5: `pgvector` compatibility
• 3: Use Cases
  • 3.1: Hybrid Search
  • 3.2: Adaptive retrieval and shortening embeddings
  • 3.3: Image Search
  • 3.4: Multi Tenancy
  • 3.5: Sparse Vector Search
  • 3.6: Generative Question-Answering
• 4: Integration
  • 4.1: LangChain
  • 4.2: LlamaIndex
  • 4.3: Citus
• 5: Administration
  • 5.1: Configuration
  • 5.2: Upgrading
  • 5.3: Logical replication
  • 5.4: Foreign data wrapper (FDW)
  • 5.5: Kubernetes
  • 5.6: Migration
• 6: FAQs
  • 6.1: General FAQ
  • 6.2: Comparison with pgvector
  • 6.3: Why not a specialty vector database?
  • 6.4: Benchmarks
• 7: Reference
  • 7.1: Schema
  • 7.2: Indexing Options
  • 7.3: Search Options
  • 7.4: Vector Types
    • 7.4.1: Half-Precision Vector: `vecf16`
    • 7.4.2: 8-Bit Integer Vector: `veci8`
    • 7.4.3: Sparse Vector: `svector`
    • 7.4.4: Binary Vector: `bvector`
• 8: Community
  • 8.1: Contributing to `pgvecto.rs`
  • 8.2: Roadmap
  • 8.3: Adopters
• 9: Developers
  • 9.1: Development

1 - Getting Started

1.1 - Overview

An introduction to the pgvecto.rs.

What is pgvecto.rs

pgvecto.rs is a Postgres extension that provides vector similarity search functions. It is written in Rust and based on pgrx. It is currently in the beta status, we invite you to try it out in production and provide us with feedback. Read more at 📝our launch blog.

Comparison with pgvector

Checkout pgvecto.rs vs pgvector for more details.

Quick start

For new users, we recommend using the Docker image to get started quickly.

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

Then you can connect to the database using the psql command line tool. The default username is postgres, and the default password is mysecretpassword.

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

Run the following SQL to ensure the extension is enabled.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

pgvecto.rs introduces a new data type vector(n) denoting an n-dimensional vector. The n within the brackets signifies the dimensions of the vector.

You could create a table with the following SQL.

-- create table with a vector column

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  embedding vector(3) NOT NULL -- 3 dimensions
);

You can then populate the table with vector data as follows.

-- insert values

INSERT INTO items (embedding)
VALUES ('[1,2,3]'), ('[4,5,6]');

-- or insert values using a casting from array to vector

INSERT INTO items (embedding)
VALUES (ARRAY[1, 2, 3]::real[]), (ARRAY[4, 5, 6]::real[]);

We support three operators to calculate the distance between two vectors.

• <->: squared Euclidean distance, defined as $\Sigma (x_i - y_i) ^ 2$.
• <#>: negative dot product, defined as $- \Sigma x_iy_i$.
• <=>: cosine distance, defined as $1 - \frac{\Sigma x_iy_i}{\sqrt{\Sigma x_i^2 \Sigma y_i^2}}$.

-- call the distance function through operators

-- squared Euclidean distance
SELECT '[1, 2, 3]'::vector <-> '[3, 2, 1]'::vector;
-- negative dot product
SELECT '[1, 2, 3]'::vector <#> '[3, 2, 1]'::vector;
-- cosine distance
SELECT '[1, 2, 3]'::vector <=> '[3, 2, 1]'::vector;

You can search for a vector simply like this.

-- query the similar embeddings
SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

A simple Question-Answering application

Please check out the Question-Answering application tutorial.

Half-precision floating-point

vecf16 type is the same with vector in anything but the scalar type. It stores 16-bit floating point numbers. If you want to reduce the memory usage to get better performance, you can try to replace vector type with vecf16 type.

For more usage of vecf16, please refer to vector types.

Sparse vector

svector type is a sparse vector type. It stores a vector in a sparse format. It is suitable for vectors with many zeros.

For more usage of svector, please refer to vector types.

Binary vector

bvector type is a binary vector type. It is a fixed-length bit string. Except for above 3 distances, we also support jaccard distance <~>, which defined as $1 - \frac{|X\cap Y|}{|X\cup Y|}$. And hamming distance is the same with squared Euclidean distance, you can use <-> operator to calculate it. We also provide binarize function to construct a bvector from a vector, which set the positive elements to 1, otherwise 0.

For more usage of bvector, please refer to vector types.

Roadmap 🗂️

Please check out ROADMAP.

Contribute 😊

We welcome all kinds of contributions from the open-source community, individuals, and partners.

• Join our discord community!
• To build from the source, please read our contributing documentation and development tutorial.

Talk with us

💬 Interested in talking with us about your experience building or managing AI/ML applications?

Set up a time to chat!

1.2 - Installation

There are four ways to install pgvecto.rs.

Docker

The easiest way to try pgvecto.rs is to run it from a ready-to use Docker image.

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

Then you can connect to the database using the psql command line tool. The default username is postgres, and the default password is mysecretpassword.

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

Run the following SQL to ensure the extension is enabled.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

To achieve full performance, please mount the volume to pg data directory by adding the option like -v $PWD/pgdata:/var/lib/postgresql/data

You can configure PostgreSQL by the reference of the parent image in https://hub.docker.com/_/postgres/.

From Debian package

Debian packages(.deb) are used in distributions based on Debian, such as Ubuntu and many others. They can be easily installed by dpkg or apt-get.

1. Download the deb package in the release page, and type sudo apt install vectors-pg15-*.deb to install the deb package.

2. Configure your PostgreSQL by modifying the shared_preload_libraries and search_path to include the extension.

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# You need restart the PostgreSQL cluster to take effects.
sudo systemctl restart postgresql.service   # for pgvecto.rs running with systemd

3. Connect to the database and enable the extension.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

From ZIP package

For systems that are not Debian based and cannot run a Docker container, please follow these steps to install:

1. Before install, make sure that you have the necessary packages installed, including PostgreSQL, pg_config, unzip, wget.

# Example for RHEL 9 dnf
# Please check your package manager
sudo dnf install -y unzip wget libpq-devel
sudo dnf module install -y postgresql:15/server
sudo postgresql-setup --initdb
sudo systemctl start postgresql.service
sudo systemctl enable postgresql.service

2. Verify whether $pkglibdir and $shardir have been set by PostgreSQL.

pg_config --pkglibdir
# Print something similar to:
# /usr/lib/postgresql/15/lib or
# /usr/lib64/pgsql

pg_config --sharedir
# Print something similar to:
# /usr/share/postgresql/15 or
# /usr/share/pgsql

3. Download the zip package in the release page and extract it to a temporary directory.

wget https://github.com/tensorchord/pgvecto.rs/releases/download/v0.2.1/vectors-pg15_x86_64-unknown-linux-gnu_0.2.1.zip -O vectors.zip
unzip vectors.zip -d vectors

4. Copy the extension files to the PostgreSQL directory.

# Copy library to `$pkglibdir`
sudo cp vectors/vectors.so $(pg_config --pkglibdir)/
# Copy schema to `$shardir`
sudo cp vectors/vectors--* $(pg_config --sharedir)/extension/
sudo cp vectors/vectors.control $(pg_config --sharedir)/extension/

5. Configure your PostgreSQL by modifying the shared_preload_libraries and search_path to include the extension.

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# You need restart the PostgreSQL cluster to take effects.
sudo systemctl restart postgresql.service   # for pgvecto.rs running with systemd

6. Connect to the database and enable the extension.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

From source

Before building from source, you could refer to the development guide to set up the development environment.

1. Then you could build and install the extension.

sed -e "s/@CARGO_VERSION@/0.0.0/g" ./vectors.control  # 0.0.0 here is the version number that you need
sudo cp -a ./sql/upgrade/. $(pg_config --sharedir)/extension
cargo pgrx install --sudo --release

2. Configure your PostgreSQL by modifying the shared_preload_libraries and search_path to include the extension.

psql -U postgres -c 'ALTER SYSTEM SET shared_preload_libraries = "vectors.so"'
psql -U postgres -c 'ALTER SYSTEM SET search_path TO "$user", public, vectors'
# You need restart the PostgreSQL cluster to take effects.
sudo systemctl restart postgresql.service   # for pgvecto.rs running with systemd
service postgresql restart  # for pgvecto.rs running in envd

3. Connect to the database and enable the extension.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

2 - Usage

2.1 - Indexing

Indexing is the process of building a data structure that allows for efficient search. pgvecto.rs supports three indexing algorithms: brute force (flat), IVF, and HNSW. The default algorithm is HNSW.

To construct an index for vectors, first create a table named items with a column named embedding of type vector(n). Then, populate the table with generated data.

CREATE TABLE items (embedding vector(3));
INSERT INTO items (embedding) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000);

You can create a vector index for squared Euclidean distance with the following SQL.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops);

The vector_l2_ops operator class calculates squared Euclidean distance. Refer to the documentation for a list of operator classes.

To switch the indexing algorithm from the default HNSW to IVF, execute the SQL command below. For additional examples of usage, refer to the documentation.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
$$);

You can now perform a KNN search with the following SQL, leveraging the vector index for your query.

SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

Examples

Here are some examples of how to create indexes using various algorithms and settings. For a complete list of options, please consult the reference guide.

-- HNSW algorithm, default settings.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops);

-- Use FP16 vector type for optimizing data storage.

CREATE TABLE items (embedding vecf16(3));
CREATE INDEX ON items USING vectors (embedding vecf16_l2_ops)

--- Or using brute force with product quantization (PQ).

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.flat]
quantization.product.ratio = "x16"
$$);

-- Or using brute force with scalar quantization (SQ).
CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = "[indexing.hnsw.quantization.scalar]");

--- Or using IVFPQ algorithm.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
quantization.product.ratio = "x16"
$$);

-- Use more threads for background building the index.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
optimizing.optimizing_threads = 16
$$);

-- Prefer smaller HNSW graph.

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
segment.max_growing_segment_size = 200000
$$);

Suggestions

Selecting the appropriate indexing algorithm and settings is crucial for vector search performance. Below are recommendations tailored to various use cases.

Brute force (flat)

If you’re conducting a limited number of searches (between 1,000 and 10,000) or require precise results, the brute force algorithm is a viable option.

This algorithm sets the standard for other indexing methods. It retains vectors in their original form without compression and avoids additional overhead. While straightforward and user-friendly, it’s not ideal for processing large datasets.

In addition, you can use quantization to reduce memory usage. For a complete list of all options, please refer to the documentation. Detailed information about quantization can also be found there.

Inverted file index (IVF)

IVF mainly uses the idea of inverted indexing to store the vectors (id, vector) under each cluster center. When querying a vector, it finds the nearest several centers and searches for the vectors under these centers respectively.

Before you begin your search, it’s necessary to train nlist cluster centers. You also have the option to set nsample for K-Means clustering; a higher value yields more precise clusters but increases training time. For additional settings, please consult the reference guide.

Before searching, train nlist cluster centers using K-Means clustering with nsample points. A higher value yields more precise clustering but requires longer training time. For additional settings, please consult the reference guide.

Hierarchical navigable small world graph (HNSW)

HNSW merges the concepts of skip list with navigable small world (NSW) graphs, employing a hierarchical structure that features longer edges on higher layers for rapid searching and shorter edges on lower layers to enhance search precision.

Navigable small world

Navigable Small World operate on the premise that each node is accessible from any other through a minimal number of “hops.” Once an entry point is chosen—be it random, heuristic-based, or algorithmically determined—the surrounding nodes are examined to identify if they’re nearer than the current one. The search progresses to the nearest neighbor and continues until no closer nodes can be found.

When constructing the graph, the m parameter determines how many connections a new node will establish with its nearest neighbors. A higher m value results in a denser and more interconnected graph but also increases memory usage and slows down insertion times. During node insertion, the algorithm identifies the closest m nodes to create two-way links with them.

The search process navigates through an NSW graph, beginning at a pre-defined entry point. The algorithm proceeds by greedily moving to adjacent vertices closer to the query vector. The process stops once it finds no closer neighbors in the current vertex’s adjacency list.

A limitation of the NSW search is that it opts for the seemingly shortest path to the nearest node, disregarding the graph’s overall structure. This “greedy search” approach can result in getting stuck at a local optimum, an issue often referred to as “early stopping.”

Skip list

Skip lists are constructed using a probabilistic approach, where nodes become less likely to appear as you move up the layers, resulting in fewer nodes at higher levels. To search for a node, the algorithm starts at the top layer’s “head” and moves to the next node that is greater than or equal to the target. If the node isn’t found, it descends to a lower (more fine-grained) layer and continues this pattern until it locates either the desired node or its nearest neighbors.

HNSW

Constructing this multi-layered graph begins with a foundational layer that encapsulates the entire dataset. Moving up the hierarchy, each successive layer provides a streamlined summary of its predecessor, featuring fewer nodes and functioning as fast tracks for making broader leaps within the graph.

Smaller m) values are better for lower-dimensional data or when you require lower recall. Larger m values are useful for higher-dimensional data or when high recall is important.

The ef_construction parameter determines the dynamic candidate list size when adding new nodes; increasing this value may enhance recall but could extend index construction time.

The HNSW index is resource-intensive, requiring additional RAM and an adjustment to the maintenance_work_mem setting for larger datasets. If you’re seeking performance that’s markedly faster than IVF, with a high recall rate and scalability that matches dataset size, HNSW is an excellent option.

2.2 - Search

Get the nearest 5 neighbors to a vector

SET vectors.hnsw_ef_search = 64;
SELECT * FROM items ORDER BY embedding <-> '[3,2,1]' LIMIT 5;

Operators

These operators are used for distance metrics:

For their definitions, see overview.

Filter

For a given category, get the nearest 10 neighbors to a vector

SELECT 1 FROM items WHERE category_id = 1 ORDER BY embedding <#> '[0.5,0.5,0.5]' limit 10

Query options

Set ivf scan lists to 1 in session:

SET vectors.ivf_nprobe=1;

Set hnsw search scope to 40 in transaction:

SET LOCAL vectors.hnsw_ef_search=40;

For all options, refer to search options.

Search modes

There are two search modes: vbase and basic.

vbase

As the default search mode, vbase is suitable for most scenarios.

For how it works, refer to the thesis VBASE: Unifying Online Vector Similarity Search and Relational Queries via Relaxed Monotonicity.

basic

basic is behaviorally consistent with vector algorithm libraries such as faiss. It will be useful if you want to align them.

Enabling basic, you must respect these restrictions:

• Execute VACUUM after updates before searches.
• Search without filter

2.3 - Monitoring

We provide a view pg_vector_index_stat to monitor the progress of indexing.

{
    "vector": {
        "dimensions": 256,
        "distance": "L2",
        "kind": "F32"
    },
    "segment": {
        "max_growing_segment_size": 20000,
        "max_sealed_segment_size": 1000000
    },
    "optimizing": {
        "sealing_secs": 60,
        "sealing_size": 1,
        "delete_threshold": 0.2,
        "optimizing_threads": 2
    },
    "indexing": {
        "hnsw": {
            "m": 12,
            "ef_construction": 300,
            "quantization": {
                "trivial": {}
            }
        }
    }
}

If there are no more insertions and you wait for all indexing to be finished, you can check the field idx_indexing every 60 seconds util it’s true.

2.4 - Quantization

Quantization is a technique to reduce the memory consumption and computation complexity with the affordable decrease in precision. In the context of vector search, there are mainly two types of quantization: production quantization and scalar quantization.

Production Quantization

Production quantization is a method that trains several clustering models on the dataset and maps each vector to a list of cluster center indexes. It is commonly used address the problem of excessive memory usage with high-dimensional data.

In pgvecto.rs, you can enable the product quantization with:

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = $$
[indexing.ivf]
quantization.product.ratio = "x16"
$$);

More configurations can be found in indexing.

Scalar Quantization

Scalar quantization maps each value to low-bit format like from 32-bit float to 8-bit integer. Usually, this won’t affect the precision a lot but can reduce the memory usage and improve the query performance. pgvecto.rs supports uint8 scalar quantization with:

CREATE INDEX ON items USING vectors (embedding vector_l2_ops)
WITH (options = "[indexing.hnsw.quantization.scalar]");

Suggestions

When quantization is required, it’s always recommended to try the scalar quantization first. It’s a good balance between precision and memory usage.

If the memory usage is still too high, you can try the product quantization. This method can reduce the memory usage significantly but may affect the precision. Since it’s highly related to the training dataset. You need to try different ratio and sample values to find the best one for your dataset.

Otherwise, you may also use some dimension reduction methods like PCA to reduce the vector dimension.

2.5 - `pgvector` compatibility

pgvecto.rs is natively compatible with pgvector at:

• CREATE TABLE commands, e.g. CREATE TABLE t (val vector(3))
• INSERT INTO commands, e.g. INSERT INTO t (val) VALUES ('[0.6,0.6,0.6]')

pgvecto.rs can be configured to be compatible with pgvector at:

• Index options, which allows you to create index by USING hnsw (val vector_ip_ops)
• Query options, which can be set by SET ivfflat.probes = 10

This feature is named pgvector compatibility. It can be enabled by SET vectors.pgvector_compatibility = on.

Index Options & Query Options

For index ivfflat and hnsw only the following options are available.

• Default value for index options is different from pgvecto.rs original, which keeps the same as in pgvector.
• Default value for query options is the same as in pgvecto.rs.

Index options for ivfflat:

Query options for ivfflat:

Index options for hnsw:

Query options for hnsw:

Examples

It’s easy to enable pgvector compatibility and start a vector query.

DROP TABLE IF EXISTS t;
SET vectors.pgvector_compatibility=on;
SET hnsw.ef_search=40;
CREATE TABLE t (val vector(3));
INSERT INTO t (val) SELECT ARRAY[random(), random(), random()]::real[] FROM generate_series(1, 1000);
CREATE INDEX hnsw_cos_index ON t USING hnsw (val vector_cosine_ops);
SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <-> '[0.5,0.5,0.5]' limit 100) t2;
DROP INDEX hnsw_cos_index;

Multiply types of indexes are accepted:

SET vectors.pgvector_compatibility=on;
SET hnsw.ef_search=40;
-- [hnsw + vector_l2_ops] index with default options
CREATE INDEX hnsw_l2_index ON t USING hnsw (val vector_l2_ops);
-- [hnsw + vector_cosine_ops] index with single ef_construction option
CREATE INDEX hnsw_cosine_index ON t USING hnsw (val vector_cosine_ops) WITH (ef_construction = 80);
-- anonymous [hnsw + vector_ip_ops] with all options
CREATE INDEX ON t USING hnsw (val vector_ip_ops) WITH (ef_construction = 80, m = 12);
SET ivfflat.probes=1;
-- [ivfflat + vector_l2_ops] index with default options
CREATE INDEX ivfflat_l2_index ON t USING ivfflat (val vector_l2_ops);
-- [ivfflat + vector_ip_ops] index with all options
CREATE INDEX ivfflat_ip_index ON t USING ivfflat (val vector_cosine_ops) WITH (nlist = 80);
-- anonymous [ivf + vector_ip_ops] with all options
CREATE INDEX ON t USING ivfflat (val vector_ip_ops) WITH (lists = 80)

3 - Use Cases

3.1 - Hybrid Search

We will show you how to use pgvecto.rs to implement hybrid search in this section.

Vector Search and Full Text Search

Postgres provides full text search, based on the match operator @@. It returns true if a tsvector (document) matches a tsquery (query). It doesn’t matter which data type is written first:

SELECT 'a fat cat sat on a mat and ate a fat rat'::tsvector @@ 'cat & rat'::tsquery;
 ?column?
----------
 t

SELECT 'cat & rat'::tsquery @@ 'a fat cat sat on a mat and ate a fat rat'::tsvector;
 ?column?
----------
 t

You could use full text search and vector search together in a single query, with the help of pgvecto.rs. In this example, we create a table named items with two columns: content and embedding. The column content is of type text, and the column embedding is of type vector(3). We insert four items into the table:

CREATE TABLE items (
  id bigserial PRIMARY KEY,
  content text NOT NULL,
  embedding vector(3) NOT NULL -- 3 dimensions
);

INSERT INTO items (content, embedding) VALUES
  ('a fat cat sat on a mat and ate a fat rat', '[1, 2, 3]'),
  ('a fat dog sat on a mat and ate a fat rat', '[4, 5, 6]'),
  ('a thin cat sat on a mat and ate a thin rat', '[7, 8, 9]'),
  ('a thin dog sat on a mat and ate a thin rat', '[10, 11, 12]');

You should get the following table:

postgres=# select * from items;
 id |                  content                   |  embedding   
----+--------------------------------------------+--------------
  2 | a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
  3 | a fat dog sat on a mat and ate a fat rat   | [4, 5, 6]
  4 | a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
  5 | a thin dog sat on a mat and ate a thin rat | [10, 11, 12]
(4 rows)

Let’s search for the items with the word cat and rat in their content first. We will use full text search to do that. The function to_tsvector converts a text to a tsvector. The operator @@ checks if a tsvector matches a tsquery.

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery;

SELECT to_tsvector('english', 'a fat  cat sat on a mat - it ate a fat rats');
                  to_tsvector
-----------------------------------------------------
 'ate':9 'cat':3 'fat':2,11 'mat':7 'rat':12 'sat':4

We should get the following result:

                  content                   | embedding 
--------------------------------------------+-----------
 a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
 a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
(2 rows)

Then we will sort the result by the distance between the embedding of each item and the embedding of the query. We will use vector search to do that:

SELECT content, embedding
FROM items
WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
ORDER BY embedding <-> '[7, 8, 9]'::vector;

We will get the item with the embedding [7, 8, 9] first, and the item with the embedding [1, 2, 3] second:

                  content                   | embedding 
--------------------------------------------+-----------
 a thin cat sat on a mat and ate a thin rat | [7, 8, 9]
 a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]
(2 rows)

Advanced search: Merge the results of full text search and vector search

You could normalize the rank of the full text search result and the vector search result, and then merge them together:

-- The query starts with a Common Table Expression (CTE), named "semantic_search".
WITH semantic_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY embedding <-> '[7, 8, 9]') AS rank
    FROM items
    ORDER BY embedding <-> '[7, 8, 9]'::vector
    ), 
-- Another CTE, named "full_text_search", is defined.
    full_text_search AS (
    SELECT id, content, embedding, 
    RANK () OVER (ORDER BY 
        ts_rank(to_tsvector('english', content), 
            'cat & rat'::tsquery) DESC) AS rank
    FROM items
    WHERE to_tsvector('english', content) @@ 'cat & rat'::tsquery
    ORDER BY ts_rank(
        to_tsvector('english', content), 'cat & rat'::tsquery) DESC
    )
-- The main query selects columns from both CTEs.
    SELECT 
        COALESCE(semantic_search.id, full_text_search.id) AS id,
        COALESCE(semantic_search.content, full_text_search.content) AS content,
        COALESCE(semantic_search.embedding, full_text_search.embedding) AS embedding,
        COALESCE(1.0 / (1 + semantic_search.rank), 0.0) + 
        COALESCE(1.0 / (1 + full_text_search.rank), 0.0) AS rank
    FROM semantic_search FULL OUTER JOIN full_text_search USING (id)
    ORDER BY rank DESC;

In this query, we first calculate the rank of the full text search result and the vector search result separately. Then we merge the two results together. We will get the following result:

 id |                  content                   |  embedding   |          rank          
----+--------------------------------------------+--------------+------------------------
  4 | a thin cat sat on a mat and ate a thin rat | [7, 8, 9]    | 1.00000000000000000000
  2 | a fat cat sat on a mat and ate a fat rat   | [1, 2, 3]    | 0.70000000000000000000
  3 | a fat dog sat on a mat and ate a fat rat   | [4, 5, 6]    | 0.33333333333333333333
  5 | a thin dog sat on a mat and ate a thin rat | [10, 11, 12] | 0.33333333333333333333
(4 rows)

3.2 - Adaptive retrieval and shortening embeddings

OpenAI latest embedding models support shortening embeddings, where developers could simply “remove some numbers from the end of a sequence” and still maintain a valid representation for text.

This enables very flexible usage. For example, when using a vector data store that only supports embeddings up to 1024 dimensions long, developers can now still use our best embedding model text-embedding-3-large and specify a value of 1024 for the dimensions API parameter, which will shorten the embedding down from 3072 dimensions, trading off some accuracy in exchange for the smaller vector size.

Shortening embeddings

OpenAI’s new embedding model text-embedding-3-large produces embeddings with 3072 dimensions by default.

But you could safely remove some numbers from the end of the sequence and still maintain a valid representation for text. For example, you could shorten the embeddings to 1024 dimensions.

embedding_3072_dim_from_openai = [-0.00611, -0.00269, ...]
embedding_1024_dim_from_openai = [-0.01520, -0.03839, ...]

embedding_1024_dim_truncated_from_3072_dim = [-0.00611, -0.00269, ...]
embedding_1024_dim_normalized = [-0.01520, -0.03839, ...]

There are lots of benefits to using shortening embeddings:

• Reduced memory usage: Shortening embeddings can reduce the memory usage to store the embeddings.
• Lower dimensions: Shortening embeddings can reduce the dimensions of the embeddings. Some vector data stores only support embeddings up to a certain number of dimensions (e.g. 2000).
• Faster indexing and search: Smaller embeddings can be indexed and searched faster than larger embeddings.
• Trade-off performance and cost: Shortening embeddings can trade off some accuracy in exchange for the smaller vector size.

However, the accuracy of the shortened embeddings is not as good as the full embeddings.

Trade-off performance and cost

We have performed an experiment to evaluate the effectiveness of reduced embeddings compared to the complete embeddings. The benchmark was conducted using the dbpedia-entities-openai3-text-embedding-3-large-3072-1M dataset. The experiment was carried out on a Google Cloud VM with specifications of n2-standard-8, which includes 8 vCPUs and 32 GB of memory.

We have established a table and populated it with 1 million embeddings, each consisting of 3072 dimensions.

CREATE TABLE openai3072 (
  id bigserial PRIMARY KEY,
  text_embedding_3_large_3072_embedding vector(3072),
  text_embedding_3_large_3072_embedding bvector(3072)
);

from pgvecto_rs.sdk import PGVectoRs, Record
from pgvecto_rs.psycopg import register_vector
import psycopg

conn = psycopg.connect(conninfo='postgresql://postgres:mysecretpassword@localhost:5433/postgres', autocommit=True)

import json
import os

import numpy as np
import random
from datasets import load_dataset
from datasets.exceptions import DatasetNotFoundError
from tqdm import tqdm

MODEL_NAME, DIMENSIONS = "text-embedding-3-large", 3072
DATASET_NAME = f"Qdrant/dbpedia-entities-openai3-{MODEL_NAME}-{DIMENSIONS}-1M"
collection_name = f"dbpedia-{MODEL_NAME}-{DIMENSIONS}"
embedding_column_name = f"{MODEL_NAME}-{DIMENSIONS}-embedding"

dataset = load_dataset(
        DATASET_NAME,
        streaming=False,
        split="train",
    )

bs = 1000
for i, record in tqdm(enumerate(dataset)):
    if i % bs == 0:
        points = []
    points.append({
        "embedding": record[embedding_column_name]
    })
    if i % bs == bs - 1:
        batch_points = f", ".join([f"('{p['embedding']}')" for p in points])
        conn.execute(f"INSERT INTO openai3072 (text_embedding_3_large_3072_embedding) VALUES %s" % (batch_points))
    print(f"Inserted {i} records")

Furthermore, we have created indexes for embeddings with dimensions of 256, 1024 and 3072, and performed an ANN (Approximate Nearest Neighbors) query for each dimension. There is no need to create the 256-dimensional column, as we can use the first 256 dimensions of the 3072-dimensional embeddings.

-- Create 256-dim index
CREATE INDEX openai_vector_index_256 ON public.openai3072 USING vectors (((text_embedding_3_large_3072_embedding[0:256])::vector(256)) vector_l2_ops);

-- Create 1024-dim index
CREATE INDEX openai_vector_index_1024 on openai3072 using vectors((text_embedding_3_large_3072_embedding[0:1024]::vector(1024)) vector_l2_ops);

-- Create 3072-dim index
CREATE INDEX openai_vector_index on openai3072 using vectors(text_embedding_3_large_3072_embedding vector_l2_ops);

Additionally, we have constructed a binary vector index for the 3072-dimensional embeddings and conducted an ANN query using this index as well.

CREATE INDEX openai_vector_index_bvector ON public.openai3072 USING vectors (text_embedding_3_large_3072_bvector bvector_l2_ops);

We run the KNN query with different limits (5, 10, 50, 100) for the number of neighbors to retrieve.

export PGHOST="localhost"
export PGUSER="postgres"
export PGPASSWORD="mysecretpassword"
export PGDATABASE="postgres"
export PGPORT=5433
export VECTOR_DIM=3072
TEST_TIME=300
for clients in 2 4 8 16 32; do
    pgbench -n -T "${TEST_TIME}" \
      --file=query.sql \
      -c "${clients}" -j "${clients}"
done

\set random_id random(1, 1000000)

select id from openai3072 order by text_embedding_3_large_3072_embedding <-> (select text_embedding_3_large_3072_embedding from openai3072 where id = :random_id) limit 100;

The results of the experiment are shown here:

As anticipated, the 3072-dimensional embeddings exhibit the highest accuracy, while the 1024-dimensional embeddings demonstrate a slightly lower accuracy at around 85%. The 256-dimensional embeddings yield the lowest accuracy, approximately 65%. On the other hand, the binary vector index achieves an accuracy of about 80%.

Regarding the Requests Per Second (RPS) metric, the binary vector index showcases the most efficient performance, followed by the 256-dimensional embeddings, the 1024-dimensional embeddings, and finally, the 3072-dimensional embeddings.

Details

The memory usage of the indexes is an important aspect to take into account. Indexing 1 million 3072-dimensional binary vectors requires only around 600MB of memory. In contrast, the 3072-dimensional vector index consumes approximately 18GB of memory. This represents a significant difference in memory usage, as the 3072-dimensional vector index utilizes approximately 30x more memory than the binary vector index.

Improve the accuracy via adaptive retrieval

Lower dimension and binary vector indexes trade off accuracy for advantages such as higher RPS and lower memory usage compared to full dimension indexes. Adaptive retrieval techniques can be used to combine the strengths of both approaches, achieving an optimized solution that balances accuracy, memory usage, and query speed.

The logic behind adaptive retrieval is really simple. Let’s take get top 100 candidates as an example. We can perform the following steps:

1. Query the lower dimensional or binary vector index first to retrieve 200 candidates from the 1 million embeddings. This is a fast operation.
2. Rerank the candidates using a KNN query to retrieve the top 100 candidates. K-NN is well-suited for situations where smaller sets and precise similarity search are necessary, making it an excellent choice for reranking in this context.

The reranking step is a bit slower than the initial query, but it is still much faster than querying the full dimension index. It could be done in the pgvecto.rs:

CREATE OR REPLACE FUNCTION match_documents_adaptive(
  query_embedding vector(3072),
  match_count int
)
RETURNS SETOF openai3072
LANGUAGE SQL
AS $$
WITH shortlist AS (
  SELECT *
  FROM openai3072
  ORDER BY (text_embedding_3_large_3072_embedding[0:256])::vector(256) <-> (query_embedding[0:256])::vector(256)
  LIMIT match_count * 2
)
SELECT *
FROM shortlist
ORDER BY text_embedding_3_large_3072_embedding <-> query_embedding
LIMIT match_count;
$$;

The match_documents_adaptive function accepts a query embedding and a match count as input parameters. It first retrieves match_count * 2 candidates from the 1 million embeddings using the 256-dimensional index. Then, it reranks the candidates using a KNN query to retrieve the top match_count candidates.

The function for binary vector and 1024-dimensional embeddings can be implemented in a similar manner.

We conduct the benchmark again with the adaptive retrieval technique. We tag the adaptive retrieval technique as Adaptive Retrieval in the following figure.

Adding a reranking step significantly improves the accuracy of the lower dimensional and binary vector indexes.

The binary vector index with adaptive retrieval achieves an accuracy of around 95%, with only a slight decrease in RPS compared to the binary vector index without adaptive retrieval. Remarkably, the memory usage of the binary vector index with adaptive retrieval is only around 600MB, which is the same as the binary vector index without adaptive retrieval.

Details

Conclusion

Shortening embeddings can reduce the memory usage to store the embeddings, and speed up indexing and search. However, the accuracy of the shortened embeddings is not as good as the full embeddings. That’s why we need to use adaptive retrieval techniques to improve the accuracy of the lower dimensional indexes.

It seems that binary vector indexes with adaptive retrieval are the best choice for the text-embedding-3-large model. You could also consider 1024-dimensional indexes with adaptive retrieval if you need a higher accuracy (e.g. 99%).

PS: All the code and data used in this experiment are available in the pgvecto.rs-matryoshka-embeddings-benchmark.

3.3 - Image Search

Image intelligent search, also known as content-based image retrieval, goes beyond traditional keyword-based search and enables users to search for images based on their visual content. For example, you could search images that contain a red mushroom (image from immich):

The feature looks like a magic because it allows you to search for images that you don’t have keywords for. The search engine understands the content of the images and videos and can find the ones that are similar to the query.

How it works

Image intelligent search leverages advanced computer vision techniques and machine learning algorithms to understand the visual characteristics and features within images, allowing users to perform freeform searches without relying on specific keywords or metadata.

The first step is to extract meaningful features from images and generate a vector representation of each image. This is typically done using deep learning models such as CLIP. The vector representation is called an embedding. It is a list of numbers that represents the image.

Then you could use vector stores like pgvecto.rs to store the embeddings and perform efficient vector search.

Example

We will leverage CLIP, a neural network that connects text and images, to generate embeddings for our images. These embeddings capture the visual concepts present in the images. Subsequently, we will utilize pgvecto.rs, a scalable vector search tool built on Postgres, to store these embeddings and conduct efficient vector searches. By combining the power of CLIP and pgvecto.rs, we can retrieve similar images based on their visual content.

Generate embeddings

You could use CLIP to generate embeddings for images:

from transformers import CLIPProcessor, CLIPModel

processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")

def compute_image_embeddings(list_of_images):
    return model.get_image_features(**processor(images=list_of_images, return_tensors="pt", padding=True))

Store embeddings

Let’s say you have a table named images with three columns: id, author and embedding. The column embedding is of type vector(512). You could insert an image into the table like this:

-- Create the table
CREATE TABLE images (
    id SERIAL PRIMARY KEY,
    author TEXT NOT NULL,
    vector vector(512) NOT NULL -- 3 dimensions
);

INSERT INTO images (author, embedding) VALUES
  ('Allen', '[...]');

Search for similar images

You could search for similar images like this:

SELECT * FROM images ORDER BY embedding <-> '[...]' LIMIT 5;

Advanced search

Assuming you want to search for the 64 nearest rows and the results don’t contain the images from the author Allen.

You may want the following SQL to work:

SET vectors.hnsw_ef_search = 64;
SELECT * FROM images ORDER BY embedding <-> '[...]' WHERE author != 'Allen' LIMIT 64;

In base mode, you may only get 32 rows since the HNSW algorithm bypasses filter conditions during the search and relies on PostgreSQL to apply filters post-search.

You could use the vbase search mode to achieve this. In vbase mode, searching results become an iterator that continuously returns results. The database can continuously retrieve results from it until the user’s query conditions are met. It’s quite different from other vector search library, such as faiss. The latter can only return a fixed-size result set, and this size needs to be determined before the search. For queries with filtering conditions, it is not possible to determine the size of the result set in advance. As a result, either insufficient results can be obtained or excessive searching will affect latency. The original idea comes from VBASE: Unifying Online Vector Similarity Search and Relational Queries via Relaxed Monotonicity.

You can enable vbase by a SQL statement SET vectors.search_mode = vbase;. Then you can get the correct behavior. This mode is used in immich’s image search feature. vbase is the default search mode in 0.2.

Real-world applications

Image intelligent search is used in immich to enable users to search for images based on their visual content. The code snippets provided here are inspired by the implementation of the image search feature in immich.

SELECT a.*, s.*, e.*
FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id
WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()
ORDER BY s.embedding <=> :embedding
LIMIT :numResults;

The query retrieves data from the asset_entity, smart_search, and exif_info tables based on specified conditions and ordering. Let’s break down the query:

SELECT a.*, s.*, e.*

This part of the query specifies the columns to be selected from the tables asset_entity, smart_search, and exif_info. The aliases a, s, and e are used to refer to each table respectively.

FROM asset_entity AS a
INNER JOIN smart_search AS s ON s.assetId = a.id
LEFT JOIN exif_info AS e ON e.assetId = a.id

Here, the asset_entity table is aliased as a, the smart_search table as s, and the exif_info table as e. The INNER JOIN is used to combine asset_entity and smart_search based on the condition s.assetId = a.id. The LEFT JOIN is used to include matching records from the exif_info table based on the condition e.assetId = a.id.

CREATE TABLE asset_entity (
  id UUID PRIMARY KEY,
  deviceAssetId VARCHAR NOT NULL,
  ownerId VARCHAR NOT NULL,
  libraryId VARCHAR NOT NULL,
  deviceId VARCHAR NOT NULL,
  type VARCHAR NOT NULL,
  originalPath VARCHAR NOT NULL,
  resizePath VARCHAR,
  webpPath VARCHAR DEFAULT '',
  thumbhash BYTEA,
  encodedVideoPath VARCHAR,
  createdAt TIMESTAMPTZ NOT NULL,
  updatedAt TIMESTAMPTZ NOT NULL,
  deletedAt TIMESTAMPTZ,
  fileCreatedAt TIMESTAMPTZ NOT NULL,
  localDateTime TIMESTAMPTZ NOT NULL,
  fileModifiedAt TIMESTAMPTZ NOT NULL,
  isFavorite BOOLEAN DEFAULT false,
  isArchived BOOLEAN DEFAULT false,
  isExternal BOOLEAN DEFAULT false,
  isReadOnly BOOLEAN DEFAULT false,
  isOffline BOOLEAN DEFAULT false,
  checksum BYTEA NOT NULL,
  duration VARCHAR,
  isVisible BOOLEAN DEFAULT true,
  livePhotoVideoId UUID,
  originalFileName VARCHAR NOT NULL,
  sidecarPath VARCHAR,
  stackId UUID,
  CONSTRAINT fk_asset_owner FOREIGN KEY (ownerId) REFERENCES user_entity (id) ON DELETE CASCADE ON UPDATE CASCADE,
  CONSTRAINT fk_asset_library FOREIGN KEY (libraryId) REFERENCES library_entity (id) ON DELETE CASCADE ON UPDATE CASCADE,
  CONSTRAINT fk_asset_live_photo_video FOREIGN KEY (livePhotoVideoId) REFERENCES asset_entity (id) ON DELETE SET NULL ON UPDATE CASCADE,
  CONSTRAINT fk_asset_stack FOREIGN KEY (stackId) REFERENCES asset_stack_entity (id) ON DELETE SET NULL ON UPDATE CASCADE
);

WHERE a.ownerId IN (:...userIds)
  AND a.isArchived = false
  AND a.isVisible = true
  AND a.fileCreatedAt < NOW()

In this part, the query applies filters to the asset_entity table. It selects rows where the ownerId column is among the specified userIds, isArchived is false, isVisible is true, and fileCreatedAt is earlier than the current time (NOW()).

ORDER BY s.embedding <=> :embedding

The query orders the results based on the embedding column of the smart_search table, comparing it to the provided :embedding value. The <=> operator is used for the comparison.

LIMIT :numResults;

Finally, limit the number of rows returned by the query to the value specified by the :numResults parameter.

This is an real-world use case to show how image intelligent search is implemented in a production system.

3.4 - Multi Tenancy

Multi-tenancy is essential for SaaS applications, allowing a single instance to serve multiple users or groups while ensuring data privacy and security. This blog outlines the design and implementation a simple multi-tenant vector search system with pgvecto.rs.

Designing the Database Schema

The success of user-based isolation in multi-tenant systems hinges on a well-structured database schema:

Users and Documents

Each user is identified by a unique user_id, which directly links them to their documents for private access. The documents table contains each user’s documents, each with a distinct id, title, and content. Additional metadata can be included in the documents table to improve search functionality. The user_id is linked to the documents instead of chunk embeddings to replicate real-world scenarios more accurately.

CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    user_id INTEGER NOT NULL,
    title TEXT,
    content TEXT,
    -- Additional metadata
);

Chunks

Embeddings store vectorized representations of chunks of documents, and use document_id to reference its document.

CREATE TABLE chunks (
    id SERIAL PRIMARY KEY,
    document_id INTEGER REFERENCES documents(id),
    embedding vector(512) NOT NULL
);

Entity-Relation Diagram

Implementing Multi-Tenancy

Create Index

Create vector index on chunks.embedding to accelerate query process. Here we use dot product as the similarity measure.

CREATE INDEX idx_chunks_embedding ON chunks USING vectors (embedding vector_dot_ops);

Run the query with user_id

SELECT chunks.id AS chunk_id, documents.id AS document_id, documents.title
FROM chunks
INNER JOIN documents ON chunks.document_id = documents.id
WHERE documents.user_id = 'alice' /* Replace with user_id */
ORDER BY chunks.embedding <-> '[3,2,1]' /* Replace with query embedding */ LIMIT 5;

In this query:

• The SELECT statement now includes chunks.id AS chunk_id, documents.id AS document_id, and documents.title to return the chunk ID, the document ID, and the document title, respectively.
• The INNER JOIN ensures that only chunks associated with documents owned by the specified user (‘alice’) are considered.
• The ORDER BY clause calculates the distance between the chunk’s embedding and the provided vector [3,2,1], sorting the results by similarity. The closest or most similar embeddings are returned first.
• Replace [3,2,1] with the query embedding and ‘specific_user_id’ with the actual ID of the user for whom you are performing the search.

Performance Tip

The vector search will scan through additional points and evaluate the conditions individually. Consequently, if there is a large number of user IDs or the filter criteria are challenging to meet, the query speed may be impacted. To enhance performance, you may want to explore utilizing PostgreSQL’s partition table function.

3.5 - Sparse Vector Search

In the past, hybrid search combined two search methods: traditional keyword-based search and vector-based similarity search. However, sparse vector can act as a substitute for keyword search, unlocking the full potential of the data pipeline with pure vector search.

This post explores how to use how BGE-M3 and pgvecto.rs to exploit the potential of hybrid search using both sparse and dense vectors.

Let’s dive in.

An Overview of Sparse and Dense Vector Embedding

Sparse vectors are high-dimensional vectors that contain few non-zero values. They are suitable for traditional information retrieval use cases. For example, a vector with 32,000 dimensions but only 2 non-zero elements is a typical sparse vector.

Dense vectors are embeddings from neural networks. They are generated by text embedding models and have most or all elements non-zero. They have fewer dimensions, such as 256 or 1536, much less than sparse vectors.

Dense vectors and sparse vectors are both used to represent features of text chunks, but they have significant differences. A single dimension in a dense vector embedding does not mean anything, as it is too abstract to determine its meaning. However, when we take all the dimensions together, they provide the semantic meaning of the input text.

Sparse vectors have dimensions equal to the vocabulary size, with each element representing the weight of a word in the text.

Generally, words with higher weights indicate greater importance in the text. Sparse vectors simplify the analysis and visualization of the relationship between text and vocabulary. To visualize dense vectors, additional dimension reduction algorithms are necessary, such as:

• Principal component analysis (PCA)
• T-distributed stochastic neighbor embedding (T-SNE)
• Uniform Manifold Approximation and Projection (UMAP)

Sparse vectors are generally less effective at capturing the relationships between words, particularly for synonyms and related terms, compared to dense vectors.

For instance, both rocket and SPACE X are related to space exploration, but Dense vectors are more likely to reveal similarities and assign a closer vector distance to corresponding texts.

How to create a vector embedding?

A vector embedding is the internal representation of input data in deep learning models, also known as embedding models. Most embedding models, such as text-embedding-3-small, only output dense embeddings.

SPLADE and BGE-M3 can generate sparse embeddings, sometimes called lexical weights. In this section, we will introduce the BGE-M3 model for generating dense and sparse vectors.

BGE-M3 is a multi-functionality model, which can simultaneously perform the three common retrieval functionalities of the embedding model: dense retrieval, multi-vector retrieval, and sparse retrieval. We will learn how to generate dense and sparse embeddings from it.

# pip install -U FlagEmbedding

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)
sentences = ["What is BM25?", "Definition of BM25"]

output = model.encode(sentences, return_dense=True, return_sparse=True)

# The dense embedding of the two texts above:
dense_embedding = output['dense_vecs']
# [[-0.02501617 -0.04525582 -0.01972856 ... -0.0099566   0.01677509
#   0.01292699]
# [-0.0106499  -0.04731942 -0.01477693 ... -0.0149255   0.01040097
#   0.00965083]

sparse_weight_id = output['lexical_weights']

# You can see the weight for each token:
# print(model.convert_id_to_token(sparse_weight_id))
# 
# [{'What': 0.10430284, 'is': 0.10090457, 'BM': 0.2635918, '25': 0.3382988, '?': 0.052101523}, 
# {'Definition': 0.14298248, 'of': 0.076763615, 'BM': 0.2577639, '25': 0.33806682}]

For all sentences and words, a text-word matrix was created:

The matrix is highly sparse, especially when all vocabularies are included. Most words only appear in a few sentences, resulting in non-zero weights.

By feeding text into sparse embedding models like BGE-M3, we can easily obtain sparse embeddings from it.

# pip install -U numpy

import numpy as np

# Sparse vector embedding of text
sparse_indices = [[int(token_id) for token_id in text] for text in sparse_weight_id]
# [[4865, 83, 90017, 2588, 32], [155455, 111, 90017, 2588]]
sparse_values = [[float(text[token_id]) for token_id in text] for text in sparse_weight_id]
# [[0.10430284, 0.10090457, 0.2635918, 0.3382988, 0.052101523], [0.14298248, 0.076763615, 0.2577639, 0.33806682]]

# Sparse indexes must be ascending for each pgvecto.rs
for i in range(len(sparse_weight_id)):
    sparse_values[i] = [x for _, x in sorted(zip(sparse_indices[i], sparse_values[i]))]
    sparse_indices[i] = sorted(sparse_indices[i])

Now we have both the sparse and dense embeddings. It’s time to insert them into the vector database.

Storing vectors in pgvecto.rs

First, a deployed pgvecto.rs instance is required for vector storage and retrieval. This can be achieved by the official pgvecto-rs Docker image:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

Now the only necessary dependency is installed. Let’s create a table to store the vectors.

pgvecto.rs supports dense or sparse vectors with different data types:

• vector for dense vector
• svector for sparse vector

The BGE-M3 model has a dense embedding dimension of 1024. The sparse embedding dimension is 250002, which matches the size of tokenizer vocabulary.

print(len(model.tokenizer))
# 250002

In this example, we create a table named documents with four columns: an ID column(id), a text column(text) and two vector columns(sparse for sparse vector and dense for dense vector).

# pip install -U psycopg

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS documents;")
    conn.execute(f"""
        CREATE TABLE documents (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            dense vector(1024) NOT NULL,
            sparse svector({len(model.tokenizer)}) NOT NULL);
    """
    )

With the table created, we can now insert embeddings into the table and create indexes for vector columns.

# pip install -U pgvecto_rs

import psycopg
from pgvecto_rs.psycopg import register_vector
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

with psycopg.connect(URL) as conn:
    register_vector(conn)
    for text, dense, sparse_ind, sparse_val in zip(sentences, dense_embedding, sparse_indices, sparse_values):
        conn.execute(
            "INSERT INTO documents (text, dense, sparse) VALUES (%s, %s, to_svector(%s, %s::real[]::int[], %s::real[]));",
            (text, dense, len(model.tokenizer), sparse_ind, sparse_val),
        )
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (sparse svector_dot_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.execute("""
        CREATE INDEX ON documents 
        USING vectors (dense vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)

For dense vectors, vector_l2_ops calculates the Squared Euclidean distance, which is the most commonly used distance metric.

For sparse vectors, vector_dot_ops calculates the dot product, which is a more efficient method.

For more information about index types in pgvecto.rs, please check out our documentation.

Retrieving and Mixing Results

With the vectors stored in the database, we can efficiently query them using the pgvecto.rs extension. To find the nearest vectors to a given input text, we need to encode the text with the same model.

This process is identical to the one described in the previous chapter:

def embedding(sentence: str) -> tuple[np.ndarray, list[list[int]], list[list[float]]]:
    """
    Create dense and sparse embeddings from text
    """
    output = model.encode(sentence, return_dense=True, return_sparse=True, return_colbert_vecs=False)
    dense_embedding = output['dense_vecs']
    sparse_weight_id = output['lexical_weights']
    sparse_indices = [int(token_id) for token_id in sparse_weight_id]
    sparse_values = [float(sparse_weight_id[token_id]) for token_id in sparse_weight_id]

    sparse_values = [x for _, x in sorted(zip(sparse_indices, sparse_values))]
    sparse_indices = sorted(sparse_indices)
    return dense_embedding, sparse_indices, sparse_values

The pgvecto.rs extension provides a simple SELECT statement to get sparse and dense vectors like this:

dense_embedding, sparse_indices, sparse_values = embedding("The text you want to search...")

with psycopg.connect(URL) as conn:
    register_vector(conn)
    cur = conn.execute(
        "SELECT text FROM documents ORDER BY sparse <#> to_svector(%s, %s::real[]::int[], %s::real[]) LIMIT 50;",
        (len(model.tokenizer), sparse_indices, sparse_values),
    )
    sparse_result = cur.fetchall()
      
    cur = conn.execute(
              "SELECT text FROM documents ORDER BY dense <-> %s LIMIT 50;",
              (dense_embedding,),
          )
    dense_result = cur.fetchall()
    
# Merge sparse and dense vector search results
mix_text = set([r[0] for r in dense_result]).union([r[0] for r in sparse_result])

To rerank the merged result, you can use popular fusion methods such as Reciprocal Ranked Fusion (RRF). However, to make better use of BGE-M3, let’s introduce the Reranker model of BGE-M3. Unlike the embedding model, Reranker uses text input and directly outputs their similarity or score.

We can initialize a Reranker model and feed it with mixed candidate texts. With this method, reranking can be as simple as sorting by the Reranker model’s estimated score.

from FlagEmbedding import FlagReranker

reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)

# Obtain relevance scores （Higher scores indicate greater relevance）
scores = reranker.compute_score([["The text you want to search...", candidate] for candidate in mix_text])
# Rerank text using combined dense and sparse embeddings
reranked_text = [t for _, t in sorted(zip(scores, mix_text))]

Congratulations! We have now completed our first step into the world of sparse vector search.

3.6 - Generative Question-Answering

Vector search can also be used to provide context to generative models such as OpenAI’s GPT series, improving the quality of the outputs.

This post will explore the use of pgvecto.rs in building and retrieving from a vector knowledge base to bolster the performance of generative models in question-answering tasks. A knowledge base serves as “long-term memory” for generative models, providing them persistent, curated, and accurate information as context to draw from when generating an answer to a question. This results in answers that are more accurate, reduces the possibility of the model producing “hallucinations” (statements with no basis in reality), and increases user trust in the outputs. This technique is known as Retrieval-Augmented Generation, or RAG: see this post or the related paper for more information.

Overview of the Question-Answering task

The generative model question-answering (QA) task is a method of Information Retrieval. The model is presented with a question, and generates an answer based on some context: typically, documents or other information sources the model can look at to retrieve the information necessary to answer the question.

This use case will involve giving the OpenAI gpt-3.5-turbo-instruct model context to answer questions about movies from the vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings Huggingface dataset.

The overall workflow will look something like this:

Question-Answering without Context

First, let’s try having the model answer some questions without providing it context.

We begin with defining a function to streamline getting a response for a given prompt:

# pip install openai

import openai

openai.api_key = "YOUR_API_KEY_HERE"
def complete(prompt):
    result = openai.Completion.create(
        engine='gpt-3.5-turbo-instruct',
        prompt=prompt,
        temperature=0,
        max_tokens=400,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    return result['choices'][0]['text'].strip()

Let’s start with asking about a relatively popular movie:

base_prompt = "What is the release year and overall plot of Skyfall?"
print(complete(base_prompt))

Accurate, if somewhat vague in the plot summary. However, what about a movie that’s much less known?

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"
print(complete(base_prompt_2))

We have found the limits of what the model “knows”: it is unable to produce the release year, and invents a potential plot based on the title. We will remedy this by creating a knowledge base for the model to draw on.

There are two options for allowing the model to better answer our question:

1. Fine-tune the model on text data in the domain of interest (movies in this case).
2. Employ retrieval augmented generation, which adds an information retrieval component to the process. This will allow us to retrieve and feed relevant information to the model as a secondary source of information.

We will be taking the second option in this post.

Building the Knowledge Base

Let’s build a knowledge base of movies to retrieve relevant information from, using the vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings Huggingface dataset. This dataset contains titles, release years, casts, Wikipedia pages, plot summaries/lengths, and vector embeddings of this information for 33,155 movies. Each movie also has text data, which encompasses title, release year, cast, and plot summary in one. We will be using the text and embeddings data for this application.

Load the Data

Start by loading the data:

# pip install -U datasets
from datasets import load_dataset
dataset = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings", split='train')

texts = dataset['train']['text']
# texts[0]
# 'Alice in Wonderland \n  starring May Clark \n  released in the year1903 \n  with the following plot  \n Alice follows a large white rabbit down a "Rabbit-hole". She finds a tiny door. When she finds a bottle labeled "Drink me", she does, and shrinks, but not enough to pass through the door. She then eats something labeled "Eat me" and grows larger. She finds a fan when enables her to shrink enough to get into the "Garden" and try to get a "Dog" to play with her. She enters the "White Rabbit\'s tiny House," but suddenly resumes her normal size. In order to get out, she has to use the "magic fan."\r\nShe enters a kitchen, in which there is a cook and a woman holding a baby. She persuades the woman to give her the child and takes the infant outside after the cook starts throwing things around. The baby then turns into a pig and squirms out of her grip. "The Duchess\'s Cheshire Cat" appears and disappears a couple of times to Alice and directs her to the Mad Hatter\'s "Mad Tea-Party." After a while, she leaves.\r\nThe Queen invites Alice to join the "ROYAL PROCESSION": a parade of marching playing cards and others headed by the White Rabbit. When Alice "unintentionally offends the Queen", the latter summons the "Executioner". Alice "boxes the ears", then flees when all the playing cards come for her. Then she wakes up and realizes it was all a dream.'

Generate Embeddings

This particular dataset comes with embeddings, so you can skip this step and use the embeddings directly.

embeddings = dataset['train']['embeddings']

If you want to generate embeddings from the text data, you can use the OpenAI, Cohere or other embeddings APIs. We provide an example using the Cohere API, you can get Cohere API key from here.

# pip install cohere
import getpass
import os
import cohere

if not os.getenv("COHERE_API_KEY"):
    os.environ["COHERE_API_KEY"] = getpass.getpass("Enter your Cohere API key: ")

co = cohere.Client()
# embedding model can refer to this documentation: https://docs.cohere.com/docs/cohere-embed
embed = co.embed(
    texts=texts,
    input_type="search_query",
    model="embed-english-v3.0"
)
embeddings = embed.embeddings

Insert Data into the Database

A deployed pgvecto.rs instance is required for vector storage and retrieval. We can use the official pgvecto-rs Docker image:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.3.0

Install the required psycopg library to interact with Postgres, and ensure the extension is running:

#pip install "psycopg[binary,pool]"
import psycopg

URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"
#Create the table
with psycopg.connect(URL) as conn:
    conn.execute("DROP EXTENSION IF EXISTS vectors;")
    conn.execute("CREATE EXTENSION vectors;")

Now, we create a table movies with a few columns: an ID column(id), a text column(texts), and a vector column(embeddings).

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Create the table
with psycopg.connect(URL) as conn:
    conn.execute("DROP TABLE IF EXISTS movies;")
    conn.execute(f"""
        CREATE TABLE movies (
            id SERIAL PRIMARY KEY, 
            text TEXT NOT NULL,
            embedding vector(768) NOT NULL);
    """
    )

With the table created, we can now insert the embeddings from the dataset and create indexes for vector columns:

# Optional install, but recommended:
# pip install tqdm
import psycopg
from tqdm import tqdm
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Insert into table
with psycopg.connect(URL) as conn:
    with conn.cursor() as cursor:    
        for text, embedding in tqdm(zip(texts, embeddings), total=len(texts)):
            cursor.execute(
                """INSERT INTO movies (text, embedding) 
                    VALUES (%s, %s::real[]);""",
                (text, embedding),
            )
        conn.commit()

Create Index

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

#Creating indexes for embedding column
with psycopg.connect(URL) as conn:
    cursor = conn.cursor()
    cursor.execute("""
        CREATE INDEX ON movies 
        USING vectors (embedding vector_l2_ops) 
        WITH (options = \"[indexing.hnsw]\");
    """)
    conn.commit()

$$

D_{L2} = \Sigma (x_i - y_i) ^ 2

$$

Question-Answering with Context

To make use of our knowledge base, queries must first be encoded to determine what texts to retrieve:

# pip install sentence-transformers

from sentence_transformers import SentenceTransformer

#some options: 'paraphrase-mpnet-base-v2', 'all-mpnet-base-v2'
encoder = SentenceTransformer('all-mpnet-base-v2')

def encode_query(query): #Query is a string
    #Preprocess the query
    query = query.lower()
    query = query.split()
    query = [word for word in query]
    query = ' '.join(query)
    
    #then encode the query
    vector = encoder.encode(query)
    return vector

Now, we define a function to streamline adding the retrieved context to the initial prompt:

import psycopg
URL = "postgresql://postgres:mysecretpassword@localhost:5432/postgres"

def add_context(prompt):
    search_vector = encode_query(prompt).tolist()
    # Construct prompt with the retrieved contexts included
    prompt_start = (
        "Answer the question based on the following context:\n\n"
    )
    prompt_end = (
        f"\n\nQuestion: {prompt}\nAnswer:"
    )
    with psycopg.connect(URL) as conn:
        cursor = conn.cursor()
        #Tune limit in accordance with how movie texts you want to give to the model
        cursor.execute("SELECT text FROM movies ORDER BY embedding <-> %s::real[] LIMIT 3;", 
                  (search_vector,),)
        result = cursor.fetchall()
    context = ""
    for r in result:
        context += r[0] + "\n\n"
    return prompt_start + context + prompt_end

Now, let’s try the prompt again, with added context:

base_prompt_2 = "What is the release year and overall plot of Cat Napping?"

context_prompt_2 = add_context(base_prompt_2)
context_response_2 = complete(context_prompt_2)
print(context_response_2)