diff options
author | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
---|---|---|
committer | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
commit | 4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch) | |
tree | ee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py | |
parent | cc961e04ba734dd72309fb548a2f97d67d578813 (diff) | |
download | gn-ai-4a52a71956a8d46fcb7294ac71734504bb09bcc2.tar.gz |
Diffstat (limited to '.venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py b/.venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py new file mode 100644 index 00000000..0f62ca98 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/pgvector/psycopg/vector.py @@ -0,0 +1,58 @@ +import psycopg +from psycopg.adapt import Loader, Dumper +from psycopg.pq import Format +from ..utils import Vector + + +class VectorDumper(Dumper): + + format = Format.TEXT + + def dump(self, obj): + return Vector._to_db(obj).encode('utf8') + + +class VectorBinaryDumper(VectorDumper): + + format = Format.BINARY + + def dump(self, obj): + return Vector._to_db_binary(obj) + + +class VectorLoader(Loader): + + format = Format.TEXT + + def load(self, data): + if isinstance(data, memoryview): + data = bytes(data) + return Vector._from_db(data.decode('utf8')) + + +class VectorBinaryLoader(VectorLoader): + + format = Format.BINARY + + def load(self, data): + if isinstance(data, memoryview): + data = bytes(data) + return Vector._from_db_binary(data) + + +def register_vector_info(context, info): + if info is None: + raise psycopg.ProgrammingError('vector type not found in the database') + info.register(context) + + # add oid to anonymous class for set_types + text_dumper = type('', (VectorDumper,), {'oid': info.oid}) + binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid}) + + adapters = context.adapters + adapters.register_dumper('numpy.ndarray', text_dumper) + adapters.register_dumper('numpy.ndarray', binary_dumper) + adapters.register_dumper(Vector, text_dumper) + adapters.register_dumper(Vector, binary_dumper) + adapters.register_loader(info.oid, VectorLoader) + adapters.register_loader(info.oid, VectorBinaryLoader) |