Added support for pg8000

Author: ankane

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.4.0 (unreleased)
 
 - Added top-level `pgvector` package
+- Added support for pg8000
 - Changed `globally` option to default to `False` for Psycopg 2
 - Changed `arrays` option to default to `True` for Psycopg 2
 - Fixed equality for `Vector`, `HalfVector`, `Bit`, and `SparseVector` classes

README.md

@@ -2,7 +2,7 @@
 
 [pgvector](https://github.com/pgvector/pgvector) support for Python
 
-Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), and [Peewee](https://github.com/coleifer/peewee)
+Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [SQLModel](https://github.com/tiangolo/sqlmodel), [Psycopg 3](https://github.com/psycopg/psycopg), [Psycopg 2](https://github.com/psycopg/psycopg2), [asyncpg](https://github.com/MagicStack/asyncpg), [pg8000](https://github.com/tlocke/pg8000), and [Peewee](https://github.com/coleifer/peewee)
 
 [![Build Status](https://github.com/pgvector/pgvector-python/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-python/actions)
 
@@ -22,6 +22,7 @@ And follow the instructions for your database library:
 - [Psycopg 3](#psycopg-3)
 - [Psycopg 2](#psycopg-2)
 - [asyncpg](#asyncpg)
+- [pg8000](#pg8000) [unreleased]
 - [Peewee](#peewee)
 
 Or check out some examples:
@@ -562,6 +563,51 @@ await conn.execute('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops
 
 Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
 
+## pg8000
+
+Enable the extension
+
+```python
+conn.run('CREATE EXTENSION IF NOT EXISTS vector')
+```
+
+Register the vector type with your connection
+
+```python
+from pgvector.pg8000 import register_vector
+
+register_vector(conn)
+```
+
+Create a table
+
+```python
+conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
+```
+
+Insert a vector
+
+```python
+embedding = np.array([1, 2, 3])
+conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
+```
+
+Get the nearest neighbors to a vector
+
+```python
+conn.run('SELECT * FROM items ORDER BY embedding <-> :embedding LIMIT 5', embedding=embedding)
+```
+
+Add an approximate index
+
+```python
+conn.run('CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)')
+# or
+conn.run('CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)')
+```
+
+Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
+
 ## Peewee
 
 Add a vector column

pgvector/pg8000/__init__.py

@@ -0,0 +1,5 @@
+from .register import register_vector
+
+__all__ = [
+    'register_vector'
+]

pgvector/pg8000/register.py

@@ -0,0 +1,23 @@
+import numpy as np
+from .. import Vector, HalfVector, SparseVector
+
+
+def register_vector(conn):
+    # use to_regtype to get first matching type in search path
+    res = conn.run("SELECT typname, oid FROM pg_type WHERE oid IN (to_regtype('vector'), to_regtype('halfvec'), to_regtype('sparsevec'))")
+    type_info = dict(res)
+
+    if 'vector' not in type_info:
+        raise RuntimeError('vector type not found in the database')
+
+    conn.register_out_adapter(Vector, Vector._to_db)
+    conn.register_out_adapter(np.ndarray, Vector._to_db)
+    conn.register_in_adapter(type_info['vector'], Vector._from_db)
+
+    if 'halfvec' in type_info:
+        conn.register_out_adapter(HalfVector, HalfVector._to_db)
+        conn.register_in_adapter(type_info['halfvec'], HalfVector._from_db)
+
+    if 'sparsevec' in type_info:
+        conn.register_out_adapter(SparseVector, SparseVector._to_db)
+        conn.register_in_adapter(type_info['sparsevec'], SparseVector._from_db)

tests/test_pg8000.py

@@ -0,0 +1,60 @@
+import numpy as np
+import os
+from pgvector import HalfVector, SparseVector, Vector
+from pgvector.pg8000 import register_vector
+from pg8000.native import Connection
+
+conn = Connection(os.environ["USER"], database='pgvector_python_test')
+
+conn.run('CREATE EXTENSION IF NOT EXISTS vector')
+conn.run('DROP TABLE IF EXISTS pg8000_items')
+conn.run('CREATE TABLE pg8000_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(conn)
+
+
+class TestPg8000:
+    def setup_method(self):
+        conn.run('DELETE FROM pg8000_items')
+
+    def test_vector(self):
+        embedding = np.array([1.5, 2, 3])
+        conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding)
+
+        res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id')
+        assert np.array_equal(res[0][0], embedding)
+        assert res[0][0].dtype == np.float32
+        assert res[1][0] is None
+
+    def test_vector_class(self):
+        embedding = Vector([1.5, 2, 3])
+        conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding)
+
+        res = conn.run('SELECT embedding FROM pg8000_items ORDER BY id')
+        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 = HalfVector([1.5, 2, 3])
+        conn.run('INSERT INTO pg8000_items (half_embedding) VALUES (:embedding), (NULL)', embedding=embedding)
+
+        res = conn.run('SELECT half_embedding FROM pg8000_items ORDER BY id')
+        assert res[0][0] == embedding
+        assert res[1][0] is None
+
+    def test_bit(self):
+        embedding = '101'
+        conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (NULL)', embedding=embedding)
+
+        res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id')
+        assert res[0][0] == '101'
+        assert res[1][0] is None
+
+    def test_sparsevec(self):
+        embedding = SparseVector([1.5, 2, 3])
+        conn.run('INSERT INTO pg8000_items (sparse_embedding) VALUES (:embedding), (NULL)', embedding=embedding)
+
+        res = conn.run('SELECT sparse_embedding FROM pg8000_items ORDER BY id')
+        assert res[0][0] == embedding
+        assert res[1][0] is None