neighbor 0.3.2 → 0.4.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0c8b5d19222742f33f51f2c30f9d03108ebd3ed99908a7e9dd5f4e49caa2e225
-  data.tar.gz: c9cfa942f2cdd8b9757c9ecfe5e89d0aced11263f8a559004ee15fa0c8adb3f4
+  metadata.gz: 8aa6de2790d94de9411b0142836b2ad181a411e299fce4b98357b96ac4161183
+  data.tar.gz: 2924d7f15f5b36bc89ee72372c1bfeb373d99481269696a9a9dcc41f90201f38
 SHA512:
-  metadata.gz: e9e0050031ce7691baa9242b3b6b5aa76afb1fe7c63575129e68b2f5c027143b3c08f68a7babfcf2a9b02f1d9327679f75e9c40b95ac2245ea7c8dd3025d3cdb
-  data.tar.gz: a9c505740cba454437617733d4025360848a16ef9a4c9c83fc16d5bc82a3e5521c77e3cba874ef3cf318cf3a1e319567958a6156481f7fd82ef72ebaa87d97eb
+  metadata.gz: 2bc1b3ee6d5b1ee0ab175b017e753cf958bd8ceb1ef2a23ba769770dfebf54eec251ac59c8f5f3b6ca56efcbad1763c34622b94924a017622c2f78fc8740f762
+  data.tar.gz: d946dda99833964582f63863b2d898fea6bf065312cf60aec873631df96195e1a54375606ad9c9cc0f767937cdb7ea38b0d9990efcbbeab15ccbb11f8a2020ef

data/CHANGELOG.md

@@ -1,3 +1,22 @@
+## 0.4.1 (2024-08-26)
+
+- Added `precision` option
+- Added support for `bit` dimensions to model generator
+- Fixed error with Numo arrays
+
+## 0.4.0 (2024-06-25)
+
+- Added support for `halfvec` and `sparsevec` types
+- Added support for `taxicab`, `hamming`, and `jaccard` distances with `vector` extension
+- Added deserialization for `cube` and `vector` columns without `has_neighbor`
+- Added support for composite primary keys
+- Changed `nearest_neighbors` to replace previous `order` scopes
+- Changed `normalize` option to use `before_save` callback
+- Changed dimensions and finite values checks to use Active Record validations
+- Fixed issue with `nearest_neighbors` scope overriding `select` values
+- Removed default attribute name
+- Dropped support for Ruby < 3.1
+
 ## 0.3.2 (2023-12-12)
 
 - Added deprecation warning for `has_neighbors` without an attribute name

data/LICENSE.txt

@@ -1,6 +1,6 @@
 The MIT License (MIT)
 
-Copyright (c) 2021-2023 Andrew Kane
+Copyright (c) 2021-2024 Andrew Kane
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

data/README.md

@@ -2,7 +2,7 @@
 
 Nearest neighbor search for Rails and Postgres
 
-[![Build Status](https://github.com/ankane/neighbor/workflows/build/badge.svg?branch=master)](https://github.com/ankane/neighbor/actions)
+[![Build Status](https://github.com/ankane/neighbor/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/neighbor/actions)
 
 ## Installation
 
@@ -35,7 +35,7 @@ rails db:migrate
 Create a migration
 
 ```ruby
-class AddEmbeddingToItems < ActiveRecord::Migration[7.1]
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
   def change
     add_column :items, :embedding, :cube
     # or
@@ -76,9 +76,11 @@ Supported values are:
 
 - `euclidean`
 - `cosine`
-- `taxicab` (cube only)
+- `taxicab`
 - `chebyshev` (cube only)
 - `inner_product` (vector only)
+- `hamming` (vector only)
+- `jaccard` (vector only)
 
 For cosine distance with cube, vectors must be normalized before being stored.
 
@@ -114,32 +116,114 @@ end
 For vector, add an approximate index to speed up queries. Create a migration with:
 
 ```ruby
-class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.1]
+class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
   def change
-    add_index :items, :embedding, using: :ivfflat, opclass: :vector_l2_ops
-    # or with pgvector 0.5.0+
     add_index :items, :embedding, using: :hnsw, opclass: :vector_l2_ops
+    # or
+    add_index :items, :embedding, using: :ivfflat, opclass: :vector_l2_ops
   end
 end
 ```
 
 Use `:vector_cosine_ops` for cosine distance and `:vector_ip_ops` for inner product.
 
-Set the number of probes with IVFFlat
+Set the size of the dynamic candidate list with HNSW
+
+```ruby
+Item.connection.execute("SET hnsw.ef_search = 100")
+```
+
+Or the number of probes with IVFFlat
 
 ```ruby
 Item.connection.execute("SET ivfflat.probes = 3")
 ```
 
-Or the size of the dynamic candidate list with HNSW
+## Half-Precision Vectors
+
+Use the `halfvec` type to store half-precision vectors
 
 ```ruby
-Item.connection.execute("SET hnsw.ef_search = 100")
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    add_column :items, :embedding, :halfvec, limit: 3 # dimensions
+  end
+end
+```
+
+## Half-Precision Indexing
+
+Index vectors at half precision for smaller indexes
+
+```ruby
+class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
+  def change
+    add_index :items, "(embedding::halfvec(3)) vector_l2_ops", using: :hnsw
+  end
+end
+```
+
+Get the nearest neighbors
+
+```ruby
+Item.nearest_neighbors(:embedding, [0.9, 1.3, 1.1], distance: "euclidean", precision: "half").first(5)
+```
+
+## Binary Vectors
+
+Use the `bit` type to store binary vectors
+
+```ruby
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    add_column :items, :embedding, :bit, limit: 3 # dimensions
+  end
+end
+```
+
+Get the nearest neighbors by Hamming distance
+
+```ruby
+Item.nearest_neighbors(:embedding, "101", distance: "hamming").first(5)
+```
+
+## Binary Quantization
+
+Use expression indexing for binary quantization
+
+```ruby
+class AddIndexToItemsEmbedding < ActiveRecord::Migration[7.2]
+  def change
+    add_index :items, "(binary_quantize(embedding)::bit(3)) bit_hamming_ops", using: :hnsw
+  end
+end
+```
+
+## Sparse Vectors
+
+Use the `sparsevec` type to store sparse vectors
+
+```ruby
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    add_column :items, :embedding, :sparsevec, limit: 3 # dimensions
+  end
+end
+```
+
+Get the nearest neighbors
+
+```ruby
+embedding = Neighbor::SparseVector.new({0 => 0.9, 1 => 1.3, 2 => 1.1}, 3)
+Item.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(5)
 ```
 
 ## Examples
 
 - [OpenAI Embeddings](#openai-embeddings)
+- [Cohere Embeddings](#cohere-embeddings)
+- [Sentence Embeddings](#sentence-embeddings)
+- [Sparse Embeddings](#sparse-embeddings)
 - [Disco Recommendations](#disco-recommendations)
 
 ### OpenAI Embeddings
@@ -170,10 +254,10 @@ def fetch_embeddings(input)
   }
   data = {
     input: input,
-    model: "text-embedding-ada-002"
+    model: "text-embedding-3-small"
   }
 
-  response = Net::HTTP.post(URI(url), data.to_json, headers)
+  response = Net::HTTP.post(URI(url), data.to_json, headers).tap(&:value)
   JSON.parse(response.body)["data"].map { |v| v["embedding"] }
 end
 ```
@@ -199,14 +283,221 @@ end
 Document.insert_all!(documents)
 ```
 
-And get similar articles
+And get similar documents
+
+```ruby
+document = Document.first
+document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:content)
+```
+
+See the [complete code](examples/openai/example.rb)
+
+### Cohere Embeddings
+
+Generate a model
+
+```sh
+rails generate model Document content:text embedding:bit{1024}
+rails db:migrate
+```
+
+And add `has_neighbors`
+
+```ruby
+class Document < ApplicationRecord
+  has_neighbors :embedding
+end
+```
+
+Create a method to call the [embed API](https://docs.cohere.com/reference/embed)
+
+```ruby
+def fetch_embeddings(input, input_type)
+  url = "https://api.cohere.com/v1/embed"
+  headers = {
+    "Authorization" => "Bearer #{ENV.fetch("CO_API_KEY")}",
+    "Content-Type" => "application/json"
+  }
+  data = {
+    texts: input,
+    model: "embed-english-v3.0",
+    input_type: input_type,
+    embedding_types: ["ubinary"]
+  }
+
+  response = Net::HTTP.post(URI(url), data.to_json, headers).tap(&:value)
+  JSON.parse(response.body)["embeddings"]["ubinary"].map { |e| e.map { |v| v.chr.unpack1("B*") }.join }
+end
+```
+
+Pass your input
+
+```ruby
+input = [
+  "The dog is barking",
+  "The cat is purring",
+  "The bear is growling"
+]
+embeddings = fetch_embeddings(input, "search_document")
+```
+
+Store the embeddings
+
+```ruby
+documents = []
+input.zip(embeddings) do |content, embedding|
+  documents << {content: content, embedding: embedding}
+end
+Document.insert_all!(documents)
+```
+
+Embed the search query
+
+```ruby
+query = "forest"
+query_embedding = fetch_embeddings([query], "search_query")[0]
+```
+
+And search the documents
+
+```ruby
+Document.nearest_neighbors(:embedding, query_embedding, distance: "hamming").first(5).map(&:content)
+```
+
+See the [complete code](examples/cohere/example.rb)
+
+### Sentence Embeddings
+
+You can generate embeddings locally with [Informers](https://github.com/ankane/informers).
+
+Generate a model
+
+```sh
+rails generate model Document content:text embedding:vector{384}
+rails db:migrate
+```
+
+And add `has_neighbors`
+
+```ruby
+class Document < ApplicationRecord
+  has_neighbors :embedding
+end
+```
+
+Load a [model](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
+
+```ruby
+model = Informers::Model.new("sentence-transformers/all-MiniLM-L6-v2")
+```
+
+Pass your input
+
+```ruby
+input = [
+  "The dog is barking",
+  "The cat is purring",
+  "The bear is growling"
+]
+embeddings = model.embed(input)
+```
+
+Store the embeddings
+
+```ruby
+documents = []
+input.zip(embeddings) do |content, embedding|
+  documents << {content: content, embedding: embedding}
+end
+Document.insert_all!(documents)
+```
+
+And get similar documents
 
 ```ruby
 document = Document.first
 document.nearest_neighbors(:embedding, distance: "cosine").first(5).map(&:content)
 ```
 
-See the [complete code](examples/openai_embeddings.rb)
+See the [complete code](examples/informers/example.rb)
+
+### Sparse Embeddings
+
+You can generate sparse embeddings locally with [Transformers.rb](https://github.com/ankane/transformers-ruby).
+
+Generate a model
+
+```sh
+rails generate model Document content:text embedding:sparsevec{30522}
+rails db:migrate
+```
+
+And add `has_neighbors`
+
+```ruby
+class Document < ApplicationRecord
+  has_neighbors :embedding
+end
+```
+
+Load a [model](https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1) to generate embeddings
+
+```ruby
+class EmbeddingModel
+  def initialize(model_id)
+    @model = Transformers::AutoModelForMaskedLM.from_pretrained(model_id)
+    @tokenizer = Transformers::AutoTokenizer.from_pretrained(model_id)
+    @special_token_ids = @tokenizer.special_tokens_map.map { |_, token| @tokenizer.vocab[token] }
+  end
+
+  def embed(input)
+    feature = @tokenizer.(input, padding: true, truncation: true, return_tensors: "pt", return_token_type_ids: false)
+    output = @model.(**feature)[0]
+    values = Torch.max(output * feature[:attention_mask].unsqueeze(-1), dim: 1)[0]
+    values = Torch.log(1 + Torch.relu(values))
+    values[0.., @special_token_ids] = 0
+    values.to_a
+  end
+end
+
+model = EmbeddingModel.new("opensearch-project/opensearch-neural-sparse-encoding-v1")
+```
+
+Pass your input
+
+```ruby
+input = [
+  "The dog is barking",
+  "The cat is purring",
+  "The bear is growling"
+]
+embeddings = model.embed(input)
+```
+
+Store the embeddings
+
+```ruby
+documents = []
+input.zip(embeddings) do |content, embedding|
+  documents << {content: content, embedding: Neighbor::SparseVector.new(embedding)}
+end
+Document.insert_all!(documents)
+```
+
+Embed the search query
+
+```ruby
+query = "forest"
+query_embedding = model.embed([query])[0]
+```
+
+And search the documents
+
+```ruby
+Document.nearest_neighbors(:embedding, Neighbor::SparseVector.new(query_embedding), distance: "inner_product").first(5).map(&:content)
+```
+
+See the [complete code](examples/sparse/example.rb)
 
 ### Disco Recommendations
 
@@ -242,7 +533,7 @@ movies = []
 recommender.item_ids.each do |item_id|
   movies << {name: item_id, factors: recommender.item_factors(item_id)}
 end
-Movie.insert_all!(movies) # use create! for Active Record < 6
+Movie.insert_all!(movies)
 ```
 
 And get similar movies
@@ -252,19 +543,7 @@ movie = Movie.find_by(name: "Star Wars (1977)")
 movie.nearest_neighbors(:factors, distance: "cosine").first(5).map(&:name)
 ```
 
-See the complete code for [cube](examples/disco_item_recs_cube.rb) and [vector](examples/disco_item_recs_vector.rb)
-
-## Upgrading
-
-### 0.2.0
-
-The `distance` option has been moved from `has_neighbors` to `nearest_neighbors`, and there is no longer a default. If you use cosine distance, set:
-
-```ruby
-class Item < ApplicationRecord
-  has_neighbors normalize: true
-end
-```
+See the complete code for [cube](examples/disco/item_recs_cube.rb) and [vector](examples/disco/item_recs_vector.rb)
 
 ## History
 
@@ -286,10 +565,5 @@ git clone https://github.com/ankane/neighbor.git
 cd neighbor
 bundle install
 createdb neighbor_test
-
-# cube
 bundle exec rake test
-
-# vector
-EXT=vector bundle exec rake test
 ```

data/lib/generators/neighbor/cube_generator.rb

@@ -1,3 +1,4 @@
+require "rails/generators"
 require "rails/generators/active_record"
 
 module Neighbor

data/lib/generators/neighbor/vector_generator.rb

@@ -1,3 +1,4 @@
+require "rails/generators"
 require "rails/generators/active_record"
 
 module Neighbor

data/lib/neighbor/model.rb

@@ -2,11 +2,9 @@ module Neighbor
   module Model
     def has_neighbors(*attribute_names, dimensions: nil, normalize: nil)
       if attribute_names.empty?
-        warn "[neighbor] has_neighbors without an attribute name is deprecated"
-        attribute_names << :neighbor_vector
-      else
-        attribute_names.map!(&:to_sym)
+        raise ArgumentError, "has_neighbors requires an attribute name"
       end
+      attribute_names.map!(&:to_sym)
 
       class_eval do
         @neighbor_attributes ||= {}
@@ -27,30 +25,46 @@ module Neighbor
         attribute_names.each do |attribute_name|
           raise Error, "has_neighbors already called for #{attribute_name.inspect}" if neighbor_attributes[attribute_name]
           @neighbor_attributes[attribute_name] = {dimensions: dimensions, normalize: normalize}
-
-          attribute attribute_name, Neighbor::Vector.new(dimensions: dimensions, normalize: normalize, model: self, attribute_name: attribute_name)
         end
 
         return if @neighbor_attributes.size != attribute_names.size
 
-        scope :nearest_neighbors, ->(attribute_name, vector = nil, options = nil) {
-          # cannot use keyword arguments with scope with Ruby 3.2 and Active Record 6.1
-          # https://github.com/rails/rails/issues/46934
-          if options.nil? && vector.is_a?(Hash)
-            options = vector
-            vector = nil
+        validate do
+          self.class.neighbor_attributes.each do |k, v|
+            value = read_attribute(k)
+            next if value.nil?
+
+            column_info = self.class.columns_hash[k.to_s]
+            dimensions = v[:dimensions] || column_info&.limit
+
+            if !Neighbor::Utils.validate_dimensions(value, column_info&.type, dimensions).nil?
+              errors.add(k, "must have #{dimensions} dimensions")
+            end
+            if !Neighbor::Utils.validate_finite(value, column_info&.type)
+              errors.add(k, "must have finite values")
+            end
           end
+        end
+
+        # TODO move to normalizes when Active Record < 7.1 no longer supported
+        before_save do
+          self.class.neighbor_attributes.each do |k, v|
+            next unless v[:normalize] && attribute_changed?(k)
+            value = read_attribute(k)
+            next if value.nil?
+            self[k] = Neighbor::Utils.normalize(value, column_info: self.class.columns_hash[k.to_s])
+          end
+        end
+
+        # cannot use keyword arguments with scope with Ruby 3.2 and Active Record 6.1
+        # https://github.com/rails/rails/issues/46934
+        scope :nearest_neighbors, ->(attribute_name, vector, options = nil) {
           raise ArgumentError, "missing keyword: :distance" unless options.is_a?(Hash) && options.key?(:distance)
           distance = options.delete(:distance)
+          precision = options.delete(:precision)
           raise ArgumentError, "unknown keywords: #{options.keys.map(&:inspect).join(", ")}" if options.any?
 
-          if vector.nil? && !attribute_name.nil? && attribute_name.respond_to?(:to_a)
-            warn "[neighbor] nearest_neighbors without an attribute name is deprecated"
-            vector = attribute_name
-            attribute_name = :neighbor_vector
-          end
           attribute_name = attribute_name.to_sym
-
           options = neighbor_attributes[attribute_name]
           raise ArgumentError, "Invalid attribute" unless options
           normalize = options[:normalize]
@@ -62,10 +76,21 @@ module Neighbor
 
           quoted_attribute = "#{connection.quote_table_name(table_name)}.#{connection.quote_column_name(attribute_name)}"
 
-          column_info = klass.type_for_attribute(attribute_name).column_info
+          column_info = columns_hash[attribute_name.to_s]
+          column_type = column_info&.type
 
           operator =
-            if column_info[:type] == :vector
+            case column_type
+            when :bit
+              case distance
+              when "hamming"
+                "<~>"
+              when "jaccard"
+                "<%>"
+              when "hamming2"
+                "#"
+              end
+            when :vector, :halfvec, :sparsevec
               case distance
               when "inner_product"
                 "<#>"
@@ -73,8 +98,10 @@ module Neighbor
                 "<=>"
               when "euclidean"
                 "<->"
+              when "taxicab"
+                "<+>"
               end
-            else
+            when :cube
               case distance
               when "taxicab"
                 "<#>"
@@ -83,27 +110,39 @@ module Neighbor
               when "euclidean", "cosine"
                 "<->"
               end
+            else
+              raise ArgumentError, "Unsupported type: #{column_type}"
             end
 
           raise ArgumentError, "Invalid distance: #{distance}" unless operator
 
           # ensure normalize set (can be true or false)
-          if distance == "cosine" && column_info[:type] == :cube && normalize.nil?
+          if distance == "cosine" && column_type == :cube && normalize.nil?
             raise Neighbor::Error, "Set normalize for cosine distance with cube"
           end
 
-          vector = Neighbor::Vector.cast(vector, dimensions: dimensions, normalize: normalize, column_info: column_info)
+          column_attribute = klass.type_for_attribute(attribute_name)
+          vector = column_attribute.cast(vector)
+          Neighbor::Utils.validate(vector, dimensions: dimensions, column_info: column_info)
+          vector = Neighbor::Utils.normalize(vector, column_info: column_info) if normalize
+
+          query = connection.quote(column_attribute.serialize(vector))
 
-          # important! neighbor_vector should already be typecast
-          # but use to_f as extra safeguard against SQL injection
-          query =
-            if column_info[:type] == :vector
-              connection.quote("[#{vector.map(&:to_f).join(", ")}]")
+          if !precision.nil?
+            case precision.to_s
+            when "half"
+              cast_dimensions = dimensions || column_info&.limit
+              raise ArgumentError, "Unknown dimensions" unless cast_dimensions
+              quoted_attribute += "::halfvec(#{connection.quote(cast_dimensions.to_i)})"
             else
-              "cube(array[#{vector.map(&:to_f).join(", ")}])"
+              raise ArgumentError, "Invalid precision"
             end
+          end
 
           order = "#{quoted_attribute} #{operator} #{query}"
+          if operator == "#"
+            order = "bit_count(#{order})"
+          end
 
           # https://stats.stackexchange.com/questions/146221/is-cosine-similarity-identical-to-l2-normalized-euclidean-distance
           # with normalized vectors:
@@ -111,31 +150,28 @@ module Neighbor
           # cosine distance = 1 - cosine similarity
           # this transformation doesn't change the order, so only needed for select
           neighbor_distance =
-            if column_info[:type] != :vector && distance == "cosine"
+            if column_type == :cube && distance == "cosine"
               "POWER(#{order}, 2) / 2.0"
-            elsif column_info[:type] == :vector && distance == "inner_product"
+            elsif [:vector, :halfvec, :sparsevec].include?(column_type) && distance == "inner_product"
               "(#{order}) * -1"
             else
               order
             end
 
           # for select, use column_names instead of * to account for ignored columns
-          select(*column_names, "#{neighbor_distance} AS neighbor_distance")
+          select_columns = select_values.any? ? [] : column_names
+          select(*select_columns, "#{neighbor_distance} AS neighbor_distance")
             .where.not(attribute_name => nil)
-            .order(Arel.sql(order))
+            .reorder(Arel.sql(order))
         }
 
-        def nearest_neighbors(attribute_name = nil, **options)
-          if attribute_name.nil?
-            warn "[neighbor] nearest_neighbors without an attribute name is deprecated"
-            attribute_name = :neighbor_vector
-          end
+        def nearest_neighbors(attribute_name, **options)
           attribute_name = attribute_name.to_sym
-          # important! check if neighbor attribute before calling send
+          # important! check if neighbor attribute before accessing
           raise ArgumentError, "Invalid attribute" unless self.class.neighbor_attributes[attribute_name]
 
           self.class
-            .where.not(self.class.primary_key => self[self.class.primary_key])
+            .where.not(Array(self.class.primary_key).to_h { |k| [k, self[k]] })
             .nearest_neighbors(attribute_name, self[attribute_name], **options)
         end
       end

data/lib/neighbor/railtie.rb

@@ -1,16 +1,16 @@
 module Neighbor
   class Railtie < Rails::Railtie
     generators do
+      require "rails/generators/generated_attribute"
+
       # rails generate model Item embedding:vector{3}
-      if defined?(Rails::Generators::GeneratedAttribute)
-        Rails::Generators::GeneratedAttribute.singleton_class.prepend(Neighbor::GeneratedAttribute)
-      end
+      Rails::Generators::GeneratedAttribute.singleton_class.prepend(Neighbor::GeneratedAttribute)
     end
   end
 
   module GeneratedAttribute
     def parse_type_and_options(type, *, **)
-      if type =~ /\A(vector)\{(\d+)\}\z/
+      if type =~ /\A(vector|halfvec|bit|sparsevec)\{(\d+)\}\z/
         return $1, limit: $2.to_i
       end
       super

data/lib/neighbor/sparse_vector.rb

@@ -0,0 +1,79 @@
+module Neighbor
+  class SparseVector
+    attr_reader :dimensions, :indices, :values
+
+    NO_DEFAULT = Object.new
+
+    def initialize(value, dimensions = NO_DEFAULT)
+      if value.is_a?(Hash)
+        if dimensions == NO_DEFAULT
+          raise ArgumentError, "missing dimensions"
+        end
+        from_hash(value, dimensions)
+      else
+        unless dimensions == NO_DEFAULT
+          raise ArgumentError, "extra argument"
+        end
+        from_array(value)
+      end
+    end
+
+    def to_s
+      "{#{@indices.zip(@values).map { |i, v| "#{i.to_i + 1}:#{v.to_f}" }.join(",")}}/#{@dimensions.to_i}"
+    end
+
+    def to_a
+      arr = Array.new(dimensions, 0.0)
+      @indices.zip(@values) do |i, v|
+        arr[i] = v
+      end
+      arr
+    end
+
+    private
+
+    def from_hash(data, dimensions)
+      elements = data.select { |_, v| v != 0 }.sort
+      @dimensions = dimensions.to_i
+      @indices = elements.map { |v| v[0].to_i }
+      @values = elements.map { |v| v[1].to_f }
+    end
+
+    def from_array(arr)
+      arr = arr.to_a
+      @dimensions = arr.size
+      @indices = []
+      @values = []
+      arr.each_with_index do |v, i|
+        if v != 0
+          @indices << i
+          @values << v.to_f
+        end
+      end
+    end
+
+    class << self
+      def from_text(string)
+        elements, dimensions = string.split("/", 2)
+        indices = []
+        values = []
+        elements[1..-2].split(",").each do |e|
+          index, value = e.split(":", 2)
+          indices << index.to_i - 1
+          values << value.to_f
+        end
+        from_parts(dimensions.to_i, indices, values)
+      end
+
+      private
+
+      def from_parts(dimensions, indices, values)
+        vec = allocate
+        vec.instance_variable_set(:@dimensions, dimensions)
+        vec.instance_variable_set(:@indices, indices)
+        vec.instance_variable_set(:@values, values)
+        vec
+      end
+    end
+  end
+end

data/lib/neighbor/type/cube.rb

@@ -1,36 +1,41 @@
 module Neighbor
   module Type
-    class Cube < ActiveRecord::Type::String
+    class Cube < ActiveRecord::Type::Value
       def type
         :cube
       end
 
-      def cast(value)
-        if value.is_a?(Array)
+      def serialize(value)
+        if value.respond_to?(:to_a)
+          value = value.to_a
           if value.first.is_a?(Array)
-            value.map { |v| cast_point(v) }.join(", ")
+            value = value.map { |v| serialize_point(v) }.join(", ")
           else
-            cast_point(value)
+            value = serialize_point(value)
           end
-        else
-          super
         end
+        super(value)
       end
 
-      # TODO uncomment in 0.4.0
-      # def deserialize(value)
-      #   if value.nil?
-      #     super
-      #   elsif value.include?("),(")
-      #     value[1..-1].split("),(").map { |v| v.split(",").map(&:to_f) }
-      #   else
-      #     value[1..-1].split(",").map(&:to_f)
-      #   end
-      # end
-
       private
 
-      def cast_point(value)
+      def cast_value(value)
+        if value.respond_to?(:to_a)
+          value.to_a
+        elsif value.is_a?(Numeric)
+          [value]
+        elsif value.is_a?(String)
+          if value.include?("),(")
+            value[1..-1].split("),(").map { |v| v.split(",").map(&:to_f) }
+          else
+            value[1..-1].split(",").map(&:to_f)
+          end
+        else
+          raise "can't cast #{value.class.name} to cube"
+        end
+      end
+
+      def serialize_point(value)
         "(#{value.map(&:to_f).join(", ")})"
       end
     end

data/lib/neighbor/type/halfvec.rb

@@ -0,0 +1,28 @@
+module Neighbor
+  module Type
+    class Halfvec < ActiveRecord::Type::Value
+      def type
+        :halfvec
+      end
+
+      def serialize(value)
+        if value.respond_to?(:to_a)
+          value = "[#{value.to_a.map(&:to_f).join(",")}]"
+        end
+        super(value)
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(String)
+          value[1..-1].split(",").map(&:to_f)
+        elsif value.respond_to?(:to_a)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to halfvec"
+        end
+      end
+    end
+  end
+end

data/lib/neighbor/type/sparsevec.rb

@@ -0,0 +1,30 @@
+module Neighbor
+  module Type
+    class Sparsevec < ActiveRecord::Type::Value
+      def type
+        :sparsevec
+      end
+
+      def serialize(value)
+        if value.is_a?(SparseVector)
+          value = "{#{value.indices.zip(value.values).map { |i, v| "#{i.to_i + 1}:#{v.to_f}" }.join(",")}}/#{value.dimensions.to_i}"
+        end
+        super(value)
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(SparseVector)
+          value
+        elsif value.is_a?(String)
+          SparseVector.from_text(value)
+        elsif value.respond_to?(:to_a)
+          value = SparseVector.new(value.to_a)
+        else
+          raise "can't cast #{value.class.name} to sparsevec"
+        end
+      end
+    end
+  end
+end

data/lib/neighbor/type/vector.rb

@@ -1,14 +1,28 @@
 module Neighbor
   module Type
-    class Vector < ActiveRecord::Type::String
+    class Vector < ActiveRecord::Type::Value
       def type
         :vector
       end
 
-      # TODO uncomment in 0.4.0
-      # def deserialize(value)
-      #   value[1..-1].split(",").map(&:to_f) unless value.nil?
-      # end
+      def serialize(value)
+        if value.respond_to?(:to_a)
+          value = "[#{value.to_a.map(&:to_f).join(",")}]"
+        end
+        super(value)
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(String)
+          value[1..-1].split(",").map(&:to_f)
+        elsif value.respond_to?(:to_a)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to vector"
+        end
+      end
     end
   end
 end

data/lib/neighbor/utils.rb

@@ -0,0 +1,42 @@
+module Neighbor
+  module Utils
+    def self.validate_dimensions(value, type, expected)
+      dimensions = type == :sparsevec ? value.dimensions : value.size
+      if expected && dimensions != expected
+        "Expected #{expected} dimensions, not #{dimensions}"
+      end
+    end
+
+    def self.validate_finite(value, type)
+      case type
+      when :bit
+        true
+      when :sparsevec
+        value.values.all?(&:finite?)
+      else
+        value.all?(&:finite?)
+      end
+    end
+
+    def self.validate(value, dimensions:, column_info:)
+      if (message = validate_dimensions(value, column_info&.type, dimensions || column_info&.limit))
+        raise Error, message
+      end
+
+      if !validate_finite(value, column_info&.type)
+        raise Error, "Values must be finite"
+      end
+    end
+
+    def self.normalize(value, column_info:)
+      raise Error, "Normalize not supported for type" unless [:cube, :vector, :halfvec].include?(column_info&.type)
+
+      norm = Math.sqrt(value.sum { |v| v * v })
+
+      # store zero vector as all zeros
+      # since NaN makes the distance always 0
+      # could also throw error
+      norm > 0 ? value.map { |v| v / norm } : value
+    end
+  end
+end

data/lib/neighbor/version.rb

@@ -1,3 +1,3 @@
 module Neighbor
-  VERSION = "0.3.2"
+  VERSION = "0.4.1"
 end

data/lib/neighbor.rb

@@ -2,6 +2,8 @@
 require "active_support"
 
 # modules
+require_relative "neighbor/sparse_vector"
+require_relative "neighbor/utils"
 require_relative "neighbor/version"
 
 module Neighbor
@@ -11,6 +13,14 @@ module Neighbor
     def initialize_type_map(m = type_map)
       super
       m.register_type "cube", Type::Cube.new
+      m.register_type "halfvec" do |_, _, sql_type|
+        limit = extract_limit(sql_type)
+        Type::Halfvec.new(limit: limit)
+      end
+      m.register_type "sparsevec" do |_, _, sql_type|
+        limit = extract_limit(sql_type)
+        Type::Sparsevec.new(limit: limit)
+      end
       m.register_type "vector" do |_, _, sql_type|
         limit = extract_limit(sql_type)
         Type::Vector.new(limit: limit)
@@ -21,8 +31,9 @@ end
 
 ActiveSupport.on_load(:active_record) do
   require_relative "neighbor/model"
-  require_relative "neighbor/vector"
   require_relative "neighbor/type/cube"
+  require_relative "neighbor/type/halfvec"
+  require_relative "neighbor/type/sparsevec"
   require_relative "neighbor/type/vector"
 
   extend Neighbor::Model
@@ -31,10 +42,12 @@ ActiveSupport.on_load(:active_record) do
 
   # ensure schema can be dumped
   ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:cube] = {name: "cube"}
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:halfvec] = {name: "halfvec"}
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:sparsevec] = {name: "sparsevec"}
   ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:vector] = {name: "vector"}
 
   # ensure schema can be loaded
-  ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :cube, :vector)
+  ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :cube, :halfvec, :sparsevec, :vector)
 
   # prevent unknown OID warning
   if ActiveRecord::VERSION::MAJOR >= 7

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: neighbor
 version: !ruby/object:Gem::Version
-  version: 0.3.2
+  version: 0.4.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-12-12 00:00:00.000000000 Z
+date: 2024-08-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -40,9 +40,12 @@ files:
 - lib/neighbor.rb
 - lib/neighbor/model.rb
 - lib/neighbor/railtie.rb
+- lib/neighbor/sparse_vector.rb
 - lib/neighbor/type/cube.rb
+- lib/neighbor/type/halfvec.rb
+- lib/neighbor/type/sparsevec.rb
 - lib/neighbor/type/vector.rb
-- lib/neighbor/vector.rb
+- lib/neighbor/utils.rb
 - lib/neighbor/version.rb
 homepage: https://github.com/ankane/neighbor
 licenses:
@@ -56,14 +59,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3'
+      version: '3.1'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.10
+rubygems_version: 3.5.11
 signing_key:
 specification_version: 4
 summary: Nearest neighbor search for Rails and Postgres

data/lib/neighbor/vector.rb

@@ -1,65 +0,0 @@
-module Neighbor
-  class Vector < ActiveRecord::Type::Value
-    def initialize(dimensions:, normalize:, model:, attribute_name:)
-      super()
-      @dimensions = dimensions
-      @normalize = normalize
-      @model = model
-      @attribute_name = attribute_name
-    end
-
-    def self.cast(value, dimensions:, normalize:, column_info:)
-      value = value.to_a.map(&:to_f)
-
-      dimensions ||= column_info[:dimensions]
-      raise Error, "Expected #{dimensions} dimensions, not #{value.size}" if dimensions && value.size != dimensions
-
-      raise Error, "Values must be finite" unless value.all?(&:finite?)
-
-      if normalize
-        norm = Math.sqrt(value.sum { |v| v * v })
-
-        # store zero vector as all zeros
-        # since NaN makes the distance always 0
-        # could also throw error
-
-        # safe to update in-place since earlier map dups
-        value.map! { |v| v / norm } if norm > 0
-      end
-
-      value
-    end
-
-    def self.column_info(model, attribute_name)
-      attribute_name = attribute_name.to_s
-      column = model.columns.detect { |c| c.name == attribute_name }
-      {
-        type: column.try(:type),
-        dimensions: column.try(:limit)
-      }
-    end
-
-    # need to be careful to avoid loading column info before needed
-    def column_info
-      @column_info ||= self.class.column_info(@model, @attribute_name)
-    end
-
-    def cast(value)
-      self.class.cast(value, dimensions: @dimensions, normalize: @normalize, column_info: column_info) unless value.nil?
-    end
-
-    def serialize(value)
-      unless value.nil?
-        if column_info[:type] == :vector
-          "[#{cast(value).join(", ")}]"
-        else
-          "(#{cast(value).join(", ")})"
-        end
-      end
-    end
-
-    def deserialize(value)
-      value[1..-1].split(",").map(&:to_f) unless value.nil?
-    end
-  end
-end