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

doc/source/whatsnew/v3.0.0.rst

@@ -1176,6 +1176,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

@@ -45,6 +45,7 @@
  InterpolateOptions,
  NpDtype,
  Scalar,
+ TakeIndexer,
  npt,
  )
 
@@ -350,6 +351,36 @@ 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.numpy_dtype in [
+ np.uint8,
+ np.uint16,
+ np.uint32,
+ np.uint64,
+ np.int8,
+ np.int16,
+ np.int32,
+ np.int64,
+ ]:
+ return type(self)(result, copy=False)
+
+ return result
+
  # ------------------------------------------------------------------------
  # Reductions
 

pandas/tests/arrays/numpy_/test_numpy.py

@@ -325,6 +325,37 @@ def test_factorize_unsigned():
  tm.assert_extension_array_equal(res_unique, NumpyExtensionArray(exp_unique))
 
 
+# TODO: Add the smaller width dtypes to the parameter sets of these tests.
+@pytest.mark.parametrize(
+ "dtype",
+ [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.
+ array = NumpyExtensionArray(np.array([1, 2, 3], dtype=dtype))
+
+ result = array.take([-1], allow_fill=True)
+
+ assert result.dtype.numpy_dtype == np.float64
+
+ result = array.take([-1], allow_fill=True, fill_value=5.0)
+
+ assert result.dtype.numpy_dtype == np.float64
+
+
+@pytest.mark.parametrize(
+ "dtype",
+ [np.uint8, np.uint16, np.uint32, np.uint64, np.int8, np.int16, np.int32, np.int64],
+)
+def test_take_assigns_integer_dtype_when_fill_disallowed(dtype):
+ # GH#62448.
+ array = NumpyExtensionArray(np.array([1, 2, 3], dtype=dtype))
+
+ result = array.take([-1], allow_fill=False)
+
+ assert result.dtype.numpy_dtype == dtype
+
+
 # ----------------------------------------------------------------------------
 # Output formatting