Manage Vector Indexes with Prisma using Atlas

AI and ML applications increasingly rely on vector data for tasks like embeddings and similarity searches. Prisma, a popular ORM, integrates with PostgreSQL extensions like pg_vector to support vector data. However, Prisma lacks native index creation for vector columns.

This guide explains how to leverage Atlas for managing schema changes, enabling efficient indexing of vector data to enhance database performance. By combining Prisma as an ORM and Atlas for schema management, developers can utilize their databases for vector-based queries.

Getting Started with Atlas and Prisma

To set up, follow the Getting Started guide for Prisma and Atlas.

Atlas Pro Feature

Atlas support for Composite Schema used in this guide is available exclusively to Pro users. To use this feature, run:

atlas login

Composite Schema

1. On your Prisma schema, we need to add the postgresqlExtensions preview feature to enable loading the extension and the vector column type:

schema.prisma

generator client {
  provider        = "prisma-client-js"
  previewFeatures = ["postgresqlExtensions"]
}

datasource db {
  provider   = "postgresql"
  url        = env("DATABASE_URL")
  extensions = [vector(schema: "public", version: "0.8.0")]
}

model Embeddings {
  id        Int    @id @default(autoincrement())
  embedding Unsupported("vector(1536)")
}

2. Create a schema.sql that defines the index on the vector column:

schema.sql

CREATE INDEX vector_idx ON "Embeddings" USING ivfflat ("embedding" vector_cosine_ops) WITH (lists = 1000);

3. Create an atlas.hcl file that composes the Prisma schema with the schema.sql file using the composite_schema data source:

atlas.hcl

data "external_schema" "prisma" {
  program = [ 
    "npx",
    "prisma",
    "migrate",
    "diff",
    "--from-empty",
    "--to-schema-datamodel",
    "prisma/schema.prisma",
    "--script"
  ]
}

data "composite_schema" "prisma" {
  schema "public" {
    url = data.external_schema.prisma.url
  }
  schema "public"{
    url = "file://schema.sql"
  }
}

env "local" {
  url = getenv("DATABASE_URL")
  dev = "docker+postgres://pgvector/pgvector:pg17/dev"
  src = data.composite_schema.prisma.url
}

note

To utilize the pg_vector extension, use the pgvector/pgvector:pg17 Docker image that has it built-in.

Usage

After setting up our composite schema, we can get its representation using the atlas schema inspect command, generate schema migrations for it, apply them to a database, and more. Below are a few commands to get you started with Atlas:

Inspect the Schema

The atlas schema inspect command is commonly used to inspect databases. However, we can also use it to inspect our composite_schema and print the SQL representation of it:

atlas schema inspect \
  --env local \
  --url env://src \
  --format '{{ sql . }}'

The command above prints the following SQL.

-- Add new schema named "public"
CREATE SCHEMA IF NOT EXISTS "public";
-- Set comment to schema: "public"
COMMENT ON SCHEMA "public" IS 'standard public schema';
-- Create extension "vector"
CREATE EXTENSION "vector" WITH SCHEMA "public" VERSION "0.8.0";
-- Create "Embeddings" table
CREATE TABLE "public"."Embeddings" ("id" serial NOT NULL, "embedding" public.vector(1536) NOT NULL, PRIMARY KEY ("id"));
-- Create index "vector_idx" to table: "Embeddings"
CREATE INDEX "vector_idx" ON "public"."Embeddings" USING ivfflat ("embedding" vector_cosine_ops);

Extensions Are Database-Level Objects

Although the SCHEMA argument is supported by the CREATE EXTENSION command, it only indicates where the extension's objects will be installed. The extension itself is installed at the database level and cannot be loaded multiple times into different schemas.

Therefore, to avoid conflicts with other schemas, when working with extensions, the scope of the migration should be set to the database, where objects are qualified with the schema name. Hence, the search_path is dropped from the dev-database URL in the atlas.hcl file.

Generate Migrations For the Schema

To generate a migration for the schema, run the following command:

atlas migrate diff --env local

Note that a new migration file is created with the following content:

migrations/20241204150848.sql

-- Create extension "vector"
CREATE EXTENSION "vector" WITH SCHEMA "public" VERSION "0.8.0";
-- Create "Embeddings" table
CREATE TABLE "public"."Embeddings" ("id" serial NOT NULL, "embedding" public.vector(1536) NOT NULL, PRIMARY KEY ("id"));
-- Create index "vector_idx" to table: "Embeddings"
CREATE INDEX "vector_idx" ON "public"."Embeddings" USING ivfflat ("embedding" vector_cosine_ops);

Apply the Migrations

To apply the migration generated above to a database, run the following command:

atlas migrate apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?sslmode=disable"