neighbor-redis 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: ca86068ccb3e4538661582c5cf25f58b522569ff2abc6fb69cb6c69f3171141e
+  data.tar.gz: 455e7ce9a5c112b3ffb2e6f5004f4cc0598033b1b1ca409fdec7fdecdd936c86
+SHA512:
+  metadata.gz: 4ff369d81326391017aa7c8c98966d44200856f7fb7be32caa8dd773f55ab01a15674ed5bb2a58aa8bf8163604d26895516fcd4c29332caf5354dc101d5d7f54
+  data.tar.gz: 5d00a85e5265d01617299daa1b0ce5a3e67f4d570c6a512eee5bc6b4023d5dedfc93a8e8583689a54fbb1cafae8401cb76acc35242b6d215cd1883aa34add26f

data/CHANGELOG.md

@@ -0,0 +1,3 @@
+## 0.1.0 (2023-02-21)
+
+- First release

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Andrew Kane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,196 @@
+# Neighbor Redis
+
+Nearest neighbor search for Ruby and Redis
+
+[![Build Status](https://github.com/ankane/neighbor-redis/workflows/build/badge.svg?branch=master)](https://github.com/ankane/neighbor-redis/actions)
+
+## Installation
+
+First, [install RediSearch](https://redis.io/docs/stack/search/quick_start/). With Docker, use:
+
+```sh
+docker run -p 6379:6379 redis/redis-stack-server:latest
+```
+
+Add this line to your application’s Gemfile:
+
+```ruby
+gem "neighbor-redis"
+```
+
+And set the Redis client:
+
+```ruby
+Neighbor::Redis.client = RedisClient.config.new_pool
+```
+
+## Getting Started
+
+Create an index
+
+```ruby
+index = Neighbor::Redis::HNSWIndex.new("items", dimensions: 3, distance: "l2")
+index.create
+```
+
+Add items
+
+```ruby
+index.add(1, [1, 1, 1])
+index.add(2, [2, 2, 2])
+index.add(3, [1, 1, 2])
+```
+
+Note: IDs are stored and returned as strings (uses less total storage)
+
+Get the nearest neighbors to an item
+
+```ruby
+index.nearest(1, count: 5)
+```
+
+Get the nearest neighbors to a vector
+
+```ruby
+index.search([1, 1, 1], count: 5)
+```
+
+## Distance
+
+Supported values are:
+
+- `l2`
+- `inner_product`
+- `cosine`
+
+## Index Types
+
+Hierarchical Navigable Small World (HNSW)
+
+```ruby
+Neighbor::Redis::HNSWIndex.new(
+  name,
+  initial_cap: nil,
+  m: 16,
+  ef_construction: 200,
+  ef_runtime: 10,
+  epsilon: 0.01
+)
+```
+
+Flat
+
+```ruby
+Neighbor::Redis::FlatIndex.new(
+  name,
+  initial_cap: nil,
+  block_size: 1024
+)
+```
+
+## Additional Options
+
+Store vectors as double precision (instead of single precision)
+
+```ruby
+Neighbor::Redis::HNSWIndex.new(name, type: "float64")
+```
+
+Store vectors as JSON (instead of a hash/blob)
+
+```ruby
+Neighbor::Redis::HNSWIndex.new(name, redis_type: "json")
+```
+
+## Changing Options
+
+Create a new index to change any index options
+
+```ruby
+Neighbor::Redis::HNSWIndex.new("items-v2", **new_options)
+```
+
+## Additional Operations
+
+Add multiple items
+
+```ruby
+index.add_all(ids, embeddings)
+```
+
+Remove an item
+
+```ruby
+index.remove(id)
+```
+
+Remove multiple items
+
+```ruby
+index.remove_all(ids)
+```
+
+Drop the index
+
+```ruby
+index.drop
+```
+
+## Example
+
+You can use Neighbor Redis for online item-based recommendations with [Disco](https://github.com/ankane/disco). We’ll use MovieLens data for this example.
+
+Create an index
+
+```ruby
+index = Neighbor::Redis::HNSWIndex.new("movies", dimensions: 20, distance: "cosine")
+index.create
+```
+
+Fit the recommender
+
+```ruby
+data = Disco.load_movielens
+recommender = Disco::Recommender.new(factors: 20)
+recommender.fit(data)
+```
+
+Store the item factors
+
+```ruby
+index.add_all(recommender.item_ids, recommender.item_factors)
+```
+
+And get similar movies
+
+```ruby
+index.nearest("Star Wars (1977)").map { |v| v[:id] }
+```
+
+See the [complete code](examples/disco_item_recs.rb)
+
+## Reference
+
+- [Vector similarity](https://redis.io/docs/stack/search/reference/vectors/)
+
+## History
+
+View the [changelog](CHANGELOG.md)
+
+## Contributing
+
+Everyone is encouraged to help improve this project. Here are a few ways you can help:
+
+- [Report bugs](https://github.com/ankane/neighbor-redis/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/neighbor-redis/pulls)
+- Write, clarify, or fix documentation
+- Suggest or add new features
+
+To get started with development:
+
+```sh
+git clone https://github.com/ankane/neighbor-redis.git
+cd neighbor-redis
+bundle install
+bundle exec rake test
+```

data/lib/neighbor/redis/flat_index.rb

@@ -0,0 +1,21 @@
+module Neighbor
+  module Redis
+    class FlatIndex < Index
+      def initialize(*args, initial_cap: nil, block_size: nil, **options)
+        super(*args, **options)
+        @algorithm = "FLAT"
+        @initial_cap = initial_cap
+        @block_size = block_size
+      end
+
+      private
+
+      def create_params
+        params = {}
+        params["INITIAL_CAP"] = @initial_cap if @initial_cap
+        params["BLOCK_SIZE"] = @block_size if @block_size
+        params
+      end
+    end
+  end
+end

data/lib/neighbor/redis/hnsw_index.rb

@@ -0,0 +1,27 @@
+module Neighbor
+  module Redis
+    class HNSWIndex < Index
+      def initialize(*args, initial_cap: nil, m: nil, ef_construction: nil, ef_runtime: nil, epsilon: nil, **options)
+        super(*args, **options)
+        @algorithm = "HNSW"
+        @initial_cap = initial_cap
+        @m = m
+        @ef_construction = ef_construction
+        @ef_runtime = ef_runtime
+        @epsilon = epsilon
+      end
+
+      private
+
+      def create_params
+        params = {}
+        params["INITIAL_CAP"] = @initial_cap if @initial_cap
+        params["M"] = @m if @m
+        params["EF_CONSTRUCTION"] = @ef_construction if @ef_construction
+        params["EF_RUNTIME"] = @ef_runtime if @ef_runtime
+        params["EPSILON"] = @epsilon if @epsilon
+        params
+      end
+    end
+  end
+end

data/lib/neighbor/redis/index.rb

@@ -0,0 +1,236 @@
+module Neighbor
+  module Redis
+    class Index
+      def initialize(name, dimensions:, distance:, type: "float32", redis_type: "hash")
+        @index_name = index_name(name)
+        @global_prefix = "neighbor:items:"
+        @prefix = "#{@global_prefix}#{name}:"
+
+        @dimensions = dimensions
+
+        unless distance.nil?
+          @distance_metric =
+            case distance.to_s
+            when "l2", "cosine"
+              distance.to_s.upcase
+            when "inner_product"
+              "IP"
+            else
+              raise ArgumentError, "invalid distance"
+            end
+        end
+
+        @float64 =
+          case type.to_s
+          when "float32"
+            false
+          when "float64"
+            true
+          else
+            raise ArgumentError, "invalid type"
+          end
+
+        @json =
+          case redis_type.to_s
+          when "hash"
+            false
+          when "json"
+            require "json"
+            true
+          else
+            raise ArgumentError, "invalid redis_type"
+          end
+      end
+
+      def self.create(...)
+        index = new(...)
+        index.create
+        index
+      end
+
+      def create
+        params = {
+          "TYPE" => @float64 ? "FLOAT64" : "FLOAT32",
+          "DIM" => @dimensions,
+          "DISTANCE_METRIC" => @distance_metric,
+        }.merge(create_params)
+
+        command = ["FT.CREATE", @index_name]
+        command.push("ON", "JSON") if @json
+        command.push("PREFIX", "1", @prefix, "SCHEMA")
+        command.push("$.v", "AS") if @json
+        command.push("v", "VECTOR", @algorithm, params.size * 2, params)
+        ft_command { redis.call(*command) }
+      end
+
+      def exists?
+        redis.call("FT.INFO", @index_name)
+        true
+      rescue ArgumentError
+        # fix for invalid value for Float(): "-nan"
+        true
+      rescue => e
+        raise unless e.message.include?("Unknown Index name")
+        false
+      end
+
+      def add(id, embedding)
+        add_all([id], [embedding])
+      end
+
+      def add_all(ids, embeddings)
+        ids = ids.to_a
+        embeddings = embeddings.to_a
+
+        raise ArgumentError, "different sizes" if ids.size != embeddings.size
+
+        embeddings.each { |e| check_dimensions(e) }
+
+        redis.pipelined do |pipeline|
+          ids.zip(embeddings).each do |id, embedding|
+            if @json
+              pipeline.call("JSON.SET", item_key(id), "$", JSON.generate({v: embedding}))
+            else
+              pipeline.call("HSET", item_key(id), {v: to_binary(embedding)})
+            end
+          end
+        end
+      end
+
+      def remove(id)
+        remove_all([id])
+      end
+
+      def remove_all(ids)
+        redis.call("DEL", ids.map { |id| item_key(id) })
+      end
+
+      def search(embedding, count: 5)
+        check_dimensions(embedding)
+
+        search_by_blob(to_binary(embedding), count)
+      end
+
+      def find(id)
+        if @json
+          s = redis.call("JSON.GET", item_key(id), "$.v")
+          JSON.parse(s)[0] if s
+        else
+          from_binary(redis.call("HGET", item_key(id), "v"))
+        end
+      end
+
+      def nearest(id, count: 5)
+        embedding =
+          if @json
+            s = redis.call("JSON.GET", item_key(id), "$.v")
+            to_binary(JSON.parse(s)[0]) if s
+          else
+            redis.call("HGET", item_key(id), "v")
+          end
+
+        unless embedding
+          raise Error, "Could not find item #{id}"
+        end
+
+        search_by_blob(embedding, count + 1).reject { |v| v[:id] == id.to_s }.first(count)
+      end
+
+      def drop
+        drop_index
+        drop_keys
+      end
+
+      def promote(alias_name)
+        redis.call("FT.ALIASUPDATE", index_name(alias_name), @index_name)
+      end
+
+      private
+
+      def index_name(name)
+        if name.include?(":")
+          raise ArgumentError, "Invalid name"
+        end
+
+        "neighbor-idx-#{name}"
+      end
+
+      def check_dimensions(embedding)
+        if embedding.size != @dimensions
+          raise ArgumentError, "expected #{@dimensions} dimensions"
+        end
+      end
+
+      def item_key(id)
+        "#{@prefix}#{id}"
+      end
+
+      def search_by_blob(blob, count)
+        resp = redis.call("FT.SEARCH", @index_name, "*=>[KNN #{count.to_i} @v $BLOB]", "PARAMS", "2", "BLOB", blob, "SORTBY", "__v_score", "DIALECT", "2")
+        len = resp.shift
+        prefix_length = nil
+        len.times.map do |i|
+          key, info = resp.shift(2)
+          info = info.each_slice(2).to_h
+          score = info["__v_score"].to_f
+          distance =
+            case @distance_metric
+            when "L2"
+              Math.sqrt(score)
+            when "IP"
+              (score * -1) + 1
+            else
+              score
+            end
+
+          prefix_length ||= find_prefix_length(key)
+
+          {
+            id: key[prefix_length..-1],
+            distance: distance
+          }
+        end
+      end
+
+      def find_prefix_length(key)
+        key[@global_prefix.length..-1].index(":") + @global_prefix.length + 1
+      end
+
+      def drop_index
+        redis.call("FT.DROPINDEX", @index_name)
+      end
+
+      def drop_keys
+        cursor = 0
+        begin
+          cursor, keys = redis.call("SCAN", cursor, "MATCH", "#{@prefix}*", "COUNT", 100)
+          redis.call("DEL", keys) if keys.any?
+        end while cursor != "0"
+      end
+
+      def to_binary(embedding)
+        embedding.to_a.pack(pack_format)
+      end
+
+      def from_binary(s)
+        s.unpack(pack_format)
+      end
+
+      def pack_format
+        @pack_format ||= @float64 ? "d#{@dimensions}" : "f#{@dimensions}"
+      end
+
+      # just use for create for now
+      def ft_command
+        yield
+      rescue => e
+        raise Error, "RediSearch not installed" if e.message.include?("ERR unknown command 'FT.")
+        raise
+      end
+
+      def redis
+        Redis.client
+      end
+    end
+  end
+end

data/lib/neighbor/redis/version.rb

@@ -0,0 +1,5 @@
+module Neighbor
+  module Redis
+    VERSION = "0.1.0"
+  end
+end

data/lib/neighbor/redis.rb

@@ -0,0 +1,18 @@
+# dependencies
+require "redis-client"
+
+# modules
+require_relative "redis/index"
+require_relative "redis/flat_index"
+require_relative "redis/hnsw_index"
+require_relative "redis/version"
+
+module Neighbor
+  module Redis
+    class Error < StandardError; end
+
+    class << self
+      attr_accessor :client
+    end
+  end
+end

data/lib/neighbor-redis.rb

@@ -0,0 +1 @@
+require_relative "neighbor/redis"

metadata

@@ -0,0 +1,65 @@
+--- !ruby/object:Gem::Specification
+name: neighbor-redis
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Andrew Kane
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2023-02-21 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: redis-client
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description:
+email: andrew@ankane.org
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- lib/neighbor-redis.rb
+- lib/neighbor/redis.rb
+- lib/neighbor/redis/flat_index.rb
+- lib/neighbor/redis/hnsw_index.rb
+- lib/neighbor/redis/index.rb
+- lib/neighbor/redis/version.rb
+homepage: https://github.com/ankane/neighbor-redis
+licenses:
+- MIT
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '2.7'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.4.6
+signing_key:
+specification_version: 4
+summary: Nearest neighbor search for Ruby and Redis
+test_files: []