Added to_binary method to Bit

Author: ankane

CHANGELOG.md

@@ -1,6 +1,6 @@
 ## 0.3.2 (unreleased)
 
-- Added `to_binary` method to `Vector` and `SparseVector`
+- Added `to_binary` method to `Vector`, `Bit`, and `SparseVector`
 
 ## 0.3.1 (2024-07-10)
 

lib/pgvector/bit.rb

@@ -24,5 +24,11 @@ def to_s
     def to_a
       @data.each_char.map { |v| v != "0" }
     end
+
+    def to_binary
+      buffer = [@data.length].pack("l>")
+      @data.split(/.{8}/).pack("B*", buffer: buffer)
+      buffer
+    end
   end
 end

lib/pgvector/pg.rb

@@ -10,7 +10,7 @@ def self.register_vector(registry)
       registry.register_type(0, "halfvec", TextEncoder::Halfvec, TextDecoder::Halfvec)
 
       registry.register_type(0, "bit", TextEncoder::Bit, TextDecoder::Bit)
-      registry.register_type(1, "bit", nil, BinaryDecoder::Bit)
+      registry.register_type(1, "bit", BinaryEncoder::Bit, BinaryDecoder::Bit)
 
       registry.register_type(0, "sparsevec", TextEncoder::Sparsevec, TextDecoder::Sparsevec)
       registry.register_type(1, "sparsevec", BinaryEncoder::Sparsevec, BinaryDecoder::Sparsevec)
@@ -41,6 +41,7 @@ module BinaryEncoder
       def self.type_map
         tm = ::PG::TypeMapByClass.new
         tm[::Pgvector::Vector] = Vector.new
+        tm[::Pgvector::Bit] = Bit.new
         tm[::Pgvector::SparseVector] = Sparsevec.new
         tm
       end
@@ -51,6 +52,12 @@ def encode(value)
         end
       end
 
+      class Bit < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_binary
+        end
+      end
+
       class Sparsevec < ::PG::SimpleEncoder
         def encode(value)
           value.to_binary

test/pg_test.rb

@@ -88,9 +88,11 @@ def test_type_map_text
 
   def test_type_map_binary
     vec = Pgvector::Vector.new([1, 2, 3])
+    binary_vec = Pgvector::Bit.new([true, false, true])
     sparse_vec = Pgvector::SparseVector.new([1, 2, 3])
     coder = PG::BinaryEncoder::CopyRow.new(type_map: Pgvector::PG::BinaryEncoder.type_map)
     assert_match vec.to_binary, coder.encode([vec])
+    assert_match binary_vec.to_binary, coder.encode([binary_vec])
     assert_match sparse_vec.to_binary, coder.encode([sparse_vec])
   end
 
@@ -112,13 +114,15 @@ def test_copy_text
 
   def test_copy_binary
     embedding = Pgvector::Vector.new([1, 2, 3])
+    binary_embedding = Pgvector::Bit.new([true, false, true])
     sparse_embedding = Pgvector::SparseVector.new([1, 2, 3])
     coder = PG::BinaryEncoder::CopyRow.new
-    conn.copy_data("COPY pg_items (embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)", coder) do
-      conn.put_copy_data([embedding.to_binary, sparse_embedding.to_binary])
+    conn.copy_data("COPY pg_items (embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)", coder) do
+      conn.put_copy_data([embedding.to_binary, binary_embedding.to_binary, sparse_embedding.to_binary])
     end
     res = conn.exec("SELECT * FROM pg_items").first
     assert_equal [1, 2, 3], res["embedding"]
+    assert_equal "101", res["binary_embedding"]
     assert_equal [1, 2, 3], res["sparse_embedding"].to_a
   end