pgx_test.go

package pgvector_test

import (
	"context"
	"math"
	"reflect"
	"testing"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pgvector/pgvector-go"
	pgxvec "github.com/pgvector/pgvector-go/pgx"
)

type PgxItem struct {
	Id              int64
	Embedding       pgvector.Vector
	HalfEmbedding   pgvector.HalfVector
	BinaryEmbedding string
	SparseEmbedding pgvector.SparseVector
	Embeddings      []pgvector.Vector
}

func CreatePgxItems(ctx context.Context, conn *pgx.Conn) {
	items := []PgxItem{
		PgxItem{
			Embedding:       pgvector.NewVector([]float32{1, 1, 1}),
			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})},
		},
		PgxItem{
			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})},
		},
		PgxItem{
			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})},
		},
	}

	for _, item := range items {
		_, err := conn.Exec(ctx, "INSERT INTO pgx_items (embedding, half_embedding, binary_embedding, sparse_embedding, embeddings) VALUES ($1, $2, $3, $4, $5)", item.Embedding, item.HalfEmbedding, item.BinaryEmbedding, item.SparseEmbedding, item.Embeddings)
		if err != nil {
			panic(err)
		}
	}
}

func TestPgx(t *testing.T) {
	ctx := context.Background()

	conn, err := pgx.Connect(ctx, "postgres://localhost/pgvector_go_test")
	if err != nil {
		panic(err)
	}
	defer conn.Close(ctx)

	_, err = conn.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector")
	if err != nil {
		panic(err)
	}

	err = pgxvec.RegisterTypes(ctx, conn)
	if err != nil {
		panic(err)
	}

	_, err = conn.Exec(ctx, "DROP TABLE IF EXISTS pgx_items")
	if err != nil {
		panic(err)
	}

	_, err = conn.Exec(ctx, "CREATE TABLE pgx_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector(3)[])")
	if err != nil {
		panic(err)
	}

	_, err = conn.Exec(ctx, "CREATE INDEX ON pgx_items USING hnsw (embedding vector_l2_ops)")
	if err != nil {
		panic(err)
	}

	CreatePgxItems(ctx, conn)

	rows, err := conn.Query(ctx, "SELECT id, embedding, half_embedding, binary_embedding, sparse_embedding, embeddings, embedding <-> $1 FROM pgx_items ORDER BY embedding <-> $1 LIMIT 5", pgvector.NewVector([]float32{1, 1, 1}))
	if err != nil {
		panic(err)
	}
	defer rows.Close()

	var items []PgxItem
	var binaryEmbeddings []pgtype.Bits
	var distances []float64
	for rows.Next() {
		var item PgxItem
		var binaryEmbedding pgtype.Bits
		var distance float64
		err = rows.Scan(&item.Id, &item.Embedding, &item.HalfEmbedding, &binaryEmbedding, &item.SparseEmbedding, &item.Embeddings, &distance)
		if err != nil {
			panic(err)
		}
		items = append(items, item)
		binaryEmbeddings = append(binaryEmbeddings, binaryEmbedding)
		distances = append(distances, distance)
	}

	if rows.Err() != nil {
		panic(rows.Err())
	}

	if items[0].Id != 1 || items[1].Id != 3 || items[2].Id != 2 {
		t.Error()
	}
	if !reflect.DeepEqual(items[1].Embedding.Slice(), []float32{1, 1, 2}) {
		t.Error()
	}
	if !reflect.DeepEqual(items[1].HalfEmbedding.Slice(), []float32{1, 1, 2}) {
		t.Error()
	}
	if binaryEmbeddings[1].Bytes[0] != (7<<5) || binaryEmbeddings[1].Len != 3 {
		t.Error()
	}
	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()
	}
	if distances[0] != 0 || distances[1] != 1 || distances[2] != math.Sqrt(3) {
		t.Error()
	}

	var item PgxItem
	row := conn.QueryRow(ctx, "SELECT embedding, half_embedding, binary_embedding, sparse_embedding FROM pgx_items ORDER BY id DESC LIMIT 1", pgx.QueryResultFormats{pgx.TextFormatCode, pgx.TextFormatCode, pgx.TextFormatCode, pgx.TextFormatCode})
	err = row.Scan(&item.Embedding, &item.HalfEmbedding, &item.BinaryEmbedding, &item.SparseEmbedding)
	if err != nil {
		panic(err)
	}
	if !reflect.DeepEqual(item.Embedding.Slice(), []float32{1, 1, 2}) {
		t.Error()
	}
	if !reflect.DeepEqual(item.HalfEmbedding.Slice(), []float32{1, 1, 2}) {
		t.Error()
	}
	if item.BinaryEmbedding != "111" {
		t.Error()
	}
	if !reflect.DeepEqual(item.SparseEmbedding.Slice(), []float32{1, 1, 2}) {
		t.Error()
	}

	_, err = conn.CopyFrom(
		ctx,
		pgx.Identifier{"pgx_items"},
		[]string{"embedding", "half_embedding", "binary_embedding", "sparse_embedding"},
		pgx.CopyFromSlice(1, func(i int) ([]any, error) {
			return []interface{}{"[1,2,3]", "[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
		}),
	)
	if err != nil {
		panic(err)
	}

	config, err := pgxpool.ParseConfig("postgres://localhost/pgvector_go_test")
	if err != nil {
		panic(err)
	}
	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
		return pgxvec.RegisterTypes(ctx, conn)
	}
	pool, err := pgxpool.NewWithConfig(ctx, config)
	if err != nil {
		panic(err)
	}
	defer pool.Close()

	_, err = pool.CopyFrom(
		ctx,
		pgx.Identifier{"pgx_items"},
		[]string{"embedding", "half_embedding", "binary_embedding", "sparse_embedding"},
		pgx.CopyFromSlice(1, func(i int) ([]any, error) {
			return []interface{}{"[1,2,3]", "[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
		}),
	)
	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()
	}
}