neighbor 0.1.0 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f346d121745bd39998267c45f77009970c014de887376022f8b6425192d37354
-  data.tar.gz: '0293bc9fd3633ca3ee811940c59995407357d4b880fcf8ea8cd9223e5df0e75b'
+  metadata.gz: 8a71d9aff5133b9551dc47cfe490d5d6a0ee63d2b77ac1d3d2f381f58aef67d6
+  data.tar.gz: d22cab3140b979ee0d13ddd2c99f064282e896d49c265bf9b1946dafbda7783a
 SHA512:
-  metadata.gz: cf080c46ad8133a460453c773faabac02e229821625fa9f2d51a320195b1e821b230c94c7fe45bc7f7d12b0b51ce6d3e5f1d4c0d74e446b0ce22dde2f5171cde
-  data.tar.gz: f2e0ae2247979bd3dd083b1869b56f7b6ed879ce8dd1c0ef31d965392d8250548acb304a78ba99a1d6cdfea116c10e0d67aff9c06185a793e576bddae4d624f1
+  metadata.gz: 3e24a71c6ace693525fc4866aaa21ee7c75f7251d5d6e0820da3b6abc0bd6fd5921b20c2edbb7fc656b75991e1fc5607357dc02870072900e7b6ef972c840280
+  data.tar.gz: 5b75200355c62c549c2d61820e74e1cc71e4ebce68a1e390af60dbe76b6e1b7a69e71d73c8121db42ea05fc296db4ecc1f5776c71ad988685f13f71fc00fec66

data/CHANGELOG.md

@@ -1,3 +1,28 @@
+## 0.2.1 (2021-12-15)
+
+- Added support for Active Record 7
+
+## 0.2.0 (2021-04-21)
+
+- Added support for pgvector
+- Added `normalize` option
+- Made `dimensions` optional
+- Raise an error if `nearest_neighbors` already defined
+- Raise an error for non-finite values
+- Fixed NaN with zero vectors and cosine distance
+
+Breaking changes
+
+- The `distance` option has been moved from `has_neighbors` to `nearest_neighbors`, and there is no longer a default
+
+## 0.1.2 (2021-02-21)
+
+- Added `nearest_neighbors` scope
+
+## 0.1.1 (2021-02-16)
+
+- Fixed `Could not dump table` error
+
 ## 0.1.0 (2021-02-15)
 
 - First release

data/README.md

@@ -12,11 +12,21 @@ Add this line to your application’s Gemfile:
 gem 'neighbor'
 ```
 
-And run:
+## Choose An Extension
+
+Neighbor supports two extensions: [cube](https://www.postgresql.org/docs/current/cube.html) and [vector](https://github.com/ankane/pgvector). cube ships with Postgres, while vector supports approximate nearest neighbor search.
+
+For cube, run:
 
 ```sh
-bundle install
-rails generate neighbor:install
+rails generate neighbor:cube
+rails db:migrate
+```
+
+For vector, [install pgvector](https://github.com/ankane/pgvector#installation) and run:
+
+```sh
+rails generate neighbor:vector
 rails db:migrate
 ```
 
@@ -28,6 +38,8 @@ Create a migration
 class AddNeighborVectorToItems < ActiveRecord::Migration[6.1]
   def change
     add_column :items, :neighbor_vector, :cube
+    # or
+    add_column :items, :neighbor_vector, :vector, limit: 3 # dimensions
   end
 end
 ```
@@ -36,7 +48,7 @@ Add to your model
 
 ```ruby
 class Item < ApplicationRecord
-  has_neighbors dimensions: 3
+  has_neighbors
 end
 ```
 
@@ -46,40 +58,80 @@ Update the vectors
 item.update(neighbor_vector: [1.0, 1.2, 0.5])
 ```
 
-> With cosine distance (the default), vectors are normalized before being stored
+Get the nearest neighbors to a record
+
+```ruby
+item.nearest_neighbors(distance: "euclidean").first(5)
+```
 
-And get the nearest neighbors
+Get the nearest neighbors to a vector
 
 ```ruby
-item.nearest_neighbors.first(5)
+Item.nearest_neighbors([0.9, 1.3, 1.1], distance: "euclidean").first(5)
 ```
 
-## Distances
+## Distance
 
-Specify the distance metric
+Supported values are:
+
+- `euclidean`
+- `cosine`
+- `taxicab` (cube only)
+- `chebyshev` (cube only)
+- `inner_product` (vector only)
+
+For cosine distance with cube, vectors must be normalized before being stored.
 
 ```ruby
 class Item < ApplicationRecord
-  has_neighbors dimensions: 20, distance: "euclidean"
+  has_neighbors normalize: true
 end
 ```
 
-Supported distances are:
+For inner product with cube, see [this example](examples/disco_user_recs_cube.rb).
 
-- `cosine` (default)
-- `euclidean`
-- `taxicab`
-- `chebyshev`
+Records returned from `nearest_neighbors` will have a `neighbor_distance` attribute
+
+```ruby
+nearest_item = item.nearest_neighbors(distance: "euclidean").first
+nearest_item.neighbor_distance
+```
+
+## Dimensions
+
+The cube data type is limited 100 dimensions by default. See the [Postgres docs](https://www.postgresql.org/docs/current/cube.html) for how to increase this. The vector data type is limited to 1024 dimensions.
+
+For cube, it’s a good idea to specify the number of dimensions to ensure all records have the same number.
+
+```ruby
+class Movie < ApplicationRecord
+  has_neighbors dimensions: 3
+end
+```
+
+## Indexing
+
+For vector, add an approximate index to speed up queries. Create a migration with:
+
+```ruby
+class AddIndexToItemsNeighborVector < ActiveRecord::Migration[6.1]
+  def change
+    add_index :items, :neighbor_vector, using: :ivfflat, opclass: :vector_l2_ops
+  end
+end
+```
 
-Returned records will have a `neighbor_distance` attribute
+Use `:vector_cosine_ops` for cosine distance and `:vector_ip_ops` for inner product.
+
+Set the number of probes
 
 ```ruby
-returned_item.neighbor_distance
+Item.connection.execute("SET ivfflat.probes = 3")
 ```
 
 ## Example
 
-You can use Neighbor for online item recommendations with [Disco](https://github.com/ankane/disco). We’ll use MovieLens data for this example.
+You can use Neighbor for online item-based recommendations with [Disco](https://github.com/ankane/disco). We’ll use MovieLens data for this example.
 
 Generate a model
 
@@ -92,7 +144,7 @@ And add `has_neighbors`
 
 ```ruby
 class Movie < ApplicationRecord
-  has_neighbors dimensions: 20
+  has_neighbors dimensions: 20, normalize: true
 end
 ```
 
@@ -107,16 +159,32 @@ recommender.fit(data)
 Use item factors for the neighbor vector
 
 ```ruby
+movies = []
 recommender.item_ids.each do |item_id|
-  Movie.create!(name: item_id, neighbor_vector: recommender.item_factors(item_id))
+  movies << {name: item_id, neighbor_vector: recommender.item_factors(item_id)}
 end
+Movie.insert_all!(movies) # use create! for Active Record < 6
 ```
 
 And get similar movies
 
 ```ruby
 movie = Movie.find_by(name: "Star Wars (1977)")
-movie.nearest_neighbors.first(5).map(&:name)
+movie.nearest_neighbors(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
 ```
 
 ## History
@@ -138,5 +206,11 @@ To get started with development:
 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/{install_generator.rb → cube_generator.rb}

@@ -2,12 +2,12 @@ require "rails/generators/active_record"
 
 module Neighbor
   module Generators
-    class InstallGenerator < Rails::Generators::Base
+    class CubeGenerator < Rails::Generators::Base
       include ActiveRecord::Generators::Migration
       source_root File.join(__dir__, "templates")
 
       def copy_migration
-        migration_template "migration.rb", "db/migrate/install_neighbor.rb", migration_version: migration_version
+        migration_template "cube.rb", "db/migrate/install_neighbor_cube.rb", migration_version: migration_version
       end
 
       def migration_version

data/lib/generators/neighbor/templates/vector.rb.tt

@@ -0,0 +1,5 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    enable_extension "vector"
+  end
+end

data/lib/generators/neighbor/vector_generator.rb

@@ -0,0 +1,18 @@
+require "rails/generators/active_record"
+
+module Neighbor
+  module Generators
+    class VectorGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+      source_root File.join(__dir__, "templates")
+
+      def copy_migration
+        migration_template "vector.rb", "db/migrate/install_neighbor_vector.rb", migration_version: migration_version
+      end
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/neighbor/model.rb

@@ -1,42 +1,89 @@
 module Neighbor
   module Model
-    def has_neighbors(dimensions:, distance: "cosine")
-      distance = distance.to_s
-      raise ArgumentError, "Invalid distance: #{distance}" unless %w(cosine euclidean taxicab chebyshev).include?(distance)
+    def has_neighbors(dimensions: nil, normalize: nil)
+      # TODO make configurable
+      # likely use argument
+      attribute_name = :neighbor_vector
 
       class_eval do
-        attribute :neighbor_vector, Neighbor::Vector.new(dimensions: dimensions, distance: distance)
+        raise Error, "nearest_neighbors already defined" if method_defined?(:nearest_neighbors)
 
-        define_method :nearest_neighbors do
-          return self.class.none if neighbor_vector.nil?
+        attribute attribute_name, Neighbor::Vector.new(dimensions: dimensions, normalize: normalize, model: self, attribute_name: attribute_name)
+
+        scope :nearest_neighbors, ->(vector, distance:) {
+          return none if vector.nil?
+
+          distance = distance.to_s
+
+          quoted_attribute = "#{connection.quote_table_name(table_name)}.#{connection.quote_column_name(attribute_name)}"
+
+          column_info = klass.type_for_attribute(attribute_name).column_info
 
           operator =
-            case distance
-            when "taxicab"
-              "<#>"
-            when "chebyshev"
-              "<=>"
+            if column_info[:type] == :vector
+              case distance
+              when "inner_product"
+                "<#>"
+              when "cosine"
+                "<=>"
+              when "euclidean"
+                "<->"
+              end
             else
-              "<->"
+              case distance
+              when "taxicab"
+                "<#>"
+              when "chebyshev"
+                "<=>"
+              when "euclidean", "cosine"
+                "<->"
+              end
             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?
+            raise Neighbor::Error, "Set normalize for cosine distance with cube"
+          end
+
+          vector = Neighbor::Vector.cast(vector, dimensions: dimensions, normalize: normalize, column_info: column_info)
+
           # important! neighbor_vector should already be typecast
           # but use to_f as extra safeguard against SQL injection
-          order = "neighbor_vector #{operator} cube(array[#{neighbor_vector.map(&:to_f).join(", ")}])"
+          query =
+            if column_info[:type] == :vector
+              connection.quote("[#{vector.map(&:to_f).join(", ")}]")
+            else
+              "cube(array[#{vector.map(&:to_f).join(", ")}])"
+            end
+
+          order = "#{quoted_attribute} #{operator} #{query}"
 
           # https://stats.stackexchange.com/questions/146221/is-cosine-similarity-identical-to-l2-normalized-euclidean-distance
           # with normalized vectors:
           # cosine similarity = 1 - (euclidean distance)**2 / 2
           # cosine distance = 1 - cosine similarity
           # this transformation doesn't change the order, so only needed for select
-          neighbor_distance = distance == "cosine" ? "POWER(#{order}, 2) / 2.0" : order
+          neighbor_distance =
+            if column_info[:type] != :vector && distance == "cosine"
+              "POWER(#{order}, 2) / 2.0"
+            elsif column_info[:type] == :vector && 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")
+            .where.not(attribute_name => nil)
+            .order(Arel.sql(order))
+        }
+
+        define_method :nearest_neighbors do |**options|
           self.class
-            .select(*self.class.column_names, "#{neighbor_distance} AS neighbor_distance")
             .where.not(self.class.primary_key => send(self.class.primary_key))
-            .where.not(neighbor_vector: nil)
-            .order(Arel.sql(order))
+            .nearest_neighbors(send(attribute_name), **options)
         end
       end
     end

data/lib/neighbor/vector.rb

@@ -1,31 +1,61 @@
 module Neighbor
   class Vector < ActiveRecord::Type::Value
-    def initialize(dimensions:, distance:)
+    def initialize(dimensions:, normalize:, model:, attribute_name:)
       super()
       @dimensions = dimensions
-      @distance = distance
+      @normalize = normalize
+      @model = model
+      @attribute_name = attribute_name
     end
 
-    def cast(value)
-      return if value.nil?
-
+    def self.cast(value, dimensions:, normalize:, column_info:)
       value = value.to_a.map(&:to_f)
-      raise Error, "Expected #{@dimensions} dimensions, not #{value.size}" unless value.size == @dimensions
 
-      if @distance == "cosine"
-        norm = 0.0
-        value.each do |v|
-          norm += v * v
-        end
-        norm = Math.sqrt(norm)
-        value.map { |v| v / norm }
-      else
-        value
+      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)
-      "(#{cast(value).join(", ")})" unless value.nil?
+      unless value.nil?
+        if column_info[:type] == :vector
+          "[#{cast(value).join(", ")}]"
+        else
+          "(#{cast(value).join(", ")})"
+        end
+      end
     end
 
     def deserialize(value)

data/lib/neighbor/version.rb

@@ -1,3 +1,3 @@
 module Neighbor
-  VERSION = "0.1.0"
+  VERSION = "0.2.1"
 end

data/lib/neighbor.rb

@@ -7,10 +7,14 @@ require "neighbor/version"
 module Neighbor
   class Error < StandardError; end
 
-  module RegisterCubeType
+  module RegisterTypes
     def initialize_type_map(m = type_map)
       super
       m.register_type "cube", ActiveRecord::ConnectionAdapters::PostgreSQL::OID::SpecializedString.new(:cube)
+      m.register_type "vector" do |_, _, sql_type|
+        limit = extract_limit(sql_type)
+        ActiveRecord::ConnectionAdapters::PostgreSQL::OID::SpecializedString.new(:vector, limit: limit)
+      end
     end
   end
 end
@@ -21,7 +25,28 @@ ActiveSupport.on_load(:active_record) do
 
   extend Neighbor::Model
 
-  # prevent unknown OID warning
   require "active_record/connection_adapters/postgresql_adapter"
-  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Neighbor::RegisterCubeType)
+
+  # ensure schema can be dumped
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:cube] = {name: "cube"}
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter::NATIVE_DATABASE_TYPES[:vector] = {name: "vector"}
+
+  # ensure schema can be loaded
+  if ActiveRecord::VERSION::MAJOR >= 6
+    ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :cube, :vector)
+  else
+    ActiveRecord::ConnectionAdapters::TableDefinition.define_method :cube do |*args, **options|
+      args.each { |name| column(name, :cube, options) }
+    end
+    ActiveRecord::ConnectionAdapters::TableDefinition.define_method :vector do |*args, **options|
+      args.each { |name| column(name, :vector, options) }
+    end
+  end
+
+  # prevent unknown OID warning
+  if ActiveRecord::VERSION::MAJOR >= 7
+    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(Neighbor::RegisterTypes)
+  else
+    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Neighbor::RegisterTypes)
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: neighbor
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-02-16 00:00:00.000000000 Z
+date: 2021-12-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '5.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '5.2'
 description:
 email: andrew@ankane.org
 executables: []
@@ -33,8 +33,10 @@ files:
 - CHANGELOG.md
 - LICENSE.txt
 - README.md
-- lib/generators/neighbor/install_generator.rb
-- lib/generators/neighbor/templates/migration.rb.tt
+- lib/generators/neighbor/cube_generator.rb
+- lib/generators/neighbor/templates/cube.rb.tt
+- lib/generators/neighbor/templates/vector.rb.tt
+- lib/generators/neighbor/vector_generator.rb
 - lib/neighbor.rb
 - lib/neighbor/model.rb
 - lib/neighbor/vector.rb
@@ -58,7 +60,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.4
+rubygems_version: 3.2.32
 signing_key:
 specification_version: 4
 summary: Nearest neighbor search for Rails and Postgres