Add __array_ufunc__ to Series / Array

pandas/core/arrays/integer.py

@@ -1,3 +1,4 @@
+import numbers
 import sys
 import warnings
 import copy
@@ -10,6 +11,7 @@
 from pandas.compat import set_function_name
 
 from pandas.core import nanops
+from pandas.core import ops
 from pandas.core.dtypes.cast import astype_nansafe
 from pandas.core.dtypes.generic import ABCSeries, ABCIndexClass
 from pandas.core.dtypes.common import (
@@ -293,6 +295,44 @@ def __array__(self, dtype=None):
         """
         return self._coerce_to_ndarray()
 
+    _HANDLED_TYPES = (np.ndarray, numbers.Number)
+
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+
+        out = kwargs.get('out', ())
+
+        for x in inputs + out:
+            if not isinstance(x, self._HANDLED_TYPES + (IntegerArray,)):
+                return NotImplemented
+
+        # for binary ops, use our custom dunder methods
+        result = ops.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs)
+        if result is not None:
+            return result
+
+        if (method == '__call__'
+                and ufunc.signature is None
+                and ufunc.nout == 1):
+            # only supports IntegerArray for now
+            args = [a._data for a in inputs]
+            masks = [a._mask for a in inputs]
+            result = ufunc(*args, **kwargs)
+            mask = np.logical_or.reduce(masks)
+            if result.dtype.kind in ('i', 'u'):
+                return IntegerArray(result, mask)
+            else:
+                result[mask] = np.nan
+                return result
+
+        # fall back to array for other ufuncs
+        inputs = tuple(
+            np.array(x) if isinstance(x, type(self)) else x
+            for x in inputs
+        )
+        return np.array(self).__array_ufunc__(
+            ufunc, method, *inputs, **kwargs)
+
     def __iter__(self):
         for i in range(len(self)):
             if self._mask[i]:

pandas/core/arrays/sparse.py

@@ -47,6 +47,7 @@
 from pandas._libs import lib
 import pandas.core.algorithms as algos
 import pandas.io.formats.printing as printing
+import pandas.core.ops as ops
 
 
 # ----------------------------------------------------------------------------
@@ -1447,44 +1448,11 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             if not isinstance(x, self._HANDLED_TYPES + (SparseArray,)):
                 return NotImplemented
 
-        special = {'add', 'sub', 'mul', 'pow', 'mod', 'floordiv', 'truediv',
-                   'divmod', 'eq', 'ne', 'lt', 'gt', 'le', 'ge', 'remainder'}
-        if compat.PY2:
-            special.add('div')
-        aliases = {
-            'subtract': 'sub',
-            'multiply': 'mul',
-            'floor_divide': 'floordiv',
-            'true_divide': 'truediv',
-            'power': 'pow',
-            'remainder': 'mod',
-            'divide': 'div',
-            'equal': 'eq',
-            'not_equal': 'ne',
-            'less': 'lt',
-            'less_equal': 'le',
-            'greater': 'gt',
-            'greater_equal': 'ge',
-        }
-
-        flipped = {
-            'lt': '__gt__',
-            'le': '__ge__',
-            'gt': '__lt__',
-            'ge': '__le__',
-            'eq': '__eq__',
-            'ne': '__ne__',
-        }
-
-        op_name = ufunc.__name__
-        op_name = aliases.get(op_name, op_name)
-
-        if op_name in special and kwargs.get('out') is None:
-            if isinstance(inputs[0], type(self)):
-                return getattr(self, '__{}__'.format(op_name))(inputs[1])
-            else:
-                name = flipped.get(op_name, '__r{}__'.format(op_name))
-                return getattr(self, name)(inputs[0])
+        # for binary ops, use our custom dunder methods
+        result = ops.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs)
+        if result is not None:
+            return result
 
         if len(inputs) == 1:
             # No alignment necessary.

pandas/core/ops.py

@@ -2146,3 +2146,46 @@ def wrapper(self, other):
 
     wrapper.__name__ = op_name
     return wrapper
+
+
+def maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs):
+
+    special = {'add', 'sub', 'mul', 'pow', 'mod', 'floordiv', 'truediv',
+               'divmod', 'eq', 'ne', 'lt', 'gt', 'le', 'ge', 'remainder'}
+    if compat.PY2:
+        special.add('div')
+    aliases = {
+        'subtract': 'sub',
+        'multiply': 'mul',
+        'floor_divide': 'floordiv',
+        'true_divide': 'truediv',
+        'power': 'pow',
+        'remainder': 'mod',
+        'divide': 'div',
+        'equal': 'eq',
+        'not_equal': 'ne',
+        'less': 'lt',
+        'less_equal': 'le',
+        'greater': 'gt',
+        'greater_equal': 'ge',
+    }
+
+    flipped = {
+        'lt': '__gt__',
+        'le': '__ge__',
+        'gt': '__lt__',
+        'ge': '__le__',
+        'eq': '__eq__',
+        'ne': '__ne__',
+    }
+
+    op_name = ufunc.__name__
+    op_name = aliases.get(op_name, op_name)
+
+    if (method == '__call__' and op_name in special
+            and kwargs.get('out') is None):
+        if isinstance(inputs[0], type(self)):
+            return getattr(self, '__{}__'.format(op_name))(inputs[1])
+        else:
+            name = flipped.get(op_name, '__r{}__'.format(op_name))
+        return getattr(self, name)(inputs[0])

pandas/core/series.py

@@ -623,6 +623,29 @@ def view(self, dtype=None):
         return self._constructor(self._values.view(dtype),
                                  index=self.index).__finalize__(self)
 
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        inputs = tuple(
+            x._values if isinstance(x, type(self)) else x
+            for x in inputs
+        )
+        # for binary ops, use our custom dunder methods
+        result = ops.maybe_dispatch_ufunc_to_dunder_op(
+            self, ufunc, method, *inputs, **kwargs)
+        if result is not None:
+            return result
+
+        if hasattr(self._values, '__array_ufunc__'):
+            result = self._values.__array_ufunc__(
+                ufunc, method, *inputs, **kwargs)
+        else:
+            result = np.array(self._values).__array_ufunc__(
+                ufunc, method, *inputs, **kwargs)
+        if result is NotImplemented:
+            raise TypeError("The '{0}' operation is not supported for "
+                            "dtype {1}.".format(ufunc.__name__, self.dtype))
+        return self._constructor(result, index=self.index,
+                                 copy=False).__finalize__(self)
+
     def __array__(self, result=None):
         """
         the array interface, return my values
@@ -640,10 +663,9 @@ def __array_prepare__(self, result, context=None):
         """
         Gets called prior to a ufunc
         """
-
         # nice error message for non-ufunc types
         if (context is not None and
-                not isinstance(self._values, (np.ndarray, ABCSparseArray))):
+                not isinstance(self._values, (np.ndarray, ExtensionArray))):
             obj = context[1][0]
             raise TypeError("{obj} with dtype {dtype} cannot perform "
                             "the numpy op {op}".format(

pandas/tests/arrays/test_integer.py

@@ -690,6 +690,70 @@ def test_astype_nansafe():
         arr.astype('uint32')
 
 
+@pytest.mark.parametrize(
+    'ufunc', [np.abs, np.positive, np.negative])
+def test_ufuncs_single_int(ufunc):
+    a = integer_array([1, 2, -3, np.nan])
+    result = ufunc(a)
+    expected = integer_array(ufunc(a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+    s = pd.Series(a)
+    result = ufunc(s)
+    expected = pd.Series(integer_array(ufunc(a.astype(float))))
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    'ufunc', [np.log, np.exp, np.sin, np.cos, np.sqrt])
+def test_ufuncs_single_float(ufunc):
+    a = integer_array([1, 2, -3, np.nan])
+    with np.errstate(invalid='ignore'):
+        result = ufunc(a)
+        expected = ufunc(a.astype(float))
+    tm.assert_numpy_array_equal(result, expected)
+
+    s = pd.Series(a)
+    with np.errstate(invalid='ignore'):
+        result = ufunc(s)
+        expected = ufunc(s.astype(float))
+    tm.assert_series_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    'ufunc', [np.add, np.subtract])
+def test_ufuncs_binary_int(ufunc):
+    # two IntegerArrays
+    a = integer_array([1, 2, -3, np.nan])
+    result = ufunc(a, a)
+    expected = integer_array(ufunc(a.astype(float), a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+    # IntegerArray with numpy array
+    arr = np.array([1, 2, 3, 4])
+    result = ufunc(a, arr)
+    expected = integer_array(ufunc(a.astype(float), arr))
+    tm.assert_extension_array_equal(result, expected)
+
+    result = ufunc(arr, a)
+    expected = integer_array(ufunc(arr, a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+    # IntegerArray with scalar
+    result = ufunc(a, 1)
+    expected = integer_array(ufunc(a.astype(float), 1))
+    tm.assert_extension_array_equal(result, expected)
+
+    result = ufunc(1, a)
+    expected = integer_array(ufunc(1, a.astype(float)))
+    tm.assert_extension_array_equal(result, expected)
+
+
+def test_ufunc_fallback():
+    a = integer_array([1, 2, -3, np.nan])
+    assert pd.isna(np.add.reduce(a))
+
+
 # TODO(jreback) - these need testing / are broken
 
 # shift

pandas/tests/extension/decimal/test_decimal.py

@@ -364,3 +364,11 @@ def test_divmod_array(reverse, expected_div, expected_mod):
 
     tm.assert_extension_array_equal(div, expected_div)
     tm.assert_extension_array_equal(mod, expected_mod)
+
+
+def test_ufunc_fallback(data):
+    a = data[:5]
+    s = pd.Series(a, index=range(3, 8))
+    result = np.abs(s)
+    expected = pd.Series(np.abs(a.astype(object)), index=range(3, 8))
+    tm.assert_series_equal(result, expected)