First commit

ankane · ankane · commit 749e8d98ac7b · 2021-06-11T21:46:35.000-07:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -0,0 +1,20 @@
+name: build
+on: [push, pull_request]
+jobs:
+  build:
+    if: "!contains(github.event.head_commit.message, '[skip ci]')"
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-python@v2
+    - run: pip install -r requirements.txt
+    - uses: ankane/setup-postgres@v1
+      with:
+        database: pgvector_python_test
+    - run: |
+        sudo apt-get update && sudo apt-get install postgresql-server-dev-13
+        git clone --branch v0.1.6 https://github.com/ankane/pgvector.git
+        cd pgvector
+        make
+        sudo make install
+    - run: pytest
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,7 @@
+build/
+dist/
+*.egg-info/
+venv/
+.cache/
+*.pyc
+__pycache__
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -0,0 +1,3 @@
+## 0.1.0 (2021-06-11)
+
+- First release
diff --git a/LICENSE.txt b/LICENSE.txt
@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2021 Andrew Kane
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
diff --git a/Makefile b/Makefile
@@ -0,0 +1,8 @@
+lint:
+	pycodestyle . --ignore=E501
+
+publish:
+	python3 setup.py bdist_wheel --universal
+	ls dist
+	# twine upload dist/*
+	rm -fr build dist pgvector.egg-info
diff --git a/README.md b/README.md
@@ -0,0 +1,97 @@
+# pgvector-python
+
+[pgvector](https://github.com/ankane/pgvector) support for Python
+
+Supports [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy) and [Psycopg 2](https://github.com/psycopg/psycopg2)
+
+[![Build Status](https://github.com/ankane/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/ankane/pgvector-python/actions)
+
+## Installation
+
+Run:
+
+```sh
+pip install pgvector
+```
+
+And follow the instructions for your database library:
+
+- [SQLAlchemy](#sqlalchemy)
+- [Psycopg 2](#psycopg-2)
+
+## SQLAlchemy
+
+Use the vector type
+
+```python
+from pgvector.sqlalchemy import Vector
+
+# Core
+item_table = Table(
+    'item',
+    metadata,
+    Column('factors', Vector(3))
+)
+
+# ORM
+class Item(Base):
+    factors = Column(Vector(3))
+```
+
+Add an index
+
+```python
+index = Index(
+    'my_index',
+    item_table.c.factors, # or Item.factors for ORM
+    postgresql_using='ivfflat',
+    postgresql_with={'lists': 100}
+)
+index.create(engine)
+```
+
+## Psycopg 2
+
+Register the vector type with your connection or cursor
+
+```python
+from pgvector.psycopg2 import register_vector
+
+register_vector(conn)
+```
+
+Insert a Numpy array as a vector
+
+```python
+factors = np.array([1, 2, 3])
+cur.execute('INSERT INTO item (factors) VALUES (%s)', (factors,))
+```
+
+Get the nearest neighbors
+
+```python
+cur.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,))
+cur.fetchall()
+```
+
+## History
+
+View the [changelog](https://github.com/ankane/pgvector-python/blob/master/CHANGELOG.md)
+
+## Contributing
+
+Everyone is encouraged to help improve this project. Here are a few ways you can help:
+
+- [Report bugs](https://github.com/ankane/pgvector-python/issues)
+- Fix bugs and [submit pull requests](https://github.com/ankane/pgvector-python/pulls)
+- Write, clarify, or fix documentation
+- Suggest or add new features
+
+To get started with development:
+
+```sh
+git clone https://github.com/ankane/pgvector-python.git
+cd pgvector-python
+pip install -r requirements.txt
+pytest
+```
diff --git a/pgvector/psycopg2/__init__.py b/pgvector/psycopg2/__init__.py
@@ -0,0 +1,33 @@
+import numpy as np
+import psycopg2
+from psycopg2.extensions import adapt, new_type, register_adapter, register_type
+from ..utils import cast_vector, quote_vector
+
+
+__all__ = ['register_vector']
+
+
+class VectorAdapter(object):
+    def __init__(self, vector):
+        self._vector = vector
+
+    def getquoted(self):
+        return adapt(quote_vector(self._vector)).getquoted()
+
+
+def cast_vector_type(value, cur):
+    return cast_vector(value)
+
+
+def register_vector(conn_or_curs=None):
+    cur = conn_or_curs.cursor() if hasattr(conn_or_curs, 'cursor') else conn_or_curs
+
+    try:
+        cur.execute('SELECT NULL::vector')
+        oid = cur.description[0][1]
+    except psycopg2.errors.UndefinedObject:
+        raise psycopg2.ProgrammingError('vector type not found in the database')
+
+    vector = new_type((oid,), 'VECTOR', cast_vector_type)
+    register_type(vector)
+    register_adapter(np.ndarray, VectorAdapter)
diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py
@@ -0,0 +1,23 @@
+from sqlalchemy.types import UserDefinedType
+from ..utils import cast_vector, quote_vector
+
+
+class Vector(UserDefinedType):
+    def __init__(self, dim=None):
+        super(UserDefinedType, self).__init__()
+        self.dim = dim
+
+    def get_col_spec(self, **kw):
+        if self.dim is None:
+            return "VECTOR"
+        return "VECTOR(%d)" % self.dim
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return quote_vector(value, self.dim)
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return cast_vector(value)
+        return process
diff --git a/pgvector/utils/__init__.py b/pgvector/utils/__init__.py
@@ -0,0 +1,28 @@
+import numpy as np
+
+
+def cast_vector(value):
+    # could be ndarray if already cast by lower-level driver
+    if value is None or isinstance(value, np.ndarray):
+        return value
+
+    return np.array(value[1:-1].split(','), dtype=np.float32)
+
+
+def quote_vector(value, dim=None):
+    if value is None:
+        return value
+
+    if isinstance(value, np.ndarray):
+        if value.ndim != 1:
+            raise ValueError('expected ndim to be 1')
+
+        if not np.issubdtype(value.dtype, np.integer) and not np.issubdtype(value.dtype, np.floating):
+            raise ValueError('dtype must be numeric')
+
+        value = value.tolist()
+
+    if dim is not None and len(value) != dim:
+        raise ValueError('expected %d dimensions, not %d' % (dim, len(value)))
+
+    return '[' + ','.join([str(float(v)) for v in value]) + ']'
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,4 @@
+numpy
+psycopg2
+pytest
+SQLAlchemy
diff --git a/setup.py b/setup.py
@@ -0,0 +1,21 @@
+from setuptools import setup
+
+setup(
+    name='pgvector',
+    version='0.1.0',
+    description='pgvector support for Python',
+    url='https://github.com/ankane/pgvector-python',
+    author='Andrew Kane',
+    author_email='andrew@ankane.org',
+    license='MIT',
+    packages=[
+        'pgvector.psycopg2',
+        'pgvector.sqlalchemy',
+        'pgvector.utils'
+    ],
+    python_requires='>=3.6',
+    install_requires=[
+        'numpy'
+    ],
+    zip_safe=False
+)
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/test_psycopg2.py b/tests/test_psycopg2.py
@@ -0,0 +1,27 @@
+import numpy as np
+from pgvector.psycopg2 import register_vector
+import psycopg2
+
+
+class TestPsycopg2(object):
+    def test_works(self):
+        conn = psycopg2.connect('dbname=pgvector_python_test')
+        conn.autocommit = True
+
+        cur = conn.cursor()
+        cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
+        cur.execute('DROP TABLE IF EXISTS item')
+        cur.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))')
+
+        register_vector(cur)
+
+        factors = np.array([1.5, 2, 3])
+        cur.execute("INSERT INTO item (factors) VALUES (%s), (NULL)", (factors,))
+
+        cur.execute("SELECT * FROM item ORDER BY id")
+        res = cur.fetchall()
+        assert res[0][0] == 1
+        assert res[1][0] == 2
+        assert np.array_equal(res[0][1], factors)
+        assert res[0][1].dtype == np.float32
+        assert res[1][1] is None
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
@@ -0,0 +1,87 @@
+import numpy as np
+from pgvector.sqlalchemy import Vector
+import pytest
+from sqlalchemy import create_engine, select, text, MetaData, Table, Column, Index, Integer
+from sqlalchemy.exc import StatementError
+from sqlalchemy.orm import declarative_base, Session
+
+engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test', future=True)
+with engine.connect() as con:
+    con.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
+    con.execute(text('DROP TABLE IF EXISTS core_item'))
+    con.execute(text('DROP TABLE IF EXISTS orm_item'))
+
+Base = declarative_base()
+
+
+class Item(Base):
+    __tablename__ = 'orm_item'
+
+    id = Column(Integer, primary_key=True)
+    factors = Column(Vector(3))
+
+
+Base.metadata.drop_all(engine)
+Base.metadata.create_all(engine)
+
+
+class TestSqlalchemy(object):
+    def test_core(self):
+        metadata = MetaData()
+
+        item_table = Table(
+            'core_item',
+            metadata,
+            Column('id', Integer, primary_key=True),
+            Column('factors', Vector(3))
+        )
+
+        metadata.drop_all(engine)
+        metadata.create_all(engine)
+
+        index = Index('my_core_index', item_table.c.factors, postgresql_using='ivfflat', postgresql_with={'lists': 1})
+        index.create(engine)
+
+    def test_orm(self):
+        item = Item(factors=np.array([1.5, 2, 3]))
+        item2 = Item(factors=[4, 5, 6])
+        item3 = Item()
+
+        session = Session(engine)
+        session.add(item)
+        session.add(item2)
+        session.add(item3)
+        session.commit()
+
+        stmt = select(Item)
+        with Session(engine) as session:
+            items = [v[0] for v in session.execute(stmt).all()]
+            assert items[0].id == 1
+            assert items[1].id == 2
+            assert items[2].id == 3
+            assert np.array_equal(items[0].factors, np.array([1.5, 2, 3]))
+            assert items[0].factors.dtype == np.float32
+            assert np.array_equal(items[1].factors, np.array([4, 5, 6]))
+            assert items[1].factors.dtype == np.float32
+            assert items[2].factors is None
+
+    def test_bad_dimensions(self):
+        item = Item(factors=[1, 2])
+        session = Session(engine)
+        session.add(item)
+        with pytest.raises(StatementError, match='expected 3 dimensions, not 2'):
+            session.commit()
+
+    def test_bad_ndim(self):
+        item = Item(factors=np.array([[1, 2, 3]]))
+        session = Session(engine)
+        session.add(item)
+        with pytest.raises(StatementError, match='expected ndim to be 1'):
+            session.commit()
+
+    def test_bad_dtype(self):
+        item = Item(factors=np.array(['one', 'two', 'three']))
+        session = Session(engine)
+        session.add(item)
+        with pytest.raises(StatementError, match='dtype must be numeric'):
+            session.commit()

-Original file line number
+Diff line change
@@ @@ -0,0 +1,7 @@ @@
 +build/
 +dist/
 +*.egg-info/
 +venv/
 +.cache/
 +*.pyc
 +__pycache__
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+## 0.1.0 (2021-06-11)`
	`2`	`+`
	`3`	`+- First release`