neighbor-s3 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 3d4c40ed4dac5051282fef9538abb3f74c7dbaa25e2acf5cbef7fdb43b72f9c5
+  data.tar.gz: cfe0700757e352d2eeb27dd869c11562fa1a31d53a9c99270c84c14ec72c9e4f
+SHA512:
+  metadata.gz: 8fc63cd4fd5030bfa9500f3f664d99a9860a60131aaff5dbd02c9072c0d43009b06394ca7ba5f70b1835d361e27c1e9a5815a352dbdf0edb373615d630dfefc3
+  data.tar.gz: 77a5d9c7cb2ca0ade3f2728ad63e05d8f024cf21fc10f925126afece477933e244dd04baf6f14a799e57aa60acfe6da1940d2feb86578533d2d4efeff4fb210b

data/CHANGELOG.md

@@ -0,0 +1,3 @@
+## 0.1.0 (2025-10-02)
+
+- First release

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2025 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,199 @@
+# Neighbor S3
+
+Nearest neighbor search for Ruby and S3 Vectors
+
+## Installation
+
+Add this line to your application’s Gemfile:
+
+```ruby
+gem "neighbor-s3"
+```
+
+## Getting Started
+
+First, create a [vector bucket](https://console.aws.amazon.com/s3/vector-buckets)
+
+Create an index
+
+```ruby
+index = Neighbor::S3::Index.new("items", bucket: "my-bucket", dimensions: 3, distance: "cosine")
+index.create
+```
+
+Add vectors
+
+```ruby
+index.add(1, [1, 1, 1])
+index.add(2, [2, 2, 2])
+index.add(3, [1, 1, 2])
+```
+
+Search for nearest neighbors to a vector
+
+```ruby
+index.search([1, 1, 1], count: 5)
+```
+
+Search for nearest neighbors to a vector in the index
+
+```ruby
+index.search_id(1, count: 5)
+```
+
+IDs are treated as strings by default, but can also be treated as integers
+
+```ruby
+Neighbor::S3::Index.new("items", id_type: "integer", ...)
+```
+
+## Operations
+
+Add or update a vector
+
+```ruby
+index.add(id, vector)
+```
+
+Add or update multiple vectors
+
+```ruby
+index.add_all([{id: 1, vector: [1, 2, 3]}, {id: 2, vector: [4, 5, 6]}])
+```
+
+Get a vector
+
+```ruby
+index.find(id)
+```
+
+Get all vectors
+
+```ruby
+index.find_in_batches do |batch|
+  # ...
+end
+```
+
+Remove a vector
+
+```ruby
+index.remove(id)
+```
+
+Remove multiple vectors
+
+```ruby
+index.remove_all(ids)
+```
+
+## Metadata
+
+Add a vector with metadata
+
+```ruby
+index.add(id, vector, metadata: {category: "A"})
+```
+
+Add multiple vectors with metadata
+
+```ruby
+index.add_all([
+  {id: 1, vector: [1, 2, 3], metadata: {category: "A"}},
+  {id: 2, vector: [4, 5, 6], metadata: {category: "B"}}
+])
+```
+
+Get metadata with search results
+
+```ruby
+index.search(vector, with_metadata: true)
+```
+
+Filter by metadata
+
+```ruby
+index.search(vector, filter: {category: "A"})
+```
+
+Supports [these operators](https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors-metadata-filtering.html#s3-vectors-metadata-filtering-filterable)
+
+Specify non-filterable metadata on index creation
+
+```ruby
+Neighbor::S3::Index.new(name, non_filterable: ["category"], ...)
+```
+
+## Example
+
+You can use Neighbor S3 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::S3::Index.new("movies", bucket: "my-bucket", dimensions: 20, distance: "cosine")
+```
+
+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.map { |v| {id: v, vector: recommender.item_factors(v)} })
+```
+
+And get similar movies
+
+```ruby
+index.search_id("Star Wars (1977)").map { |v| v[:id] }
+```
+
+See the [complete code](examples/disco_item_recs.rb)
+
+## Reference
+
+Get index info
+
+```ruby
+index.info
+```
+
+Check if an index exists
+
+```ruby
+index.exists?
+```
+
+Drop an index
+
+```ruby
+index.drop
+```
+
+## History
+
+View the [changelog](https://github.com/ankane/neighbor-s3/blob/master/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-s3/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/neighbor-s3/pulls)
+- Write, clarify, or fix documentation
+- Suggest or add new features
+
+To get started with development:
+
+```sh
+git clone https://github.com/ankane/neighbor-s3.git
+cd neighbor-s3
+bundle install
+bundle exec rake test
+```

data/lib/neighbor/s3/index.rb

@@ -0,0 +1,237 @@
+module Neighbor
+  module S3
+    class Index
+      def initialize(name, bucket:, dimensions:, distance:, id_type: "string", non_filterable: nil)
+        @name = name
+        @bucket = bucket
+        @dimensions = dimensions.to_i
+
+        @distance_metric =
+          case distance.to_s
+          when "euclidean"
+            "euclidean"
+          when "cosine"
+            "cosine"
+          else
+            raise ArgumentError, "invalid distance"
+          end
+
+        @int_ids =
+          case id_type.to_s
+          when "string"
+            false
+          when "integer"
+            true
+          else
+            raise ArgumentError, "invalid id_type"
+          end
+
+        @non_filterable = non_filterable.to_a
+      end
+
+      def self.create(*args, **options)
+        index = new(*args, **options)
+        index.create
+        index
+      end
+
+      def create
+        options = {
+          vector_bucket_name: @bucket,
+          index_name: @name,
+          data_type: "float32",
+          dimension: @dimensions,
+          distance_metric: @distance_metric
+        }
+        if @non_filterable.any?
+          options[:metadata_configuration] = {
+            non_filterable_metadata_keys: @non_filterable
+          }
+        end
+        client.create_index(options)
+        nil
+      end
+
+      def exists?
+        client.get_index({
+          vector_bucket_name: @bucket,
+          index_name: @name
+        })
+        true
+      rescue Aws::S3Vectors::Errors::NotFoundException
+        false
+      end
+
+      def info
+        client.get_index({
+          vector_bucket_name: @bucket,
+          index_name: @name
+        }).index.to_h
+      end
+
+      def add(id, vector, metadata: nil)
+        add_all([{id: id, vector: vector, metadata: metadata}])
+      end
+
+      def add_all(items)
+        # perform checks first to reduce chance of non-atomic updates
+        vectors =
+          items.map do |item|
+            vector = item.fetch(:vector).to_a
+            check_dimensions(vector)
+
+            {
+              key: item_id(item.fetch(:id)).to_s,
+              data: {float32: vector},
+              metadata: item[:metadata]
+            }
+          end
+
+        vectors.each_slice(500) do |batch|
+          client.put_vectors({
+            vector_bucket_name: @bucket,
+            index_name: @name,
+            vectors: batch
+          })
+        end
+        nil
+      end
+
+      def member?(id)
+        id = item_id(id)
+
+        client.get_vectors({
+          vector_bucket_name: @bucket,
+          index_name: @name,
+          keys: [id.to_s],
+          return_data: false,
+          return_metadata: false
+        }).vectors.any?
+      end
+      alias_method :include?, :member?
+
+      def remove(id)
+        remove_all([id])
+      end
+
+      def remove_all(ids)
+        ids = ids.to_a.map { |id| item_id(id) }
+
+        ids.each_slice(500) do |batch|
+          client.delete_vectors({
+            vector_bucket_name: @bucket,
+            index_name: @name,
+            keys: batch.map(&:to_s)
+          })
+        end
+        nil
+      end
+
+      def find(id, with_metadata: true)
+        id = item_id(id)
+
+        v =
+          client.get_vectors({
+            vector_bucket_name: @bucket,
+            index_name: @name,
+            keys: [id.to_s],
+            return_data: true,
+            return_metadata: with_metadata
+          }).vectors.first
+
+        if v
+          item = {
+            id: item_id(v.key),
+            vector: v.data.float32
+          }
+          item[:metadata] = v.metadata if with_metadata
+          item
+        end
+      end
+
+      def find_in_batches(batch_size: 1000, with_metadata: true)
+        options = {
+          vector_bucket_name: @bucket,
+          index_name: @name,
+          max_results: batch_size,
+          return_data: true,
+          return_metadata: with_metadata
+        }
+
+        begin
+          resp = client.list_vectors(options)
+          batch =
+            resp.vectors.map do |v|
+              item = {
+                id: item_id(v.key),
+                vector: v.data.float32
+              }
+              item[:metadata] = v.metadata if with_metadata
+              item
+            end
+          yield batch
+          options[:next_token] = resp.next_token
+        end while resp.next_token
+      end
+
+      def search(vector, count: 5, with_metadata: false, filter: nil)
+        check_dimensions(vector)
+
+        client.query_vectors({
+          vector_bucket_name: @bucket,
+          index_name: @name,
+          top_k: count,
+          query_vector: {
+            float32: vector,
+          },
+          filter: filter,
+          return_metadata: with_metadata,
+          return_distance: true
+        }).vectors.map do |v|
+          item = {
+            id: item_id(v.key),
+            distance: @distance_metric == "euclidean" ? Math.sqrt(v.distance) : v.distance
+          }
+          item[:metadata] = v.metadata if with_metadata
+          item
+        end
+      end
+
+      def search_id(id, count: 5, with_metadata: false, filter: nil)
+        id = item_id(id)
+
+        item = find(id)
+        unless item
+          raise Error, "Could not find item #{id}"
+        end
+
+        result = search(item[:vector], count: count + 1, with_metadata:, filter:)
+        result.reject { |v| v[:id] == id }.first(count)
+      end
+
+      def drop
+        client.delete_index({
+          vector_bucket_name: @bucket,
+          index_name: @name
+        })
+        nil
+      end
+
+      private
+
+      def check_dimensions(vector)
+        if vector.size != @dimensions
+          raise ArgumentError, "expected #{@dimensions} dimensions"
+        end
+      end
+
+      def item_id(id)
+        @int_ids ? Integer(id) : id.to_s
+      end
+
+      def client
+        S3.client
+      end
+    end
+  end
+end

data/lib/neighbor/s3/version.rb

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

data/lib/neighbor/s3.rb

@@ -0,0 +1,20 @@
+# dependencies
+require "aws-sdk-s3vectors"
+
+# modules
+require_relative "s3/index"
+require_relative "s3/version"
+
+module Neighbor
+  module S3
+    class Error < StandardError; end
+
+    class << self
+      attr_writer :client
+
+      def client
+        @client ||= Aws::S3Vectors::Client.new
+      end
+    end
+  end
+end

data/lib/neighbor-s3.rb

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

metadata

@@ -0,0 +1,59 @@
+--- !ruby/object:Gem::Specification
+name: neighbor-s3
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Andrew Kane
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-s3vectors
+  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'
+email: andrew@ankane.org
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- lib/neighbor-s3.rb
+- lib/neighbor/s3.rb
+- lib/neighbor/s3/index.rb
+- lib/neighbor/s3/version.rb
+homepage: https://github.com/ankane/neighbor-s3
+licenses:
+- MIT
+metadata: {}
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.2'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.9
+specification_version: 4
+summary: Nearest neighbor search for Ruby and S3 Vectors
+test_files: []