neighbor 0.3.2 → 0.5.0

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/sqlite_generator.rb

@@ -0,0 +1,13 @@
+require "rails/generators"
+
+module Neighbor
+  module Generators
+    class SqliteGenerator < Rails::Generators::Base
+      source_root File.join(__dir__, "templates")
+
+      def copy_templates
+        template "sqlite.rb", "config/initializers/neighbor.rb"
+      end
+    end
+  end
+end

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

@@ -0,0 +1,2 @@
+# Load the sqlite-vec extension
+Neighbor::SQLite.initialize!

data/lib/generators/neighbor/vector_generator.rb

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

data/lib/neighbor/attribute.rb

@@ -0,0 +1,48 @@
+module Neighbor
+  class Attribute < ActiveRecord::Type::Value
+    delegate :type, :serialize, :deserialize, :cast, to: :new_cast_type
+
+    def initialize(cast_type:, model:, type:, attribute_name:)
+      @cast_type = cast_type
+      @model = model
+      @type = type
+      @attribute_name = attribute_name
+    end
+
+    private
+
+    def cast_value(...)
+      new_cast_type.send(:cast_value, ...)
+    end
+
+    def new_cast_type
+      @new_cast_type ||= begin
+        if @cast_type.is_a?(ActiveModel::Type::Value)
+          case Utils.adapter(@model)
+          when :sqlite
+            case @type&.to_sym
+            when :int8
+              Type::SqliteInt8Vector.new
+            when :bit
+              @cast_type
+            when :float32, nil
+              Type::SqliteVector.new
+            else
+              raise ArgumentError, "Unsupported type"
+            end
+          when :mariadb
+            if @model.columns_hash[@attribute_name.to_s]&.type == :integer
+              @cast_type
+            else
+              Type::MysqlVector.new
+            end
+          else
+            @cast_type
+          end
+        else
+          @cast_type
+        end
+      end
+    end
+  end
+end

data/lib/neighbor/model.rb

@@ -1,12 +1,10 @@
 module Neighbor
   module Model
-    def has_neighbors(*attribute_names, dimensions: nil, normalize: nil)
+    def has_neighbors(*attribute_names, dimensions: nil, normalize: nil, type: 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 ||= {}
@@ -26,84 +24,116 @@ 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}
+          @neighbor_attributes[attribute_name] = {dimensions: dimensions, normalize: normalize, type: type&.to_sym}
+        end
 
-          attribute attribute_name, Neighbor::Vector.new(dimensions: dimensions, normalize: normalize, model: self, attribute_name: attribute_name)
+        if ActiveRecord::VERSION::STRING.to_f >= 7.2
+          decorate_attributes(attribute_names) do |name, cast_type|
+            Neighbor::Attribute.new(cast_type: cast_type, model: self, type: type, attribute_name: name)
+          end
+        else
+          attribute_names.each do |attribute_name|
+            attribute attribute_name do |cast_type|
+              Neighbor::Attribute.new(cast_type: cast_type, model: self, type: type, attribute_name: attribute_name)
+            end
+          end
+        end
+
+        if normalize
+          if ActiveRecord::VERSION::STRING.to_f >= 7.1
+            attribute_names.each do |attribute_name|
+              normalizes attribute_name, with: ->(v) { Neighbor::Utils.normalize(v, column_info: columns_hash[attribute_name.to_s]) }
+            end
+          else
+            attribute_names.each do |attribute_name|
+              attribute attribute_name do |cast_type|
+                Neighbor::NormalizedAttribute.new(cast_type: cast_type, model: self, attribute_name: attribute_name)
+              end
+            end
+          end
         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
-          end
-          raise ArgumentError, "missing keyword: :distance" unless options.is_a?(Hash) && options.key?(:distance)
-          distance = options.delete(:distance)
-          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
+        validate do
+          adapter = Utils.adapter(self.class)
+
+          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]
+            dimensions ||= column_info&.limit unless column_info&.type == :binary
+            type = v[:type] || Utils.type(adapter, column_info&.type)
+
+            if !Neighbor::Utils.validate_dimensions(value, type, dimensions, adapter).nil?
+              errors.add(k, "must have #{dimensions} dimensions")
+            end
+            if !Neighbor::Utils.validate_finite(value, type)
+              errors.add(k, "must have finite values")
+            end
           end
-          attribute_name = attribute_name.to_sym
+        end
 
+        scope :nearest_neighbors, ->(attribute_name, vector, distance:, precision: nil) {
+          attribute_name = attribute_name.to_sym
           options = neighbor_attributes[attribute_name]
           raise ArgumentError, "Invalid attribute" unless options
           normalize = options[:normalize]
           dimensions = options[:dimensions]
+          type = options[:type]
 
           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 = columns_hash[attribute_name.to_s]
+          column_type = column_info&.type
 
-          column_info = klass.type_for_attribute(attribute_name).column_info
-
-          operator =
-            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
+          adapter = Neighbor::Utils.adapter(klass)
+          if type && adapter != :sqlite
+            raise ArgumentError, "type only works with SQLite"
+          end
 
+          operator = Neighbor::Utils.operator(adapter, column_type, distance)
           raise ArgumentError, "Invalid distance: #{distance}" unless operator
 
           # ensure normalize set (can be true or false)
-          if distance == "cosine" && column_info[:type] == :cube && normalize.nil?
+          normalize_required = Utils.normalize_required?(adapter, column_type)
+          if distance == "cosine" && normalize_required && 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)
+          dimensions ||= column_info&.limit unless column_info&.type == :binary
+          Neighbor::Utils.validate(vector, dimensions: dimensions, type: type || Utils.type(adapter, column_info&.type), adapter: adapter)
+          vector = Neighbor::Utils.normalize(vector, column_info: column_info) if normalize
+
+          quoted_attribute = nil
+          query = nil
+          connection_pool.with_connection do |c|
+            quoted_attribute = "#{c.quote_table_name(table_name)}.#{c.quote_column_name(attribute_name)}"
+            query = c.quote(column_attribute.serialize(vector))
+          end
 
-          # 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?
+            if adapter != :postgresql || column_type != :vector
+              raise ArgumentError, "Precision not supported for this type"
+            end
+
+            case precision.to_s
+            when "half"
+              cast_dimensions = dimensions || column_info&.limit
+              raise ArgumentError, "Unknown dimensions" unless cast_dimensions
+              quoted_attribute += "::halfvec(#{connection_pool.with_connection { |c| c.quote(cast_dimensions.to_i) }})"
             else
-              "cube(array[#{vector.map(&:to_f).join(", ")}])"
+              raise ArgumentError, "Invalid precision"
             end
+          end
 
-          order = "#{quoted_attribute} #{operator} #{query}"
+          order = Utils.order(adapter, type, operator, quoted_attribute, query)
 
           # https://stats.stackexchange.com/questions/146221/is-cosine-similarity-identical-to-l2-normalized-euclidean-distance
           # with normalized vectors:
@@ -111,31 +141,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 distance == "cosine" && normalize_required
               "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/mysql.rb

@@ -0,0 +1,37 @@
+module Neighbor
+  module MySQL
+    def self.initialize!
+      require_relative "type/mysql_vector"
+
+      require "active_record/connection_adapters/abstract_mysql_adapter"
+
+      # ensure schema can be dumped
+      ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter::NATIVE_DATABASE_TYPES[:vector] = {name: "vector"}
+
+      # ensure schema can be loaded
+      unless ActiveRecord::ConnectionAdapters::TableDefinition.method_defined?(:vector)
+        ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :vector)
+      end
+
+      # prevent unknown OID warning
+      ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter.singleton_class.prepend(RegisterTypes)
+      if ActiveRecord::VERSION::STRING.to_f < 7.1
+        ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter.register_vector_type(ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter::TYPE_MAP)
+      end
+    end
+
+    module RegisterTypes
+      def initialize_type_map(m)
+        super
+        register_vector_type(m)
+      end
+
+      def register_vector_type(m)
+        m.register_type %r(^vector)i do |sql_type|
+          limit = extract_limit(sql_type)
+          Type::MysqlVector.new(limit: limit)
+        end
+      end
+    end
+  end
+end

data/lib/neighbor/normalized_attribute.rb

@@ -0,0 +1,21 @@
+module Neighbor
+  class NormalizedAttribute < ActiveRecord::Type::Value
+    delegate :type, :serialize, :deserialize, to: :@cast_type
+
+    def initialize(cast_type:, model:, attribute_name:)
+      @cast_type = cast_type
+      @model = model
+      @attribute_name = attribute_name.to_s
+    end
+
+    def cast(...)
+      Neighbor::Utils.normalize(@cast_type.cast(...), column_info: @model.columns_hash[@attribute_name])
+    end
+
+    private
+
+    def cast_value(...)
+      @cast_type.send(:cast_value, ...)
+    end
+  end
+end

data/lib/neighbor/postgresql.rb

@@ -0,0 +1,43 @@
+module Neighbor
+  module PostgreSQL
+    def self.initialize!
+      require_relative "type/cube"
+      require_relative "type/halfvec"
+      require_relative "type/sparsevec"
+      require_relative "type/vector"
+
+      require "active_record/connection_adapters/postgresql_adapter"
+
+      # 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, :halfvec, :sparsevec, :vector)
+
+      # prevent unknown OID warning
+      ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(RegisterTypes)
+    end
+
+    module RegisterTypes
+      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)
+        end
+      end
+    end
+  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/reranking.rb

@@ -0,0 +1,27 @@
+module Neighbor
+  module Reranking
+    def self.rrf(first_ranking, *rankings, k: 60)
+      rankings.unshift(first_ranking)
+
+      ranks = []
+      results = []
+      rankings.each do |ranking|
+        ranks << ranking.map.with_index.to_h { |v, i| [v, i + 1] }
+        results.concat(ranking)
+      end
+
+      results =
+        results.uniq.map do |result|
+          score =
+            ranks.sum do |rank|
+              r = rank[result]
+              r ? 1.0 / (k + r) : 0.0
+            end
+
+          {result: result, score: score}
+        end
+
+      results.sort_by { |v| -v[:score] }
+    end
+  end
+end

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/sqlite.rb

@@ -0,0 +1,28 @@
+module Neighbor
+  module SQLite
+    # note: this is a public API (unlike PostgreSQL and MySQL)
+    def self.initialize!
+      return if defined?(@initialized)
+
+      require_relative "type/sqlite_vector"
+      require_relative "type/sqlite_int8_vector"
+
+      require "sqlite_vec"
+      require "active_record/connection_adapters/sqlite3_adapter"
+
+      ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(InstanceMethods)
+
+      @initialized = true
+    end
+
+    module InstanceMethods
+      def configure_connection
+        super
+        db = ActiveRecord::VERSION::STRING.to_f >= 7.1 ? @raw_connection : @connection
+        db.enable_load_extension(1)
+        SqliteVec.load(db)
+        db.enable_load_extension(0)
+      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 Utils.array?(value)
+          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 Utils.array?(value)
+          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 Utils.array?(value)
+          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 Utils.array?(value)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to halfvec"
+        end
+      end
+    end
+  end
+end