diff --git a/pgvector/pg8000/register.py b/pgvector/pg8000/register.py index f3978d7..efad043 100644 --- a/pgvector/pg8000/register.py +++ b/pgvector/pg8000/register.py @@ -1,6 +1,6 @@ from pg8000.native import Connection from typing import cast -from .. import Vector, HalfVector, SparseVector +from .. import Bit, Vector, HalfVector, SparseVector def register_vector(conn: Connection, /) -> None: @@ -14,6 +14,8 @@ def register_vector(conn: Connection, /) -> None: conn.register_out_adapter(Vector, lambda v: v.to_text()) conn.register_in_adapter(type_info['vector'], Vector.from_text) + conn.register_out_adapter(Bit, lambda v: v.to_text()) + try: import numpy as np conn.register_out_adapter(np.ndarray, lambda v: Vector(v).to_text()) diff --git a/pgvector/psycopg2/bit.py b/pgvector/psycopg2/bit.py new file mode 100644 index 0000000..75f1fb0 --- /dev/null +++ b/pgvector/psycopg2/bit.py @@ -0,0 +1,15 @@ +from psycopg2.extensions import adapt, register_adapter +from typing import Any +from .. import Bit + + +class BitAdapter: + def __init__(self, value: Bit) -> None: + self._value = value + + def getquoted(self) -> Any: + return adapt(self._value.to_text()).getquoted() + + +def register_bit_info() -> None: + register_adapter(Bit, BitAdapter) diff --git a/pgvector/psycopg2/register.py b/pgvector/psycopg2/register.py index 2a5af49..958371d 100644 --- a/pgvector/psycopg2/register.py +++ b/pgvector/psycopg2/register.py @@ -1,5 +1,6 @@ import psycopg2 from psycopg2.extensions import connection, cursor +from .bit import register_bit_info from .halfvec import register_halfvec_info from .sparsevec import register_sparsevec_info from .vector import register_vector_info @@ -20,6 +21,8 @@ def register_vector(conn_or_curs: connection | cursor, /, *, globally: bool = Fa register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope) + register_bit_info() + if 'halfvec' in type_info: register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope) diff --git a/tests/test_pg8000.py b/tests/test_pg8000.py index 8c8e02c..66da911 100644 --- a/tests/test_pg8000.py +++ b/tests/test_pg8000.py @@ -1,5 +1,5 @@ from getpass import getuser -from pgvector import HalfVector, SparseVector, Vector +from pgvector import Bit, HalfVector, SparseVector, Vector from pgvector.pg8000 import register_vector from pg8000.native import Connection from .conftest import numpy as np @@ -42,6 +42,14 @@ def test_bit(self) -> None: res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') assert res == [['101'], [None]] + def test_bit_class(self) -> None: + embedding = Bit([True, False, True]) + embedding2 = None + conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2) + + res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id') + assert res == [['101'], [None]] + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) embedding2 = None diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py index badb7c2..008c505 100644 --- a/tests/test_psycopg2.py +++ b/tests/test_psycopg2.py @@ -1,4 +1,4 @@ -from pgvector import HalfVector, SparseVector, Vector +from pgvector import Bit, HalfVector, SparseVector, Vector from pgvector.psycopg2 import register_vector import psycopg2 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor @@ -61,6 +61,15 @@ def test_bit(self) -> None: assert res[0][0] == embedding assert res[1][0] is None + def test_bit_class(self) -> None: + embedding = Bit([True, False, True]) + cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,)) + + cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id') + res = cur.fetchall() + assert res[0][0] == '101' + assert res[1][0] is None + def test_sparsevec(self) -> None: embedding = SparseVector([1.5, 2, 3]) cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,))