hnswlib 0.1.0

data/ext/hnswlib/src/visited_list_pool.h

@@ -0,0 +1,78 @@
+#pragma once
+
+#include <mutex>
+#include <string.h>
+
+namespace hnswlib {
+    typedef unsigned short int vl_type;
+
+    class VisitedList {
+    public:
+        vl_type curV;
+        vl_type *mass;
+        unsigned int numelements;
+
+        VisitedList(int numelements1) {
+            curV = -1;
+            numelements = numelements1;
+            mass = new vl_type[numelements];
+        }
+
+        void reset() {
+            curV++;
+            if (curV == 0) {
+                memset(mass, 0, sizeof(vl_type) * numelements);
+                curV++;
+            }
+        };
+
+        ~VisitedList() { delete[] mass; }
+    };
+///////////////////////////////////////////////////////////
+//
+// Class for multi-threaded pool-management of VisitedLists
+//
+/////////////////////////////////////////////////////////
+
+    class VisitedListPool {
+        std::deque<VisitedList *> pool;
+        std::mutex poolguard;
+        int numelements;
+
+    public:
+        VisitedListPool(int initmaxpools, int numelements1) {
+            numelements = numelements1;
+            for (int i = 0; i < initmaxpools; i++)
+                pool.push_front(new VisitedList(numelements));
+        }
+
+        VisitedList *getFreeVisitedList() {
+            VisitedList *rez;
+            {
+                std::unique_lock <std::mutex> lock(poolguard);
+                if (pool.size() > 0) {
+                    rez = pool.front();
+                    pool.pop_front();
+                } else {
+                    rez = new VisitedList(numelements);
+                }
+            }
+            rez->reset();
+            return rez;
+        };
+
+        void releaseVisitedList(VisitedList *vl) {
+            std::unique_lock <std::mutex> lock(poolguard);
+            pool.push_front(vl);
+        };
+
+        ~VisitedListPool() {
+            while (pool.size()) {
+                VisitedList *rez = pool.front();
+                pool.pop_front();
+                delete rez;
+            }
+        };
+    };
+}
+

data/hnswlib.gemspec

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+require_relative 'lib/hnswlib/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = 'hnswlib'
+  spec.version       = Hnswlib::VERSION
+  spec.authors       = ['yoshoku']
+  spec.email         = ['yoshoku@outlook.com']
+
+  spec.summary       = 'Ruby bindings for the Hnswlib.'
+  spec.description   = 'Hnswlib.rb provides Ruby bindings for the Hnswlib.'
+  spec.homepage      = 'https://github.com/yoshoku/hnswlib.rb'
+  spec.license       = 'Apache-2.0'
+
+  spec.metadata['homepage_uri'] = spec.homepage
+  spec.metadata['source_code_uri'] = spec.homepage
+  spec.metadata['changelog_uri'] = 'https://github.com/yoshoku/hnswlib.rb/blob/main/CHANGELOG.md'
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files = Dir.chdir(File.expand_path(__dir__)) do
+    `git ls-files -z`.split("\x0").reject { |f| f.match(%r{\A(?:test|spec|features)/}) }
+  end
+  spec.bindir        = 'exe'
+  spec.executables   = spec.files.grep(%r{\Aexe/}) { |f| File.basename(f) }
+  spec.require_paths = ['lib']
+  spec.extensions    = ['ext/hnswlib/extconf.rb']
+
+  # Uncomment to register a new dependency of your gem
+  # spec.add_dependency "example-gem", "~> 1.0"
+
+  # For more information and examples about making a new gem, checkout our
+  # guide at: https://bundler.io/guides/creating_gem.html
+end

data/lib/hnswlib.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+require_relative 'hnswlib/version'
+require_relative 'hnswlib/hnswlibext'
+
+module Hnswlib
+  # Index is a class that provides functions for k-nearest eighbors search.
+  #
+  # @example
+  #   require 'hnswlib'
+  #
+  #   index = Hnswlib::Index.new(n_features: 100, max_item: 10000)
+  #
+  #   5000.times do |item_id|
+  #     item_vec = Array.new(100) { rand - 0.5 }
+  #     index.add_item(item_id, item_vec)
+  #   end
+  #
+  #   index.get_nns_by_item(0, 100)
+  #
+  class Index
+    # Returns the metric of index.
+    # @return [String]
+    attr_reader :metric
+
+    # Create a new search index.
+    #
+    # @param n_features [Integer] The number of features (dimensions) of stored vector.
+    # @param max_item [Integer] The maximum number of items.
+    # @param metric [String] The distance metric between vectors ('l2' or 'dot').
+    # @param m [Integer] The maximum number of outgoing connections in the graph
+    # @param ef_construction [Integer] The size of the dynamic list for the nearest neighbors. It controls the index time/accuracy trade-off.
+    # @param random_seed [Integer] The seed value using to initialize the random generator.
+    def initialize(n_features:, max_item:, metric: 'l2', m: 16, ef_construction: 200, random_seed: 100)
+      @metric = metric
+      space = if @metric == 'dot'
+                Hnswlib::InnerProductSpace.new(n_features)
+              else
+                Hnswlib::L2Space.new(n_features)
+              end
+      @index = Hnswlib::HierarchicalNSW.new(
+        space: space, max_elements: max_item, m: m, ef_construction: ef_construction, random_seed: random_seed
+      )
+    end
+
+    # Add item to be indexed.
+    #
+    # @param i [Integer] The ID of item.
+    # @param v [Array] The vector of item.
+    # @return [Boolean]
+    def add_item(i, v)
+      @index.add_point(v, i)
+    end
+
+    # Return the item vector.
+    #
+    # @param i [Integer] The ID of item.
+    # @return [Array]
+    def get_item(i)
+      @index.get_point(i)
+    end
+
+    # Remove the item vector.
+    #
+    # @param i [Integer] The ID of item.
+    # @return [Array]
+    def remove_item(i)
+      @index.mark_deleted(i)
+    end
+
+    # Search the n closest items.
+    #
+    # @param i [Integer] The ID of query item.
+    # @param n [Integer] The number of nearest neighbors.
+    # @param include_distances [Boolean] The flag indicating whether to returns all corresponding distances.
+    # @return [Array<Integer> or Array<Array<Integer>, Array<Float>>]
+    def get_nns_by_item(i, n, include_distances: false)
+      v = @index.get_point(i)
+      ids, dists = @index.search_knn(v, n)
+      include_distances ? [ids, dists] : ids
+    end
+
+    # Search the n closest items.
+    #
+    # @param v [Array] The vector of query item.
+    # @param n [Integer] The number of nearest neighbors.
+    # @param include_distances [Boolean] The flag indicating whether to returns all corresponding distances.
+    # @return [Array<Integer> or Array<Array<Integer>, Array<Float>>]
+    def get_nns_by_vector(v, n, include_distances: false)
+      ids, dists = @index.search_knn(v, n)
+      include_distances ? [ids, dists] : ids
+    end
+
+    # Reize the search index.
+    #
+    # @param new_max_item [Integer] The maximum number of items.
+    def resize_index(new_max_item)
+      @index.reisze_index(new_max_item)
+    end
+
+    # Set the size of the dynamic list for the nearest neighbors.
+    #
+    # @param ef [Integer] The size of the dynamic list.
+    def set_ef(ef)
+      @index.set_ef(ef)
+    end
+
+    # Save the search index to disk.
+    #
+    # @param filename [String] The filename of search index.
+    def save(filename)
+      @index.save_index(filename)
+    end
+
+    # Load a search index from disk.
+    #
+    # @param filename [String] The filename of search index.
+    def load(filename)
+      @index.load_index(filename)
+    end
+
+    # Calculate the distances between items.
+    #
+    # @param i [Integer] The ID of item.
+    # @param j [Integer] The ID of item.
+    # @return [Float or Integer]
+    def get_distance(i, j)
+      vi = @index.get_point(i)
+      vj = @index.get_point(j)
+      @index.space.distance(vi, vj)
+    end
+
+    # Return the number of items in the search index.
+    #
+    # @return [Integer]
+    def n_items
+      @index.current_count
+    end
+
+    # Returns the number of features of indexed item.
+    #
+    # @return [Integer]
+    def n_features
+      @index.space.dim
+    end
+
+    # Return the maximum number of items.
+    #
+    # @return [Integer]
+    def max_item
+      @index.max_elements
+    end
+  end
+end

data/lib/hnswlib/version.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+# Hnswlib.rb provides Ruby bindings for the Hnswlib.
+module Hnswlib
+  # The version of Hnswlib.rb you install.
+  VERSION = '0.1.0'
+  # The version of Hnswlib included with gem.
+  HSWLIB_VERSION = '0.5.2'
+end

metadata

@@ -0,0 +1,69 @@
+--- !ruby/object:Gem::Specification
+name: hnswlib
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- yoshoku
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2021-07-24 00:00:00.000000000 Z
+dependencies: []
+description: Hnswlib.rb provides Ruby bindings for the Hnswlib.
+email:
+- yoshoku@outlook.com
+executables: []
+extensions:
+- ext/hnswlib/extconf.rb
+extra_rdoc_files: []
+files:
+- ".github/workflows/build.yml"
+- ".gitignore"
+- ".rspec"
+- CHANGELOG.md
+- CODE_OF_CONDUCT.md
+- Gemfile
+- LICENSE.txt
+- README.md
+- Rakefile
+- ext/hnswlib/extconf.rb
+- ext/hnswlib/hnswlibext.cpp
+- ext/hnswlib/hnswlibext.hpp
+- ext/hnswlib/src/LICENSE
+- ext/hnswlib/src/bruteforce.h
+- ext/hnswlib/src/hnswalg.h
+- ext/hnswlib/src/hnswlib.h
+- ext/hnswlib/src/space_ip.h
+- ext/hnswlib/src/space_l2.h
+- ext/hnswlib/src/visited_list_pool.h
+- hnswlib.gemspec
+- lib/hnswlib.rb
+- lib/hnswlib/version.rb
+homepage: https://github.com/yoshoku/hnswlib.rb
+licenses:
+- Apache-2.0
+metadata:
+  homepage_uri: https://github.com/yoshoku/hnswlib.rb
+  source_code_uri: https://github.com/yoshoku/hnswlib.rb
+  changelog_uri: https://github.com/yoshoku/hnswlib.rb/blob/main/CHANGELOG.md
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.1.6
+signing_key:
+specification_version: 4
+summary: Ruby bindings for the Hnswlib.
+test_files: []