hnswlib 0.6.2 → 0.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fc8a40172623a3439b17c8b1854017c845bf4dc23f4e1714e9427e4b6b701c9b
-  data.tar.gz: 95fa48d791f000bf48809e7067efaba9a654ab0dd517c626f689c66af3668d0e
+  metadata.gz: a564835f983f6c07a04d62e198aac0b2a80eef8eaa784a14d3d7bdc5dadaa962
+  data.tar.gz: 98cc90158fbe92a012a6e0f945a1c82cd3bbfc1ae973cf1c4c1eccf95c249fed
 SHA512:
-  metadata.gz: 07c449031cc7afbf80803ae15c8094589a146b59b1852ea72e88fcd39067def572bbdd58277e3a85c47c537e8dc6eef645f9f404a2b874e4e217b559adc38054
-  data.tar.gz: 4b71a7f4fe0ffb3c15124e32c9c8fbba9eb6f82e6f0b95f5d91dd35c6b10be8aeba25d7313eabeaa5078b051b3410195415843ecd600b58ec1b0b36532b7ae1d
+  metadata.gz: ada8080314124e768aba55385f92c98a8e1206932c25a3d93cff2eaf8659d44375622a22226bb96379c7377055d8840703aa8dac25b49440f1a862f5c9b45444
+  data.tar.gz: 6ce0ee55fd1d0174bd1d06377eb1a5ed3cf875ea1225da0320ba6e4b003bdd18d73d00f126fc17484dc8c2b45621a2cfe86c74134f6c227b4b1364b631977d69

data/CHANGELOG.md

@@ -1,3 +1,9 @@
+## [0.7.0] - 2023-03-04
+
+- Update bundled hnswlib version to 0.7.0.
+- Add support for replacing an element marked for deletion with a new element.
+- Add support filtering function by label in search_knn method of BruteforeceSearch and HierarchicalNSW.
+
 ## [0.6.2] - 2022-06-25
 
 - Refactor codes and configs with RuboCop and clang-format.

data/ext/hnswlib/hnswlibext.cpp

@@ -1,7 +1,7 @@
 /**
  * hnswlib.rb is a Ruby binding for the Hnswlib.
  *
- * Copyright (c) 2021-2022 Atsushi Tatsuma
+ * Copyright (c) 2021-2023 Atsushi Tatsuma
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

data/ext/hnswlib/hnswlibext.hpp

@@ -1,7 +1,7 @@
 /**
  * hnswlib.rb is a Ruby binding for the Hnswlib.
  *
- * Copyright (c) 2021-2022 Atsushi Tatsuma
+ * Copyright (c) 2021-2023 Atsushi Tatsuma
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,9 @@
 
 #include <hnswlib.h>
 
+#include <new>
+#include <vector>
+
 VALUE rb_cHnswlibL2Space;
 VALUE rb_cHnswlibInnerProductSpace;
 VALUE rb_cHnswlibHierarchicalNSW;
@@ -64,12 +67,12 @@ private:
   static VALUE _hnsw_l2space_init(VALUE self, VALUE dim) {
     rb_iv_set(self, "@dim", dim);
     hnswlib::L2Space* ptr = get_hnsw_l2space(self);
-    new (ptr) hnswlib::L2Space(NUM2INT(rb_iv_get(self, "@dim")));
+    new (ptr) hnswlib::L2Space(NUM2SIZET(rb_iv_get(self, "@dim")));
     return Qnil;
   };
 
   static VALUE _hnsw_l2space_distance(VALUE self, VALUE arr_a, VALUE arr_b) {
-    const int dim = NUM2INT(rb_iv_get(self, "@dim"));
+    const size_t dim = NUM2SIZET(rb_iv_get(self, "@dim"));
     if (!RB_TYPE_P(arr_a, T_ARRAY) || !RB_TYPE_P(arr_b, T_ARRAY)) {
       rb_raise(rb_eArgError, "Expect input vector to be Ruby Array.");
       return Qnil;
@@ -79,9 +82,9 @@ private:
       return Qnil;
     }
     float* vec_a = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) vec_a[i] = (float)NUM2DBL(rb_ary_entry(arr_a, i));
+    for (size_t i = 0; i < dim; i++) vec_a[i] = (float)NUM2DBL(rb_ary_entry(arr_a, i));
     float* vec_b = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) vec_b[i] = (float)NUM2DBL(rb_ary_entry(arr_b, i));
+    for (size_t i = 0; i < dim; i++) vec_b[i] = (float)NUM2DBL(rb_ary_entry(arr_b, i));
     hnswlib::DISTFUNC<float> dist_func = get_hnsw_l2space(self)->get_dist_func();
     const float dist = dist_func(vec_a, vec_b, get_hnsw_l2space(self)->get_dist_func_param());
     ruby_xfree(vec_a);
@@ -140,12 +143,12 @@ private:
   static VALUE _hnsw_ipspace_init(VALUE self, VALUE dim) {
     rb_iv_set(self, "@dim", dim);
     hnswlib::InnerProductSpace* ptr = get_hnsw_ipspace(self);
-    new (ptr) hnswlib::InnerProductSpace(NUM2INT(rb_iv_get(self, "@dim")));
+    new (ptr) hnswlib::InnerProductSpace(NUM2SIZET(rb_iv_get(self, "@dim")));
     return Qnil;
   };
 
   static VALUE _hnsw_ipspace_distance(VALUE self, VALUE arr_a, VALUE arr_b) {
-    const int dim = NUM2INT(rb_iv_get(self, "@dim"));
+    const size_t dim = NUM2SIZET(rb_iv_get(self, "@dim"));
     if (!RB_TYPE_P(arr_a, T_ARRAY) || !RB_TYPE_P(arr_b, T_ARRAY)) {
       rb_raise(rb_eArgError, "Expect input vector to be Ruby Array.");
       return Qnil;
@@ -155,9 +158,9 @@ private:
       return Qnil;
     }
     float* vec_a = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) vec_a[i] = (float)NUM2DBL(rb_ary_entry(arr_a, i));
+    for (size_t i = 0; i < dim; i++) vec_a[i] = (float)NUM2DBL(rb_ary_entry(arr_a, i));
     float* vec_b = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) vec_b[i] = (float)NUM2DBL(rb_ary_entry(arr_b, i));
+    for (size_t i = 0; i < dim; i++) vec_b[i] = (float)NUM2DBL(rb_ary_entry(arr_b, i));
     hnswlib::DISTFUNC<float> dist_func = get_hnsw_ipspace(self)->get_dist_func();
     const float dist = dist_func(vec_a, vec_b, get_hnsw_ipspace(self)->get_dist_func_param());
     ruby_xfree(vec_a);
@@ -180,6 +183,19 @@ const rb_data_type_t RbHnswlibInnerProductSpace::hnsw_ipspace_type = {
 };
 // clang-format on
 
+class CustomFilterFunctor : public hnswlib::BaseFilterFunctor {
+public:
+  CustomFilterFunctor(const VALUE& callback) : callback_(callback) {}
+
+  bool operator()(hnswlib::labeltype id) {
+    VALUE result = rb_funcall(callback_, rb_intern("call"), 1, SIZET2NUM(id));
+    return result == Qtrue ? true : false;
+  }
+
+private:
+  VALUE callback_;
+};
+
 class RbHnswlibHierarchicalNSW {
 public:
   static VALUE hnsw_hierarchicalnsw_alloc(VALUE self) {
@@ -206,17 +222,21 @@ public:
     rb_cHnswlibHierarchicalNSW = rb_define_class_under(rb_mHnswlib, "HierarchicalNSW", rb_cObject);
     rb_define_alloc_func(rb_cHnswlibHierarchicalNSW, hnsw_hierarchicalnsw_alloc);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "initialize", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_init), -1);
-    rb_define_method(rb_cHnswlibHierarchicalNSW, "add_point", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_add_point), 2);
-    rb_define_method(rb_cHnswlibHierarchicalNSW, "search_knn", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_search_knn), 2);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "add_point", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_add_point), -1);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "search_knn", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_search_knn), -1);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "save_index", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_save_index), 1);
-    rb_define_method(rb_cHnswlibHierarchicalNSW, "load_index", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_load_index), 1);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "load_index", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_load_index), -1);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "get_point", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_get_point), 1);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "get_ids", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_get_ids), 0);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "mark_deleted", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_mark_deleted), 1);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "unmark_deleted", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_unmark_deleted), 1);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "resize_index", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_resize_index), 1);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "set_ef", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_set_ef), 1);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "get_ef", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_get_ef), 0);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "max_elements", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_max_elements), 0);
     rb_define_method(rb_cHnswlibHierarchicalNSW, "current_count", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_current_count), 0);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "ef_construction", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_ef_construction), 0);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "m", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_m), 0);
     rb_define_attr(rb_cHnswlibHierarchicalNSW, "space", 1, 0);
     return rb_cHnswlibHierarchicalNSW;
   };
@@ -226,14 +246,15 @@ private:
 
   static VALUE _hnsw_hierarchicalnsw_init(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[5] = {rb_intern("space"), rb_intern("max_elements"), rb_intern("m"), rb_intern("ef_construction"),
-                      rb_intern("random_seed")};
-    VALUE kw_values[5] = {Qundef, Qundef, Qundef, Qundef, Qundef};
+    ID kw_table[6] = {rb_intern("space"),           rb_intern("max_elements"), rb_intern("m"),
+                      rb_intern("ef_construction"), rb_intern("random_seed"),  rb_intern("allow_replace_deleted")};
+    VALUE kw_values[6] = {Qundef, Qundef, Qundef, Qundef, Qundef, Qundef};
     rb_scan_args(argc, argv, ":", &kw_args);
-    rb_get_kwargs(kw_args, kw_table, 2, 3, kw_values);
-    if (kw_values[2] == Qundef) kw_values[2] = INT2NUM(16);
-    if (kw_values[3] == Qundef) kw_values[3] = INT2NUM(200);
-    if (kw_values[4] == Qundef) kw_values[4] = INT2NUM(100);
+    rb_get_kwargs(kw_args, kw_table, 2, 4, kw_values);
+    if (kw_values[2] == Qundef) kw_values[2] = SIZET2NUM(16);
+    if (kw_values[3] == Qundef) kw_values[3] = SIZET2NUM(200);
+    if (kw_values[4] == Qundef) kw_values[4] = SIZET2NUM(100);
+    if (kw_values[5] == Qundef) kw_values[5] = Qfalse;
 
     if (!(rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space) ||
           rb_obj_is_instance_of(kw_values[0], rb_cHnswlibInnerProductSpace))) {
@@ -256,6 +277,10 @@ private:
       rb_raise(rb_eTypeError, "expected random_seed, Integer");
       return Qnil;
     }
+    if (!RB_TYPE_P(kw_values[5], T_TRUE) && !RB_TYPE_P(kw_values[5], T_FALSE)) {
+      rb_raise(rb_eTypeError, "expected allow_replace_deleted, Boolean");
+      return Qnil;
+    }
 
     rb_iv_set(self, "@space", kw_values[0]);
     hnswlib::SpaceInterface<float>* space;
@@ -264,14 +289,15 @@ private:
     } else {
       space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(kw_values[0]);
     }
-    const size_t max_elements = (size_t)NUM2INT(kw_values[1]);
-    const size_t m = (size_t)NUM2INT(kw_values[2]);
-    const size_t ef_construction = (size_t)NUM2INT(kw_values[3]);
-    const size_t random_seed = (size_t)NUM2INT(kw_values[4]);
+    const size_t max_elements = NUM2SIZET(kw_values[1]);
+    const size_t m = NUM2SIZET(kw_values[2]);
+    const size_t ef_construction = NUM2SIZET(kw_values[3]);
+    const size_t random_seed = NUM2SIZET(kw_values[4]);
+    const bool allow_replace_deleted = kw_values[5] == Qtrue ? true : false;
 
     hnswlib::HierarchicalNSW<float>* ptr = get_hnsw_hierarchicalnsw(self);
     try {
-      new (ptr) hnswlib::HierarchicalNSW<float>(space, max_elements, m, ef_construction, random_seed);
+      new (ptr) hnswlib::HierarchicalNSW<float>(space, max_elements, m, ef_construction, random_seed, allow_replace_deleted);
     } catch (const std::runtime_error& e) {
       rb_raise(rb_eRuntimeError, "%s", e.what());
       return Qnil;
@@ -280,33 +306,63 @@ private:
     return Qnil;
   };
 
-  static VALUE _hnsw_hierarchicalnsw_add_point(VALUE self, VALUE arr, VALUE idx) {
-    const int dim = NUM2INT(rb_iv_get(rb_iv_get(self, "@space"), "@dim"));
+  static VALUE _hnsw_hierarchicalnsw_add_point(int argc, VALUE* argv, VALUE self) {
+    VALUE _arr, _idx, _replace_deleted;
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = {rb_intern("replace_deleted")};
+    VALUE kw_values[1] = {Qundef};
+
+    rb_scan_args(argc, argv, "2:", &_arr, &_idx, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 0, 1, kw_values);
+    _replace_deleted = kw_values[0] != Qundef ? kw_values[0] : Qfalse;
 
-    if (!RB_TYPE_P(arr, T_ARRAY)) {
+    const size_t dim = NUM2SIZET(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)) {
+    if (!RB_INTEGER_TYPE_P(_idx)) {
       rb_raise(rb_eArgError, "Expect index to be Ruby Integer.");
       return Qfalse;
     }
-    if (dim != RARRAY_LEN(arr)) {
+    if (dim != RARRAY_LEN(_arr)) {
       rb_raise(rb_eArgError, "Array size does not match to index dimensionality.");
       return Qfalse;
     }
+    if (!RB_TYPE_P(_replace_deleted, T_TRUE) && !RB_TYPE_P(_replace_deleted, T_FALSE)) {
+      rb_raise(rb_eArgError, "Expect replace_deleted to be Boolean.");
+      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));
+    float* arr = (float*)ruby_xmalloc(dim * sizeof(float));
+    for (size_t i = 0; i < dim; i++) arr[i] = (float)NUM2DBL(rb_ary_entry(_arr, i));
+    const size_t idx = NUM2SIZET(_idx);
+    const bool replace_deleted = _replace_deleted == Qtrue ? true : false;
 
-    get_hnsw_hierarchicalnsw(self)->addPoint((void*)vec, (size_t)NUM2INT(idx));
+    try {
+      get_hnsw_hierarchicalnsw(self)->addPoint((void*)arr, idx, replace_deleted);
+    } catch (const std::runtime_error& e) {
+      ruby_xfree(arr);
+      rb_raise(rb_eRuntimeError, "%s", e.what());
+      return Qfalse;
+    }
 
-    ruby_xfree(vec);
+    ruby_xfree(arr);
     return Qtrue;
   };
 
-  static VALUE _hnsw_hierarchicalnsw_search_knn(VALUE self, VALUE arr, VALUE k) {
-    const int dim = NUM2INT(rb_iv_get(rb_iv_get(self, "@space"), "@dim"));
+  static VALUE _hnsw_hierarchicalnsw_search_knn(int argc, VALUE* argv, VALUE self) {
+    VALUE arr, k, filter;
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = {rb_intern("filter")};
+    VALUE kw_values[1] = {Qundef};
+
+    rb_scan_args(argc, argv, "2:", &arr, &k, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 0, 1, kw_values);
+    filter = kw_values[0] != Qundef ? kw_values[0] : Qnil;
+
+    const size_t dim = NUM2SIZET(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.");
@@ -321,14 +377,24 @@ private:
       return Qnil;
     }
 
+    CustomFilterFunctor* filter_func = nullptr;
+    if (!NIL_P(filter)) {
+      try {
+        filter_func = new CustomFilterFunctor(filter);
+      } catch (const std::bad_alloc& e) {
+        rb_raise(rb_eRuntimeError, "%s", e.what());
+        return Qnil;
+      }
+    }
+
     float* vec = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) {
+    for (size_t i = 0; i < dim; i++) {
       vec[i] = (float)NUM2DBL(rb_ary_entry(arr, i));
     }
 
     std::priority_queue<std::pair<float, size_t>> result;
     try {
-      result = get_hnsw_hierarchicalnsw(self)->searchKnn((void*)vec, (size_t)NUM2INT(k));
+      result = get_hnsw_hierarchicalnsw(self)->searchKnn((void*)vec, NUM2SIZET(k), filter_func);
     } catch (const std::runtime_error& e) {
       ruby_xfree(vec);
       rb_raise(rb_eRuntimeError, "%s", e.what());
@@ -336,8 +402,9 @@ private:
     }
 
     ruby_xfree(vec);
+    if (filter_func) delete filter_func;
 
-    if (result.size() != (size_t)NUM2INT(k)) {
+    if (result.size() != NUM2SIZET(k)) {
       rb_warning("Cannot return as many search results as the requested number of neighbors. Probably ef or M is too small.");
     }
 
@@ -347,7 +414,7 @@ private:
     while (!result.empty()) {
       const std::pair<float, size_t>& result_tuple = result.top();
       rb_ary_unshift(distances_arr, DBL2NUM((double)result_tuple.first));
-      rb_ary_unshift(neighbors_arr, INT2NUM((int)result_tuple.second));
+      rb_ary_unshift(neighbors_arr, SIZET2NUM(result_tuple.second));
       result.pop();
     }
 
@@ -364,8 +431,28 @@ private:
     return Qnil;
   };
 
-  static VALUE _hnsw_hierarchicalnsw_load_index(VALUE self, VALUE _filename) {
+  static VALUE _hnsw_hierarchicalnsw_load_index(int argc, VALUE* argv, VALUE self) {
+    VALUE _filename, _allow_replace_deleted;
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = {rb_intern("allow_replace_deleted")};
+    VALUE kw_values[1] = {Qundef};
+
+    rb_scan_args(argc, argv, "1:", &_filename, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 0, 1, kw_values);
+    _allow_replace_deleted = kw_values[0] != Qundef ? kw_values[0] : Qfalse;
+
+    if (!RB_TYPE_P(_filename, T_STRING)) {
+      rb_raise(rb_eArgError, "Expect filename to be Ruby Array.");
+      return Qnil;
+    }
+    if (!NIL_P(_allow_replace_deleted) && !RB_TYPE_P(_allow_replace_deleted, T_TRUE) &&
+        !RB_TYPE_P(_allow_replace_deleted, T_FALSE)) {
+      rb_raise(rb_eArgError, "Expect replace_deleted to be Boolean.");
+      return Qnil;
+    }
+
     std::string filename(StringValuePtr(_filename));
+    const bool allow_replace_deleted = _allow_replace_deleted == Qtrue ? true : false;
     VALUE ivspace = rb_iv_get(self, "@space");
     hnswlib::SpaceInterface<float>* space;
     if (rb_obj_is_instance_of(ivspace, rb_cHnswlibL2Space)) {
@@ -373,6 +460,7 @@ private:
     } else {
       space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(ivspace);
     }
+
     hnswlib::HierarchicalNSW<float>* index = get_hnsw_hierarchicalnsw(self);
     if (index->data_level0_memory_) {
       free(index->data_level0_memory_);
@@ -392,12 +480,15 @@ private:
       delete index->visited_list_pool_;
       index->visited_list_pool_ = nullptr;
     }
+
     try {
       index->loadIndex(filename, space);
+      index->allow_replace_deleted_ = allow_replace_deleted;
     } catch (const std::runtime_error& e) {
       rb_raise(rb_eRuntimeError, "%s", e.what());
       return Qnil;
     }
+
     RB_GC_GUARD(_filename);
     return Qnil;
   };
@@ -405,7 +496,7 @@ private:
   static VALUE _hnsw_hierarchicalnsw_get_point(VALUE self, VALUE idx) {
     VALUE ret = Qnil;
     try {
-      std::vector<float> vec = get_hnsw_hierarchicalnsw(self)->template getDataByLabel<float>((size_t)NUM2INT(idx));
+      std::vector<float> vec = get_hnsw_hierarchicalnsw(self)->template getDataByLabel<float>(NUM2SIZET(idx));
       ret = rb_ary_new2(vec.size());
       for (size_t i = 0; i < vec.size(); i++) rb_ary_store(ret, i, DBL2NUM((double)vec[i]));
     } catch (const std::runtime_error& e) {
@@ -417,13 +508,23 @@ private:
 
   static VALUE _hnsw_hierarchicalnsw_get_ids(VALUE self) {
     VALUE ret = rb_ary_new();
-    for (auto kv : get_hnsw_hierarchicalnsw(self)->label_lookup_) rb_ary_push(ret, INT2NUM((int)kv.first));
+    for (auto kv : get_hnsw_hierarchicalnsw(self)->label_lookup_) rb_ary_push(ret, SIZET2NUM(kv.first));
     return ret;
   };
 
   static VALUE _hnsw_hierarchicalnsw_mark_deleted(VALUE self, VALUE idx) {
     try {
-      get_hnsw_hierarchicalnsw(self)->markDelete((size_t)NUM2INT(idx));
+      get_hnsw_hierarchicalnsw(self)->markDelete(NUM2SIZET(idx));
+    } catch (const std::runtime_error& e) {
+      rb_raise(rb_eRuntimeError, "%s", e.what());
+      return Qnil;
+    }
+    return Qnil;
+  };
+
+  static VALUE _hnsw_hierarchicalnsw_unmark_deleted(VALUE self, VALUE idx) {
+    try {
+      get_hnsw_hierarchicalnsw(self)->unmarkDelete(NUM2SIZET(idx));
     } catch (const std::runtime_error& e) {
       rb_raise(rb_eRuntimeError, "%s", e.what());
       return Qnil;
@@ -432,31 +533,42 @@ private:
   };
 
   static VALUE _hnsw_hierarchicalnsw_resize_index(VALUE self, VALUE new_max_elements) {
-    if ((size_t)NUM2INT(new_max_elements) < get_hnsw_hierarchicalnsw(self)->cur_element_count) {
+    if (NUM2SIZET(new_max_elements) < get_hnsw_hierarchicalnsw(self)->cur_element_count) {
       rb_raise(rb_eArgError, "Cannot resize, max element is less than the current number of elements.");
       return Qnil;
     }
     try {
-      get_hnsw_hierarchicalnsw(self)->resizeIndex((size_t)NUM2INT(new_max_elements));
+      get_hnsw_hierarchicalnsw(self)->resizeIndex(NUM2SIZET(new_max_elements));
     } catch (const std::runtime_error& e) {
       rb_raise(rb_eRuntimeError, "%s", e.what());
       return Qnil;
+    } catch (const std::bad_alloc& e) {
+      rb_raise(rb_eRuntimeError, "%s", e.what());
+      return Qnil;
     }
     return Qnil;
   };
 
   static VALUE _hnsw_hierarchicalnsw_set_ef(VALUE self, VALUE ef) {
-    get_hnsw_hierarchicalnsw(self)->ef_ = (size_t)NUM2INT(ef);
+    get_hnsw_hierarchicalnsw(self)->setEf(NUM2SIZET(ef));
     return Qnil;
   };
 
+  static VALUE _hnsw_hierarchicalnsw_get_ef(VALUE self) { return SIZET2NUM(get_hnsw_hierarchicalnsw(self)->ef_); };
+
   static VALUE _hnsw_hierarchicalnsw_max_elements(VALUE self) {
-    return INT2NUM((int)(get_hnsw_hierarchicalnsw(self)->max_elements_));
+    return SIZET2NUM(get_hnsw_hierarchicalnsw(self)->max_elements_);
   };
 
   static VALUE _hnsw_hierarchicalnsw_current_count(VALUE self) {
-    return INT2NUM((int)(get_hnsw_hierarchicalnsw(self)->cur_element_count));
+    return SIZET2NUM(get_hnsw_hierarchicalnsw(self)->cur_element_count);
+  };
+
+  static VALUE _hnsw_hierarchicalnsw_ef_construction(VALUE self) {
+    return SIZET2NUM(get_hnsw_hierarchicalnsw(self)->ef_construction_);
   };
+
+  static VALUE _hnsw_hierarchicalnsw_m(VALUE self) { return SIZET2NUM(get_hnsw_hierarchicalnsw(self)->M_); };
 };
 
 // clang-format off
@@ -500,7 +612,7 @@ public:
     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, "search_knn", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_search_knn), -1);
     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);
@@ -537,7 +649,7 @@ private:
     } else {
       space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(kw_values[0]);
     }
-    const size_t max_elements = (size_t)NUM2INT(kw_values[1]);
+    const size_t max_elements = NUM2SIZET(kw_values[1]);
 
     hnswlib::BruteforceSearch<float>* ptr = get_hnsw_bruteforcesearch(self);
     try {
@@ -551,7 +663,7 @@ private:
   };
 
   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"));
+    const size_t dim = NUM2SIZET(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.");
@@ -567,10 +679,10 @@ private:
     }
 
     float* vec = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) vec[i] = (float)NUM2DBL(rb_ary_entry(arr, i));
+    for (size_t i = 0; i < dim; i++) vec[i] = (float)NUM2DBL(rb_ary_entry(arr, i));
 
     try {
-      get_hnsw_bruteforcesearch(self)->addPoint((void*)vec, (size_t)NUM2INT(idx));
+      get_hnsw_bruteforcesearch(self)->addPoint((void*)vec, NUM2SIZET(idx));
     } catch (const std::runtime_error& e) {
       ruby_xfree(vec);
       rb_raise(rb_eRuntimeError, "%s", e.what());
@@ -581,8 +693,17 @@ private:
     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"));
+  static VALUE _hnsw_bruteforcesearch_search_knn(int argc, VALUE* argv, VALUE self) {
+    VALUE arr, k, filter;
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = {rb_intern("filter")};
+    VALUE kw_values[1] = {Qundef};
+
+    rb_scan_args(argc, argv, "2:", &arr, &k, &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 0, 1, kw_values);
+    filter = kw_values[0] != Qundef ? kw_values[0] : Qnil;
+
+    const size_t dim = NUM2SIZET(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.");
@@ -597,17 +718,28 @@ private:
       return Qnil;
     }
 
+    CustomFilterFunctor* filter_func = nullptr;
+    if (!NIL_P(filter)) {
+      try {
+        filter_func = new CustomFilterFunctor(filter);
+      } catch (const std::bad_alloc& e) {
+        rb_raise(rb_eRuntimeError, "%s", e.what());
+        return Qnil;
+      }
+    }
+
     float* vec = (float*)ruby_xmalloc(dim * sizeof(float));
-    for (int i = 0; i < dim; i++) {
+    for (size_t 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));
+        get_hnsw_bruteforcesearch(self)->searchKnn((void*)vec, NUM2SIZET(k), filter_func);
 
     ruby_xfree(vec);
+    if (filter_func) delete filter_func;
 
-    if (result.size() != (size_t)NUM2INT(k)) {
+    if (result.size() != NUM2SIZET(k)) {
       rb_warning("Cannot return as many search results as the requested number of neighbors.");
     }
 
@@ -617,7 +749,7 @@ private:
     while (!result.empty()) {
       const std::pair<float, size_t>& result_tuple = result.top();
       rb_ary_unshift(distances_arr, DBL2NUM((double)result_tuple.first));
-      rb_ary_unshift(neighbors_arr, INT2NUM((int)result_tuple.second));
+      rb_ary_unshift(neighbors_arr, SIZET2NUM(result_tuple.second));
       result.pop();
     }
 
@@ -659,16 +791,16 @@ private:
   };
 
   static VALUE _hnsw_bruteforcesearch_remove_point(VALUE self, VALUE idx) {
-    get_hnsw_bruteforcesearch(self)->removePoint((size_t)NUM2INT(idx));
+    get_hnsw_bruteforcesearch(self)->removePoint(NUM2SIZET(idx));
     return Qnil;
   };
 
   static VALUE _hnsw_bruteforcesearch_max_elements(VALUE self) {
-    return INT2NUM((int)(get_hnsw_bruteforcesearch(self)->maxelements_));
+    return SIZET2NUM(get_hnsw_bruteforcesearch(self)->maxelements_);
   };
 
   static VALUE _hnsw_bruteforcesearch_current_count(VALUE self) {
-    return INT2NUM((int)(get_hnsw_bruteforcesearch(self)->cur_element_count));
+    return SIZET2NUM(get_hnsw_bruteforcesearch(self)->cur_element_count);
   };
 };