neighbor 0.4.3 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e8d611fd277cd48d309b2a087fdeb22f39f43d8ef81fcab57763bd5b4b2e48b3
-  data.tar.gz: fe0a5f7e4aa1ebd81f8c5849be67dc0f3c948d53d313b80259def04b9b7e9e84
+  metadata.gz: 5d7036a69b1c57161eaeb11e38feee92c1d5082ddbe8907a83ac3126adf9ae56
+  data.tar.gz: c88e4400b75d2a87f766f7e0b7ff6c5311d9c7866d32e7c2b9aa7601f60c474f
 SHA512:
-  metadata.gz: 0d9d9d0be9f2929f1eab7e5df52a0606ca8c220ea74e6ef349c2ab77e5da44bd4b98ed4d2271345b2d679882f1790ac54b8373c13a69f75607935322bdb68754
-  data.tar.gz: 834c7d6e26be9b6d8fc262280048e5104cd022587fed30d030f6c0edd849966a49a08bc793f1b49764b1fc5afc014c1463dd8acacfb5739bc507f4a77d3281a1
+  metadata.gz: d3c4c25404fb64f324fbba70edcf06d827d3708905ed4a84404a6c9ce39f27b6890d449b285ce302e24495a666c34f0bf3050270b54ef8d283f53ebeb19e4e91
+  data.tar.gz: 63927a8801a88edd48f74ce85d056d7112fa526b37d473b232256b5f2d47e5254b34b25e0312afc01bfedcd6a9d7826496f208e0ab0d3f57c3307fa298b8984e

data/CHANGELOG.md

@@ -1,3 +1,12 @@
+## 0.5.0 (2024-10-07)
+
+- Added experimental support for SQLite (sqlite-vec)
+- Added experimental support for MariaDB 11.6 Vector
+- Added experimental support for MySQL 9
+- Changed `normalize` option to use Active Record normalization
+- Fixed connection leasing for Active Record 7.2
+- Dropped support for Active Record < 7
+
 ## 0.4.3 (2024-09-02)
 
 - Added `rrf` method

data/README.md

@@ -1,6 +1,13 @@
 # Neighbor
 
-Nearest neighbor search for Rails and Postgres
+Nearest neighbor search for Rails
+
+Supports:
+
+- Postgres (cube and pgvector)
+- SQLite (sqlite-vec) - experimental
+- MariaDB 11.6 Vector - experimental
+- MySQL 9 (searching requires HeatWave) - experimental
 
 [![Build Status](https://github.com/ankane/neighbor/actions/workflows/build.yml/badge.svg)](https://github.com/ankane/neighbor/actions)
 
@@ -12,7 +19,7 @@ Add this line to your application’s Gemfile:
 gem "neighbor"
 ```
 
-## Choose An Extension
+### For Postgres
 
 Neighbor supports two extensions: [cube](https://www.postgresql.org/docs/current/cube.html) and [pgvector](https://github.com/pgvector/pgvector). cube ships with Postgres, while pgvector supports more dimensions and approximate nearest neighbor search.
 
@@ -30,6 +37,20 @@ rails generate neighbor:vector
 rails db:migrate
 ```
 
+### For SQLite
+
+Add this line to your application’s Gemfile:
+
+```ruby
+gem "sqlite-vec"
+```
+
+And run:
+
+```sh
+rails generate neighbor:sqlite
+```
+
 ## Getting Started
 
 Create a migration
@@ -37,9 +58,14 @@ Create a migration
 ```ruby
 class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
   def change
+    # cube
     add_column :items, :embedding, :cube
-    # or
+
+    # pgvector and MySQL
     add_column :items, :embedding, :vector, limit: 3 # dimensions
+
+    # sqlite-vec and MariaDB
+    add_column :items, :embedding, :binary
   end
 end
 ```
@@ -81,6 +107,9 @@ See the additional docs for:
 
 - [cube](#cube)
 - [pgvector](#pgvector)
+- [sqlite-vec](#sqlite-vec)
+- [MariaDB](#mariadb)
+- [MySQL](#mysql)
 
 Or check out some [examples](#examples)
 
@@ -134,6 +163,12 @@ Supported values are:
 
 The `vector` type can have up to 16,000 dimensions, and vectors with up to 2,000 dimensions can be indexed.
 
+The `halfvec` type can have up to 16,000 dimensions, and half vectors with up to 4,000 dimensions can be indexed.
+
+The `bit` type can have up to 83 million dimensions, and bit vectors with up to 64,000 dimensions can be indexed.
+
+The `sparsevec` type can have up to 16,000 non-zero elements, and sparse vectors with up to 1,000 non-zero elements can be indexed.
+
 ### Indexing
 
 Add an approximate index to speed up queries. Create a migration with:
@@ -241,6 +276,190 @@ embedding = Neighbor::SparseVector.new({0 => 0.9, 1 => 1.3, 2 => 1.1}, 3)
 Item.nearest_neighbors(:embedding, embedding, distance: "euclidean").first(5)
 ```
 
+## sqlite-vec
+
+### Distance
+
+Supported values are:
+
+- `euclidean`
+- `cosine`
+- `taxicab`
+- `hamming`
+
+### Dimensions
+
+For sqlite-vec, it’s a good idea to specify the number of dimensions to ensure all records have the same number.
+
+```ruby
+class Item < ApplicationRecord
+  has_neighbors :embedding, dimensions: 3
+end
+```
+
+### Virtual Tables
+
+You can also use [virtual tables](https://alexgarcia.xyz/sqlite-vec/features/knn.html)
+
+```ruby
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    # Rails < 8
+    execute <<~SQL
+      CREATE VIRTUAL TABLE items USING vec0(
+        embedding float[3] distance_metric=L2
+      )
+    SQL
+
+    # Rails 8+
+    create_virtual_table :items, :vec0, [
+      "embedding float[3] distance_metric=L2"
+    ]
+  end
+end
+```
+
+Use `distance_metric=cosine` for cosine distance
+
+You can optionally ignore any shadow tables that are created
+
+```ruby
+ActiveRecord::SchemaDumper.ignore_tables += [
+  "items_chunks", "items_rowids", "items_vector_chunks00"
+]
+```
+
+Create a model with `rowid` as the primary key
+
+```ruby
+class Item < ApplicationRecord
+  self.primary_key = "rowid"
+
+  has_neighbors :embedding, dimensions: 3
+end
+```
+
+Get the `k` nearest neighbors
+
+```ruby
+Item.where("embedding MATCH ?", [1, 2, 3].to_s).where(k: 5).order(:distance)
+```
+
+Filter by primary key
+
+```ruby
+Item.where(rowid: [2, 3]).where("embedding MATCH ?", [1, 2, 3].to_s).where(k: 5).order(:distance)
+```
+
+### Int8 Vectors
+
+Use the `type` option for int8 vectors
+
+```ruby
+class Item < ApplicationRecord
+  has_neighbors :embedding, dimensions: 3, type: :int8
+end
+```
+
+### Binary Vectors
+
+Use the `type` option for binary vectors
+
+```ruby
+class Item < ApplicationRecord
+  has_neighbors :embedding, dimensions: 8, type: :bit
+end
+```
+
+Get the nearest neighbors by Hamming distance
+
+```ruby
+Item.nearest_neighbors(:embedding, "\x05", distance: "hamming").first(5)
+```
+
+## MariaDB
+
+### Distance
+
+Supported values are:
+
+- `euclidean`
+- `cosine`
+- `hamming`
+
+For cosine distance with MariaDB, vectors must be normalized before being stored.
+
+```ruby
+class Item < ApplicationRecord
+  has_neighbors :embedding, normalize: true
+end
+```
+
+### Indexing
+
+Vector columns must use `null: false` to add a vector index
+
+```ruby
+class CreateItems < ActiveRecord::Migration[7.2]
+  def change
+    create_table :items do |t|
+      t.binary :embedding, null: false
+      t.index :embedding, type: :vector
+    end
+  end
+end
+```
+
+### Binary Vectors
+
+Use the `bigint` type to store binary vectors
+
+```ruby
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    add_column :items, :embedding, :bigint
+  end
+end
+```
+
+Note: Binary vectors can have up to 64 dimensions
+
+Get the nearest neighbors by Hamming distance
+
+```ruby
+Item.nearest_neighbors(:embedding, 5, distance: "hamming").first(5)
+```
+
+## MySQL
+
+### Distance
+
+Supported values are:
+
+- `euclidean`
+- `cosine`
+- `hamming`
+
+Note: The `DISTANCE()` function is [only available on HeatWave](https://dev.mysql.com/doc/refman/9.0/en/vector-functions.html)
+
+### Binary Vectors
+
+Use the `binary` type to store binary vectors
+
+```ruby
+class AddEmbeddingToItems < ActiveRecord::Migration[7.2]
+  def change
+    add_column :items, :embedding, :binary
+  end
+end
+```
+
+Get the nearest neighbors by Hamming distance
+
+```ruby
+Item.nearest_neighbors(:embedding, "\x05", distance: "hamming").first(5)
+```
+
 ## Examples
 
 - [Embeddings](#openai-embeddings) with OpenAI
@@ -472,12 +691,9 @@ end
 Create some documents
 
 ```ruby
-texts = [
-  "The dog is barking",
-  "The cat is purring",
-  "The bear is growling"
-]
-documents = Document.create!(texts.map { |v| {content: v} })
+Document.create!(content: "The dog is barking")
+Document.create!(content: "The cat is purring")
+Document.create!(content: "The bear is growling")
 ```
 
 Generate an embedding for each document
@@ -485,9 +701,9 @@ Generate an embedding for each document
 ```ruby
 embed = Informers.pipeline("embedding", "Snowflake/snowflake-arctic-embed-m-v1.5")
 embed_options = {model_output: "sentence_embedding", pooling: "none"} # specific to embedding model
-embeddings = embed.(documents.map(&:content), **embed_options)
 
-documents.zip(embeddings) do |document, embedding|
+Document.find_each do |document|
+  embedding = embed.(document.content, **embed_options)
   document.update!(embedding: embedding)
 end
 ```
@@ -511,7 +727,7 @@ semantic_results =
 To combine the results, use Reciprocal Rank Fusion (RRF)
 
 ```ruby
-Neighbor::Reranking.rrf(keyword_results, semantic_results)
+Neighbor::Reranking.rrf(keyword_results, semantic_results).first(5)
 ```
 
 Or a reranking model
@@ -519,7 +735,7 @@ Or a reranking model
 ```ruby
 rerank = Informers.pipeline("reranking", "mixedbread-ai/mxbai-rerank-xsmall-v1")
 results = (keyword_results + semantic_results).uniq
-rerank.(query, results.map(&:content), top_k: 5).map { |v| results[v[:doc_id]] }
+rerank.(query, results.map(&:content)).first(5).map { |v| results[v[:doc_id]] }
 ```
 
 See the [complete code](examples/hybrid/example.rb)
@@ -667,6 +883,19 @@ To get started with development:
 git clone https://github.com/ankane/neighbor.git
 cd neighbor
 bundle install
+
+# Postgres
 createdb neighbor_test
-bundle exec rake test
+bundle exec rake test:postgresql
+
+# SQLite
+bundle exec rake test:sqlite
+
+# MariaDB
+docker run -e MARIADB_ALLOW_EMPTY_ROOT_PASSWORD=1 -e MARIADB_DATABASE=neighbor_test -p 3307:3306 quay.io/mariadb-foundation/mariadb-devel:11.6-vector-preview
+bundle exec rake test:mariadb
+
+# MySQL
+docker run -e MYSQL_ALLOW_EMPTY_PASSWORD=1 -e MYSQL_DATABASE=neighbor_test -p 3306:3306 mysql:9
+bundle exec rake test:mysql
 ```

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/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,6 +1,6 @@
 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?
         raise ArgumentError, "has_neighbors requires an attribute name"
       end
@@ -24,125 +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
+
+        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
 
         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] || column_info&.limit
+            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, column_info&.type, dimensions).nil?
+            if !Neighbor::Utils.validate_dimensions(value, type, dimensions, adapter).nil?
               errors.add(k, "must have #{dimensions} dimensions")
             end
-            if !Neighbor::Utils.validate_finite(value, column_info&.type)
+            if !Neighbor::Utils.validate_finite(value, 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?
-
+        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
 
-          operator =
-            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
+          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_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
 
           column_attribute = klass.type_for_attribute(attribute_name)
           vector = column_attribute.cast(vector)
-          Neighbor::Utils.validate(vector, dimensions: dimensions, column_info: column_info)
+          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
 
-          query = connection.quote(column_attribute.serialize(vector))
+          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
 
           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.quote(cast_dimensions.to_i)})"
+              quoted_attribute += "::halfvec(#{connection_pool.with_connection { |c| c.quote(cast_dimensions.to_i) }})"
             else
               raise ArgumentError, "Invalid precision"
             end
           end
 
-          order = "#{quoted_attribute} #{operator} #{query}"
-          if operator == "#"
-            order = "bit_count(#{order})"
-          end
+          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:
@@ -150,7 +141,7 @@ module Neighbor
           # cosine distance = 1 - cosine similarity
           # this transformation doesn't change the order, so only needed for select
           neighbor_distance =
-            if column_type == :cube && distance == "cosine"
+            if distance == "cosine" && normalize_required
               "POWER(#{order}, 2) / 2.0"
             elsif [:vector, :halfvec, :sparsevec].include?(column_type) && distance == "inner_product"
               "(#{order}) * -1"

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

@@ -1,7 +1,9 @@
 module Neighbor
   module Utils
-    def self.validate_dimensions(value, type, expected)
+    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
@@ -9,7 +11,7 @@ module Neighbor
 
     def self.validate_finite(value, type)
       case type
-      when :bit
+      when :bit, :integer
         true
       when :sparsevec
         value.values.all?(&:finite?)
@@ -18,18 +20,20 @@ module Neighbor
       end
     end
 
-    def self.validate(value, dimensions:, column_info:)
-      if (message = validate_dimensions(value, column_info&.type, dimensions || column_info&.limit))
+    def self.validate(value, dimensions:, type:, adapter:)
+      if (message = validate_dimensions(value, type, dimensions, adapter))
         raise Error, message
       end
 
-      if !validate_finite(value, column_info&.type)
+      if !validate_finite(value, 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)
+      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 })
 
@@ -42,5 +46,156 @@ module Neighbor
     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.4.3"
+  VERSION = "0.5.0"
 end

data/lib/neighbor.rb

@@ -1,6 +1,11 @@
 # 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"
@@ -9,53 +14,22 @@ 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 "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
 
 ActiveSupport.on_load(:active_record) do
+  require_relative "neighbor/attribute"
   require_relative "neighbor/model"
-  require_relative "neighbor/type/cube"
-  require_relative "neighbor/type/halfvec"
-  require_relative "neighbor/type/sparsevec"
-  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[: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
-  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.4.3
+  version: 0.5.0
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-09-02 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,17 +34,27 @@ 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/utils.rb
 - lib/neighbor/version.rb
@@ -67,8 +77,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.11
+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: []