diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 952be69..5a01635 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -6,10 +6,10 @@ jobs: strategy: fail-fast: false matrix: - python: [3.14, "3.10"] + python: [3.15, "3.10"] steps: - uses: actions/checkout@v7 - - uses: actions/setup-python@v6 + - uses: actions/setup-python@v7 with: python-version: ${{ matrix.python }} - run: pip install --group dev @@ -19,7 +19,7 @@ jobs: dev-files: true - run: | cd /tmp - git clone --branch v0.8.4 https://github.com/pgvector/pgvector.git + git clone --branch v0.8.7 https://github.com/pgvector/pgvector.git cd pgvector make sudo make install diff --git a/CHANGELOG.md b/CHANGELOG.md index cd577c8..f033cf1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,10 @@ -## 0.5.1 (unreleased) +## 0.5.2 (unreleased) + +- Fixed memory leak with `register_vector` and `register_vector_async` for Psycopg 3 +- Fixed error with `Bit` constructor and empty strings +- Fixed equality for `HalfVector` with signed zeros + +## 0.5.1 (2026-10-08) - Fixed error with non-canonical SciPy sparse arrays - Fixed error with SciPy sparse arrays with explicit zeros diff --git a/examples/cohere/example.py b/examples/cohere/example.py index 5cf8d69..9b32073 100644 --- a/examples/cohere/example.py +++ b/examples/cohere/example.py @@ -9,12 +9,12 @@ register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') -conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(1536))') +conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding bit(2048))') def embed(input, input_type): co = cohere.ClientV2() - response = co.embed(texts=input, model='embed-v4.0', input_type=input_type, embedding_types=['ubinary']) + response = co.embed(texts=input, model='embed-v5.0-pro', input_type=input_type, embedding_types=['ubinary']) return [bytes(embedding) for embedding in response.embeddings.ubinary] diff --git a/pgvector/bit.py b/pgvector/bit.py index a157b30..a9370e4 100644 --- a/pgvector/bit.py +++ b/pgvector/bit.py @@ -34,7 +34,10 @@ def __init__( value += '0' * (8 - (length % 8)) self._length = length - self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + if length > 0: + self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big') + else: + self._data = b'' elif is_ndarray(value): import numpy as np @@ -85,7 +88,7 @@ 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: + if length < 0 or len(data) != (length + 7) // 8: raise ValueError('invalid length') bit = cls.__new__(cls) diff --git a/pgvector/halfvec.py b/pgvector/halfvec.py index 1a2c8c7..2b3b325 100644 --- a/pgvector/halfvec.py +++ b/pgvector/halfvec.py @@ -37,7 +37,7 @@ def __repr__(self) -> str: def __eq__(self, other: object, /) -> bool: if not isinstance(other, self.__class__): return NotImplemented - return self._value == other._value + return self.to_list() == other.to_list() def dimensions(self) -> int: return len(self._value) @@ -59,6 +59,7 @@ def to_binary(self) -> bytes: else: value = array('H', self._value) value.byteswap() + # TODO change to >hh in 0.6.0 return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod @@ -67,7 +68,7 @@ def from_text(cls, value: str, /) -> HalfVector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: - dim, unused = struct.unpack_from('>HH', value) + dim, unused = struct.unpack_from('>hh', value) data = memoryview(value)[4:] if len(data) != 2 * dim: @@ -87,6 +88,8 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> HalfVector: @classmethod def _from_text(cls, value: str, /) -> list[float]: + if len(value) == 2: + return [] return [float(v) for v in value[1:-1].split(',')] @classmethod diff --git a/pgvector/psycopg/bit.py b/pgvector/psycopg/bit.py index f899497..4cb502f 100644 --- a/pgvector/psycopg/bit.py +++ b/pgvector/psycopg/bit.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Dumper @@ -26,9 +27,19 @@ def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) -> info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (BitDumper,), {'oid': info.oid}) - binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid}) + text_dumper = bit_text_dumper(info.oid) + binary_dumper = bit_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(Bit, text_dumper) adapters.register_dumper(Bit, binary_dumper) + + +@cache +def bit_text_dumper(oid: int) -> type[BitDumper]: + return type('', (BitDumper,), {'oid': oid}) + + +@cache +def bit_binary_dumper(oid: int) -> type[BitBinaryDumper]: + return type('', (BitBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/halfvec.py b/pgvector/psycopg/halfvec.py index 9ed0453..05e7994 100644 --- a/pgvector/psycopg/halfvec.py +++ b/pgvector/psycopg/halfvec.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper @@ -41,11 +42,21 @@ def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> No info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid}) + text_dumper = halfvec_text_dumper(info.oid) + binary_dumper = halfvec_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(HalfVector, text_dumper) adapters.register_dumper(HalfVector, binary_dumper) adapters.register_loader(info.oid, HalfVectorLoader) adapters.register_loader(info.oid, HalfVectorBinaryLoader) + + +@cache +def halfvec_text_dumper(oid: int) -> type[HalfVectorDumper]: + return type('', (HalfVectorDumper,), {'oid': oid}) + + +@cache +def halfvec_binary_dumper(oid: int) -> type[HalfVectorBinaryDumper]: + return type('', (HalfVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/sparsevec.py b/pgvector/psycopg/sparsevec.py index 57fd19a..0c953bf 100644 --- a/pgvector/psycopg/sparsevec.py +++ b/pgvector/psycopg/sparsevec.py @@ -1,3 +1,4 @@ +from functools import cache from psycopg import BaseConnection from psycopg.abc import Buffer from psycopg.adapt import Loader, Dumper @@ -41,11 +42,21 @@ def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) -> info.register(context) # add oid to anonymous class for set_types - text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid}) - binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid}) + text_dumper = sparsevec_text_dumper(info.oid) + binary_dumper = sparsevec_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(SparseVector, text_dumper) adapters.register_dumper(SparseVector, binary_dumper) adapters.register_loader(info.oid, SparseVectorLoader) adapters.register_loader(info.oid, SparseVectorBinaryLoader) + + +@cache +def sparsevec_text_dumper(oid: int) -> type[SparseVectorDumper]: + return type('', (SparseVectorDumper,), {'oid': oid}) + + +@cache +def sparsevec_binary_dumper(oid: int) -> type[SparseVectorBinaryDumper]: + return type('', (SparseVectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/psycopg/vector.py b/pgvector/psycopg/vector.py index 7ef16b0..f4b66e4 100644 --- a/pgvector/psycopg/vector.py +++ b/pgvector/psycopg/vector.py @@ -1,4 +1,5 @@ from __future__ import annotations +from functools import cache import psycopg from psycopg import BaseConnection from psycopg.abc import Buffer @@ -50,8 +51,8 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) 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}) + text_dumper = vector_text_dumper(info.oid) + binary_dumper = vector_binary_dumper(info.oid) adapters = context.adapters adapters.register_dumper(Vector, text_dumper) @@ -60,3 +61,13 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /) adapters.register_dumper('numpy.ndarray', binary_dumper) adapters.register_loader(info.oid, VectorLoader) adapters.register_loader(info.oid, VectorBinaryLoader) + + +@cache +def vector_text_dumper(oid: int) -> type[VectorDumper]: + return type('', (VectorDumper,), {'oid': oid}) + + +@cache +def vector_binary_dumper(oid: int) -> type[VectorBinaryDumper]: + return type('', (VectorBinaryDumper,), {'oid': oid}) diff --git a/pgvector/sparsevec.py b/pgvector/sparsevec.py index 0d3434c..76e1356 100644 --- a/pgvector/sparsevec.py +++ b/pgvector/sparsevec.py @@ -101,6 +101,8 @@ def _from_dict(self, d: dict[int, float], dim: int) -> None: self._dim = int(dim) self._indices = [int(v[0]) for v in elements] self._values = [float(v[1]) for v in elements] + # TODO uncomment in 0.6.0 + # self._validate() def _from_sparse(self, arr: sparray | spmatrix, /) -> None: value: coo_array | coo_matrix = arr.tocoo(copy=False) # type: ignore @@ -135,8 +137,25 @@ def _from_sparse(self, arr: sparray | spmatrix, /) -> None: def _from_dense(self, value: list[float] | ndarray, /) -> None: self._dim = len(value) - self._indices = [i for i, v in enumerate(value) if v != 0] - self._values = [float(value[i]) for i in self._indices] + self._indices = [] + self._values = [] + for i, v in enumerate(value): + v = float(v) + if v != 0: + self._indices.append(i) + self._values.append(v) + + def _validate(self) -> None: + if self._dim < 0: + raise ValueError('sparsevec cannot have negative dimensions') + + if any(i < 0 or i >= self._dim for i in self._indices): + raise ValueError('sparsevec index out of bounds') + + def _validate_sorted(self) -> None: + for i in range(len(self._indices) - 1): + if self._indices[i] >= self._indices[i + 1]: + raise ValueError('sparsevec indices must be in ascending order') @classmethod def from_text(cls, value: str, /) -> SparseVector: @@ -171,6 +190,9 @@ def _from_parts(cls, dim: int, indices: list[int], values: list[float], /) -> Sp vec._dim = dim vec._indices = indices vec._values = values + # TODO uncomment in 0.6.0 + # vec._validate() + # vec._validate_sorted() return vec @classmethod diff --git a/pgvector/sqlalchemy/bit.py b/pgvector/sqlalchemy/bit.py index 7c861ef..a5af3b0 100644 --- a/pgvector/sqlalchemy/bit.py +++ b/pgvector/sqlalchemy/bit.py @@ -14,7 +14,7 @@ def process_bind_param(self, value: Any, dialect: Any) -> Any: return asyncpg.BitString(value) # type: ignore return value - class Comparator(TypeDecorator.Comparator[Any]): + class Comparator(TypeDecorator.Comparator[Any]): # type: ignore def hamming_distance(self, other: object, /) -> Operators: return self.op('<~>', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/halfvec.py b/pgvector/sqlalchemy/halfvec.py index 9d717c3..d82c5ac 100644 --- a/pgvector/sqlalchemy/halfvec.py +++ b/pgvector/sqlalchemy/halfvec.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return HalfVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/sparsevec.py b/pgvector/sqlalchemy/sparsevec.py index 9bd1478..5cc910d 100644 --- a/pgvector/sqlalchemy/sparsevec.py +++ b/pgvector/sqlalchemy/sparsevec.py @@ -35,7 +35,7 @@ def process(value: Any) -> SparseVector | None: return SparseVector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py index 83acb09..4e7e4b3 100644 --- a/pgvector/sqlalchemy/vector.py +++ b/pgvector/sqlalchemy/vector.py @@ -35,7 +35,7 @@ def process(value: Any) -> list[float] | None: return Vector._from_db(value) return process - class Comparator(TypeEngine.Comparator[Any]): + class Comparator(TypeEngine.Comparator[Any]): # type: ignore def l2_distance(self, other: object, /) -> Operators: return self.op('<->', return_type=Float)(other) diff --git a/pgvector/vector.py b/pgvector/vector.py index d357113..d86cba0 100644 --- a/pgvector/vector.py +++ b/pgvector/vector.py @@ -57,6 +57,7 @@ def to_binary(self) -> bytes: else: value = array('f', self._value) value.byteswap() + # TODO change to >hh in 0.6.0 return struct.pack('>HH', len(value), 0) + memoryview(value) @classmethod @@ -65,7 +66,7 @@ def from_text(cls, value: str, /) -> Vector: @classmethod def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: - dim, unused = struct.unpack_from('>HH', value) + dim, unused = struct.unpack_from('>hh', value) data = memoryview(value)[4:] if len(data) != 4 * dim: @@ -85,6 +86,8 @@ def from_binary(cls, value: bytes | bytearray | memoryview, /) -> Vector: @classmethod def _from_text(cls, value: str, /) -> list[float]: + if len(value) == 2: + return [] return [float(v) for v in value[1:-1].split(',')] @classmethod diff --git a/pyproject.toml b/pyproject.toml index ad130b5..3ff09d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pgvector" -version = "0.5.0" +version = "0.5.1" description = "pgvector support for Python" readme = "README.md" authors = [ diff --git a/tests/test_bit.py b/tests/test_bit.py index ed13374..9179cd3 100644 --- a/tests/test_bit.py +++ b/tests/test_bit.py @@ -1,42 +1,51 @@ from pgvector import Bit import pytest import random -from struct import pack +import struct from .conftest import numpy as np class TestBit: - def test_list(self) -> None: + def test_init_list(self) -> None: assert Bit([True, False, True]).to_list() == [True, False, True] - def test_list_none(self) -> None: + def test_init_list_empty(self) -> None: + assert Bit([]).to_list() == [] + + def test_init_list_none(self) -> None: with pytest.raises(ValueError) as error: Bit([True, None, True]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_int(self) -> None: + def test_init_list_int(self) -> None: with pytest.raises(ValueError) as error: Bit([254, 7, 0]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: Bit([[True, False], [True, False]]) # type: ignore assert str(error.value) == 'expected list[bool]' - def test_str(self) -> None: + def test_init_str(self) -> None: assert Bit('101').to_list() == [True, False, True] - def test_str_two(self) -> None: + def test_init_str_empty(self) -> None: + assert Bit('').to_list() == [] + + def test_init_str_two(self) -> None: with pytest.raises(ValueError) as error: Bit('201') assert str(error.value) == 'expected bit string' - def test_bytes(self) -> None: + def test_init_bytes(self) -> None: assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000' assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000' - def test_ndarray(self) -> None: + def test_init_bytes_empty(self) -> None: + assert Bit(b'').to_text() == '' + + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -44,14 +53,14 @@ def test_ndarray(self) -> None: assert Bit(arr).to_list() == [True, False, True] assert np.array_equal(Bit(arr).to_numpy(), arr) - def test_ndarray_unpackbits(self) -> None: + def test_init_ndarray_unpackbits(self) -> None: if np is None: pytest.skip('NumPy required') arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8)) assert Bit(arr).to_text() == '111111100000011100000000' - def test_ndarray_uint8(self) -> None: + def test_init_ndarray_uint8(self) -> None: if np is None: pytest.skip('NumPy required') @@ -60,7 +69,7 @@ def test_ndarray_uint8(self) -> None: Bit(arr) assert str(error.value) == 'expected elements to be boolean' - def test_ndarray_uint16(self) -> None: + def test_init_ndarray_uint16(self) -> None: if np is None: pytest.skip('NumPy required') @@ -69,12 +78,12 @@ def test_ndarray_uint16(self) -> None: Bit(arr) # type: ignore assert str(error.value) == 'expected elements to be boolean' - def test_bool(self) -> None: + def test_init_bool(self) -> None: with pytest.raises(ValueError) as error: Bit(True) # type: ignore assert str(error.value) == 'expected bytes, str, list, or ndarray' - def test_random(self) -> None: + def test_init_random(self) -> None: value = ''.join(random.choices(['0', '1'], k=random.randint(1024, 2048))) assert Bit(value).to_text() == value @@ -82,22 +91,54 @@ def test_repr(self) -> None: assert repr(Bit([True, False, True])) == 'Bit(101)' assert str(Bit([True, False, True])) == 'Bit(101)' - def test_equality(self) -> None: + def test_eq(self) -> None: assert Bit([True, False, True]) == Bit([True, False, True]) assert Bit([True, False, True]) != Bit([True, False, False]) assert Bit([True, False, True]) != 1 + def test_to_list(self) -> None: + assert Bit('101').to_list() == [True, False, True] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = Bit('101').to_numpy() + assert np.array_equal(arr, np.array([True, False, True])) + assert arr.dtype == np.bool_ + + def test_to_text(self) -> None: + assert Bit('101').to_text() == '101' + + def test_to_binary(self) -> None: + assert Bit('101').to_binary() == struct.pack('>iB', 3, 5 << 5) + def test_from_text(self) -> None: - vec = Bit.from_text('101') - assert vec.to_list() == [True, False, True] - if np is not None: - assert np.array_equal(vec.to_numpy(), [True, False, True]) - assert vec.to_text() == '101' + assert Bit.from_text('101') == Bit('101') + + def test_from_text_zero_length(self) -> None: + assert Bit.from_text('') == Bit([]) + + def test_from_text_invalid_character(self) -> None: + with pytest.raises(ValueError, match='expected bit string'): + Bit.from_text('2') def test_from_binary(self) -> None: - data = pack('>iB', 3, 5 << 5) - vec = Bit.from_binary(data) - assert vec.to_list() == [True, False, True] - if np is not None: - assert np.array_equal(vec.to_numpy(), [True, False, True]) - assert vec.to_binary() == data + assert Bit.from_binary(struct.pack('>iB', 3, 5 << 5)) == Bit('101') + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + Bit.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>iB', 9, 0)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>iBBB', 16, 0, 0, 0)) + + def test_from_binary_negative_length(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Bit.from_binary(struct.pack('>i', -1)) diff --git a/tests/test_half_vector.py b/tests/test_half_vector.py index 26c5ac7..a907c94 100644 --- a/tests/test_half_vector.py +++ b/tests/test_half_vector.py @@ -1,29 +1,29 @@ from pgvector import HalfVector import pytest -from struct import pack +import struct from .conftest import numpy as np class TestHalfVector: - def test_list(self) -> None: + def test_init_list(self) -> None: arr = [1.0, 2.0, 3.0] assert HalfVector(arr).to_list() == arr assert HalfVector(arr).to_list() is not arr - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert HalfVector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError) as error: HalfVector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: HalfVector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -31,18 +31,32 @@ def test_ndarray(self) -> None: assert HalfVector(arr).to_list() == [1, 2, 3] assert HalfVector(arr).to_numpy() is not arr assert HalfVector(arr).to_numpy().dtype == np.float16 - # non-contiguous + + def test_init_ndarray_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(np.flip(arr)).to_list() == [3, 2, 1] assert HalfVector(np.flip(arr)).to_binary() == HalfVector([3, 2, 1]).to_binary() - # big endian + + def test_init_ndarray_big_endian(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float16) assert HalfVector(arr.astype('>f2')).to_list() == [1, 2, 3] assert HalfVector(arr.astype('>f2')).to_binary() == HalfVector([1, 2, 3]).to_binary() + def test_init_ndarray_str(self) -> None: + if np is None: + pytest.skip('NumPy required') + with pytest.raises(ValueError) as error: HalfVector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) - def test_int(self) -> None: + def test_init_int(self) -> None: with pytest.raises(ValueError) as error: HalfVector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' @@ -51,25 +65,74 @@ def test_repr(self) -> None: assert repr(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' assert str(HalfVector([1, 2, 3])) == 'HalfVector([1.0, 2.0, 3.0])' - def test_equality(self) -> None: + def test_eq(self) -> None: assert HalfVector([1, 2, 3]) == HalfVector([1, 2, 3]) assert HalfVector([1, 2, 3]) != HalfVector([1, 2, 4]) assert HalfVector([1, 2, 3]) != 1 + assert HalfVector([0.0]) == HalfVector([-0.0]) def test_dimensions(self) -> None: assert HalfVector([1, 2, 3]).dimensions() == 3 + def test_to_list(self) -> None: + assert HalfVector([1.5, 2, 3]).to_list() == [1.5, 2, 3] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = HalfVector([1.5, 2, 3]).to_numpy() + assert np.array_equal(arr, np.array([1.5, 2, 3])) + assert arr.dtype == np.float16 + + def test_to_text(self) -> None: + assert HalfVector([1.5, 2, 3]).to_text() == '[1.5,2.0,3.0]' + + def test_to_binary(self) -> None: + assert HalfVector([1.5, 2, 3]).to_binary() == struct.pack('>hh3e', 3, 0, 1.5, 2, 3) + def test_from_text(self) -> None: - vec = HalfVector.from_text('[1.5,2,3]') - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_text() == '[1.5,2.0,3.0]' + assert HalfVector.from_text('[1.5,2,3]') == HalfVector([1.5, 2, 3]) + + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + HalfVector.from_text('') + + def test_from_text_zero_dimensions(self) -> None: + assert HalfVector.from_text('[]') == HalfVector([]) + + def test_from_text_open_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert HalfVector.from_text('{1]') + + def test_from_text_close_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert HalfVector.from_text('[1}') + + def test_from_text_non_number(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert HalfVector.from_text('[a]') def test_from_binary(self) -> None: - data = pack('>HH3e', 3, 0, 1.5, 2, 3) - vec = HalfVector.from_binary(data) - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_binary() == data + assert HalfVector.from_binary(struct.pack('>hh3e', 3, 0, 1.5, 2, 3)) == HalfVector([1.5, 2, 3]) + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + HalfVector.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hhe', 2, 0, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hheee', 2, 0, 1, 2, 3)) + + def test_from_binary_negative_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + HalfVector.from_binary(struct.pack('>hh', -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + HalfVector.from_binary(struct.pack('>hh', 0, 1)) diff --git a/tests/test_sparse_vector.py b/tests/test_sparse_vector.py index 262e23a..e3bb8cc 100644 --- a/tests/test_sparse_vector.py +++ b/tests/test_sparse_vector.py @@ -1,6 +1,6 @@ from pgvector import SparseVector import pytest -from struct import pack +import struct from .conftest import numpy as np try: @@ -10,26 +10,26 @@ class TestSparseVector: - def test_list(self) -> None: + def test_init_list(self) -> None: vec = SparseVector([1, 0, 2, 0, 3, 0]) assert vec.to_list() == [1, 0, 2, 0, 3, 0] if np is not None: assert np.array_equal(vec.to_numpy(), [1, 0, 2, 0, 3, 0]) assert vec.indices() == [0, 2, 4] - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert SparseVector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError, match='could not convert string to float'): SparseVector([1, 'two', 3]) # type: ignore - def test_list_dimensions(self) -> None: + def test_init_list_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector([1, 0, 2, 0, 3, 0], 6) # type: ignore assert str(error.value) == 'extra argument' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -37,20 +37,35 @@ def test_ndarray(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict(self) -> None: + def test_init_dict(self) -> None: vec = SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dict_empty(self) -> None: + def test_init_dict_empty(self) -> None: assert SparseVector({}, 0).to_list() == [] - def test_dict_no_dimensions(self) -> None: + # TODO uncomment in 0.6.0 + # def test_dict_negative_index(self) -> None: + # with pytest.raises(ValueError, match='sparsevec index out of bounds'): + # SparseVector({-1: 1}, 3) + + # TODO uncomment in 0.6.0 + # def test_dict_higher_index(self) -> None: + # with pytest.raises(ValueError, match='sparsevec index out of bounds'): + # SparseVector({3: 1}, 3) + + # TODO uncomment in 0.6.0 + # def test_dict_negative_dimensions(self) -> None: + # with pytest.raises(ValueError, match='sparsevec cannot have negative dimensions'): + # SparseVector({}, -1) + + def test_init_dict_no_dimensions(self) -> None: with pytest.raises(ValueError) as error: SparseVector({0: 1, 2: 2, 4: 3}) # type: ignore assert str(error.value) == 'missing dimensions' - def test_coo_array(self) -> None: + def test_init_coo_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -60,7 +75,7 @@ def test_coo_array(self) -> None: assert vec.indices() == [0, 2, 4] assert isinstance(vec.values()[0], float) - def test_coo_array_duplicates(self) -> None: + def test_init_coo_array_duplicates(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -70,7 +85,7 @@ def test_coo_array_duplicates(self) -> None: assert vec.to_list() == [0, 3, 0] assert vec.indices() == [1] - def test_coo_array_dimensions(self) -> None: + def test_init_coo_array_dimensions(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -78,7 +93,7 @@ def test_coo_array_dimensions(self) -> None: SparseVector(sparse.coo_array(np.array([1, 0, 2, 0, 3, 0])), 6) # type: ignore assert str(error.value) == 'extra argument' - def test_coo_matrix(self) -> None: + def test_init_coo_matrix(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -87,7 +102,7 @@ def test_coo_matrix(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_dok_array(self) -> None: + def test_init_dok_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -96,7 +111,7 @@ def test_dok_array(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_csr_array(self) -> None: + def test_init_csr_array(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -105,7 +120,7 @@ def test_csr_array(self) -> None: assert vec.to_list() == [1, 0, 2, 0, 3, 0] assert vec.indices() == [0, 2, 4] - def test_csr_matrix(self) -> None: + def test_init_csr_matrix(self) -> None: if np is None or sparse is None: pytest.skip('NumPy and SciPy required') @@ -118,7 +133,7 @@ def test_repr(self) -> None: assert repr(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' assert str(SparseVector([1, 0, 2, 0, 3, 0])) == 'SparseVector({0: 1.0, 2: 2.0, 4: 3.0}, 6)' - def test_equality(self) -> None: + def test_eq(self) -> None: assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector([1, 0, 2, 0, 3, 0]) assert SparseVector([1, 0, 2, 0, 3, 0]) != SparseVector([1, 0, 2, 0, 3, 1]) assert SparseVector([1, 0, 2, 0, 3, 0]) == SparseVector({2: 2, 4: 3, 0: 1, 3: 0}, 6) @@ -140,27 +155,72 @@ def test_to_coo(self) -> None: assert np.array_equal(SparseVector([1, 0, 2, 0, 3, 0]).to_coo().toarray(), [[1, 0, 2, 0, 3, 0]]) - def test_zero_vector_text(self) -> None: - vec = SparseVector({}, 3) - assert vec.to_list() == SparseVector.from_text(vec.to_text()).to_list() + def test_to_list(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_list() == [1.5, 0, 2, 0, 3, 0] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_numpy() + assert np.array_equal(arr, np.array([1.5, 0, 2, 0, 3, 0])) + assert arr.dtype == np.float32 + + def test_to_text(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_text() == '{1:1.5,3:2.0,5:3.0}/6' + + def test_to_binary(self) -> None: + assert SparseVector({2: 2, 4: 3, 0: 1.5, 3: 0}, 6).to_binary() == struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) def test_from_text(self) -> None: - vec = SparseVector.from_text('{1:1.5,3:2,5:3}/6') - assert vec.dimensions() == 6 - assert vec.indices() == [0, 2, 4] - assert vec.values() == [1.5, 2, 3] - assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) - assert vec.to_text() == '{1:1.5,3:2.0,5:3.0}/6' + assert SparseVector.from_text('{1:1.5,3:2,5:3}/6') == SparseVector([1.5, 0, 2, 0, 3, 0]) + + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='not enough values to unpack'): + SparseVector.from_text('') + + def test_from_text_zero_dimensions(self) -> None: + assert SparseVector.from_text('{}/0') == SparseVector([]) + + def test_from_text_non_number_index(self) -> None: + with pytest.raises(ValueError, match='invalid literal for int'): + assert SparseVector.from_text('{a:1}/1') + + def test_from_text_non_number_value(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert SparseVector.from_text('{1:a}/1') + + # TODO uncomment in 0.6.0 + # def test_from_text_unsorted(self) -> None: + # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): + # SparseVector.from_text('{2:1,1:1}/3') + + # TODO uncomment in 0.6.0 + # def test_from_text_duplicate(self) -> None: + # with pytest.raises(ValueError, match='sparsevec indices must be in ascending order'): + # SparseVector.from_text('{2:1,2:1}/3') def test_from_binary(self) -> None: - data = pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) - vec = SparseVector.from_binary(data) - assert vec.dimensions() == 6 - assert vec.indices() == [0, 2, 4] - assert vec.values() == [1.5, 2, 3] - assert vec.to_list() == [1.5, 0, 2, 0, 3, 0] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 0, 2, 0, 3, 0]) - assert vec.to_binary() == data + data = struct.pack('>iii3i3f', 6, 3, 0, 0, 2, 4, 1.5, 2, 3) + assert SparseVector.from_binary(data) == SparseVector([1.5, 0, 2, 0, 3, 0]) + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 12 bytes'): + SparseVector.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii2if', 2, 2, 0, 0, 1, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii2ifff', 2, 2, 0, 0, 1, 1, 2, 3)) + + def test_from_binary_negative_nnz(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + SparseVector.from_binary(struct.pack('>iii', 0, -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + SparseVector.from_binary(struct.pack('>iii', 0, 0, 1)) diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py index 3dd7bac..dc58aeb 100644 --- a/tests/test_sqlalchemy.py +++ b/tests/test_sqlalchemy.py @@ -4,6 +4,7 @@ from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum import pytest from sqlalchemy import create_engine, event, insert, inspect, literal, select, text, MetaData, Table, Column, Index, Integer, ARRAY, Engine +from sqlalchemy.dialects import postgresql from sqlalchemy.exc import StatementError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncEngine from sqlalchemy.ext.automap import automap_base @@ -12,6 +13,8 @@ from typing import Any from .conftest import numpy as np +sqlalchemy_version = 2.1 if hasattr(postgresql, 'BitString') else 2 + psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') psycopg2_type_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test') @@ -420,7 +423,7 @@ def test_select(self, engine: Engine) -> None: def test_select_orm(self, engine: Engine) -> None: with Session(engine) as session: session.add(Item(embedding=[2, 3, 3])) - items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() + items = session.scalars(select(Item.embedding.l2_distance([1, 1, 1]))).all() # type: ignore assert items == [3] def test_avg(self, engine: Engine) -> None: @@ -634,9 +637,15 @@ async def test_bit(self, engine: AsyncEngine) -> None: async with async_session() as session: async with session.begin(): - # typing issue - # https://github.com/MagicStack/py-pgproto/pull/32 - embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101' # type: ignore + embedding: Any + if sqlalchemy_version >= 2.1: + embedding = postgresql.BitString('101') # type: ignore + elif engine == asyncpg_engine: + # typing issue + # https://github.com/MagicStack/py-pgproto/pull/32 + embedding = asyncpg.BitString('101') # type: ignore + else: + embedding = '101' session.add(Item(id=1, binary_embedding=embedding)) item = await session.get_one(Item, 1) assert item.binary_embedding == embedding diff --git a/tests/test_vector.py b/tests/test_vector.py index a08370e..13a9cc7 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -1,29 +1,29 @@ from pgvector import Vector import pytest -from struct import pack +import struct from .conftest import numpy as np class TestVector: - def test_list(self) -> None: + def test_init_list(self) -> None: arr = [1.0, 2.0, 3.0] assert Vector(arr).to_list() == arr assert Vector(arr).to_list() is not arr - def test_list_empty(self) -> None: + def test_init_list_empty(self) -> None: assert Vector([]).to_list() == [] - def test_list_str(self) -> None: + def test_init_list_str(self) -> None: with pytest.raises(ValueError) as error: Vector([1, 'two', 3]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_list_list(self) -> None: + def test_init_list_list(self) -> None: with pytest.raises(ValueError) as error: Vector([[1, 2], [3, 4]]) # type: ignore assert str(error.value) == 'expected list[float]' - def test_ndarray(self) -> None: + def test_init_ndarray(self) -> None: if np is None: pytest.skip('NumPy required') @@ -31,18 +31,32 @@ def test_ndarray(self) -> None: assert Vector(arr).to_list() == [1, 2, 3] assert Vector(arr).to_numpy() is not arr assert Vector(arr).to_numpy().dtype == np.float32 - # non-contiguous + + def test_init_ndarray_non_contiguous(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(np.flip(arr)).to_list() == [3, 2, 1] assert Vector(np.flip(arr)).to_binary() == Vector([3, 2, 1]).to_binary() - # big endian + + def test_init_ndarray_big_endian(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = np.array([1, 2, 3], dtype=np.float32) assert Vector(arr.astype('>f4')).to_list() == [1, 2, 3] assert Vector(arr.astype('>f4')).to_binary() == Vector([1, 2, 3]).to_binary() + def test_init_ndarray_str(self) -> None: + if np is None: + pytest.skip('NumPy required') + with pytest.raises(ValueError) as error: Vector(np.array(['one', 'two', 'three'])) assert 'could not convert string to float' in str(error.value) - def test_int(self) -> None: + def test_init_int(self) -> None: with pytest.raises(ValueError) as error: Vector(1) # type: ignore assert str(error.value) == 'expected list or ndarray' @@ -51,25 +65,74 @@ def test_repr(self) -> None: assert repr(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' assert str(Vector([1, 2, 3])) == 'Vector([1.0, 2.0, 3.0])' - def test_equality(self) -> None: + def test_eq(self) -> None: assert Vector([1, 2, 3]) == Vector([1, 2, 3]) assert Vector([1, 2, 3]) != Vector([1, 2, 4]) assert Vector([1, 2, 3]) != 1 + assert Vector([0.0]) == Vector([-0.0]) def test_dimensions(self) -> None: assert Vector([1, 2, 3]).dimensions() == 3 + def test_to_list(self) -> None: + assert Vector([1.5, 2, 3]).to_list() == [1.5, 2, 3] + + def test_to_numpy(self) -> None: + if np is None: + pytest.skip('NumPy required') + + arr = Vector([1.5, 2, 3]).to_numpy() + assert np.array_equal(arr, np.array([1.5, 2, 3])) + assert arr.dtype == np.float32 + + def test_to_text(self) -> None: + assert Vector([1.5, 2, 3]).to_text() == '[1.5,2.0,3.0]' + + def test_to_binary(self) -> None: + assert Vector([1.5, 2, 3]).to_binary() == struct.pack('>hh3f', 3, 0, 1.5, 2, 3) + def test_from_text(self) -> None: - vec = Vector.from_text('[1.5,2,3]') - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_text() == '[1.5,2.0,3.0]' + assert Vector.from_text('[1.5,2,3]') == Vector([1.5, 2, 3]) + + def test_from_text_zero_length(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + Vector.from_text('') + + def test_from_text_zero_dimensions(self) -> None: + assert Vector.from_text('[]') == Vector([]) + + def test_from_text_open_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert Vector.from_text('{1]') + + def test_from_text_close_bracket(self) -> None: + # TODO raise ValueError in 0.6.0 + assert Vector.from_text('[1}') + + def test_from_text_non_number(self) -> None: + with pytest.raises(ValueError, match='could not convert string to float'): + assert Vector.from_text('[a]') def test_from_binary(self) -> None: - data = pack('>HH3f', 3, 0, 1.5, 2, 3) - vec = Vector.from_binary(data) - assert vec.to_list() == [1.5, 2, 3] - if np is not None: - assert np.array_equal(vec.to_numpy(), [1.5, 2, 3]) - assert vec.to_binary() == data + assert Vector.from_binary(struct.pack('>hh3f', 3, 0, 1.5, 2, 3)) == Vector([1.5, 2, 3]) + + def test_from_binary_zero_length(self) -> None: + # TODO change to ValueError('invalid length') in 0.6.0 + with pytest.raises(struct.error, match='unpack_from requires a buffer of at least 4 bytes'): + Vector.from_binary(b'') + + def test_from_binary_missing_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hhf', 2, 0, 1)) + + def test_from_binary_extra_data(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hhfff', 2, 0, 1, 2, 3)) + + def test_from_binary_negative_dimensions(self) -> None: + with pytest.raises(ValueError, match='invalid length'): + Vector.from_binary(struct.pack('>hh', -1, 0)) + + def test_from_binary_non_zero_unused(self) -> None: + with pytest.raises(ValueError, match='expected unused to be 0'): + Vector.from_binary(struct.pack('>hh', 0, 1))