faiss 0.5.2 → 0.5.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8ab397b95262fa87d28a4f08098de35f487b44e6240682ec8e98c4efc823c4df
-  data.tar.gz: 719490d3f20bad4531f588e9469d3fa9ec2362cc18f800e913725ffb6b0aeca0
+  metadata.gz: 46b01025b31b0b052236ffb42c9d046de719f2c6ffbdff50299e51f339def739
+  data.tar.gz: 1d50d77f839ad49dd6827776dca2f7eeea25eb333ccac6d86077a4366a085065
 SHA512:
-  metadata.gz: 56e16fad10a1ba95c9727f2e356470e984ca29b44406cf125b1ab529c83fee2c9bf15efa564a543cda331489dbdc56c9c44f4eda4ad53178979370987b6fbf9e
-  data.tar.gz: 691e9a843eaac7e699dafc15d5505f6febf783ae6d474ee0409cb01d1b1f1575477db13cb7290a4d2cc639692503099fa7d4ff767d3c3034838d05ea0476521a
+  metadata.gz: ffd3ca8a5726dacc053da5c489f19212399c780fd41048793ebd9785716e934ac3e158a5a1fa09c11723fa12ab172d65166bb12216f245dd01864cef83b4608f
+  data.tar.gz: c485d9b425913d52f3a53c7e229dffa04fe1535de151b2abb5ac2e09a56d4ba4cfcfd8838846eab5b9a18c898115f56fa107640c210001b69ea8fd8a7eac8428

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## 0.5.3 (2026-02-12)
+
+- Added `IndexIDMap2` class
+- Added `remove_ids` method
+- Added `reconstruct_batch` and `reconstruct_n` methods
+- Fixed error when searching frozen indexes
+- Fixed memory leak with frozen exceptions
+
 ## 0.5.2 (2026-01-02)
 
 - Fixed error with Rice 4.8

data/LICENSE.txt

@@ -1,7 +1,7 @@
 MIT License
 
 Copyright (c) Facebook, Inc. and its affiliates.
-Copyright (c) 2020-2025 Andrew Kane
+Copyright (c) 2020-2026 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

data/ext/faiss/index.cpp

@@ -1,4 +1,6 @@
+#include <algorithm>
 #include <string>
+#include <vector>
 
 #include <faiss/AutoTune.h>
 #include <faiss/Index.h>
@@ -108,7 +110,7 @@ void init_index(Rice::Module& m) {
     .define_method(
       "train",
       [](Rice::Object rb_self, numo::SFloat objects) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::Index>{rb_self};
         auto n = check_shape(objects, self.d);
@@ -117,7 +119,7 @@ void init_index(Rice::Module& m) {
     .define_method(
       "add",
       [](Rice::Object rb_self, numo::SFloat objects) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::Index>{rb_self};
         auto n = check_shape(objects, self.d);
@@ -126,7 +128,7 @@ void init_index(Rice::Module& m) {
     .define_method(
       "add_with_ids",
       [](Rice::Object rb_self, numo::SFloat objects, numo::Int64 ids) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::Index>{rb_self};
         auto n = check_shape(objects, self.d);
@@ -135,6 +137,19 @@ void init_index(Rice::Module& m) {
         }
         self.add_with_ids(n, objects.read_ptr(), ids.read_ptr());
       })
+    .define_method(
+      "remove_ids",
+      [](Rice::Object rb_self, numo::Int64 ids) {
+        check_frozen(rb_self);
+
+        auto &self = *Rice::Data_Object<faiss::Index>{rb_self};
+        if (ids.ndim() != 1) {
+          throw Rice::Exception(rb_eArgError, "expected ids to be 1d array");
+        }
+        auto n = ids.shape()[0];
+        faiss::IDSelectorBatch sel(n, ids.read_ptr());
+        return self.remove_ids(sel);
+      })
     .define_method(
       "search",
       [](Rice::Object rb_self, numo::SFloat objects, size_t k) {
@@ -145,9 +160,18 @@ void init_index(Rice::Module& m) {
         auto labels = numo::Int64({n, k});
 
         if (rb_self.is_frozen()) {
+          // Don't mess with Ruby-owned memory while the GVL is released
+          auto objects_ptr = objects.read_ptr();
+          std::vector<float> objects_vec(objects_ptr, objects_ptr + n * self.d);
+          std::vector<float> distances_vec(n * k);
+          std::vector<int64_t> labels_vec(n * k);
+
           Rice::detail::no_gvl([&] {
-            self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
+            self.search(n, objects_vec.data(), k, distances_vec.data(), labels_vec.data());
           });
+
+          std::copy(distances_vec.begin(), distances_vec.end(), distances.write_ptr());
+          std::copy(labels_vec.begin(), labels_vec.end(), labels.write_ptr());
         } else {
           self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
         }
@@ -160,7 +184,7 @@ void init_index(Rice::Module& m) {
     .define_method(
       "nprobe=",
       [](Rice::Object rb_self, double val) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::Index>{rb_self};
         faiss::ParameterSpace().set_index_parameter(&self, "nprobe", val);
@@ -173,6 +197,34 @@ void init_index(Rice::Module& m) {
         self.reconstruct(key, recons.write_ptr());
         return recons;
       })
+    .define_method(
+      "reconstruct_batch",
+      [](faiss::Index &self, numo::Int64 ids) {
+        if (ids.ndim() != 1) {
+          throw Rice::Exception(rb_eArgError, "expected ids to be 1d array");
+        }
+        auto n = static_cast<std::size_t>(ids.shape()[0]);
+        auto d = static_cast<std::size_t>(self.d);
+        auto recons = numo::SFloat({n, d});
+        self.reconstruct_batch(n, ids.read_ptr(), recons.write_ptr());
+        return recons;
+      })
+    .define_method(
+      "reconstruct_n",
+      [](faiss::Index &self, int64_t i0, int64_t ni) {
+        if (ni < 0) {
+          throw Rice::Exception(rb_eArgError, "expected n to be non-negative");
+        }
+        // second condition written this way to avoid overflow
+        if (i0 < 0 || i0 > self.ntotal - ni) {
+          throw Rice::Exception(rb_eIndexError, "index out of range");
+        }
+        auto d = static_cast<std::size_t>(self.d);
+        auto n = static_cast<std::size_t>(ni);
+        auto recons = numo::SFloat({n, d});
+        self.reconstruct_n(i0, ni, recons.write_ptr());
+        return recons;
+      })
     .define_method(
       "save",
       [](faiss::Index &self, Rice::String fname) {
@@ -227,4 +279,7 @@ void init_index(Rice::Module& m) {
 
   Rice::define_class_under<faiss::IndexIDMap, faiss::Index>(m, "IndexIDMap")
     .define_constructor(Rice::Constructor<faiss::IndexIDMap, faiss::Index*>());
+
+  Rice::define_class_under<faiss::IndexIDMap2, faiss::Index>(m, "IndexIDMap2")
+    .define_constructor(Rice::Constructor<faiss::IndexIDMap2, faiss::Index*>());
 }

data/ext/faiss/index_binary.cpp

@@ -1,3 +1,6 @@
+#include <algorithm>
+#include <vector>
+
 #include <faiss/IndexBinary.h>
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexBinaryIVF.h>
@@ -28,7 +31,7 @@ void init_index_binary(Rice::Module& m) {
     .define_method(
       "train",
       [](Rice::Object rb_self, numo::UInt8 objects) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::IndexBinary>{rb_self};
         auto n = check_shape(objects, self.d / 8);
@@ -37,12 +40,25 @@ void init_index_binary(Rice::Module& m) {
     .define_method(
       "add",
       [](Rice::Object rb_self, numo::UInt8 objects) {
-        rb_check_frozen(rb_self.value());
+        check_frozen(rb_self);
 
         auto &self = *Rice::Data_Object<faiss::IndexBinary>{rb_self};
         auto n = check_shape(objects, self.d / 8);
         self.add(n, objects.read_ptr());
       })
+    .define_method(
+      "remove_ids",
+      [](Rice::Object rb_self, numo::Int64 ids) {
+        check_frozen(rb_self);
+
+        auto &self = *Rice::Data_Object<faiss::IndexBinary>{rb_self};
+        if (ids.ndim() != 1) {
+          throw Rice::Exception(rb_eArgError, "expected ids to be 1d array");
+        }
+        auto n = ids.shape()[0];
+        faiss::IDSelectorBatch sel(n, ids.read_ptr());
+        return self.remove_ids(sel);
+      })
     .define_method(
       "search",
       [](Rice::Object rb_self, numo::UInt8 objects, size_t k) {
@@ -53,9 +69,18 @@ void init_index_binary(Rice::Module& m) {
         auto labels = numo::Int64({n, k});
 
         if (rb_self.is_frozen()) {
+          // Don't mess with Ruby-owned memory while the GVL is released
+          auto objects_ptr = objects.read_ptr();
+          std::vector<uint8_t> objects_vec(objects_ptr, objects_ptr + n * (self.d / 8));
+          std::vector<int32_t> distances_vec(n * k);
+          std::vector<int64_t> labels_vec(n * k);
+
           Rice::detail::no_gvl([&] {
-            self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
+            self.search(n, objects_vec.data(), k, distances_vec.data(), labels_vec.data());
           });
+
+          std::copy(distances_vec.begin(), distances_vec.end(), distances.write_ptr());
+          std::copy(labels_vec.begin(), labels_vec.end(), labels.write_ptr());
         } else {
           self.search(n, objects.read_ptr(), k, distances.write_ptr(), labels.write_ptr());
         }
@@ -65,6 +90,30 @@ void init_index_binary(Rice::Module& m) {
         ret.push(std::move(labels), false);
         return ret;
       })
+    .define_method(
+      "reconstruct",
+      [](faiss::IndexBinary &self, int64_t key) {
+        auto d = static_cast<std::size_t>(self.d / 8);
+        auto recons = numo::UInt8({d});
+        self.reconstruct(key, recons.write_ptr());
+        return recons;
+      })
+    .define_method(
+      "reconstruct_n",
+      [](faiss::IndexBinary &self, int64_t i0, int64_t ni) {
+        if (ni < 0) {
+          throw Rice::Exception(rb_eArgError, "expected n to be non-negative");
+        }
+        // second condition written this way to avoid overflow
+        if (i0 < 0 || i0 > self.ntotal - ni) {
+          throw Rice::Exception(rb_eIndexError, "index out of range");
+        }
+        auto d = static_cast<std::size_t>(self.d / 8);
+        auto n = static_cast<std::size_t>(ni);
+        auto recons = numo::UInt8({n, d});
+        self.reconstruct_n(i0, ni, recons.write_ptr());
+        return recons;
+      })
     .define_method(
       "save",
       [](faiss::IndexBinary &self, Rice::String fname) {

data/ext/faiss/kmeans.cpp

@@ -1,3 +1,5 @@
+#include <algorithm>
+
 #include <faiss/Clustering.h>
 #include <faiss/IndexFlat.h>
 #include <rice/rice.hpp>
@@ -27,7 +29,7 @@ void init_kmeans(Rice::Module& m) {
       "centroids",
       [](faiss::Clustering &self) {
         auto centroids = numo::SFloat({self.k, self.d});
-        memcpy(centroids.write_ptr(), self.centroids.data(), self.centroids.size() * sizeof(float));
+        std::copy(self.centroids.begin(), self.centroids.end(), centroids.write_ptr());
         return centroids;
       })
     .define_method(

data/ext/faiss/numo.hpp

@@ -1,5 +1,5 @@
 /*!
- * Numo.hpp v0.3.0
+ * Numo.hpp v0.3.1
  * https://github.com/ankane/numo.hpp
  * BSD-2-Clause License
  */
@@ -38,7 +38,7 @@ public:
   }
 
   bool is_contiguous() const {
-    return nary_check_contiguous(this->_value) == Qtrue;
+    return Rice::detail::protect(nary_check_contiguous, this->_value) == Qtrue;
   }
 
   operator Rice::Object() const {
@@ -46,26 +46,28 @@ public:
   }
 
   const void* read_ptr() {
-    if (!is_contiguous()) {
-      this->_value = nary_dup(this->_value);
-    }
-    return nary_get_pointer_for_read(this->_value) + nary_get_offset(this->_value);
+    return Rice::detail::protect([&]() {
+      if (!nary_check_contiguous(this->_value)) {
+        this->_value = nary_dup(this->_value);
+      }
+      return nary_get_pointer_for_read(this->_value) + nary_get_offset(this->_value);
+    });
   }
 
   void* write_ptr() {
-    return nary_get_pointer_for_write(this->_value);
+    return Rice::detail::protect(nary_get_pointer_for_write, this->_value);
   }
 
 protected:
   NArray() { }
 
   void construct_value(VALUE dtype, VALUE v) {
-    this->_value = rb_funcall(dtype, rb_intern("cast"), 1, v);
+    this->_value = Rice::detail::protect(rb_funcall, dtype, rb_intern("cast"), 1, v);
   }
 
   void construct_shape(VALUE dtype, std::initializer_list<size_t> shape) {
     // rb_narray_new doesn't modify shape, but not marked as const
-    this->_value = rb_narray_new(dtype, shape.size(), const_cast<size_t*>(shape.begin()));
+    this->_value = Rice::detail::protect(rb_narray_new, dtype, shape.size(), const_cast<size_t*>(shape.begin()));
   }
 
   VALUE _value;

data/ext/faiss/utils.cpp

@@ -14,3 +14,8 @@ size_t check_shape(const numo::NArray& objects, size_t k) {
   }
   return shape[0];
 }
+
+void check_frozen(Rice::Object obj) {
+  // use protect to avoid skipping destructors and leaking memory
+  Rice::detail::protect(rb_check_frozen, obj.value());
+}

data/ext/faiss/utils.h

@@ -3,3 +3,5 @@
 #include "numo.hpp"
 
 size_t check_shape(const numo::NArray& objects, size_t k);
+
+void check_frozen(Rice::Object obj);

data/lib/faiss/version.rb

@@ -1,3 +1,3 @@
 module Faiss
-  VERSION = "0.5.2"
+  VERSION = "0.5.3"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: faiss
 version: !ruby/object:Gem::Version
-  version: 0.5.2
+  version: 0.5.3
 platform: ruby
 authors:
 - Andrew Kane