Added support for Postgres arrays for pgx - closes #20

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.2.3 (unreleased)
+
+- Added support for Postgres arrays for pgx
+
 ## 0.2.2 (2024-08-10)
 
 - Added support for `CopyFrom` with `string` values

pgx/register.go

@@ -10,9 +10,12 @@ import (
 
 func RegisterTypes(ctx context.Context, conn *pgx.Conn) error {
 	var vectorOid *uint32
+	var vectorArrayOid *uint32
 	var halfvecOid *uint32
+	var halfvecArrayOid *uint32
 	var sparsevecOid *uint32
-	err := conn.QueryRow(ctx, "SELECT to_regtype('vector')::oid, to_regtype('halfvec')::oid, to_regtype('sparsevec')::oid").Scan(&vectorOid, &halfvecOid, &sparsevecOid)
+	var sparsevecArrayOid *uint32
+	err := conn.QueryRow(ctx, "SELECT to_regtype('vector')::oid, to_regtype('_vector')::oid, to_regtype('halfvec')::oid, to_regtype('_halfvec')::oid, to_regtype('sparsevec')::oid, to_regtype('_sparsevec')::oid").Scan(&vectorOid, &vectorArrayOid, &halfvecOid, &halfvecArrayOid, &sparsevecOid, &sparsevecArrayOid)
 	if err != nil {
 		return err
 	}
@@ -22,14 +25,20 @@ func RegisterTypes(ctx context.Context, conn *pgx.Conn) error {
 	}
 
 	tm := conn.TypeMap()
-	tm.RegisterType(&pgtype.Type{Name: "vector", OID: *vectorOid, Codec: &VectorCodec{}})
+	vectorType := pgtype.Type{Name: "vector", OID: *vectorOid, Codec: &VectorCodec{}}
+	tm.RegisterType(&vectorType)
+	tm.RegisterType(&pgtype.Type{Name: "_vector", OID: *vectorArrayOid, Codec: &pgtype.ArrayCodec{ElementType: &vectorType}})
 
 	if halfvecOid != nil {
-		tm.RegisterType(&pgtype.Type{Name: "halfvec", OID: *halfvecOid, Codec: &HalfVectorCodec{}})
+		halfvecType := pgtype.Type{Name: "halfvec", OID: *halfvecOid, Codec: &HalfVectorCodec{}}
+		tm.RegisterType(&halfvecType)
+		tm.RegisterType(&pgtype.Type{Name: "_halfvec", OID: *halfvecArrayOid, Codec: &pgtype.ArrayCodec{ElementType: &halfvecType}})
 	}
 
 	if sparsevecOid != nil {
-		tm.RegisterType(&pgtype.Type{Name: "sparsevec", OID: *sparsevecOid, Codec: &SparseVectorCodec{}})
+		sparsevecType := pgtype.Type{Name: "sparsevec", OID: *sparsevecOid, Codec: &SparseVectorCodec{}}
+		tm.RegisterType(&sparsevecType)
+		tm.RegisterType(&pgtype.Type{Name: "_sparsevec", OID: *sparsevecArrayOid, Codec: &pgtype.ArrayCodec{ElementType: &sparsevecType}})
 	}
 
 	return nil

pgx_test.go

@@ -169,4 +169,25 @@ func TestPgx(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
+
+	embeddings := []pgvector.Vector{pgvector.NewVector([]float32{1, 2, 3}), pgvector.NewVector([]float32{4, 5, 6})}
+	halfEmbeddings := []pgvector.HalfVector{pgvector.NewHalfVector([]float32{1, 2, 3}), pgvector.NewHalfVector([]float32{4, 5, 6})}
+	sparseEmbeddings := []pgvector.SparseVector{pgvector.NewSparseVector([]float32{1, 2, 3}), pgvector.NewSparseVector([]float32{4, 5, 6})}
+	row = conn.QueryRow(ctx, "SELECT $1::vector[], $2::halfvec[], $3::sparsevec[]", embeddings, halfEmbeddings, sparseEmbeddings)
+	var scanEmbeddings []pgvector.Vector
+	var scanHalfEmbeddings []pgvector.HalfVector
+	var scanSparseEmbeddings []pgvector.SparseVector
+	err = row.Scan(&scanEmbeddings, &scanHalfEmbeddings, &scanSparseEmbeddings)
+	if err != nil {
+		panic(err)
+	}
+	if !reflect.DeepEqual(scanEmbeddings, embeddings) {
+		t.Error()
+	}
+	if !reflect.DeepEqual(scanHalfEmbeddings, halfEmbeddings) {
+		t.Error()
+	}
+	if !reflect.DeepEqual(scanSparseEmbeddings, sparseEmbeddings) {
+		t.Error()
+	}
 }