REF: move __array_ufunc__ to base NumericArray (#38412)

Author: jorisvandenbossche

pandas/core/arrays/floating.py

@@ -1,4 +1,3 @@
-import numbers
 from typing import List, Optional, Tuple, Type
 import warnings
 
@@ -22,7 +21,6 @@
 from pandas.core.dtypes.dtypes import ExtensionDtype, register_extension_dtype
 from pandas.core.dtypes.missing import isna
 
-from pandas.core import ops
 from pandas.core.ops import invalid_comparison
 from pandas.core.tools.numeric import to_numeric
 
@@ -255,54 +253,6 @@ def _from_sequence_of_strings(
  scalars = to_numeric(strings, errors="raise")
  return cls._from_sequence(scalars, dtype=dtype, copy=copy)
 
- _HANDLED_TYPES = (np.ndarray, numbers.Number)
-
- def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):
- # For FloatingArray inputs, we apply the ufunc to ._data
- # and mask the result.
- if method == "reduce":
- # Not clear how to handle missing values in reductions. Raise.
- raise NotImplementedError("The 'reduce' method is not supported.")
- out = kwargs.get("out", ())
-
- for x in inputs + out:
- if not isinstance(x, self._HANDLED_TYPES + (FloatingArray,)):
- 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 NotImplemented:
- return result
-
- mask = np.zeros(len(self), dtype=bool)
- inputs2 = []
- for x in inputs:
- if isinstance(x, FloatingArray):
- mask |= x._mask
- inputs2.append(x._data)
- else:
- inputs2.append(x)
-
- def reconstruct(x):
- # we don't worry about scalar `x` here, since we
- # raise for reduce up above.
-
- # TODO
- if is_float_dtype(x.dtype):
- m = mask.copy()
- return FloatingArray(x, m)
- else:
- x[mask] = np.nan
- return x
-
- result = getattr(ufunc, method)(*inputs2, **kwargs)
- if isinstance(result, tuple):
- tuple(reconstruct(x) for x in result)
- else:
- return reconstruct(result)
-
  def _coerce_to_array(self, value) -> Tuple[np.ndarray, np.ndarray]:
  return coerce_to_array(value, dtype=self.dtype)
 

pandas/core/arrays/integer.py

@@ -1,4 +1,3 @@
-import numbers
 from typing import Dict, List, Optional, Tuple, Type
 import warnings
 
@@ -22,7 +21,6 @@
 )
 from pandas.core.dtypes.missing import isna
 
-from pandas.core import ops
 from pandas.core.ops import invalid_comparison
 from pandas.core.tools.numeric import to_numeric
 
@@ -316,53 +314,6 @@ def _from_sequence_of_strings(
  scalars = to_numeric(strings, errors="raise")
  return cls._from_sequence(scalars, dtype=dtype, copy=copy)
 
- _HANDLED_TYPES = (np.ndarray, numbers.Number)
-
- def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):
- # For IntegerArray inputs, we apply the ufunc to ._data
- # and mask the result.
- if method == "reduce":
- # Not clear how to handle missing values in reductions. Raise.
- raise NotImplementedError("The 'reduce' method is not supported.")
- 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 NotImplemented:
- return result
-
- mask = np.zeros(len(self), dtype=bool)
- inputs2 = []
- for x in inputs:
- if isinstance(x, IntegerArray):
- mask |= x._mask
- inputs2.append(x._data)
- else:
- inputs2.append(x)
-
- def reconstruct(x):
- # we don't worry about scalar `x` here, since we
- # raise for reduce up above.
-
- if is_integer_dtype(x.dtype):
- m = mask.copy()
- return IntegerArray(x, m)
- else:
- x[mask] = np.nan
- return x
-
- result = getattr(ufunc, method)(*inputs2, **kwargs)
- if isinstance(result, tuple):
- return tuple(reconstruct(x) for x in result)
- else:
- return reconstruct(result)
-
  def _coerce_to_array(self, value) -> Tuple[np.ndarray, np.ndarray]:
  return coerce_to_array(value, dtype=self.dtype)
 

pandas/core/arrays/numeric.py

@@ -1,5 +1,6 @@
 import datetime
-from typing import TYPE_CHECKING, Union
+import numbers
+from typing import TYPE_CHECKING, Any, List, Union
 
 import numpy as np
 
@@ -14,6 +15,8 @@
  is_list_like,
 )
 
+from pandas.core import ops
+
 from .masked import BaseMaskedArray, BaseMaskedDtype
 
 if TYPE_CHECKING:
@@ -130,3 +133,57 @@ def _arith_method(self, other, op):
  )
 
  return self._maybe_mask_result(result, mask, other, op_name)
+
+ _HANDLED_TYPES = (np.ndarray, numbers.Number)
+
+ def __array_ufunc__(self, ufunc, method: str, *inputs, **kwargs):
+ # For NumericArray inputs, we apply the ufunc to ._data
+ # and mask the result.
+ if method == "reduce":
+ # Not clear how to handle missing values in reductions. Raise.
+ raise NotImplementedError("The 'reduce' method is not supported.")
+ out = kwargs.get("out", ())
+
+ for x in inputs + out:
+ if not isinstance(x, self._HANDLED_TYPES + (NumericArray,)):
+ 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 NotImplemented:
+ return result
+
+ mask = np.zeros(len(self), dtype=bool)
+ inputs2: List[Any] = []
+ for x in inputs:
+ if isinstance(x, NumericArray):
+ mask |= x._mask
+ inputs2.append(x._data)
+ else:
+ inputs2.append(x)
+
+ def reconstruct(x):
+ # we don't worry about scalar `x` here, since we
+ # raise for reduce up above.
+
+ if is_integer_dtype(x.dtype):
+ from pandas.core.arrays import IntegerArray
+
+ m = mask.copy()
+ return IntegerArray(x, m)
+ elif is_float_dtype(x.dtype):
+ from pandas.core.arrays import FloatingArray
+
+ m = mask.copy()
+ return FloatingArray(x, m)
+ else:
+ x[mask] = np.nan
+ return x
+
+ result = getattr(ufunc, method)(*inputs2, **kwargs)
+ if isinstance(result, tuple):
+ return tuple(reconstruct(x) for x in result)
+ else:
+ return reconstruct(result)

pandas/tests/arrays/integer/test_function.py

@@ -3,6 +3,7 @@
 
 import pandas as pd
 import pandas._testing as tm
+from pandas.core.arrays import FloatingArray
 
 
 @pytest.mark.parametrize("ufunc", [np.abs, np.sign])
@@ -25,13 +26,13 @@ def test_ufuncs_single_float(ufunc):
  a = pd.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)
+ expected = FloatingArray(ufunc(a.astype(float)), mask=a._mask)
+ tm.assert_extension_array_equal(result, expected)
 
  s = pd.Series(a)
  with np.errstate(invalid="ignore"):
  result = ufunc(s)
-  expected = ufunc(s.astype(float))
+ expected = pd.Series(expected)
  tm.assert_series_equal(result, expected)
 
 
@@ -67,14 +68,13 @@ def test_ufunc_binary_output():
  a = pd.array([1, 2, np.nan])
  result = np.modf(a)
  expected = np.modf(a.to_numpy(na_value=np.nan, dtype="float"))
+ expected = (pd.array(expected[0]), pd.array(expected[1]))
 
  assert isinstance(result, tuple)
  assert len(result) == 2
 
  for x, y in zip(result, expected):
- # TODO(FloatArray): This will return an extension array.
- # y = pd.array(y)
- tm.assert_numpy_array_equal(x, y)
+ tm.assert_extension_array_equal(x, y)
 
 
 @pytest.mark.parametrize("values", [[0, 1], [0, None]])