Skip to content

Commit

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

func CreatePgItems(db *pg.DB) {
Expand All @@ -28,18 +29,21 @@ func CreatePgItems(db *pg.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})},
},
PgItem{
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})},
},
PgItem{
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 @@ -93,6 +97,9 @@ func TestPg(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.Model(&items).ColumnExpr("embedding <-> ?", pgvector.NewVector([]float32{1, 1, 1})).Order("id").Select(&distances)
Expand Down

0 comments on commit 48fdcf2

Please sign in to comment.