pgvector 0.4.1__tar.gz → 0.5.0__tar.gz

{pgvector-0.4.1 → pgvector-0.5.0}/LICENSE.txt

@@ -1,6 +1,6 @@
 The MIT License (MIT)
 
-Copyright (c) 2021-2025 Andrew Kane
+Copyright (c) 2021-2026 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

{pgvector-0.4.1/pgvector.egg-info → pgvector-0.5.0}/PKG-INFO

@@ -1,21 +1,20 @@
 Metadata-Version: 2.4
 Name: pgvector
-Version: 0.4.1
+Version: 0.5.0
 Summary: pgvector support for Python
 Author-email: Andrew Kane <andrew@ankane.org>
-License: MIT
+License-Expression: MIT
 Project-URL: Homepage, https://github.com/pgvector/pgvector-python
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: numpy
 Dynamic: license-file
 
 # pgvector-python
 
 [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), [pg8000](https://github.com/tlocke/pg8000), 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://codeberg.org/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)
 
@@ -190,10 +189,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
 Add a vector column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 
 class Item(Base):
-    embedding = mapped_column(Vector(3))
+    embedding: Mapped[list[float]] = mapped_column(VECTOR(3))
 ```
 
 Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -272,7 +271,6 @@ index = Index(
     'my_index',
     func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
     postgresql_using='hnsw',
-    postgresql_with={'m': 16, 'ef_construction': 64},
     postgresql_ops={'embedding': 'halfvec_l2_ops'}
 )
 ```
@@ -284,16 +282,47 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
 session.scalars(select(Item).order_by(order).limit(5))
 ```
 
+#### Binary Quantization
+
+Use expression indexing for binary quantization
+
+```python
+from pgvector.sqlalchemy import BIT
+from sqlalchemy.sql import func
+
+index = Index(
+    'my_index',
+    func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
+    postgresql_using='hnsw',
+    postgresql_ops={'embedding': 'bit_hamming_ops'}
+)
+```
+
+Get the nearest neighbors by Hamming distance
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+session.scalars(select(Item).order_by(order).limit(5))
+```
+
+Re-rank by the original vectors for better recall
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+subquery = session.query(Item).order_by(order).limit(20).subquery()
+session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
+```
+
 #### Arrays
 
 Add an array column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 from sqlalchemy import ARRAY
 
 class Item(Base):
-    embeddings = mapped_column(ARRAY(Vector(3)))
+    embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1))
 ```
 
 And register the types with the underlying driver
@@ -328,7 +357,7 @@ from sqlalchemy import event
 
 @event.listens_for(engine, "connect")
 def connect(dbapi_connection, connection_record):
-    register_vector(dbapi_connection, arrays=True)
+    register_vector(dbapi_connection)
 ```
 
 ## SQLModel
@@ -342,10 +371,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
 Add a vector column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 
 class Item(SQLModel, table=True):
-    embedding: Any = Field(sa_type=Vector(3))
+    embedding: list[float] = Field(sa_type=VECTOR(3))
 ```
 
 Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -422,7 +451,7 @@ Enable the extension
 conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.psycopg import register_vector
@@ -456,7 +485,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
 ```
 
@@ -485,7 +516,7 @@ cur = conn.cursor()
 cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection or cursor
+Register the types with your connection or cursor
 
 ```python
 from pgvector.psycopg2 import register_vector
@@ -502,7 +533,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))'
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
 ```
 
@@ -531,7 +564,7 @@ Enable the extension
 await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.asyncpg import register_vector
@@ -557,7 +590,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)
 ```
 
@@ -585,7 +620,7 @@ Enable the extension
 conn.run('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.pg8000 import register_vector
@@ -602,7 +637,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
 ```
 
@@ -681,6 +718,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc
 
 ## Reference
 
+### Vectors
+
+Create a vector from a list
+
+```python
+vec = Vector([1, 2, 3])
+```
+
+Or a NumPy array
+
+```python
+vec = Vector(np.array([1, 2, 3]))
+```
+
+Get a list
+
+```python
+lst = vec.to_list()
+```
+
+Get a NumPy array
+
+```python
+arr = vec.to_numpy()
+```
+
 ### Half Vectors
 
 Create a half vector from a list
@@ -790,7 +853,7 @@ To get started with development:
 ```sh
 git clone https://github.com/pgvector/pgvector-python.git
 cd pgvector-python
-pip install -r requirements.txt
+pip install --group dev
 createdb pgvector_python_test
 pytest
 ```
@@ -799,7 +862,7 @@ To run an example:
 
 ```sh
 cd examples/loading
-pip install -r requirements.txt
+pip install --group dev
 createdb pgvector_example
 python3 example.py
 ```

pgvector-0.4.1/PKG-INFO → pgvector-0.5.0/README.md

@@ -1,21 +1,8 @@
-Metadata-Version: 2.4
-Name: pgvector
-Version: 0.4.1
-Summary: pgvector support for Python
-Author-email: Andrew Kane <andrew@ankane.org>
-License: MIT
-Project-URL: Homepage, https://github.com/pgvector/pgvector-python
-Requires-Python: >=3.9
-Description-Content-Type: text/markdown
-License-File: LICENSE.txt
-Requires-Dist: numpy
-Dynamic: license-file
-
 # pgvector-python
 
 [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), [pg8000](https://github.com/tlocke/pg8000), 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://codeberg.org/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)
 
@@ -190,10 +177,10 @@ session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
 Add a vector column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 
 class Item(Base):
-    embedding = mapped_column(Vector(3))
+    embedding: Mapped[list[float]] = mapped_column(VECTOR(3))
 ```
 
 Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -272,7 +259,6 @@ index = Index(
     'my_index',
     func.cast(Item.embedding, HALFVEC(3)).label('embedding'),
     postgresql_using='hnsw',
-    postgresql_with={'m': 16, 'ef_construction': 64},
     postgresql_ops={'embedding': 'halfvec_l2_ops'}
 )
 ```
@@ -284,16 +270,47 @@ order = func.cast(Item.embedding, HALFVEC(3)).l2_distance([3, 1, 2])
 session.scalars(select(Item).order_by(order).limit(5))
 ```
 
+#### Binary Quantization
+
+Use expression indexing for binary quantization
+
+```python
+from pgvector.sqlalchemy import BIT
+from sqlalchemy.sql import func
+
+index = Index(
+    'my_index',
+    func.cast(func.binary_quantize(Item.embedding), BIT(3)).label('embedding'),
+    postgresql_using='hnsw',
+    postgresql_ops={'embedding': 'bit_hamming_ops'}
+)
+```
+
+Get the nearest neighbors by Hamming distance
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+session.scalars(select(Item).order_by(order).limit(5))
+```
+
+Re-rank by the original vectors for better recall
+
+```python
+order = func.cast(func.binary_quantize(Item.embedding), BIT(3)).hamming_distance(func.binary_quantize(func.cast([3, -1, 2], VECTOR(3))))
+subquery = session.query(Item).order_by(order).limit(20).subquery()
+session.scalars(select(subquery).order_by(subquery.c.embedding.cosine_distance([3, -1, 2])).limit(5))
+```
+
 #### Arrays
 
 Add an array column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 from sqlalchemy import ARRAY
 
 class Item(Base):
-    embeddings = mapped_column(ARRAY(Vector(3)))
+    embeddings: Mapped[list[list[float]]] = mapped_column(ARRAY(VECTOR(3), dimensions=1))
 ```
 
 And register the types with the underlying driver
@@ -328,7 +345,7 @@ from sqlalchemy import event
 
 @event.listens_for(engine, "connect")
 def connect(dbapi_connection, connection_record):
-    register_vector(dbapi_connection, arrays=True)
+    register_vector(dbapi_connection)
 ```
 
 ## SQLModel
@@ -342,10 +359,10 @@ session.exec(text('CREATE EXTENSION IF NOT EXISTS vector'))
 Add a vector column
 
 ```python
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import VECTOR
 
 class Item(SQLModel, table=True):
-    embedding: Any = Field(sa_type=Vector(3))
+    embedding: list[float] = Field(sa_type=VECTOR(3))
 ```
 
 Also supports `HALFVEC`, `BIT`, and `SPARSEVEC`
@@ -422,7 +439,7 @@ Enable the extension
 conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.psycopg import register_vector
@@ -456,7 +473,9 @@ conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 conn.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
 ```
 
@@ -485,7 +504,7 @@ cur = conn.cursor()
 cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection or cursor
+Register the types with your connection or cursor
 
 ```python
 from pgvector.psycopg2 import register_vector
@@ -502,7 +521,9 @@ cur.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))'
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 cur.execute('INSERT INTO items (embedding) VALUES (%s)', (embedding,))
 ```
 
@@ -531,7 +552,7 @@ Enable the extension
 await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.asyncpg import register_vector
@@ -557,7 +578,9 @@ await conn.execute('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vect
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 await conn.execute('INSERT INTO items (embedding) VALUES ($1)', embedding)
 ```
 
@@ -585,7 +608,7 @@ Enable the extension
 conn.run('CREATE EXTENSION IF NOT EXISTS vector')
 ```
 
-Register the vector type with your connection
+Register the types with your connection
 
 ```python
 from pgvector.pg8000 import register_vector
@@ -602,7 +625,9 @@ conn.run('CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))')
 Insert a vector
 
 ```python
-embedding = np.array([1, 2, 3])
+from pgvector import Vector
+
+embedding = Vector([1, 2, 3])
 conn.run('INSERT INTO items (embedding) VALUES (:embedding)', embedding=embedding)
 ```
 
@@ -681,6 +706,32 @@ Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distanc
 
 ## Reference
 
+### Vectors
+
+Create a vector from a list
+
+```python
+vec = Vector([1, 2, 3])
+```
+
+Or a NumPy array
+
+```python
+vec = Vector(np.array([1, 2, 3]))
+```
+
+Get a list
+
+```python
+lst = vec.to_list()
+```
+
+Get a NumPy array
+
+```python
+arr = vec.to_numpy()
+```
+
 ### Half Vectors
 
 Create a half vector from a list
@@ -790,7 +841,7 @@ To get started with development:
 ```sh
 git clone https://github.com/pgvector/pgvector-python.git
 cd pgvector-python
-pip install -r requirements.txt
+pip install --group dev
 createdb pgvector_python_test
 pytest
 ```
@@ -799,7 +850,7 @@ To run an example:
 
 ```sh
 cd examples/loading
-pip install -r requirements.txt
+pip install --group dev
 createdb pgvector_example
 python3 example.py
 ```

pgvector-0.5.0/pgvector/_utils.py

@@ -0,0 +1,23 @@
+import sys
+from typing import TYPE_CHECKING, TypeAlias
+
+if TYPE_CHECKING:
+    import numpy as np
+
+    ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]]
+else:
+    # any value works since not type checking
+    # TODO use Never when Python 3.10 no longer supported
+    ndarray = None
+
+
+def is_ndarray(value: object, /) -> bool:
+    if (numpy := sys.modules.get('numpy')):
+        return isinstance(value, numpy.ndarray)
+    return False
+
+
+def is_sparse_array(value: object, /) -> bool:
+    if (sparse := sys.modules.get('scipy.sparse')):
+        return isinstance(value, (sparse.sparray, sparse.spmatrix))
+    return False

pgvector-0.5.0/pgvector/asyncpg/register.py

@@ -0,0 +1,32 @@
+from asyncpg import Connection
+from .. import Vector, HalfVector, SparseVector
+
+
+async def register_vector(conn: Connection, /, *, schema: str = 'public') -> None:
+    await conn.set_type_codec(
+        'vector',
+        schema=schema,
+        encoder=lambda v: (v if isinstance(v, Vector) else Vector(v)).to_binary(),
+        decoder=Vector.from_binary,
+        format='binary'
+    )
+
+    try:
+        await conn.set_type_codec(
+            'halfvec',
+            schema=schema,
+            encoder=lambda v: (v if isinstance(v, HalfVector) else HalfVector(v)).to_binary(),
+            decoder=HalfVector.from_binary,
+            format='binary'
+        )
+
+        await conn.set_type_codec(
+            'sparsevec',
+            schema=schema,
+            encoder=lambda v: (v if isinstance(v, SparseVector) else SparseVector(v)).to_binary(),
+            decoder=SparseVector.from_binary,
+            format='binary'
+        )
+    except ValueError as e:
+        if not str(e).startswith('unknown type:'):
+            raise e

pgvector-0.5.0/pgvector/bit.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+from struct import pack, unpack_from
+from typing import TYPE_CHECKING
+from ._utils import is_ndarray
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+class Bit:
+    _length: int
+    _data: bytes
+
+    def __init__(
+        self,
+        value: bytes | str | list[bool] | np.ndarray[tuple[int, ...], np.dtype[np.bool | np.uint8]],
+        /
+    ) -> None:
+        if isinstance(value, bytes):
+            self._length = 8 * len(value)
+            self._data = value
+        elif isinstance(value, (list, str)):
+            if isinstance(value, list):
+                bits = {True: '1', False: '0'}
+                try:
+                    value = ''.join([bits[v] for v in value])
+                except (KeyError, TypeError):
+                    raise ValueError('expected list[bool]')
+
+            length = len(value)
+            if length % 8 != 0:
+                value += '0' * (8 - (length % 8))
+
+            self._length = length
+            try:
+                self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
+            except ValueError:
+                raise ValueError('expected bit string')
+        elif is_ndarray(value):
+            import numpy as np
+
+            if value.dtype != np.bool:
+                # skip error for result of np.unpackbits
+                if value.dtype != np.uint8 or np.any(value > 1):
+                    raise ValueError('expected elements to be boolean')
+                value = value.astype(bool)
+
+            if value.ndim != 1:
+                raise ValueError('expected ndim to be 1')
+
+            self._length = len(value)
+            self._data = np.packbits(value).tobytes()  # type: ignore
+        else:
+            raise ValueError('expected bytes, str, list, or ndarray')
+
+    def __repr__(self) -> str:
+        return f'Bit({self.to_text()})'
+
+    def __eq__(self, other: object, /) -> bool:
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return self._length == other._length and self._data == other._data
+
+    def to_list(self) -> list[bool]:
+        # TODO improve
+        return [v != '0' for v in self.to_text()]
+
+    def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]:
+        import numpy as np
+
+        return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool)
+
+    def to_text(self) -> str:
+        return ''.join(format(v, '08b') for v in self._data)[:self._length]
+
+    def to_binary(self) -> bytes:
+        return pack('>i', self._length) + self._data
+
+    @classmethod
+    def from_text(cls, value: str, /) -> Bit:
+        # cast to ensure always uses str constructor
+        return cls(str(value))
+
+    @classmethod
+    def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Bit:
+        length, = unpack_from('>i', value)
+        data = memoryview(value)[4:].tobytes()
+
+        if len(data) != (length + 7) // 8:
+            raise ValueError('invalid length')
+
+        bit = cls.__new__(cls)
+        bit._length = length
+        bit._data = data
+        return bit

{pgvector-0.4.1 → pgvector-0.5.0}/pgvector/django/__init__.py

@@ -6,9 +6,6 @@ from .indexes import IvfflatIndex, HnswIndex
 from .sparsevec import SparseVectorField
 from .vector import VectorField
 
-# TODO remove
-from .. import HalfVector, SparseVector
-
 __all__ = [
     'VectorExtension',
     'VectorField',
@@ -22,7 +19,5 @@ __all__ = [
     'CosineDistance',
     'L1Distance',
     'HammingDistance',
-    'JaccardDistance',
-    'HalfVector',
-    'SparseVector'
+    'JaccardDistance'
 ]