neighbor 0.3.2 → 0.5.0

data/lib/neighbor/type/mysql_vector.rb

@@ -0,0 +1,33 @@
+module Neighbor
+  module Type
+    class MysqlVector < ActiveRecord::Type::Binary
+      def type
+        :vector
+      end
+
+      def serialize(value)
+        if Utils.array?(value)
+          value = value.to_a.pack("e*")
+        end
+        super(value)
+      end
+
+      def deserialize(value)
+        value = super
+        cast_value(value) unless value.nil?
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(String)
+          value.unpack("e*")
+        elsif Utils.array?(value)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to vector"
+        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 Utils.array?(value)
+          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/sqlite_int8_vector.rb

@@ -0,0 +1,29 @@
+module Neighbor
+  module Type
+    class SqliteInt8Vector < ActiveRecord::Type::Binary
+      def serialize(value)
+        if Utils.array?(value)
+          value = value.to_a.pack("c*")
+        end
+        super(value)
+      end
+
+      def deserialize(value)
+        value = super
+        cast_value(value) unless value.nil?
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(String)
+          value.unpack("c*")
+        elsif Utils.array?(value)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to vector"
+        end
+      end
+    end
+  end
+end

data/lib/neighbor/type/sqlite_vector.rb

@@ -0,0 +1,29 @@
+module Neighbor
+  module Type
+    class SqliteVector < ActiveRecord::Type::Binary
+      def serialize(value)
+        if Utils.array?(value)
+          value = value.to_a.pack("f*")
+        end
+        super(value)
+      end
+
+      def deserialize(value)
+        value = super
+        cast_value(value) unless value.nil?
+      end
+
+      private
+
+      def cast_value(value)
+        if value.is_a?(String)
+          value.unpack("f*")
+        elsif Utils.array?(value)
+          value.to_a
+        else
+          raise "can't cast #{value.class.name} to vector"
+        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 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 vector"
+        end
+      end
     end
   end
 end

data/lib/neighbor/utils.rb

@@ -0,0 +1,201 @@
+module Neighbor
+  module Utils
+    def self.validate_dimensions(value, type, expected, adapter)
+      dimensions = type == :sparsevec ? value.dimensions : value.size
+      dimensions *= 8 if type == :bit && [:sqlite, :mysql].include?(adapter)
+
+      if expected && dimensions != expected
+        "Expected #{expected} dimensions, not #{dimensions}"
+      end
+    end
+
+    def self.validate_finite(value, type)
+      case type
+      when :bit, :integer
+        true
+      when :sparsevec
+        value.values.all?(&:finite?)
+      else
+        value.all?(&:finite?)
+      end
+    end
+
+    def self.validate(value, dimensions:, type:, adapter:)
+      if (message = validate_dimensions(value, type, dimensions, adapter))
+        raise Error, message
+      end
+
+      if !validate_finite(value, type)
+        raise Error, "Values must be finite"
+      end
+    end
+
+    def self.normalize(value, column_info:)
+      return nil if value.nil?
+
+      raise Error, "Normalize not supported for type" unless [:cube, :vector, :halfvec, :binary].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
+
+    def self.array?(value)
+      !value.nil? && value.respond_to?(:to_a)
+    end
+
+    def self.adapter(model)
+      case model.connection_db_config.adapter
+      when /sqlite/i
+        :sqlite
+      when /mysql|trilogy/i
+        model.connection_pool.with_connection { |c| c.try(:mariadb?) } ? :mariadb : :mysql
+      else
+        :postgresql
+      end
+    end
+
+    def self.type(adapter, column_type)
+      case adapter
+      when :mysql
+        if column_type == :binary
+          :bit
+        else
+          column_type
+        end
+      else
+        column_type
+      end
+    end
+
+    def self.operator(adapter, column_type, distance)
+      case adapter
+      when :sqlite
+        case distance
+        when "euclidean"
+          "vec_distance_L2"
+        when "cosine"
+          "vec_distance_cosine"
+        when "taxicab"
+          "vec_distance_L1"
+        when "hamming"
+          "vec_distance_hamming"
+        end
+      when :mariadb
+        case column_type
+        when :binary
+          case distance
+          when "euclidean", "cosine"
+            "VEC_DISTANCE"
+          end
+        when :integer
+          case distance
+          when "hamming"
+            "BIT_COUNT"
+          end
+        else
+          raise ArgumentError, "Unsupported type: #{column_type}"
+        end
+      when :mysql
+        case column_type
+        when :vector
+          case distance
+          when "cosine"
+            "COSINE"
+          when "euclidean"
+            "EUCLIDEAN"
+          end
+        when :binary
+          case distance
+          when "hamming"
+            "BIT_COUNT"
+          end
+        else
+          raise ArgumentError, "Unsupported type: #{column_type}"
+        end
+      else
+        case column_type
+        when :bit
+          case distance
+          when "hamming"
+            "<~>"
+          when "jaccard"
+            "<%>"
+          when "hamming2"
+            "#"
+          end
+        when :vector, :halfvec, :sparsevec
+          case distance
+          when "inner_product"
+            "<#>"
+          when "cosine"
+            "<=>"
+          when "euclidean"
+            "<->"
+          when "taxicab"
+            "<+>"
+          end
+        when :cube
+          case distance
+          when "taxicab"
+            "<#>"
+          when "chebyshev"
+            "<=>"
+          when "euclidean", "cosine"
+            "<->"
+          end
+        else
+          raise ArgumentError, "Unsupported type: #{column_type}"
+        end
+      end
+    end
+
+    def self.order(adapter, type, operator, quoted_attribute, query)
+      case adapter
+      when :sqlite
+        case type
+        when :int8
+          "#{operator}(vec_int8(#{quoted_attribute}), vec_int8(#{query}))"
+        when :bit
+          "#{operator}(vec_bit(#{quoted_attribute}), vec_bit(#{query}))"
+        else
+          "#{operator}(#{quoted_attribute}, #{query})"
+        end
+      when :mariadb
+        if operator == "BIT_COUNT"
+          "BIT_COUNT(#{quoted_attribute} ^ #{query})"
+        else
+          "VEC_DISTANCE(#{quoted_attribute}, #{query})"
+        end
+      when :mysql
+        if operator == "BIT_COUNT"
+          "BIT_COUNT(#{quoted_attribute} ^ #{query})"
+        elsif operator == "COSINE"
+          "DISTANCE(#{quoted_attribute}, #{query}, 'COSINE')"
+        else
+          "DISTANCE(#{quoted_attribute}, #{query}, 'EUCLIDEAN')"
+        end
+      else
+        if operator == "#"
+          "bit_count(#{quoted_attribute} # #{query})"
+        else
+          "#{quoted_attribute} #{operator} #{query}"
+        end
+      end
+    end
+
+    def self.normalize_required?(adapter, column_type)
+      case adapter
+      when :postgresql
+        column_type == :cube
+      when :mariadb
+        true
+      else
+        false
+      end
+    end
+  end
+end

data/lib/neighbor/version.rb

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

data/lib/neighbor.rb

@@ -1,47 +1,35 @@
 # dependencies
 require "active_support"
 
+# adapter hooks
+require_relative "neighbor/mysql"
+require_relative "neighbor/postgresql"
+require_relative "neighbor/sqlite"
+
 # modules
+require_relative "neighbor/reranking"
+require_relative "neighbor/sparse_vector"
+require_relative "neighbor/utils"
 require_relative "neighbor/version"
 
 module Neighbor
   class Error < StandardError; end
-
-  module RegisterTypes
-    def initialize_type_map(m = type_map)
-      super
-      m.register_type "cube", Type::Cube.new
-      m.register_type "vector" do |_, _, sql_type|
-        limit = extract_limit(sql_type)
-        Type::Vector.new(limit: limit)
-      end
-    end
-  end
 end
 
 ActiveSupport.on_load(:active_record) do
+  require_relative "neighbor/attribute"
   require_relative "neighbor/model"
-  require_relative "neighbor/vector"
-  require_relative "neighbor/type/cube"
-  require_relative "neighbor/type/vector"
+  require_relative "neighbor/normalized_attribute"
 
   extend Neighbor::Model
 
-  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[:vector] = {name: "vector"}
-
-  # ensure schema can be loaded
-  ActiveRecord::ConnectionAdapters::TableDefinition.send(:define_column_methods, :cube, :vector)
-
-  # prevent unknown OID warning
-  if ActiveRecord::VERSION::MAJOR >= 7
-    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(Neighbor::RegisterTypes)
-  else
-    ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Neighbor::RegisterTypes)
+  begin
+    Neighbor::PostgreSQL.initialize!
+  rescue Gem::LoadError
+    # tries to load pg gem, which may not be available
   end
+
+  Neighbor::MySQL.initialize!
 end
 
 require_relative "neighbor/railtie" if defined?(Rails::Railtie)

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: neighbor
 version: !ruby/object:Gem::Version
-  version: 0.3.2
+  version: 0.5.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-12-12 00:00:00.000000000 Z
+date: 2024-10-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '7'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '7'
 description:
 email: andrew@ankane.org
 executables: []
@@ -34,15 +34,29 @@ files:
 - LICENSE.txt
 - README.md
 - lib/generators/neighbor/cube_generator.rb
+- lib/generators/neighbor/sqlite_generator.rb
 - lib/generators/neighbor/templates/cube.rb.tt
+- lib/generators/neighbor/templates/sqlite.rb.tt
 - lib/generators/neighbor/templates/vector.rb.tt
 - lib/generators/neighbor/vector_generator.rb
 - lib/neighbor.rb
+- lib/neighbor/attribute.rb
 - lib/neighbor/model.rb
+- lib/neighbor/mysql.rb
+- lib/neighbor/normalized_attribute.rb
+- lib/neighbor/postgresql.rb
 - lib/neighbor/railtie.rb
+- lib/neighbor/reranking.rb
+- lib/neighbor/sparse_vector.rb
+- lib/neighbor/sqlite.rb
 - lib/neighbor/type/cube.rb
+- lib/neighbor/type/halfvec.rb
+- lib/neighbor/type/mysql_vector.rb
+- lib/neighbor/type/sparsevec.rb
+- lib/neighbor/type/sqlite_int8_vector.rb
+- lib/neighbor/type/sqlite_vector.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,15 +70,15 @@ 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.16
 signing_key:
 specification_version: 4
-summary: Nearest neighbor search for Rails and Postgres
+summary: Nearest neighbor search for Rails
 test_files: []

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