Improved numpy ffi for boolean arrays (#540)

python/core/python/geoarrow/rust/core/_rust.pyi

@@ -459,11 +459,13 @@ class RectArray:
     def total_bounds(self) -> Tuple[float, float, float, float]: ...
 
 class BooleanArray:
+    def __array__(self) -> NDArray[np.bool_]: ...
     def __eq__(self, other: Self) -> bool: ...
     def __len__(self) -> int: ...
     def __arrow_c_array__(
         self, requested_schema: object | None = None
     ) -> Tuple[object, object]: ...
+    def to_numpy(self) -> NDArray[np.bool_]: ...
 
 # class Float16Array:
 #     def __array__(self) -> NDArray[np.float16]: ...
@@ -899,6 +901,7 @@ class ChunkedRectArray:
     def total_bounds(self) -> Tuple[float, float, float, float]: ...
 
 class ChunkedBooleanArray:
+    def __array__(self) -> NDArray[np.bool_]: ...
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: ...
     def __eq__(self, other: Self) -> bool: ...
     def __len__(self) -> int: ...
@@ -907,6 +910,7 @@ class ChunkedBooleanArray:
     @classmethod
     def from_arrow_arrays(cls, input: Sequence[ArrowArrayExportable]) -> Self: ...
     def num_chunks(self) -> int: ...
+    def to_numpy(self) -> NDArray[np.bool_]: ...
 
 # class ChunkedFloat16Array:
 #     def __array__(self) -> NDArray[np.float16]: ...

python/core/src/interop/numpy/to_numpy.rs

@@ -48,6 +48,25 @@ impl_array!(Int16Array);
 impl_array!(Int32Array);
 impl_array!(Int64Array);
 
+#[pymethods]
+impl BooleanArray {
+    pub fn __array__<'py>(&'py self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        if self.0.null_count() > 0 {
+            return Err(PyValueError::new_err(
+                "Cannot create numpy array from pyarrow array with nulls.",
+            ));
+        }
+
+        let bools = self.0.values().iter().collect::<Vec<_>>();
+        Ok(bools.to_pyarray(py))
+    }
+
+    /// Copy this array to a `numpy` NDArray
+    pub fn to_numpy<'py>(&'py self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        self.__array__(py)
+    }
+}
+
 macro_rules! impl_chunked {
     ($struct_name:ty) => {
         #[pymethods]
@@ -67,6 +86,7 @@ macro_rules! impl_chunked {
                 let numpy_mod = py.import(intern!(py, "numpy"))?;
                 numpy_mod.call_method1(intern!(py, "concatenate"), (py_arrays,))
             }
+
             /// Copy this array to a `numpy` NDArray
             pub fn to_numpy<'py>(&'py self, py: Python<'py>) -> PyResult<&'py PyAny> {
                 self.__array__(py)
@@ -88,6 +108,31 @@ impl_chunked!(ChunkedInt16Array);
 impl_chunked!(ChunkedInt32Array);
 impl_chunked!(ChunkedInt64Array);
 
+#[pymethods]
+impl ChunkedBooleanArray {
+    pub fn __array__<'py>(&'py self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        if self.0.null_count() > 0 {
+            return Err(PyValueError::new_err(
+                "Cannot create numpy array from pyarrow array with nulls.",
+            ));
+        }
+
+        let np_chunks = self
+            .0
+            .chunks()
+            .iter()
+            .map(|chunk| Ok(BooleanArray(chunk.clone()).__array__(py)?.to_object(py)))
+            .collect::<PyResult<Vec<_>>>()?;
+
+        let numpy_mod = py.import(intern!(py, "numpy"))?;
+        numpy_mod.call_method1(intern!(py, "concatenate"), (np_chunks,))
+    }
+
+    /// Copy this array to a `numpy` NDArray
+    pub fn to_numpy<'py>(&'py self, py: Python<'py>) -> PyResult<&'py PyAny> {
+        self.__array__(py)
+    }
+}
 #[pymethods]
 impl WKBArray {
     /// An implementation of the Array interface, for interoperability with numpy and other

src/chunked_array/chunked_array.rs

@@ -54,6 +54,12 @@ impl<A: Array> ChunkedArray<A> {
         self.chunks.first().unwrap().data_type()
     }
 
+    pub fn null_count(&self) -> usize {
+        self.chunks()
+            .iter()
+            .fold(0, |acc, chunk| acc + chunk.null_count())
+    }
+
     pub fn chunks(&self) -> &[A] {
         self.chunks.as_slice()
     }