neighbor 0.1.2 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e64a3292445759187b7c08fd9cc54fe0da340ea3e3ab5de909620de6395b4984
-  data.tar.gz: 274d73ed464a0503973c5549022dca8820257fe470c9746b5599da9a153d46e3
+  metadata.gz: 9a8839eeaeeef5b3ff8deb63165791929a81f9a9c1c807dd3d094e9bf770e881
+  data.tar.gz: 14035d37686f7035bc0e874de7933259dc0b8be9d046fca1d31867cacaf1c922
 SHA512:
-  metadata.gz: 17df9f8a5848337fc570c6fa685f54f4aa9d49bea8e52e74fbec098e0bc9e450b655c7d96cb66c6e8b6aaae4824f2c831ea26a14f7f46b50ef64c955cfb9839b
-  data.tar.gz: a47165d174ec86ebfdcd128e62c50b85d097aac535b6b6d424ea31988428a57237cdade2829d8e22a15e566ac7597371c95fabafad42cdbc5d56f063abab55e4
+  metadata.gz: bfbf7bca1b4290abf2ec91dd5ea032dd1878ab9d7a0938dfed4e6e63441aa557397345090713222a735beb32bc9f1b6cd7971536de9ca427db8650f053917040
+  data.tar.gz: 1ba36cba82aafbb2bb003e016f0817884b1129a7d0a8b2be7385e438999e6e7728224f7e146d4de611e171fc851ae64ebea3aa7bfa65f55868cb0c5477d8ecff

data/CHANGELOG.md

@@ -1,3 +1,16 @@
+## 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

data/README.md

@@ -12,15 +12,23 @@ 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
 ```
 
-This enables the [cube extension](https://www.postgresql.org/docs/current/cube.html) in Postgres
+For vector, install [pgvector](https://github.com/ankane/pgvector#installation) and run:
+
+```sh
+rails generate neighbor:vector
+rails db:migrate
+```
 
 ## Getting Started
 
@@ -30,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
   end
 end
 ```
@@ -38,7 +48,7 @@ Add to your model
 
 ```ruby
 class Item < ApplicationRecord
-  has_neighbors dimensions: 3
+  has_neighbors
 end
 ```
 
@@ -48,49 +58,76 @@ 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.first(5)
+item.nearest_neighbors(distance: "euclidean").first(5)
 ```
 
 Get the nearest neighbors to a vector
 
 ```ruby
-Item.nearest_neighbors([1, 2, 3])
+Item.nearest_neighbors([0.9, 1.3, 1.1], distance: "euclidean").first(5)
 ```
 
 ## 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 values are:
-
-- `cosine` (default)
-- `euclidean`
-- `taxicab`
-- `chebyshev`
-
-For inner product, see [this example](examples/disco_user_recs.rb)
+For inner product with cube, see [this example](examples/disco_user_recs.rb).
 
 Records returned from `nearest_neighbors` will have a `neighbor_distance` attribute
 
 ```ruby
-nearest_item = item.nearest_neighbors.first
+nearest_item = item.nearest_neighbors(distance: "euclidean").first
 nearest_item.neighbor_distance
 ```
 
 ## Dimensions
 
-By default, Postgres limits the `cube` data type to 100 dimensions. See the [Postgres docs](https://www.postgresql.org/docs/current/cube.html) for how to increase this.
+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
+  end
+end
+```
+
+Add `opclass: :vector_cosine_ops` for cosine distance and `opclass: :vector_ip_ops` for inner product.
+
+Set the number of probes
+
+```ruby
+Item.connection.execute("SET ivfflat.probes = 3")
+```
 
 ## Example
 
@@ -107,7 +144,7 @@ And add `has_neighbors`
 
 ```ruby
 class Movie < ApplicationRecord
-  has_neighbors dimensions: 20
+  has_neighbors dimensions: 20, normalize: true
 end
 ```
 
@@ -131,11 +168,23 @@ 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)
 ```
 
 [Complete code](examples/disco_item_recs.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
 
 View the [changelog](https://github.com/ankane/neighbor/blob/master/CHANGELOG.md)

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.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
@@ -25,16 +29,20 @@ 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[:vector] = {name: "vector"}
 
   # ensure schema can be loaded
   if ActiveRecord::VERSION::MAJOR >= 6
-    ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :cube)
+    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
-  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Neighbor::RegisterCubeType)
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Neighbor::RegisterTypes)
 end

data/lib/neighbor/model.rb

@@ -1,42 +1,78 @@
 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 attribute_name, Neighbor::Vector.new(dimensions: dimensions, distance: distance)
+        raise Error, "nearest_neighbors already defined" if method_defined?(:nearest_neighbors)
+
+        attribute attribute_name, Neighbor::Vector.new(dimensions: dimensions, normalize: normalize, model: self, attribute_name: attribute_name)
 
-        scope :nearest_neighbors, ->(vector) {
+        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
-          vector = Neighbor::Vector.cast(vector, dimensions: dimensions, distance: distance)
-          order = "#{quoted_attribute} #{operator} cube(array[#{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")
@@ -44,10 +80,10 @@ module Neighbor
             .order(Arel.sql(order))
         }
 
-        define_method :nearest_neighbors do
+        define_method :nearest_neighbors do |**options|
           self.class
             .where.not(self.class.primary_key => send(self.class.primary_key))
-            .nearest_neighbors(send(attribute_name))
+            .nearest_neighbors(send(attribute_name), **options)
         end
       end
     end

data/lib/neighbor/vector.rb

@@ -1,29 +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 self.cast(value, dimensions:, distance:)
+    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"
+      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 })
-        value.map { |v| v / norm }
-      else
-        value
+
+        # 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, distance: @distance) unless value.nil?
+      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.2"
+  VERSION = "0.2.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: neighbor
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.2.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-02-22 00:00:00.000000000 Z
+date: 2021-04-22 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