pgvector-go

pgvector support for Go

Supports pgx, pg, Bun, and Ent

Getting Started

Run:

go get github.com/pgvector/pgvector-go

Import the package:

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

And follow the instructions for your database library:

• pgx
• pg
• Bun
• Ent

Or check out some examples:

• Embeddings with OpenAI
• Recommendations with Disco

pgx

Enable the extension

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

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 ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
// or
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

pg

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 ivfflat (embedding vector_l2_ops) WITH (lists = 100)")
// or
_, err := conn.Exec(ctx, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")

Use vector_ip_ops for inner product and vector_cosine_ops for cosine distance

See a full example

Bun

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

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

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