Skip to content

Commit

Permalink
Added test for Postgres arrays with Bun
Browse files Browse the repository at this point in the history
  • Loading branch information
ankane committed Jan 16, 2025
1 parent a4fefe2 commit 73016ab
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions bun_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ type BunItem struct {
HalfEmbedding pgvector.HalfVector `bun:"type:halfvec(3)"`
BinaryEmbedding string `bun:"type:bit(3)"`
SparseEmbedding pgvector.SparseVector `bun:"type:sparsevec(3)"`
Embeddings []pgvector.Vector `bun:"type:vector(3)[]"`
}

var _ bun.AfterCreateTableHook = (*BunItem)(nil)
Expand All @@ -43,18 +44,21 @@ func CreateBunItems(ctx context.Context, db *bun.DB) {
HalfEmbedding: pgvector.NewHalfVector([]float32{1, 1, 1}),
BinaryEmbedding: "000",
SparseEmbedding: pgvector.NewSparseVector([]float32{1, 1, 1}),
Embeddings: []pgvector.Vector{pgvector.NewVector([]float32{1, 1, 1})},
},
BunItem{
Embedding: pgvector.NewVector([]float32{2, 2, 2}),
HalfEmbedding: pgvector.NewHalfVector([]float32{2, 2, 2}),
BinaryEmbedding: "101",
SparseEmbedding: pgvector.NewSparseVector([]float32{2, 2, 2}),
Embeddings: []pgvector.Vector{pgvector.NewVector([]float32{2, 2, 2})},
},
BunItem{
Embedding: pgvector.NewVector([]float32{1, 1, 2}),
HalfEmbedding: pgvector.NewHalfVector([]float32{1, 1, 2}),
BinaryEmbedding: "111",
SparseEmbedding: pgvector.NewSparseVector([]float32{1, 1, 2}),
Embeddings: []pgvector.Vector{pgvector.NewVector([]float32{1, 1, 2})},
},
}

Expand Down Expand Up @@ -105,6 +109,9 @@ func TestBun(t *testing.T) {
if !reflect.DeepEqual(items[1].SparseEmbedding.Slice(), []float32{1, 1, 2}) {
t.Error()
}
if !reflect.DeepEqual(items[1].Embeddings, []pgvector.Vector{pgvector.NewVector([]float32{1, 1, 2})}) {
t.Error()
}

var distances []float64
err = db.NewSelect().Model(&items).ColumnExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 1, 1})).Order("id").Scan(ctx, &distances)
Expand Down

0 comments on commit 73016ab

Please sign in to comment.