Improved error message for negative k

Author: ankane

ext/faiss/index_binary.cpp

@@ -62,19 +62,22 @@ void init_index_binary(Rice::Module& m) {
       })
     .define_method(
       "search",
-      [](Rice::Object rb_self, numo::UInt8 objects, size_t k) {
+      [](Rice::Object rb_self, numo::UInt8 objects, int64_t k) {
         auto& self = *Rice::Data_Object<faiss::IndexBinary>{rb_self};
         size_t n = check_shape(objects, self.d / 8);
+        if (k < 0) {
+          throw Rice::Exception(rb_eArgError, "expected k to be non-negative");
+        }
 
-        numo::Int32 distances{{n, k}};
-        numo::Int64 labels{{n, k}};
+        numo::Int32 distances{{n, static_cast<size_t>(k)}};
+        numo::Int64 labels{{n, static_cast<size_t>(k)}};
 
         if (rb_self.is_frozen()) {
           // Don't mess with Ruby-owned memory while the GVL is released
           const 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);
+          std::vector<int32_t> distances_vec(n * static_cast<size_t>(k));
+          std::vector<int64_t> labels_vec(n * static_cast<size_t>(k));
 
           Rice::detail::no_gvl([&] {
             self.search(n, objects_vec.data(), k, distances_vec.data(), labels_vec.data());

ext/faiss/index_rb.cpp

@@ -153,19 +153,22 @@ void init_index(Rice::Module& m) {
       })
     .define_method(
       "search",
-      [](Rice::Object rb_self, numo::SFloat objects, size_t k) {
+      [](Rice::Object rb_self, numo::SFloat objects, int64_t k) {
         auto& self = *Rice::Data_Object<faiss::Index>{rb_self};
         size_t n = check_shape(objects, self.d);
+        if (k < 0) {
+          throw Rice::Exception(rb_eArgError, "expected k to be non-negative");
+        }
 
-        numo::SFloat distances{{n, k}};
-        numo::Int64 labels{{n, k}};
+        numo::SFloat distances{{n, static_cast<size_t>(k)}};
+        numo::Int64 labels{{n, static_cast<size_t>(k)}};
 
         if (rb_self.is_frozen()) {
           // Don't mess with Ruby-owned memory while the GVL is released
           const 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);
+          std::vector<float> distances_vec(n * static_cast<size_t>(k));
+          std::vector<int64_t> labels_vec(n * static_cast<size_t>(k));
 
           Rice::detail::no_gvl([&] {
             self.search(n, objects_vec.data(), k, distances_vec.data(), labels_vec.data());

test/index_binary_test.rb

@@ -56,6 +56,21 @@ def test_remove_ids
     assert_equal [1, 0, -1], ids[0, true].to_a
   end
 
+  def test_search_negative_k
+    objects = [
+      [1, 1, 2, 1],
+      [5, 4, 6, 5],
+      [1, 2, 1, 2]
+    ]
+    index = Faiss::IndexBinaryFlat.new(32)
+    index.add(objects)
+
+    error = assert_raises(ArgumentError) do
+      index.search(objects, -1)
+    end
+    assert_equal "expected k to be non-negative", error.message
+  end
+
   def test_add_frozen
     index = Faiss::IndexBinaryFlat.new(32)
     index.freeze

test/index_test.rb

@@ -336,6 +336,21 @@ def test_remove_ids_empty
     assert_equal 0, index.remove_ids([])
   end
 
+  def test_search_negative_k
+    objects = [
+      [1, 1, 2, 1],
+      [5, 4, 6, 5],
+      [1, 2, 1, 2]
+    ]
+    index = Faiss::IndexFlatL2.new(4)
+    index.add(objects)
+
+    error = assert_raises(ArgumentError) do
+      index.search(objects, -1)
+    end
+    assert_equal "expected k to be non-negative", error.message
+  end
+
   def test_add_frozen
     index = Faiss::IndexFlatL2.new(4)
     index.freeze