Merge pull request #214 from xhochy/allow-numpy-outputs-where-clear

Allow NumPy ufuncs to work with `np.ndarray` outputs where operations are clearly defined

Author: xhochy

CHANGELOG.md

@@ -7,6 +7,11 @@ Starting with 0.5, we will follow the following versioning scheme:
 * We bump MINOR on breaking changes.
 * We increase PATCH otherwise.
 
+0.7.2 (2021-01-17)
+------------------
+
+* Allow NumPy ufuncs to work with `np.ndarray` outputs where operations are clearly defined (i.e. the fletcher array has no nulls).
+
 0.7.1 (2020-12-29)
 ------------------
 

fletcher/base.py

@@ -209,6 +209,10 @@ def name(self) -> str:
         """
         return str(self)
 
+    @property
+    def itemsize(self) -> int:
+        return self.arrow_dtype.bit_width
+
     @property
     def _is_boolean(self):
         return pa.types.is_boolean(self.arrow_dtype)
@@ -547,6 +551,24 @@ def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):
                 )
         if len(inputs) != 2:
             raise NotImplementedError("Only ufuncs with a second input are supported")
+        if "out" in kwargs:
+            out = kwargs.pop("out")
+            if len(out) == 1:
+                out_array = out[0]
+                if isinstance(out_array, np.ndarray) and self.data.null_count == 0:
+                    self_as_np = np.asarray(self)
+                    mapped_inputs = [self_as_np if i is self else i for i in inputs]
+                    return self_as_np.__array_ufunc__(
+                        ufunc, method, *mapped_inputs, **kwargs
+                    )
+                else:
+                    raise NotImplementedError(
+                        "Currently ufuncs with outputs are only supported for arrays without missings"
+                    )
+            else:
+                raise NotImplementedError(
+                    "Currently ufuncs only support a single output"
+                )
         if len(kwargs) > 0:
             raise NotImplementedError("ufuncs with kwargs aren't supported")
         if isinstance(inputs[0], FletcherBaseArray):

tests/test_pandas_integration.py

@@ -278,3 +278,15 @@ def test_setitem_chunked_int_index(indices, test_array_chunked):
     ser[integer_index] = ["int", "index"]
     assert ser[indices[0]] == "int"
     assert ser[indices[1]] == "index"
+
+
+def test_numpy_ufunc_with_keyword(fletcher_array):
+    # https://github.com/xhochy/fletcher/issues/213
+
+    x = fletcher_array([0])
+    y = np.array([False])
+
+    # Test without in-place storage
+    y | (x == 1)
+    # Test with in-place storage
+    y |= x == 1