hnswlib 0.1.1 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1cd9cbb9fec801aa832f424f382c7efe7818c4909d9f0909e39afbe6d095311a
-  data.tar.gz: 557b1bed74b3c53814d6b68c85d588583546a83bcc1998f4dbfd6c265c188797
+  metadata.gz: 84f3ca262eddd8b331cbedc163e55e28d77eba76d2184332fef6afd0e2ddcce5
+  data.tar.gz: 248f8ca655ae83e2f56b26910dde6e96a23771040d61bac7ee7204e294063c2a
 SHA512:
-  metadata.gz: b764cd9bc82036ba04865a287caf78691e87b44eac24d74719d3956a26a5cc586f2c37f7dd44ad90bdc43d2bffbbe4902dd52f5599fe647e997b9a35e7bd5659
-  data.tar.gz: b20823f47c3a40ce39bd1edebe6e72b04084be07d16d55eef761d4c97fdda764bd3792691a57b94edf8c479a7ee5b786d02d66cbf7a861de63741d4a6cbc1b29
+  metadata.gz: 56d528a7ce9af7f96b291e74fc6a4c8ff1d3da3d66e1dbeed0dc8e884f5f5822b04e35d0fb12f83155a97ac1505a9a1db7b245cdcb2b08ba96ae6657743ccfeb
+  data.tar.gz: 47ca9c4db92798e07576505bc0c6b2a4f3ffb8afdcfdd020226232d8fca31d4259571509ce08ccde8eeecd06bcfa7cc2a5512c29d2568fa9b5d0ec0736aaa2ea

data/.gitignore

@@ -15,4 +15,5 @@ mkmf.log
 # rspec failure tracking
 .rspec_status
 
-spec/test.ann
+*.ann
+/bin/

data/CHANGELOG.md

@@ -1,3 +1,9 @@
+## [0.2.0] - 2021-08-02
+
+- Add binding class for the BruteforceSearch.
+- Add type check for arguments of initialize methods of BruteforceSearch and HierarchicalNSW.
+- Add dummy constructor call at memory allocation for class to prevent occuring segment fault on GC when initialize method is failed.
+
 ## [0.1.1] - 2021-07-25
 
 - Fix to use `rb_obj_is_isntance_of` for klass comparison due to type error when loading search index on irb 1.3.x: [issue #1](https://github.com/yoshoku/hnswlib.rb/issues/1)

data/README.md

@@ -49,6 +49,10 @@ u.load('test.ann')
 p u.get_nns_by_item(0, 100) # will find the 100 nearest neighbors.
 ```
 
+## License
+
+The gem is available as open source under the terms of the [Apache-2.0 License](https://www.apache.org/licenses/LICENSE-2.0).
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/yoshoku/hnswlib.rb.

data/ext/hnswlib/hnswlibext.cpp

@@ -26,4 +26,5 @@ void Init_hnswlibext(void) {
   RbHnswlibL2Space::define_class(rb_mHnswlib);
   RbHnswlibInnerProductSpace::define_class(rb_mHnswlib);
   RbHnswlibHierarchicalNSW::define_class(rb_mHnswlib);
+  RbHnswlibBruteforceSearch::define_class(rb_mHnswlib);
 }

data/ext/hnswlib/hnswlibext.hpp

@@ -26,11 +26,13 @@
 VALUE rb_cHnswlibL2Space;
 VALUE rb_cHnswlibInnerProductSpace;
 VALUE rb_cHnswlibHierarchicalNSW;
+VALUE rb_cHnswlibBruteforceSearch;
 
 class RbHnswlibL2Space {
   public:
     static VALUE hnsw_l2space_alloc(VALUE self) {
       hnswlib::L2Space* ptr = (hnswlib::L2Space*)ruby_xmalloc(sizeof(hnswlib::L2Space));
+      new (ptr) hnswlib::L2Space(); // dummy call to constructor for GC.
       return TypedData_Wrap_Struct(self, &hnsw_l2space_type, ptr);
     };
 
@@ -106,6 +108,7 @@ class RbHnswlibInnerProductSpace {
   public:
     static VALUE hnsw_ipspace_alloc(VALUE self) {
       hnswlib::InnerProductSpace* ptr = (hnswlib::InnerProductSpace*)ruby_xmalloc(sizeof(hnswlib::InnerProductSpace));
+      new (ptr) hnswlib::InnerProductSpace(); // dummy call to constructor for GC.
       return TypedData_Wrap_Struct(self, &hnsw_ipspace_type, ptr);
     };
 
@@ -181,6 +184,7 @@ class RbHnswlibHierarchicalNSW {
   public:
     static VALUE hnsw_hierarchicalnsw_alloc(VALUE self) {
       hnswlib::HierarchicalNSW<float>* ptr = (hnswlib::HierarchicalNSW<float>*)ruby_xmalloc(sizeof(hnswlib::HierarchicalNSW<float>));
+      new (ptr) hnswlib::HierarchicalNSW<float>(); // dummy call to constructor for GC.
       return TypedData_Wrap_Struct(self, &hnsw_hierarchicalnsw_type, ptr);
     };
 
@@ -239,6 +243,27 @@ class RbHnswlibHierarchicalNSW {
       if (kw_values[3] == Qundef) kw_values[3] = INT2NUM(200);
       if (kw_values[4] == Qundef) kw_values[4] = INT2NUM(100);
 
+      if (!(rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space) || rb_obj_is_instance_of(kw_values[0], rb_cHnswlibInnerProductSpace))) {
+        rb_raise(rb_eTypeError, "expected space, Hnswlib::L2Space or Hnswlib::InnerProductSpace");
+        return Qnil;
+      }
+      if (!RB_INTEGER_TYPE_P(kw_values[1])) {
+        rb_raise(rb_eTypeError, "expected max_elements, Integer");
+        return Qnil;
+      }
+      if (!RB_INTEGER_TYPE_P(kw_values[2])) {
+        rb_raise(rb_eTypeError, "expected m, Integer");
+        return Qnil;
+      }
+      if (!RB_INTEGER_TYPE_P(kw_values[3])) {
+        rb_raise(rb_eTypeError, "expected ef_construction, Integer");
+        return Qnil;
+      }
+      if (!RB_INTEGER_TYPE_P(kw_values[4])) {
+        rb_raise(rb_eTypeError, "expected random_seed, Integer");
+        return Qnil;
+      }
+
       rb_iv_set(self, "@space", kw_values[0]);
       hnswlib::SpaceInterface<float>* space;
       if (rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space)) {
@@ -418,4 +443,201 @@ const rb_data_type_t RbHnswlibHierarchicalNSW::hnsw_hierarchicalnsw_type = {
   RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+class RbHnswlibBruteforceSearch {
+  public:
+    static VALUE hnsw_bruteforcesearch_alloc(VALUE self) {
+      hnswlib::BruteforceSearch<float>* ptr = (hnswlib::BruteforceSearch<float>*)ruby_xmalloc(sizeof(hnswlib::BruteforceSearch<float>));
+      new (ptr) hnswlib::BruteforceSearch<float>(); // dummy call to constructor for GC.
+      return TypedData_Wrap_Struct(self, &hnsw_bruteforcesearch_type, ptr);
+    };
+
+    static void hnsw_bruteforcesearch_free(void* ptr) {
+      ((hnswlib::BruteforceSearch<float>*)ptr)->~BruteforceSearch();
+      ruby_xfree(ptr);
+    };
+
+    static size_t hnsw_bruteforcesearch_size(const void* ptr) {
+      return sizeof(*((hnswlib::BruteforceSearch<float>*)ptr));
+    };
+
+    static hnswlib::BruteforceSearch<float>* get_hnsw_bruteforcesearch(VALUE self) {
+      hnswlib::BruteforceSearch<float>* ptr;
+      TypedData_Get_Struct(self, hnswlib::BruteforceSearch<float>, &hnsw_bruteforcesearch_type, ptr);
+      return ptr;
+    };
+
+    static VALUE define_class(VALUE rb_mHnswlib) {
+      rb_cHnswlibBruteforceSearch = rb_define_class_under(rb_mHnswlib, "BruteforceSearch", rb_cObject);
+      rb_define_alloc_func(rb_cHnswlibBruteforceSearch, hnsw_bruteforcesearch_alloc);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "initialize", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_init), -1);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "add_point", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_add_point), 2);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "search_knn", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_search_knn), 2);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "save_index", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_save_index), 1);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "load_index", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_load_index), 1);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "remove_point", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_remove_point), 1);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "max_elements", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_max_elements), 0);
+      rb_define_method(rb_cHnswlibBruteforceSearch, "current_count", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_current_count), 0);
+      rb_define_attr(rb_cHnswlibBruteforceSearch, "space", 1, 0);
+      return rb_cHnswlibBruteforceSearch;
+    };
+
+  private:
+    static const rb_data_type_t hnsw_bruteforcesearch_type;
+
+    static VALUE _hnsw_bruteforcesearch_init(int argc, VALUE* argv, VALUE self) {
+      VALUE kw_args = Qnil;
+      ID kw_table[2] = { rb_intern("space"), rb_intern("max_elements") };
+      VALUE kw_values[2] = { Qundef, Qundef };
+      rb_scan_args(argc, argv, ":", &kw_args);
+      rb_get_kwargs(kw_args, kw_table, 2, 0, kw_values);
+
+      if (!(rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space) || rb_obj_is_instance_of(kw_values[0], rb_cHnswlibInnerProductSpace))) {
+        rb_raise(rb_eTypeError, "expected space, Hnswlib::L2Space or Hnswlib::InnerProductSpace");
+        return Qnil;
+      }
+      if (!RB_INTEGER_TYPE_P(kw_values[1])) {
+        rb_raise(rb_eTypeError, "expected max_elements, Integer");
+        return Qnil;
+      }
+
+      rb_iv_set(self, "@space", kw_values[0]);
+      hnswlib::SpaceInterface<float>* space;
+      if (rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space)) {
+        space = RbHnswlibL2Space::get_hnsw_l2space(kw_values[0]);
+      } else {
+        space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(kw_values[0]);
+      }
+      const size_t max_elements = (size_t)NUM2INT(kw_values[1]);
+
+      hnswlib::BruteforceSearch<float>* ptr = get_hnsw_bruteforcesearch(self);
+      new (ptr) hnswlib::BruteforceSearch<float>(space, max_elements);
+
+      return Qnil;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_add_point(VALUE self, VALUE arr, VALUE idx) {
+      const int dim = NUM2INT(rb_iv_get(rb_iv_get(self, "@space"), "@dim"));
+
+      if (!RB_TYPE_P(arr, T_ARRAY)) {
+        rb_raise(rb_eArgError, "Expect point vector to be Ruby Array.");
+        return Qfalse;
+      }
+
+      if (!RB_INTEGER_TYPE_P(idx)) {
+        rb_raise(rb_eArgError, "Expect index to be Ruby Integer.");
+        return Qfalse;
+      }
+
+      if (dim != RARRAY_LEN(arr)) {
+        rb_raise(rb_eArgError, "Array size does not match to index dimensionality.");
+        return Qfalse;
+      }
+
+      float* vec = (float*)ruby_xmalloc(dim * sizeof(float));
+      for (int i = 0; i < dim; i++) {
+        vec[i] = (float)NUM2DBL(rb_ary_entry(arr, i));
+      }
+
+      get_hnsw_bruteforcesearch(self)->addPoint((void *)vec, (size_t)NUM2INT(idx));
+
+      ruby_xfree(vec);
+      return Qtrue;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_search_knn(VALUE self, VALUE arr, VALUE k) {
+      const int dim = NUM2INT(rb_iv_get(rb_iv_get(self, "@space"), "@dim"));
+
+      if (!RB_TYPE_P(arr, T_ARRAY)) {
+        rb_raise(rb_eArgError, "Expect query vector to be Ruby Array.");
+        return Qnil;
+      }
+
+      if (!RB_INTEGER_TYPE_P(k)) {
+        rb_raise(rb_eArgError, "Expect the number of nearest neighbors to be Ruby Integer.");
+        return Qnil;
+      }
+
+      if (dim != RARRAY_LEN(arr)) {
+        rb_raise(rb_eArgError, "Array size does not match to index dimensionality.");
+        return Qnil;
+      }
+
+      float* vec = (float*)ruby_xmalloc(dim * sizeof(float));
+      for (int i = 0; i < dim; i++) {
+        vec[i] = (float)NUM2DBL(rb_ary_entry(arr, i));
+      }
+
+      std::priority_queue<std::pair<float, size_t>> result =
+        get_hnsw_bruteforcesearch(self)->searchKnn((void *)vec, (size_t)NUM2INT(k));
+
+      ruby_xfree(vec);
+
+      if (result.size() != (size_t)NUM2INT(k)) {
+        rb_raise(rb_eRuntimeError, "Cannot return the results in a contigious 2D array. Probably ef or M is too small.");
+        return Qnil;
+      }
+
+      VALUE distances_arr = rb_ary_new2(result.size());
+      VALUE neighbors_arr = rb_ary_new2(result.size());
+
+      for (int i = NUM2INT(k) - 1; i >= 0; i--) {
+        const std::pair<float, size_t>& result_tuple = result.top();
+        rb_ary_store(distances_arr, i, DBL2NUM((double)result_tuple.first));
+        rb_ary_store(neighbors_arr, i, INT2NUM((int)result_tuple.second));
+        result.pop();
+      }
+
+      VALUE ret = rb_ary_new2(2);
+      rb_ary_store(ret, 0, neighbors_arr);
+      rb_ary_store(ret, 1, distances_arr);
+      return ret;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_save_index(VALUE self, VALUE _filename) {
+      std::string filename(StringValuePtr(_filename));
+      get_hnsw_bruteforcesearch(self)->saveIndex(filename);
+      RB_GC_GUARD(_filename);
+      return Qnil;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_load_index(VALUE self, VALUE _filename) {
+      std::string filename(StringValuePtr(_filename));
+      VALUE ivspace = rb_iv_get(self, "@space");
+      hnswlib::SpaceInterface<float>* space;
+      if (rb_obj_is_instance_of(ivspace, rb_cHnswlibL2Space)) {
+        space = RbHnswlibL2Space::get_hnsw_l2space(ivspace);
+      } else {
+        space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(ivspace);
+      }
+      get_hnsw_bruteforcesearch(self)->loadIndex(filename, space);
+      RB_GC_GUARD(_filename);
+      return Qnil;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_remove_point(VALUE self, VALUE idx) {
+      get_hnsw_bruteforcesearch(self)->removePoint((size_t)NUM2INT(idx));
+      return Qnil;
+    };
+
+    static VALUE _hnsw_bruteforcesearch_max_elements(VALUE self) {
+      return INT2NUM((int)(get_hnsw_bruteforcesearch(self)->maxelements_));
+    };
+
+    static VALUE _hnsw_bruteforcesearch_current_count(VALUE self) {
+      return INT2NUM((int)(get_hnsw_bruteforcesearch(self)->cur_element_count));
+    };
+};
+
+const rb_data_type_t RbHnswlibBruteforceSearch::hnsw_bruteforcesearch_type = {
+  "RbHnswlibBruteforceSearch",
+  {
+    NULL,
+    RbHnswlibBruteforceSearch::hnsw_bruteforcesearch_free,
+    RbHnswlibBruteforceSearch::hnsw_bruteforcesearch_size
+  },
+  NULL,
+  NULL,
+  RUBY_TYPED_FREE_IMMEDIATELY
+};
+
 #endif /* HNSWLIBEXT_HPP */

data/ext/hnswlib/src/bruteforce.h

@@ -8,6 +8,7 @@ namespace hnswlib {
     template<typename dist_t>
     class BruteforceSearch : public AlgorithmInterface<dist_t> {
     public:
+        BruteforceSearch() : data_(nullptr) { }
         BruteforceSearch(SpaceInterface <dist_t> *s) {
 
         }
@@ -91,13 +92,13 @@ namespace hnswlib {
         searchKnn(const void *query_data, size_t k) const {
             std::priority_queue<std::pair<dist_t, labeltype >> topResults;
             if (cur_element_count == 0) return topResults;
-            for (int i = 0; i < k; i++) {
+            for (size_t i = 0; i < k; i++) {
                 dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_);
                 topResults.push(std::pair<dist_t, labeltype>(dist, *((labeltype *) (data_ + size_per_element_ * i +
                                                                                     data_size_))));
             }
             dist_t lastdist = topResults.top().first;
-            for (int i = k; i < cur_element_count; i++) {
+            for (size_t i = k; i < cur_element_count; i++) {
                 dist_t dist = fstdistfunc_(query_data, data_ + size_per_element_ * i, dist_func_param_);
                 if (dist <= lastdist) {
                     topResults.push(std::pair<dist_t, labeltype>(dist, *((labeltype *) (data_ + size_per_element_ * i +

data/ext/hnswlib/src/hnswalg.h

@@ -17,6 +17,7 @@ namespace hnswlib {
     class HierarchicalNSW : public AlgorithmInterface<dist_t> {
     public:
         static const tableint max_update_element_locks = 65536;
+        HierarchicalNSW() : visited_list_pool_(nullptr), data_level0_memory_(nullptr), linkLists_(nullptr), cur_element_count(0) { }
         HierarchicalNSW(SpaceInterface<dist_t> *s) {
 
         }
@@ -757,7 +758,7 @@ namespace hnswlib {
             size_t dim = *((size_t *) dist_func_param_);
             std::vector<data_t> data;
             data_t* data_ptr = (data_t*) data_ptrv;
-            for (int i = 0; i < dim; i++) {
+            for (size_t i = 0; i < dim; i++) {
                 data.push_back(*data_ptr);
                 data_ptr += 1;
             }

data/ext/hnswlib/src/space_ip.h

@@ -247,6 +247,7 @@ namespace hnswlib {
         size_t data_size_;
         size_t dim_;
     public:
+        InnerProductSpace() : data_size_(0), dim_(0) { }
         InnerProductSpace(size_t dim) {
             fstdistfunc_ = InnerProduct;
     #if defined(USE_AVX) || defined(USE_SSE)

data/ext/hnswlib/src/space_l2.h

@@ -172,6 +172,7 @@ namespace hnswlib {
         size_t data_size_;
         size_t dim_;
     public:
+        L2Space() : data_size_(0), dim_(0) { }
         L2Space(size_t dim) {
             fstdistfunc_ = L2Sqr;
         #if defined(USE_SSE) || defined(USE_AVX)

data/lib/hnswlib/version.rb

@@ -3,7 +3,7 @@
 # Hnswlib.rb provides Ruby bindings for the Hnswlib.
 module Hnswlib
   # The version of Hnswlib.rb you install.
-  VERSION = '0.1.1'
+  VERSION = '0.2.0'
 
   # The version of Hnswlib included with gem.
   HSWLIB_VERSION = '0.5.2'

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: hnswlib
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.2.0
 platform: ruby
 authors:
 - yoshoku
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2021-07-25 00:00:00.000000000 Z
+date: 2021-08-02 00:00:00.000000000 Z
 dependencies: []
 description: Hnswlib.rb provides Ruby bindings for the Hnswlib.
 email:
@@ -62,7 +62,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.6
+rubygems_version: 3.2.22
 signing_key:
 specification_version: 4
 summary: Ruby bindings for the Hnswlib.