Added support for halfvec binary format for pgx - #24

Author: ankane

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.3.0 (unreleased)
 
 - Added distance functions for Ent
+- Added support for `halfvec` binary format for pgx
 - Dropped support for Go < 1.23
 
 ## 0.2.3 (2025-01-15)

go.mod

@@ -38,6 +38,7 @@ require (
 	github.com/vmihailenco/msgpack/v5 v5.3.5 // indirect
 	github.com/vmihailenco/tagparser v0.1.2 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
+	github.com/x448/float16 v0.8.4 // indirect
 	github.com/zclconf/go-cty v1.16.2 // indirect
 	github.com/zclconf/go-cty-yaml v1.1.0 // indirect
 	golang.org/x/crypto v0.36.0 // indirect

halfvec.go

@@ -71,6 +71,9 @@ func (v *HalfVector) Scan(src interface{}) (err error) {
 		return v.Parse(string(src))
 	case string:
 		return v.Parse(src)
+	case []float32:
+		v.vec = src
+		return nil
 	default:
 		return fmt.Errorf("unsupported data type: %T", src)
 	}

pgx/halfvec.go

@@ -2,21 +2,24 @@ package pgx
 
 import (
 	"database/sql/driver"
+	"encoding/binary"
 	"fmt"
+	"slices"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/pgvector/pgvector-go"
+	"github.com/x448/float16"
 )
 
 type HalfVectorCodec struct{}
 
 func (HalfVectorCodec) FormatSupported(format int16) bool {
-	return format == pgx.TextFormatCode
+	return format == pgx.BinaryFormatCode || format == pgx.TextFormatCode
 }
 
 func (HalfVectorCodec) PreferredFormat() int16 {
-	return pgx.TextFormatCode
+	return pgx.BinaryFormatCode
 }
 
 func (HalfVectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan {
@@ -25,35 +28,76 @@ func (HalfVectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value
 		return nil
 	}
 
-	if format == pgx.TextFormatCode {
+	switch format {
+	case pgx.BinaryFormatCode:
+		return encodePlanHalfVectorCodecBinary{}
+	case pgx.TextFormatCode:
 		return encodePlanHalfVectorCodecText{}
 	}
 
 	return nil
 }
 
+type encodePlanHalfVectorCodecBinary struct{}
+
+func (encodePlanHalfVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	v := value.(pgvector.HalfVector)
+	vec := v.Slice()
+	dim := len(vec)
+	buf = slices.Grow(buf, 4+2*dim)
+	buf = binary.BigEndian.AppendUint16(buf, uint16(dim))
+	buf = binary.BigEndian.AppendUint16(buf, 0)
+	for _, v := range vec {
+		buf = binary.BigEndian.AppendUint16(buf, float16.Fromfloat32(v).Bits())
+	}
+	return buf, nil
+}
+
 type encodePlanHalfVectorCodecText struct{}
 
 func (encodePlanHalfVectorCodecText) Encode(value any, buf []byte) (newBuf []byte, err error) {
 	v := value.(pgvector.HalfVector)
 	return v.EncodeText(buf)
 }
 
-type scanPlanHalfVectorCodecText struct{}
-
 func (HalfVectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan {
 	_, ok := target.(*pgvector.HalfVector)
 	if !ok {
 		return nil
 	}
 
-	if format == pgx.TextFormatCode {
+	switch format {
+	case pgx.BinaryFormatCode:
+		return scanPlanHalfVectorCodecBinary{}
+	case pgx.TextFormatCode:
 		return scanPlanHalfVectorCodecText{}
 	}
 
 	return nil
 }
 
+type scanPlanHalfVectorCodecBinary struct{}
+
+func (scanPlanHalfVectorCodecBinary) Scan(src []byte, dst any) error {
+	v := (dst).(*pgvector.HalfVector)
+	buf := src
+	dim := int(binary.BigEndian.Uint16(buf[0:2]))
+	unused := binary.BigEndian.Uint16(buf[2:4])
+	if unused != 0 {
+		return fmt.Errorf("expected unused to be 0")
+	}
+
+	vec := make([]float32, 0, dim)
+	offset := 4
+	for i := 0; i < dim; i++ {
+		vec = append(vec, float16.Frombits(binary.BigEndian.Uint16(buf[offset:offset+2])).Float32())
+		offset += 2
+	}
+	return v.Scan(vec)
+}
+
+type scanPlanHalfVectorCodecText struct{}
+
 func (scanPlanHalfVectorCodecText) Scan(src []byte, dst any) error {
 	v := (dst).(*pgvector.HalfVector)
 	return v.Scan(src)

pgx_test.go

@@ -161,9 +161,9 @@ func TestPgx(t *testing.T) {
 	_, err = conn.CopyFrom(
 		ctx,
 		pgx.Identifier{"pgx_items"},
-		[]string{"embedding", "binary_embedding", "sparse_embedding"},
+		[]string{"embedding", "half_embedding", "binary_embedding", "sparse_embedding"},
 		pgx.CopyFromSlice(1, func(i int) ([]any, error) {
-			return []interface{}{"[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
+			return []interface{}{"[1,2,3]", "[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
 		}),
 	)
 	if err != nil {
@@ -186,9 +186,9 @@ func TestPgx(t *testing.T) {
 	_, err = pool.CopyFrom(
 		ctx,
 		pgx.Identifier{"pgx_items"},
-		[]string{"embedding", "binary_embedding", "sparse_embedding"},
+		[]string{"embedding", "half_embedding", "binary_embedding", "sparse_embedding"},
 		pgx.CopyFromSlice(1, func(i int) ([]any, error) {
-			return []interface{}{"[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
+			return []interface{}{"[1,2,3]", "[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
 		}),
 	)
 	if err != nil {