Added arrays option for Psycopg 2

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.3.6 (unreleased)
+
+- Added `arrays` option for Psycopg 2
+
 ## 0.3.5 (2024-10-05)
 
 - Added `avg` function with type casting to SQLAlchemy

pgvector/psycopg2/halfvec.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import HalfVector
 
 
@@ -14,7 +14,12 @@ def cast_halfvec(value, cur):
     return HalfVector._from_db(value)
 
 
-def register_halfvec_info(oid, scope):
+def register_halfvec_info(oid, array_oid, scope):
     halfvec = new_type((oid,), 'HALFVEC', cast_halfvec)
     register_type(halfvec, scope)
+
+    if array_oid is not None:
+        halfvecarray = new_array_type((array_oid,), 'HALFVECARRAY', halfvec)
+        register_type(halfvecarray, scope)
+
     register_adapter(HalfVector, HalfvecAdapter)

pgvector/psycopg2/register.py

@@ -7,22 +7,23 @@
 
 # TODO make globally False by default in 0.4.0
 # note: register_adapter is always global
-def register_vector(conn_or_curs=None, globally=True):
+# TODO make arrays True by defalt in 0.4.0
+def register_vector(conn_or_curs=None, globally=True, arrays=False):
     conn = conn_or_curs if hasattr(conn_or_curs, 'cursor') else conn_or_curs.connection
     cur = conn.cursor(cursor_factory=cursor)
     scope = None if globally else conn_or_curs
 
     # use to_regtype to get first matching type in search path
-    cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")
+    cur.execute("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('_vector'), to_regtype('halfvec'), to_regtype('_halfvec'), to_regtype('sparsevec'), to_regtype('_sparsevec'))")
     type_info = dict(cur.fetchall())
 
     if 'vector' not in type_info:
         raise psycopg2.ProgrammingError('vector type not found in the database')
 
-    register_vector_info(type_info['vector'], scope)
+    register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope)
 
     if 'halfvec' in type_info:
-        register_halfvec_info(type_info['halfvec'], scope)
+        register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope)
 
     if 'sparsevec' in type_info:
-        register_sparsevec_info(type_info['sparsevec'], scope)
+        register_sparsevec_info(type_info['sparsevec'], type_info['_sparsevec'] if arrays else None, scope)

pgvector/psycopg2/sparsevec.py

@@ -1,4 +1,4 @@
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import SparseVector
 
 
@@ -14,7 +14,12 @@ def cast_sparsevec(value, cur):
     return SparseVector._from_db(value)
 
 
-def register_sparsevec_info(oid, scope):
+def register_sparsevec_info(oid, array_oid, scope):
     sparsevec = new_type((oid,), 'SPARSEVEC', cast_sparsevec)
     register_type(sparsevec, scope)
+
+    if array_oid is not None:
+        sparsevecarray = new_array_type((array_oid,), 'SPARSEVECARRAY', sparsevec)
+        register_type(sparsevecarray, scope)
+
     register_adapter(SparseVector, SparsevecAdapter)

pgvector/psycopg2/vector.py

@@ -1,5 +1,5 @@
 import numpy as np
-from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from psycopg2.extensions import adapt, new_array_type, new_type, register_adapter, register_type
 from ..utils import Vector
 
 
@@ -15,7 +15,12 @@ def cast_vector(value, cur):
     return Vector._from_db(value)
 
 
-def register_vector_info(oid, scope):
+def register_vector_info(oid, array_oid, scope):
     vector = new_type((oid,), 'VECTOR', cast_vector)
     register_type(vector, scope)
+
+    if array_oid is not None:
+        vectorarray = new_array_type((array_oid,), 'VECTORARRAY', vector)
+        register_type(vectorarray, scope)
+
     register_adapter(np.ndarray, VectorAdapter)

tests/test_psycopg2.py

@@ -1,5 +1,5 @@
 import numpy as np
-from pgvector.psycopg2 import register_vector, SparseVector
+from pgvector.psycopg2 import register_vector, HalfVector, SparseVector
 import psycopg2
 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
 
@@ -9,9 +9,9 @@
 cur = conn.cursor()
 cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
 cur.execute('DROP TABLE IF EXISTS psycopg2_items')
-cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3))')
+cur.execute('CREATE TABLE psycopg2_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(3), sparse_embedding sparsevec(3), embeddings vector[], half_embeddings halfvec[], sparse_embeddings sparsevec[])')
 
-register_vector(cur, globally=False)
+register_vector(cur, globally=False, arrays=True)
 
 
 class TestPsycopg2:
@@ -55,6 +55,33 @@ def test_sparsevec(self):
         assert res[0][0].to_list() == [1.5, 2, 3]
         assert res[1][0] is None
 
+    def test_vector_array(self):
+        embeddings = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (embeddings) VALUES (%s::vector[])', (embeddings,))
+
+        cur.execute('SELECT embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert np.array_equal(res[0][0], embeddings[0])
+        assert np.array_equal(res[0][1], embeddings[1])
+
+    def test_halfvec_array(self):
+        embeddings = [HalfVector([1.5, 2, 3]), HalfVector([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (half_embeddings) VALUES (%s::halfvec[])', (embeddings,))
+
+        cur.execute('SELECT half_embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert res[0][0].to_list() == [1.5, 2, 3]
+        assert res[0][1].to_list() == [4.5, 5, 6]
+
+    def test_sparsevec_array(self):
+        embeddings = [SparseVector([1.5, 2, 3]), SparseVector([4.5, 5, 6])]
+        cur.execute('INSERT INTO psycopg2_items (sparse_embeddings) VALUES (%s::sparsevec[])', (embeddings,))
+
+        cur.execute('SELECT sparse_embeddings FROM psycopg2_items ORDER BY id')
+        res = cur.fetchone()
+        assert res[0][0].to_list() == [1.5, 2, 3]
+        assert res[0][1].to_list() == [4.5, 5, 6]
+
     def test_cursor_factory(self):
         for cursor_factory in [DictCursor, RealDictCursor, NamedTupleCursor]:
             conn = psycopg2.connect(dbname='pgvector_python_test')