-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
6 changed files
with
206 additions
and
58 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
package pgx | ||
|
||
import ( | ||
"context" | ||
"fmt" | ||
|
||
"github.com/jackc/pgx/v5" | ||
"github.com/jackc/pgx/v5/pgtype" | ||
) | ||
|
||
func RegisterTypes(ctx context.Context, conn *pgx.Conn) error { | ||
var vectorOid *uint32 | ||
var sparsevecOid *uint32 | ||
err := conn.QueryRow(ctx, "SELECT to_regtype('vector')::oid, to_regtype('sparsevec')::oid").Scan(&vectorOid, &sparsevecOid) | ||
if err != nil { | ||
return err | ||
} | ||
|
||
if vectorOid == nil { | ||
return fmt.Errorf("vector type not found in the database") | ||
} | ||
|
||
tm := conn.TypeMap() | ||
tm.RegisterType(&pgtype.Type{Name: "vector", OID: *vectorOid, Codec: &VectorCodec{}}) | ||
|
||
if sparsevecOid != nil { | ||
tm.RegisterType(&pgtype.Type{Name: "sparsevec", OID: *sparsevecOid, Codec: &SparseVectorCodec{}}) | ||
} | ||
|
||
return nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,83 @@ | ||
package pgx | ||
|
||
import ( | ||
"database/sql/driver" | ||
"fmt" | ||
|
||
"github.com/jackc/pgx/v5" | ||
"github.com/jackc/pgx/v5/pgtype" | ||
"github.com/pgvector/pgvector-go" | ||
) | ||
|
||
type SparseVectorCodec struct{} | ||
|
||
func (SparseVectorCodec) FormatSupported(format int16) bool { | ||
return format == pgx.BinaryFormatCode | ||
} | ||
|
||
func (SparseVectorCodec) PreferredFormat() int16 { | ||
return pgx.BinaryFormatCode | ||
} | ||
|
||
func (SparseVectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan { | ||
_, ok := value.(pgvector.SparseVector) | ||
if !ok { | ||
return nil | ||
} | ||
|
||
if format == pgx.BinaryFormatCode { | ||
return encodePlanSparseVectorCodecBinary{} | ||
} | ||
|
||
return nil | ||
} | ||
|
||
type encodePlanSparseVectorCodecBinary struct{} | ||
|
||
func (encodePlanSparseVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { | ||
v := value.(pgvector.SparseVector) | ||
return v.EncodeBinary(buf) | ||
} | ||
|
||
type scanPlanSparseVectorCodecBinary struct{} | ||
|
||
func (SparseVectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan { | ||
_, ok := target.(*pgvector.SparseVector) | ||
if !ok { | ||
return nil | ||
} | ||
|
||
if format == pgx.BinaryFormatCode { | ||
return scanPlanSparseVectorCodecBinary{} | ||
} | ||
|
||
return nil | ||
} | ||
|
||
func (scanPlanSparseVectorCodecBinary) Scan(src []byte, dst any) error { | ||
v := (dst).(*pgvector.SparseVector) | ||
return v.DecodeBinary(src) | ||
} | ||
|
||
func (c SparseVectorCodec) DecodeDatabaseSQLValue(m *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) { | ||
return c.DecodeValue(m, oid, format, src) | ||
} | ||
|
||
func (c SparseVectorCodec) DecodeValue(m *pgtype.Map, oid uint32, format int16, src []byte) (any, error) { | ||
if src == nil { | ||
return nil, nil | ||
} | ||
|
||
var vec pgvector.SparseVector | ||
scanPlan := c.PlanScan(m, oid, format, &vec) | ||
if scanPlan == nil { | ||
return nil, fmt.Errorf("Unable to decode sparsevec type") | ||
} | ||
|
||
err := scanPlan.Scan(src, &vec) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
return vec, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,83 @@ | ||
package pgx | ||
|
||
import ( | ||
"database/sql/driver" | ||
"fmt" | ||
|
||
"github.com/jackc/pgx/v5" | ||
"github.com/jackc/pgx/v5/pgtype" | ||
"github.com/pgvector/pgvector-go" | ||
) | ||
|
||
type VectorCodec struct{} | ||
|
||
func (VectorCodec) FormatSupported(format int16) bool { | ||
return format == pgx.BinaryFormatCode | ||
} | ||
|
||
func (VectorCodec) PreferredFormat() int16 { | ||
return pgx.BinaryFormatCode | ||
} | ||
|
||
func (VectorCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan { | ||
_, ok := value.(pgvector.Vector) | ||
if !ok { | ||
return nil | ||
} | ||
|
||
if format == pgx.BinaryFormatCode { | ||
return encodePlanVectorCodecBinary{} | ||
} | ||
|
||
return nil | ||
} | ||
|
||
type encodePlanVectorCodecBinary struct{} | ||
|
||
func (encodePlanVectorCodecBinary) Encode(value any, buf []byte) (newBuf []byte, err error) { | ||
v := value.(pgvector.Vector) | ||
return v.EncodeBinary(buf) | ||
} | ||
|
||
type scanPlanVectorCodecBinary struct{} | ||
|
||
func (VectorCodec) PlanScan(m *pgtype.Map, oid uint32, format int16, target any) pgtype.ScanPlan { | ||
_, ok := target.(*pgvector.Vector) | ||
if !ok { | ||
return nil | ||
} | ||
|
||
if format == pgx.BinaryFormatCode { | ||
return scanPlanVectorCodecBinary{} | ||
} | ||
|
||
return nil | ||
} | ||
|
||
func (scanPlanVectorCodecBinary) Scan(src []byte, dst any) error { | ||
v := (dst).(*pgvector.Vector) | ||
return v.DecodeBinary(src) | ||
} | ||
|
||
func (c VectorCodec) DecodeDatabaseSQLValue(m *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) { | ||
return c.DecodeValue(m, oid, format, src) | ||
} | ||
|
||
func (c VectorCodec) DecodeValue(m *pgtype.Map, oid uint32, format int16, src []byte) (any, error) { | ||
if src == nil { | ||
return nil, nil | ||
} | ||
|
||
var vec pgvector.Vector | ||
scanPlan := c.PlanScan(m, oid, format, &vec) | ||
if scanPlan == nil { | ||
return nil, fmt.Errorf("Unable to decode vector type") | ||
} | ||
|
||
err := scanPlan.Scan(src, &vec) | ||
if err != nil { | ||
return nil, err | ||
} | ||
|
||
return vec, nil | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters