Added support for halfvec

ankane · Dec 31, 2024 · 4e12728 · 4e12728
1 parent e4a764b
commit 4e12728
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/lib/neighbor/postgresql.rb b/lib/neighbor/postgresql.rb
@@ -46,7 +46,7 @@ def initialize_type_map(m = type_map)
     module ArrayMethods
       def type_cast_array(value, method)
         # TODO support halfvec
-        if subtype.is_a?(Neighbor::Type::Vector) && method != :deserialize && value.is_a?(::Array) && value.all? { |v| v.is_a?(::Numeric) }
+        if (subtype.is_a?(Neighbor::Type::Vector) || subtype.is_a?(Neighbor::Type::Halfvec)) && method != :deserialize && value.is_a?(::Array) && value.all? { |v| v.is_a?(::Numeric) }
           @subtype.public_send(method, value)
         else
           super

diff --git a/test/halfvec_test.rb b/test/halfvec_test.rb
@@ -61,4 +61,16 @@ def test_nan
     end
     assert_equal "Validation failed: Half embedding must have finite values", error.message
   end
+
+  def test_array
+    item = Item.create!(half_embeddings: [[1, 2, 3], [4, 5, 6]])
+    assert_equal [[1, 2, 3], [4, 5, 6]], item.half_embeddings
+    assert_equal [[1, 2, 3], [4, 5, 6]], Item.last.half_embeddings
+  end
+
+  def test_array_2d
+    item = Item.create!(half_embeddings: [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]])
+    assert_equal [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], item.half_embeddings
+    assert_equal [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], Item.last.half_embeddings
+  end
 end
diff --git a/test/support/postgresql.rb b/test/support/postgresql.rb
@@ -19,6 +19,7 @@ class PostgresRecord < ActiveRecord::Base
     t.sparsevec :sparse_embedding, limit: 3
     t.sparsevec :sparse_factors, limit: 5
     t.vector :embeddings, limit: 3, array: true
+    t.halfvec :half_embeddings, limit: 3, array: true
   end
   add_index :items, :cube_embedding, using: :gist
   add_index :items, :embedding, using: :hnsw, opclass: :vector_cosine_ops