Improved support for Postgres arrays

ankane · Dec 31, 2024 · 1648762 · 1648762
1 parent 884db5e
commit 1648762
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.5.2 (unreleased)
+
+- Improved support for Postgres arrays
+
 ## 0.5.1 (2024-12-03)
 
 - Added experimental support for MariaDB 11.7

diff --git a/lib/neighbor/postgresql.rb b/lib/neighbor/postgresql.rb
@@ -19,6 +19,9 @@ def self.initialize!
 
       # prevent unknown OID warning
       ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(RegisterTypes)
+
+      # patch arrays
+      ActiveRecord::ConnectionAdapters::PostgreSQL::OID::Array.prepend(ArrayMethods)
     end
 
     module RegisterTypes
@@ -39,5 +42,16 @@ def initialize_type_map(m = type_map)
         end
       end
     end
+
+    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) }
+          @subtype.public_send(method, value)
+        else
+          super
+        end
+      end
+    end
   end
 end
diff --git a/test/vector_test.rb b/test/vector_test.rb
@@ -99,8 +99,8 @@ def test_nan
   end
 
   def test_array
-    Item.connection.execute("INSERT INTO items (embeddings) VALUES (ARRAY['[1,2,3]', '[4,5,6]']::vector[])")
-    item = Item.last
+    item = Item.create!(embeddings: [[1, 2, 3], [4, 5, 6]])
     assert_equal [[1, 2, 3], [4, 5, 6]], item.embeddings
+    assert_equal [[1, 2, 3], [4, 5, 6]], Item.last.embeddings
   end
 end