hnswlib 0.6.2 → 0.8.0

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,11 @@
 
 #include <hnswlib.h>
 
+#include <cmath>
+#include <new>
+#include <vector>
+
+VALUE rb_mHnswlib;
 VALUE rb_cHnswlibL2Space;
 VALUE rb_cHnswlibInnerProductSpace;
 VALUE rb_cHnswlibHierarchicalNSW;
@@ -49,8 +54,8 @@ public:
     return ptr;
   };
 
-  static VALUE define_class(VALUE rb_mHnswlib) {
-    rb_cHnswlibL2Space = rb_define_class_under(rb_mHnswlib, "L2Space", rb_cObject);
+  static VALUE define_class(VALUE outer) {
+    rb_cHnswlibL2Space = rb_define_class_under(outer, "L2Space", rb_cObject);
     rb_define_alloc_func(rb_cHnswlibL2Space, hnsw_l2space_alloc);
     rb_define_method(rb_cHnswlibL2Space, "initialize", RUBY_METHOD_FUNC(_hnsw_l2space_init), 1);
     rb_define_method(rb_cHnswlibL2Space, "distance", RUBY_METHOD_FUNC(_hnsw_l2space_distance), 2);
@@ -64,12 +69,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 +84,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);
@@ -125,8 +130,8 @@ public:
     return ptr;
   };
 
-  static VALUE define_class(VALUE rb_mHnswlib) {
-    rb_cHnswlibInnerProductSpace = rb_define_class_under(rb_mHnswlib, "InnerProductSpace", rb_cObject);
+  static VALUE define_class(VALUE outer) {
+    rb_cHnswlibInnerProductSpace = rb_define_class_under(outer, "InnerProductSpace", rb_cObject);
     rb_define_alloc_func(rb_cHnswlibInnerProductSpace, hnsw_ipspace_alloc);
     rb_define_method(rb_cHnswlibInnerProductSpace, "initialize", RUBY_METHOD_FUNC(_hnsw_ipspace_init), 1);
     rb_define_method(rb_cHnswlibInnerProductSpace, "distance", RUBY_METHOD_FUNC(_hnsw_ipspace_distance), 2);
@@ -140,12 +145,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 +160,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 +185,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) {
@@ -202,21 +220,26 @@ public:
     return ptr;
   };
 
-  static VALUE define_class(VALUE rb_mHnswlib) {
-    rb_cHnswlibHierarchicalNSW = rb_define_class_under(rb_mHnswlib, "HierarchicalNSW", rb_cObject);
+  static VALUE define_class(VALUE outer) {
+    rb_cHnswlibHierarchicalNSW = rb_define_class_under(outer, "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, "initialize", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_initialize), -1);
+    rb_define_method(rb_cHnswlibHierarchicalNSW, "init_index", RUBY_METHOD_FUNC(_hnsw_hierarchicalnsw_init_index), -1);
+    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;
   };
@@ -224,54 +247,91 @@ public:
 private:
   static const rb_data_type_t hnsw_hierarchicalnsw_type;
 
-  static VALUE _hnsw_hierarchicalnsw_init(int argc, VALUE* argv, VALUE self) {
+  static VALUE _hnsw_hierarchicalnsw_initialize(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[2] = {rb_intern("space"), rb_intern("dim")};
+    VALUE kw_values[2] = {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, 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");
+    if (!RB_TYPE_P(kw_values[0], T_STRING)) {
+      rb_raise(rb_eTypeError, "expected space, String");
+      return Qnil;
+    }
+    if (strcmp(StringValueCStr(kw_values[0]), "l2") != 0 && strcmp(StringValueCStr(kw_values[0]), "ip") != 0 &&
+        strcmp(StringValueCStr(kw_values[0]), "cosine") != 0) {
+      rb_raise(rb_eArgError, "expected space, 'l2', 'ip', or 'cosine' only");
       return Qnil;
     }
     if (!RB_INTEGER_TYPE_P(kw_values[1])) {
+      rb_raise(rb_eTypeError, "expected dim, Integer");
+      return Qnil;
+    }
+
+    if (strcmp(StringValueCStr(kw_values[0]), "l2") == 0) {
+      rb_iv_set(self, "@space", rb_funcall(rb_const_get(rb_mHnswlib, rb_intern("L2Space")), rb_intern("new"), 1, kw_values[1]));
+    } else {
+      rb_iv_set(self, "@space",
+                rb_funcall(rb_const_get(rb_mHnswlib, rb_intern("InnerProductSpace")), rb_intern("new"), 1, kw_values[1]));
+    }
+
+    rb_iv_set(self, "@normalize", Qfalse);
+    if (strcmp(StringValueCStr(kw_values[0]), "cosine") == 0) rb_iv_set(self, "@normalize", Qtrue);
+
+    return Qnil;
+  };
+
+  static VALUE _hnsw_hierarchicalnsw_init_index(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[5] = {rb_intern("max_elements"), rb_intern("m"), rb_intern("ef_construction"), rb_intern("random_seed"),
+                      rb_intern("allow_replace_deleted")};
+    VALUE kw_values[5] = {Qundef, Qundef, Qundef, Qundef, Qundef};
+    rb_scan_args(argc, argv, ":", &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 4, kw_values);
+    if (kw_values[1] == Qundef) kw_values[1] = SIZET2NUM(16);
+    if (kw_values[2] == Qundef) kw_values[2] = SIZET2NUM(200);
+    if (kw_values[3] == Qundef) kw_values[3] = SIZET2NUM(100);
+    if (kw_values[4] == Qundef) kw_values[4] = Qfalse;
+
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
       rb_raise(rb_eTypeError, "expected max_elements, Integer");
       return Qnil;
     }
-    if (!RB_INTEGER_TYPE_P(kw_values[2])) {
+    if (!RB_INTEGER_TYPE_P(kw_values[1])) {
       rb_raise(rb_eTypeError, "expected m, Integer");
       return Qnil;
     }
-    if (!RB_INTEGER_TYPE_P(kw_values[3])) {
+    if (!RB_INTEGER_TYPE_P(kw_values[2])) {
       rb_raise(rb_eTypeError, "expected ef_construction, Integer");
       return Qnil;
     }
-    if (!RB_INTEGER_TYPE_P(kw_values[4])) {
+    if (!RB_INTEGER_TYPE_P(kw_values[3])) {
       rb_raise(rb_eTypeError, "expected random_seed, Integer");
       return Qnil;
     }
+    if (!RB_TYPE_P(kw_values[4], T_TRUE) && !RB_TYPE_P(kw_values[4], 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;
-    if (rb_obj_is_instance_of(kw_values[0], rb_cHnswlibL2Space)) {
-      space = RbHnswlibL2Space::get_hnsw_l2space(kw_values[0]);
+    hnswlib::SpaceInterface<float>* space = nullptr;
+    VALUE ivspace = rb_iv_get(self, "@space");
+    if (rb_obj_is_instance_of(ivspace, rb_cHnswlibL2Space)) {
+      space = RbHnswlibL2Space::get_hnsw_l2space(ivspace);
     } else {
-      space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(kw_values[0]);
+      space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(ivspace);
     }
-    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[0]);
+    const size_t m = NUM2SIZET(kw_values[1]);
+    const size_t ef_construction = NUM2SIZET(kw_values[2]);
+    const size_t random_seed = NUM2SIZET(kw_values[3]);
+    const bool allow_replace_deleted = kw_values[4] == 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);
+      ptr->~HierarchicalNSW();
+      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 +340,72 @@ 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};
 
-    if (!RB_TYPE_P(arr, T_ARRAY)) {
+    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;
+
+    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));
+    for (size_t i = 0; i < dim; i++) vec[i] = (float)NUM2DBL(rb_ary_entry(_arr, i));
+    const size_t idx = NUM2SIZET(_idx);
+    const bool replace_deleted = _replace_deleted == Qtrue ? true : false;
+
+    if (rb_iv_get(self, "@normalize") == Qtrue) {
+      float norm = 0.0;
+      for (size_t i = 0; i < dim; i++) norm += vec[i] * vec[i];
+      norm = std::sqrt(std::fabs(norm));
+      if (norm >= 0.0) {
+        for (size_t i = 0; i < dim; i++) vec[i] /= norm;
+      }
+    }
 
-    get_hnsw_hierarchicalnsw(self)->addPoint((void*)vec, (size_t)NUM2INT(idx));
+    try {
+      get_hnsw_hierarchicalnsw(self)->addPoint((void*)vec, idx, replace_deleted);
+    } catch (const std::runtime_error& e) {
+      ruby_xfree(vec);
+      rb_raise(rb_eRuntimeError, "%s", e.what());
+      return Qfalse;
+    }
 
     ruby_xfree(vec);
     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 +420,33 @@ 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));
     }
 
+    if (rb_iv_get(self, "@normalize") == Qtrue) {
+      float norm = 0.0;
+      for (size_t i = 0; i < dim; i++) norm += vec[i] * vec[i];
+      norm = std::sqrt(std::fabs(norm));
+      if (norm >= 0.0) {
+        for (size_t i = 0; i < dim; i++) vec[i] /= norm;
+      }
+    }
+
     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 +454,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 +466,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 +483,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 +512,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 +532,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 +548,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 +560,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 +585,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
@@ -495,12 +659,13 @@ public:
     return ptr;
   };
 
-  static VALUE define_class(VALUE rb_mHnswlib) {
-    rb_cHnswlibBruteforceSearch = rb_define_class_under(rb_mHnswlib, "BruteforceSearch", rb_cObject);
+  static VALUE define_class(VALUE outer) {
+    rb_cHnswlibBruteforceSearch = rb_define_class_under(outer, "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, "initialize", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_initialize), -1);
+    rb_define_method(rb_cHnswlibBruteforceSearch, "init_index", RUBY_METHOD_FUNC(_hnsw_bruteforcesearch_init_index), -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);
@@ -513,34 +678,66 @@ public:
 private:
   static const rb_data_type_t hnsw_bruteforcesearch_type;
 
-  static VALUE _hnsw_bruteforcesearch_init(int argc, VALUE* argv, VALUE self) {
+  static VALUE _hnsw_bruteforcesearch_initialize(int argc, VALUE* argv, VALUE self) {
     VALUE kw_args = Qnil;
-    ID kw_table[2] = {rb_intern("space"), rb_intern("max_elements")};
+    ID kw_table[2] = {rb_intern("space"), rb_intern("dim")};
     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");
+    if (!RB_TYPE_P(kw_values[0], T_STRING)) {
+      rb_raise(rb_eTypeError, "expected space, String");
+      return Qnil;
+    }
+    if (strcmp(StringValueCStr(kw_values[0]), "l2") != 0 && strcmp(StringValueCStr(kw_values[0]), "ip") != 0 &&
+        strcmp(StringValueCStr(kw_values[0]), "cosine") != 0) {
+      rb_raise(rb_eArgError, "expected space, 'l2', 'ip', or 'cosine' only");
       return Qnil;
     }
     if (!RB_INTEGER_TYPE_P(kw_values[1])) {
-      rb_raise(rb_eTypeError, "expected max_elements, Integer");
+      rb_raise(rb_eTypeError, "expected dim, 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]);
+    if (strcmp(StringValueCStr(kw_values[0]), "l2") == 0) {
+      rb_iv_set(self, "@space", rb_funcall(rb_const_get(rb_mHnswlib, rb_intern("L2Space")), rb_intern("new"), 1, kw_values[1]));
+    } else {
+      rb_iv_set(self, "@space",
+                rb_funcall(rb_const_get(rb_mHnswlib, rb_intern("InnerProductSpace")), rb_intern("new"), 1, kw_values[1]));
+    }
+
+    rb_iv_set(self, "@normalize", Qfalse);
+    if (strcmp(StringValueCStr(kw_values[0]), "cosine") == 0) rb_iv_set(self, "@normalize", Qtrue);
+
+    return Qnil;
+  };
+
+  static VALUE _hnsw_bruteforcesearch_init_index(int argc, VALUE* argv, VALUE self) {
+    VALUE kw_args = Qnil;
+    ID kw_table[1] = {rb_intern("max_elements")};
+    VALUE kw_values[1] = {Qundef};
+    rb_scan_args(argc, argv, ":", &kw_args);
+    rb_get_kwargs(kw_args, kw_table, 1, 0, kw_values);
+
+    if (!RB_INTEGER_TYPE_P(kw_values[0])) {
+      rb_raise(rb_eTypeError, "expected max_elements, Integer");
+      return Qnil;
+    }
+
+    hnswlib::SpaceInterface<float>* space = nullptr;
+    VALUE ivspace = rb_iv_get(self, "@space");
+    if (rb_obj_is_instance_of(ivspace, rb_cHnswlibL2Space)) {
+      space = RbHnswlibL2Space::get_hnsw_l2space(ivspace);
     } else {
-      space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(kw_values[0]);
+      space = RbHnswlibInnerProductSpace::get_hnsw_ipspace(ivspace);
     }
-    const size_t max_elements = (size_t)NUM2INT(kw_values[1]);
+
+    const size_t max_elements = NUM2SIZET(kw_values[0]);
 
     hnswlib::BruteforceSearch<float>* ptr = get_hnsw_bruteforcesearch(self);
     try {
+      ptr->~BruteforceSearch();
       new (ptr) hnswlib::BruteforceSearch<float>(space, max_elements);
     } catch (const std::runtime_error& e) {
       rb_raise(rb_eRuntimeError, "%s", e.what());
@@ -551,7 +748,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 +764,19 @@ 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));
+
+    if (rb_iv_get(self, "@normalize") == Qtrue) {
+      float norm = 0.0;
+      for (size_t i = 0; i < dim; i++) norm += vec[i] * vec[i];
+      norm = std::sqrt(std::fabs(norm));
+      if (norm >= 0.0) {
+        for (size_t i = 0; i < dim; i++) vec[i] /= norm;
+      }
+    }
 
     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 +787,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 +812,35 @@ 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++) {
-      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));
+
+    if (rb_iv_get(self, "@normalize") == Qtrue) {
+      float norm = 0.0;
+      for (size_t i = 0; i < dim; i++) norm += vec[i] * vec[i];
+      norm = std::sqrt(std::fabs(norm));
+      if (norm >= 0.0) {
+        for (size_t i = 0; i < dim; i++) vec[i] /= norm;
+      }
     }
 
     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 +850,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 +892,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);
   };
 };