Added support for Vector class to Psycopg 2 [skip ci]

Author: ankane

pgvector/psycopg2/__init__.py

@@ -1,8 +1,9 @@
 from .register import register_vector
-from ..utils import HalfVector, SparseVector
+from ..utils import HalfVector, SparseVector, Vector
 
 __all__ = [
     'register_vector',
+    'Vector',
     'HalfVector',
     'SparseVector'
 ]

pgvector/psycopg2/vector.py

@@ -24,3 +24,4 @@ def register_vector_info(oid, array_oid, scope):
         register_type(vectorarray, scope)
 
     register_adapter(np.ndarray, VectorAdapter)
+    register_adapter(Vector, VectorAdapter)

tests/test_psycopg2.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.psycopg2 import register_vector, HalfVector, SparseVector
+from pgvector.psycopg2 import register_vector, Vector, HalfVector, SparseVector
 import psycopg2
 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
 from psycopg2.pool import ThreadedConnectionPool
@@ -29,6 +29,16 @@ def test_vector(self):
         assert res[0][0].dtype == np.float32
         assert res[1][0] is None
 
+    def test_vector_class(self):
+        embedding = Vector([1.5, 2, 3])
+        cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,))
+
+        cur.execute('SELECT embedding FROM psycopg2_items ORDER BY id')
+        res = cur.fetchall()
+        assert np.array_equal(res[0][0], embedding.to_numpy())
+        assert res[0][0].dtype == np.float32
+        assert res[1][0] is None
+
     def test_halfvec(self):
         embedding = [1.5, 2, 3]
         cur.execute('INSERT INTO psycopg2_items (half_embedding) VALUES (%s), (NULL)', (embedding,))