Easily deploy a pgvector-enabled PostgreSQL server to Azure

The pgvector extension was first released in 2021 but has shot up in popularity in the last year, as developers discover what is possible with vector embeddings, vector similarity, and vector search. Once we start storing embedding vectors in database rows, we can make queries like "which movies are more similar to each other, based on their synopsis?" and "which retail item's descriptions most closely match this user's query?"

What is pgvector?

The pgvector extension allows PostgreSQL users to store columns of a vector type, create an index (HNSW or IVF) to efficiently index the vector fields, and query using vector distance operators (cosine, Euclidean, or inner product). Or, to put that in SQL form:

CREATE EXTENSION IF NOT EXISTS vector; CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3)); CREATE INDEX ON items USING hnsw (embedding vector_l2_ops); INSERT INTO items (embedding) VALUES ([-1, 2, 1]); ... SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 2;

Note: While the extension is typically referred to as "pgvector", the actual extension name is "vector", so that is what's used in the CREATE EXTENSION statement above.

Using pgvector on Azure PaaS

Azure now has multiple PaaS offerings for PostgreSQL that support the pgvector extension: PostgreSQL Flexible Server and Azure Cosmos DB PostgreSQL server.

If you already have an existing Azure PostgreSQL server, you can enable the extension manually in the Portal, as shown below and described in the linked tutorials.

A PostgreSQL flexible server with pgvector extension enabled.

An easy template for pgvector deployment

To make it even easier to get started with pgvector on PostgreSQL Flexible Server, we've created a template project that contains infrastructure-as-code (Bicep files) and support for the Azure Developer CLI (azd). Clone or download the project here:

https://github.com/Azure-Samples/azure-postgres-pgvector-python

With a few commands, you'll have a pgvector-enabled PostgreSQL server provisioned in your Azure account. We've also added keyless authentication to the template, so you can authenticate with your Azure credential instead of a secret.

The template project includes multiple Python scripts showing you how to connect to your PostgreSQL server and use the pgvector extension in the most common SQL packages: psycopg2, asyncpg, SQLAlchemy, and SQLModel.

For example, here's a selection of the code from a SQLAlchemy example for storing movie titles and their embeddings:

class Movie(Base): __tablename__ = "movies" id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) title: Mapped[str] = mapped_column() title_vector = mapped_column(Vector(1536)) # ada-002 is 1536-dimensional index = Index( "hnsw_index", Movie.title_vector, postgresql_using="hnsw", postgresql_with={"m": 16, "ef_construction": 64}, postgresql_ops={"title_vector": "vector_cosine_ops"}, ) Base.metadata.create_all(engine) # (Insert rows from a JSON) most_similars = session.scalars( select(Movie).order_by( Movie.title_vector.cosine_distance(target_movie.title_vector) ).limit(5))

See the full code in the repository.

There are so many ways that you can use the pgvector extension once you've gotten started with it, both as a tool in building generative AI applications (especially RAG apps), but also in any situation where similarity is a heuristic, like recommendations, fraud detection, and more. Start bringing vectors into your apps today and let us know what you build!

Published on: April 10, 2024