BUG: Take method of NumpyExtensionArray now returns another extension array with the correct dtype. (#62502)

Author: aijams

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`)

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
 

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