diff --git a/lib/pgvector/pg.rb b/lib/pgvector/pg.rb index b115b05..8752dc5 100644 --- a/lib/pgvector/pg.rb +++ b/lib/pgvector/pg.rb @@ -36,6 +36,21 @@ def decode(string, tuple = nil, field = nil) end end + module BinaryEncoder + # experimental + def self.type_map + tm = ::PG::TypeMapByClass.new + tm[::Pgvector::Vector] = Vector.new + tm + end + + class Vector < ::PG::SimpleEncoder + def encode(value) + value.to_binary + end + end + end + module TextDecoder class Vector < ::PG::SimpleDecoder def decode(string, tuple = nil, field = nil) @@ -61,5 +76,20 @@ def decode(string, tuple = nil, field = nil) end end end + + module TextEncoder + # experimental + def self.type_map + tm = ::PG::TypeMapByClass.new + tm[::Pgvector::Vector] = Vector.new + tm + end + + class Vector < ::PG::SimpleEncoder + def encode(value) + value.to_s + end + end + end end end diff --git a/test/pg_test.rb b/test/pg_test.rb index 9a11729..ca8bdd5 100644 --- a/test/pg_test.rb +++ b/test/pg_test.rb @@ -70,6 +70,17 @@ def test_sparsevec_binary assert_nil res[1]["sparse_embedding"] end + def test_type_map_binary + vec = Pgvector::Vector.new([1, 2, 3]) + coder = PG::BinaryEncoder::CopyRow.new(type_map: Pgvector::PG::BinaryEncoder.type_map) + assert_include vec.to_binary, coder.encode([vec]) + end + + def test_type_map_binary + coder = PG::TextEncoder::CopyRow.new(type_map: Pgvector::PG::TextEncoder.type_map) + assert_equal "[1.0,2.0,3.0]\n", coder.encode([Pgvector::Vector.new([1, 2, 3])]) + end + def conn @@conn ||= begin conn = PG.connect(dbname: "pgvector_ruby_test")