diff --git a/doc/source/whatsnew/v3.0.0.rst b/doc/source/whatsnew/v3.0.0.rst index 293f1cb6f5e79..42d845edd5c2f 100644 --- a/doc/source/whatsnew/v3.0.0.rst +++ b/doc/source/whatsnew/v3.0.0.rst @@ -1171,6 +1171,7 @@ ExtensionArray - Bug in :class:`Categorical` when constructing with an :class:`Index` with :class:`ArrowDtype` (:issue:`60563`) - Bug in :meth:`.arrays.ArrowExtensionArray.__setitem__` which caused wrong behavior when using an integer array with repeated values as a key (:issue:`58530`) - Bug in :meth:`ArrowExtensionArray.factorize` where NA values were dropped when input was dictionary-encoded even when dropna was set to False(:issue:`60567`) +- Bug in :meth:`NDArrayBackedExtensionArray.take` which produced arrays whose dtypes didn't match their underlying data, when called with integer arrays (:issue:`62448`) - Bug in :meth:`api.types.is_datetime64_any_dtype` where a custom :class:`ExtensionDtype` would return ``False`` for array-likes (:issue:`57055`) - Bug in comparison between object with :class:`ArrowDtype` and incompatible-dtyped (e.g. string vs bool) incorrectly raising instead of returning all-``False`` (for ``==``) or all-``True`` (for ``!=``) (:issue:`59505`) - Bug in constructing pandas data structures when passing into ``dtype`` a string of the type followed by ``[pyarrow]`` while PyArrow is not installed would raise ``NameError`` rather than ``ImportError`` (:issue:`57928`) diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 84aca81420fe1..1245beb4c22d0 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -40,6 +40,7 @@ from pandas.core.dtypes.dtypes import ( DatetimeTZDtype, ExtensionDtype, + NumpyEADtype, PeriodDtype, ) from pandas.core.dtypes.missing import array_equivalent @@ -173,6 +174,17 @@ def take( fill_value=fill_value, axis=axis, ) + if self.dtype in [ + NumpyEADtype(np.uint8), + NumpyEADtype(np.uint16), + NumpyEADtype(np.uint32), + NumpyEADtype(np.uint64), + NumpyEADtype(np.int8), + NumpyEADtype(np.int16), + NumpyEADtype(np.int32), + NumpyEADtype(np.int64), + ]: + return type(self)(new_data) return self._from_backing_data(new_data) # ------------------------------------------------------------------------ diff --git a/pandas/tests/arrays/numpy_/test_numpy.py b/pandas/tests/arrays/numpy_/test_numpy.py index 620a553d5a731..8337c26e4ba10 100644 --- a/pandas/tests/arrays/numpy_/test_numpy.py +++ b/pandas/tests/arrays/numpy_/test_numpy.py @@ -325,6 +325,15 @@ def test_factorize_unsigned(): tm.assert_extension_array_equal(res_unique, NumpyExtensionArray(exp_unique)) +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_take_assigns_correct_dtype(dtype): + array = NumpyExtensionArray(np.array([1, 2, 3], dtype=dtype)) + + result = array.take([-1], allow_fill=True) + + assert result.dtype == NumpyEADtype(np.float64) + + # ---------------------------------------------------------------------------- # Output formatting