pgvector support for Go
Supports pgx, pg, Bun, Ent, GORM, and sqlx
Run:
go get github.com/pgvector/pgvector-go
And follow the instructions for your database library:
Or check out some examples:
- Embeddings with OpenAI
- Binary embeddings with Cohere
- Recommendations with Disco
- Horizontal scaling with Citus
- Bulk loading with
COPY
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
err := 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
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
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
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
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
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
View the changelog
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