Removed dependency on NumPy - closes #146

Author: ankane

CHANGELOG.md

@@ -4,6 +4,7 @@
 - Changed `vector` type to return `Vector` object instead of NumPy array
 - Removed `utils` package (use top-level `pgvector` package instead)
 - Removed re-exported classes (use top-level `pgvector` package instead)
+- Removed dependency on NumPy
 - Dropped support for Python < 3.10
 - Dropped support for SQLAlchemy < 2
 

pgvector/bit.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
-import numpy as np
 from struct import pack, unpack_from
 from warnings import warn
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class Bit:
     def __init__(self, value: bytes | str | list[bool] | np.ndarray[tuple[int], np.dtype[np.bool | np.uint8]]) -> None:
@@ -30,7 +34,7 @@ def bit_value(v: bool) -> str:
                     data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
                 except ValueError:
                     raise ValueError('expected bit string')
-            elif isinstance(value, np.ndarray):
+            elif np is not None and isinstance(value, np.ndarray):
                 if value.dtype != np.bool:
                     # skip warning for result of np.unpackbits
                     if value.dtype != np.uint8 or np.any(value > 1):
@@ -63,7 +67,8 @@ def _length(self):
         return length
 
     def to_list(self) -> list[bool]:
-        return self.to_numpy().tolist()
+        # TODO improve
+        return [v != '0' for v in self.to_text()]
 
     def to_numpy(self) -> np.ndarray[tuple[int], np.dtype[np.bool]]:
         return np.unpackbits(np.frombuffer(self._value[4:], dtype=np.uint8), count=self._length()).astype(bool)

pgvector/halfvec.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
-import numpy as np
 import struct
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class HalfVector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
@@ -11,7 +15,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float
                 self._value = struct.pack(f'>HH{dim}e', dim, 0, *value)
             except struct.error:
                 raise ValueError('expected list[float]')
-        elif isinstance(value, np.ndarray):
+        elif np is not None and isinstance(value, np.ndarray):
             if value.ndim != 1:
                 raise ValueError('expected ndim to be 1')
 

pgvector/sparsevec.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
-import numpy as np
 from struct import pack, unpack_from
 from typing import Any, overload
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 NO_DEFAULT = object()
 
 

pgvector/vector.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
-import numpy as np
 import struct
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class Vector:
     def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.floating]]) -> None:
@@ -11,7 +15,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float
                 self._value = struct.pack(f'>HH{dim}f', dim, 0, *value)
             except struct.error:
                 raise ValueError('expected list[float]')
-        elif isinstance(value, np.ndarray):
+        elif np is not None and isinstance(value, np.ndarray):
             if value.ndim != 1:
                 raise ValueError('expected ndim to be 1')
 

pyproject.toml

@@ -13,7 +13,6 @@ authors = [
 license = "MIT"
 requires-python = ">= 3.10"
 dependencies = [
-    "numpy"
 ]
 
 [project.urls]
@@ -29,10 +28,13 @@ dev = [
     "psycopg2-binary",
     "pytest",
     "pytest-asyncio",
-    "scipy",
     "SQLAlchemy[asyncio]>=2",
     "sqlmodel>=0.0.12"
 ]
+dev-optional = [
+    "numpy",
+    "scipy"
+]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"

tests/test_asyncpg.py

@@ -1,5 +1,4 @@
 import asyncpg
-import numpy as np
 from pgvector import HalfVector, SparseVector, Vector
 from pgvector.asyncpg import register_vector
 import pytest
@@ -16,7 +15,7 @@ async def test_vector(self):
         await register_vector(conn)
 
         embedding = Vector([1.5, 2, 3])
-        embedding2 = np.array([4.5, 5, 6])
+        embedding2 = [4.5, 5, 6]
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
@@ -111,7 +110,7 @@ async def test_vector_array(self):
         embeddings = [Vector([1.5, 2, 3]), Vector([4.5, 5, 6])]
         await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES ($1)", embeddings)
 
-        embeddings2 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
+        embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]]
         await conn.execute("INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])", embeddings2[0], embeddings2[1])
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")
@@ -133,7 +132,7 @@ async def init(conn):
             await conn.execute('CREATE TABLE asyncpg_items (id bigserial PRIMARY KEY, embedding vector(3))')
 
             embedding = Vector([1.5, 2, 3])
-            embedding2 = np.array([1.5, 2, 3])
+            embedding2 = [1.5, 2, 3]
             await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), (NULL)", embedding, embedding2)
 
             res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")

tests/test_bit.py

@@ -1,7 +1,11 @@
-import numpy as np
 from pgvector import Bit
 import pytest
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class TestBit:
     def test_list(self):
@@ -29,20 +33,24 @@ def test_bytes(self):
         assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000'
         assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000'
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndarray(self):
         arr = np.array([True, False, True])
         assert Bit(arr).to_list() == [True, False, True]
         assert np.array_equal(Bit(arr).to_numpy(), arr)
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndarray_unpackbits(self):
         arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8))
         assert Bit(arr).to_text() == '111111100000011100000000'
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndarray_uint8(self):
         arr = np.array([254, 7, 0], dtype=np.uint8)
         with pytest.warns(UserWarning, match='expected elements to be boolean'):
             assert Bit(arr).to_text() == '110'
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndarray_uint16(self):
         arr = np.array([254, 7, 0], dtype=np.uint16)
         with pytest.warns(UserWarning, match='expected elements to be boolean'):

tests/test_django.py

@@ -11,13 +11,17 @@
 from django.db.migrations.loader import MigrationLoader
 from django.forms import ModelForm
 from math import sqrt
-import numpy as np
 import os
 import pgvector.django
 from pgvector import HalfVector, SparseVector, Vector
 from pgvector.django import VectorExtension, VectorField, HalfVectorField, BitField, SparseVectorField, IvfflatIndex, HnswIndex, L2Distance, MaxInnerProduct, CosineDistance, L1Distance, HammingDistance, JaccardDistance
 from unittest import mock
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 settings.configure(
     DATABASES={
         'default': {
@@ -458,7 +462,7 @@ def test_missing(self):
         assert Item.objects.first().sparse_embedding is None
 
     def test_vector_array(self):
-        Item(id=1, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]).save()
+        Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]).save()
 
         with connection.cursor() as cursor:
             from pgvector.psycopg import register_vector

tests/test_half_vector.py

@@ -1,8 +1,12 @@
-import numpy as np
 from pgvector import HalfVector
 import pytest
 from struct import pack
 
+try:
+    import numpy as np
+except ImportError:
+    np = None
+
 
 class TestHalfVector:
     def test_list(self):
@@ -21,6 +25,7 @@ def test_list_list(self):
             HalfVector([[1, 2], [3, 4]])  # ty: ignore[invalid-argument-type]
         assert str(error.value) == 'expected list[float]'
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_ndarray(self):
         arr = np.array([1, 2, 3])
         assert HalfVector(arr).to_list() == [1, 2, 3]
@@ -42,6 +47,7 @@ def test_equality(self):
     def test_dimensions(self):
         assert HalfVector([1, 2, 3]).dimensions() == 3
 
+    @pytest.mark.skipif(np is None, reason='NumPy required')
     def test_to_numpy_readonly(self):
         arr = HalfVector([1, 2, 3]).to_numpy()
         with pytest.raises(ValueError) as error:
@@ -51,11 +57,13 @@ def test_to_numpy_readonly(self):
     def test_from_text(self):
         vec = HalfVector.from_text('[1.5,2,3]')
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
+        if np is not None:
+            assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
 
     def test_from_binary(self):
         data = pack('>HH3e', 3, 0, 1.5, 2, 3)
         vec = HalfVector.from_binary(data)
         assert vec.to_list() == [1.5, 2, 3]
-        assert np.array_equal(vec.to_numpy(), [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