pgvector-go

pgvector support for Go

Supports pgx, pg, Bun, Ent, GORM, and sqlx

Getting Started

Run:

go get github.com/pgvector/pgvector-go

And follow the instructions for your database library:

• pgx
• pg
• Bun
• Ent
• GORM
• sqlx

Or check out some examples:

• Embeddings with OpenAI
• Binary embeddings with Cohere
• Recommendations with Disco
• Horizontal scaling with Citus
• Bulk loading with COPY

pgx

Import the packages

import (
    "github.com/pgvector/pgvector-go"
    pgxvector "github.com/pgvector/pgvector-go/pgx"
)

Enable the extension

_, err := conn.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector")

Register the types with the connection

err := pgxvector.RegisterTypes(ctx, conn)

or the pool

config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
    return pgxvector.RegisterTypes(ctx, conn)
}

Create a table

_, err := conn.Exec(ctx, "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))")

Insert a vector

_, err := conn.Exec(ctx, "INSERT INTO items (embedding) VALUES ($1)", pgvector.NewVector([]float32{1, 2, 3}))

Get the nearest neighbors to a vector

rows, err := conn.Query(ctx, "SELECT id FROM items ORDER BY embedding <-> $1 LIMIT 5", pgvector.NewVector([]float32{1, 2, 3}))

Add an approximate index

_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
// or
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

pg

Import the package

import "github.com/pgvector/pgvector-go"

Enable the extension

_, err := db.Exec("CREATE EXTENSION IF NOT EXISTS vector")

Add a vector column

type Item struct {
    Embedding pgvector.Vector `pg:"type:vector(3)"`
}

Insert a vector

item := Item{
    Embedding: pgvector.NewVector([]float32{1, 2, 3}),
}
_, err := db.Model(&item).Insert()

Get the nearest neighbors to a vector

var items []Item
err := db.Model(&items).
    OrderExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 2, 3})).
    Limit(5).
    Select()

Add an approximate index

_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
// or
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

Bun

Import the package

import "github.com/pgvector/pgvector-go"

Enable the extension

_, err := db.Exec("CREATE EXTENSION IF NOT EXISTS vector")

Add a vector column

type Item struct {
    Embedding pgvector.Vector `bun:"type:vector(3)"`
}

Insert a vector

item := Item{
    Embedding: pgvector.NewVector([]float32{1, 2, 3}),
}
_, err := db.NewInsert().Model(&item).Exec(ctx)

Get the nearest neighbors to a vector

var items []Item
err := db.NewSelect().
    Model(&items).
    OrderExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 2, 3})).
    Limit(5).
    Scan(ctx)

Add an approximate index

var _ bun.AfterCreateTableHook = (*Item)(nil)

func (*Item) AfterCreateTable(ctx context.Context, query *bun.CreateTableQuery) error {
    _, err := query.DB().NewCreateIndex().
        Model((*Item)(nil)).
        Index("items_embedding_idx").
        ColumnExpr("embedding vector_l2_ops").
        Using("hnsw").
        Exec(ctx)
    return err
}

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

Ent

Import the package

import "github.com/pgvector/pgvector-go"

Enable the extension (requires the sql/execquery feature)

_, err := client.ExecContext(ctx, "CREATE EXTENSION IF NOT EXISTS vector")

Add a vector column

func (Item) Fields() []ent.Field {
    return []ent.Field{
        field.Other("embedding", pgvector.Vector{}).
            SchemaType(map[string]string{
                dialect.Postgres: "vector(3)",
            }),
    }
}

Insert a vector

_, err := client.Item.
    Create().
    SetEmbedding(pgvector.NewVector([]float32{1, 2, 3})).
    Save(ctx)

Get the nearest neighbors to a vector

items, err := client.Item.
    Query().
    Order(func(s *sql.Selector) {
        s.OrderExpr(sql.ExprP("embedding <-> $1", pgvector.NewVector([]float32{1, 2, 3})))
    }).
    Limit(5).
    All(ctx)

Add an approximate index

func (Item) Indexes() []ent.Index {
    return []ent.Index{
        index.Fields("embedding").
            Annotations(
                entsql.IndexType("hnsw"),
                entsql.OpClass("vector_l2_ops"),
            ),
    }
}

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

GORM

Import the package

import "github.com/pgvector/pgvector-go"

Enable the extension

db.Exec("CREATE EXTENSION IF NOT EXISTS vector")

Add a vector column

type Item struct {
    Embedding pgvector.Vector `gorm:"type:vector(3)"`
}

Insert a vector

item := Item{
    Embedding: pgvector.NewVector([]float32{1, 2, 3}),
}
result := db.Create(&item)

Get the nearest neighbors to a vector

var items []Item
db.Clauses(clause.OrderBy{
    Expression: clause.Expr{SQL: "embedding <-> ?", Vars: []interface{}{pgvector.NewVector([]float32{1, 1, 1})}},
}).Limit(5).Find(&items)

Add an approximate index

db.Exec("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
// or
db.Exec("CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

sqlx

Import the package

import "github.com/pgvector/pgvector-go"

Enable the extension

db.MustExec("CREATE EXTENSION IF NOT EXISTS vector")

Add a vector column

type Item struct {
    Embedding pgvector.Vector
}

Insert a vector

item := Item{
    Embedding: pgvector.NewVector([]float32{1, 2, 3}),
}
_, err := db.NamedExec(`INSERT INTO items (embedding) VALUES (:embedding)`, item)

Get the nearest neighbors to a vector

var items []Item
db.Select(&items, "SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5", pgvector.NewVector([]float32{1, 1, 1}))

Add an approximate index

db.MustExec("CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
// or
db.MustExec("CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

History

View the changelog

Contributing

Everyone is encouraged to help improve this project. Here are a few ways you can help:

• Report bugs
• Fix bugs and submit pull requests
• Write, clarify, or fix documentation
• Suggest or add new features

To get started with development:

git clone https://github.com/pgvector/pgvector-go.git
cd pgvector-go
go mod tidy
createdb pgvector_go_test
go generate ./ent
go test -v

To run an example:

createdb pgvector_example
go run ./examples/loading