diff --git a/pyproject.toml b/pyproject.toml index 23a98558..ee62fcb5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -225,6 +225,13 @@ module = [ ] ignore_missing_imports = true +[[tool.mypy.overrides]] +# https://github.com/apache/arrow/issues/50123 +module = [ + "pyarrow.*", +] +follow_imports = "skip" + [tool.ruff] line-length = 79 extend-exclude = [ diff --git a/src/neo4j/_optional_deps.py b/src/neo4j/_optional_deps.py index a8bb4f6a..70f6246a 100644 --- a/src/neo4j/_optional_deps.py +++ b/src/neo4j/_optional_deps.py @@ -32,13 +32,16 @@ import pandas as pd # type: ignore[no-redef] pa: t.Any = None +pa_compute: t.Any = None with suppress(ImportError): import pyarrow as pa # type: ignore[no-redef] + import pyarrow.compute as pa_compute # type: ignore[no-redef] __all__ = [ "np", "pa", + "pa_compute", "pd", ] diff --git a/src/neo4j/vector.py b/src/neo4j/vector.py index ca11568c..b3b09066 100644 --- a/src/neo4j/vector.py +++ b/src/neo4j/vector.py @@ -32,10 +32,12 @@ # This beautiful construct helps sphinx to properly resolve the type hints. import numpy as _np import pyarrow as _pa + import pyarrow.compute as _pa_compute else: from ._optional_deps import ( np as _np, pa as _pa, + pa_compute as _pa_compute, ) @@ -795,7 +797,7 @@ def to_numpy(self) -> _np.ndarray: ... def from_pyarrow(cls, data: _pa.Array, /) -> _t.Self: width = data.type.byte_width assert cls.size == width - if _pa.compute.count(data, mode="only_null").as_py(): + if _pa_compute.count(data, mode="only_null").as_py(): raise ValueError("PyArrow array must not contain any null values.") _, buffer = data.buffers() buffer = buffer[ diff --git a/tests/unit/common/vector/test_vector.py b/tests/unit/common/vector/test_vector.py index 05aacaca..219b60e9 100644 --- a/tests/unit/common/vector/test_vector.py +++ b/tests/unit/common/vector/test_vector.py @@ -28,6 +28,7 @@ from neo4j._optional_deps import ( np, pa, + pa_compute, ) from neo4j.vector import ( _swap_endian, @@ -1044,7 +1045,7 @@ def test_to_pyarrow_random( v = _vector_from_data(data_be, dtype, endian) array = v.to_pyarrow() assert array.type == pa_type - assert pa.compute.count(array, mode="only_null").as_py() == 0 + assert pa_compute.count(array, mode="only_null").as_py() == 0 buffers = array.buffers() assert len(buffers) == 2 assert buffers[0] is None @@ -1076,7 +1077,7 @@ def test_to_pyarrow_special_values( v = _vector_from_data(data_be, dtype, endian) array = v.to_pyarrow() assert array.type == pa_type - assert pa.compute.count(array, mode="only_null").as_py() == 0 + assert pa_compute.count(array, mode="only_null").as_py() == 0 buffers = array.buffers() assert len(buffers) == 2 assert buffers[0] is None