pgvector 0.3.1 → 0.3.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 3d3ef6a53417383ff7f8a2a514df78513dfe9029e524f0f624df8828fb99dba0
-  data.tar.gz: f326a21d3942079cdadc22d9a61367e7a6651ae37c5eb07a3955182f9f569ad0
+  metadata.gz: fbdde0af357aaae0f727dff71e8678c705b1b347ecca7b49a8ca097189b42db0
+  data.tar.gz: d9d8bee13760ca165905f888e79507350565425ae819c452bf723ff775899235
 SHA512:
-  metadata.gz: fa4d6519685d179e3d5712cbcf1e6989ca4f98e5b0902f5061eca2be55451162f199fcc75f19841dbd8ada8c18de69ac6fa39cad545d4b5a6c542e9cdd0109ea
-  data.tar.gz: fd71245492b2ff8a06a0af60dbd12e57e44025f2662d4ea77eb7176aaaa4c4cb3b39c1e159edce92944da5b54fddbb19d1b504acd76a12c4efdb9a0fb6e797da
+  metadata.gz: '04629b78419ba4fd5325792d2b47f60a6a42fe2e5f8c681bcbd53d2f769b0c8db74ba7e53b9ef84c579b1c98ded45b51e3fe58347394b5f7b707efc81a6500f3'
+  data.tar.gz: 995963de5eb7bd7587bff1f70a709cb9b4e7f0791ef6f23b6930104c5a0ba61514a684b2dbf784b2e5b62464e84f5c6edbf569681578d75c5db860500c46bbae

data/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.3.2 (2024-07-17)
+
+- Added `to_binary` method to `Vector`, `Bit`, and `SparseVector`
+
 ## 0.3.1 (2024-07-10)
 
 - Added support for `bit` type to pg

data/README.md

@@ -24,6 +24,7 @@ And follow the instructions for your database library:
 Or check out some examples:
 
 - [Embeddings](examples/openai_embeddings.rb) with OpenAI
+- [Binary embeddings](examples/cohere_embeddings.rb) with Cohere
 - [User-based recommendations](examples/disco_user_recs.rb) with Disco
 - [Item-based recommendations](examples/disco_item_recs.rb) with Disco
 - [Bulk loading](examples/bulk_loading.rb) with `COPY`

data/lib/pgvector/bit.rb

@@ -13,8 +13,8 @@ module Pgvector
     end
 
     def self.from_binary(string)
-      length = string[..3].unpack1("l>")
-      Bit.new(string[4..].unpack("B*").join[...length])
+      length, data = string.unpack("l>B*")
+      Bit.new(data[...length])
     end
 
     def to_s
@@ -24,5 +24,9 @@ module Pgvector
     def to_a
       @data.each_char.map { |v| v != "0" }
     end
+
+    def to_binary
+      [@data.length, @data].pack("l>B*")
+    end
   end
 end

data/lib/pgvector/pg.rb

@@ -3,17 +3,17 @@ require "pg"
 module Pgvector
   module PG
     def self.register_vector(registry)
-      registry.register_type(0, "vector", nil, TextDecoder::Vector)
-      registry.register_type(1, "vector", nil, BinaryDecoder::Vector)
+      registry.register_type(0, "vector", TextEncoder::Vector, TextDecoder::Vector)
+      registry.register_type(1, "vector", BinaryEncoder::Vector, BinaryDecoder::Vector)
 
       # no binary decoder for halfvec since unpack does not have directive for half-precision
-      registry.register_type(0, "halfvec", nil, TextDecoder::Halfvec)
+      registry.register_type(0, "halfvec", TextEncoder::Halfvec, TextDecoder::Halfvec)
 
-      registry.register_type(0, "bit", nil, TextDecoder::Bit)
-      registry.register_type(1, "bit", nil, BinaryDecoder::Bit)
+      registry.register_type(0, "bit", TextEncoder::Bit, TextDecoder::Bit)
+      registry.register_type(1, "bit", BinaryEncoder::Bit, BinaryDecoder::Bit)
 
-      registry.register_type(0, "sparsevec", nil, TextDecoder::Sparsevec)
-      registry.register_type(1, "sparsevec", nil, BinaryDecoder::Sparsevec)
+      registry.register_type(0, "sparsevec", TextEncoder::Sparsevec, TextDecoder::Sparsevec)
+      registry.register_type(1, "sparsevec", BinaryEncoder::Sparsevec, BinaryDecoder::Sparsevec)
     end
 
     module BinaryDecoder
@@ -31,7 +31,36 @@ module Pgvector
 
       class Sparsevec < ::PG::SimpleDecoder
         def decode(string, tuple = nil, field = nil)
-          SparseVector.from_binary(string)
+          ::Pgvector::SparseVector.from_binary(string)
+        end
+      end
+    end
+
+    module BinaryEncoder
+      # experimental
+      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
+
+      class Vector < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_binary
+        end
+      end
+
+      class Bit < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_binary
+        end
+      end
+
+      class Sparsevec < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_binary
         end
       end
     end
@@ -45,7 +74,7 @@ module Pgvector
 
       class Halfvec < ::PG::SimpleDecoder
         def decode(string, tuple = nil, field = nil)
-          HalfVector.from_text(string).to_a
+          ::Pgvector::HalfVector.from_text(string).to_a
         end
       end
 
@@ -57,7 +86,43 @@ module Pgvector
 
       class Sparsevec < ::PG::SimpleDecoder
         def decode(string, tuple = nil, field = nil)
-          SparseVector.from_text(string)
+          ::Pgvector::SparseVector.from_text(string)
+        end
+      end
+    end
+
+    module TextEncoder
+      # experimental
+      def self.type_map
+        tm = ::PG::TypeMapByClass.new
+        tm[::Pgvector::Vector] = Vector.new
+        tm[::Pgvector::HalfVector] = Halfvec.new
+        tm[::Pgvector::Bit] = Bit.new
+        tm[::Pgvector::SparseVector] = Sparsevec.new
+        tm
+      end
+
+      class Vector < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_s
+        end
+      end
+
+      class Halfvec < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_s
+        end
+      end
+
+      class Bit < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_s
+        end
+      end
+
+      class Sparsevec < ::PG::SimpleEncoder
+        def encode(value)
+          value.to_s
         end
       end
     end

data/lib/pgvector/sparse_vector.rb

@@ -30,6 +30,14 @@ module Pgvector
       arr
     end
 
+    def to_binary
+      nnz = @indices.size
+      buffer = [dimensions, nnz, 0].pack("l>l>l>")
+      @indices.pack("l>#{nnz}", buffer: buffer)
+      @values.pack("g#{nnz}", buffer: buffer)
+      buffer
+    end
+
     private
 
     def from_hash(data, dimensions)

data/lib/pgvector/vector.rb

@@ -1,7 +1,13 @@
 module Pgvector
   class Vector
     def initialize(data)
-      @data = data.to_a.map(&:to_f)
+      # keep as NArray when possible for performance
+      @data =
+        if numo?(data)
+          data.cast_to(Numo::SFloat)
+        else
+          data.to_a.map(&:to_f)
+        end
     end
 
     def self.from_text(string)
@@ -19,7 +25,23 @@ module Pgvector
     end
 
     def to_a
-      @data
+      @data.to_a
+    end
+
+    def to_binary
+      if numo?(@data)
+        [@data.shape[0], 0].pack("s>s>") + @data.to_network.to_binary
+      else
+        buffer = [@data.size, 0].pack("s>s>")
+        @data.pack("g*", buffer: buffer)
+        buffer
+      end
+    end
+
+    private
+
+    def numo?(data)
+      defined?(Numo::NArray) && data.is_a?(Numo::NArray)
     end
   end
 end

data/lib/pgvector/version.rb

@@ -1,3 +1,3 @@
 module Pgvector
-  VERSION = "0.3.1"
+  VERSION = "0.3.2"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pgvector
 version: !ruby/object:Gem::Version
-  version: 0.3.1
+  version: 0.3.2
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-07-11 00:00:00.000000000 Z
+date: 2024-07-17 00:00:00.000000000 Z
 dependencies: []
 description:
 email: andrew@ankane.org