diff --git a/doc/source/whatsnew/v3.0.0.rst b/doc/source/whatsnew/v3.0.0.rst index 0f7be8cfbcb68..666a00f971d96 100644 --- a/doc/source/whatsnew/v3.0.0.rst +++ b/doc/source/whatsnew/v3.0.0.rst @@ -1299,6 +1299,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/numpy_.py b/pandas/core/arrays/numpy_.py index 9d174a49b5db1..57336cd4e15f1 100644 --- a/pandas/core/arrays/numpy_.py +++ b/pandas/core/arrays/numpy_.py @@ -48,6 +48,7 @@ InterpolateOptions, NpDtype, Scalar, + TakeIndexer, npt, ) @@ -365,6 +366,27 @@ def interpolate( return self return type(self)._simple_new(out_data, dtype=self.dtype) + def take( + self, + indices: TakeIndexer, + *, + allow_fill: bool = False, + fill_value: Any = None, + axis: AxisInt = 0, + ) -> Self: + """ + Take entries from this array at each index in a list of indices, + producing an array containing only those entries. + """ + result = super().take( + indices, allow_fill=allow_fill, fill_value=fill_value, axis=axis + ) + # See GH#62448. + if self.dtype.kind in "iub": + return type(self)(result._ndarray, copy=False) + + return result + # ------------------------------------------------------------------------ # Reductions diff --git a/pandas/tests/arrays/numpy_/test_numpy.py b/pandas/tests/arrays/numpy_/test_numpy.py index f0e495cac4b3d..b43b6a42ae084 100644 --- a/pandas/tests/arrays/numpy_/test_numpy.py +++ b/pandas/tests/arrays/numpy_/test_numpy.py @@ -348,6 +348,42 @@ def test_factorize_unsigned(): tm.assert_extension_array_equal(res_unique, NumpyExtensionArray(exp_unique)) +@pytest.mark.parametrize( + "dtype", + [ + np.bool_, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.int8, + np.int16, + np.int32, + np.int64, + ], +) +def test_take_assigns_floating_point_dtype(dtype): + # GH#62448. + if dtype == np.bool_: + array = NumpyExtensionArray(np.array([False, True, False], dtype=dtype)) + expected = np.dtype(object) + else: + array = NumpyExtensionArray(np.array([1, 2, 3], dtype=dtype)) + expected = np.float64 + + result = array.take([-1], allow_fill=True) + assert result.dtype.numpy_dtype == expected + + result = array.take([-1], allow_fill=True, fill_value=5.0) + assert result.dtype.numpy_dtype == expected + + +def test_take_preserves_boolean_arrays(): + array = NumpyExtensionArray(np.array([False, True, False], dtype=np.bool_)) + result = array.take([-1], allow_fill=False) + assert result.dtype.numpy_dtype == np.bool_ + + # ---------------------------------------------------------------------------- # Output formatting