Closes #5219: improve astype in ak.pandas.extension

Author: ajpotts

arkouda/numpy/_typing/_typing.py

@@ -28,7 +28,9 @@
     None,
 ]
 
-StringDTypeTypes: TypeAlias = _Union[Literal["str", "str_"], type[str_], type[str], type[Strings]]
+StringDTypeTypes: TypeAlias = _Union[
+    Literal["str", "str_", "string"], type[str_], type[str], type[Strings]
+]
 
 _ArrayLikeNum: TypeAlias = _Union[
     np.ndarray,  # keeps it simple; or list your NDArray[...]
@@ -88,4 +90,9 @@
 
 def is_string_dtype_hint(x: object) -> TypeGuard["_StringDType"]:
     # accept the spellings you want to map to Arkouda Strings
-    return x in ("str", "str_") or x is str_ or x is str_ or x is Strings
+    return (
+        (isinstance(x, str) and x.lower() in ("str", "str_", "string", "trings"))
+        or x is str_
+        or x is str
+        or x is Strings
+    )

arkouda/numpy/dtypes.py

@@ -258,7 +258,7 @@ def dtype(x):
         return bigint()
 
     # ---- String dtype spellings ----
-    if isinstance(x, str) and x.lower() in {"str", "str_", "Strings", "strings"}:
+    if isinstance(x, str) and x.lower() in {"str", "str_", "strings", "string"}:
         return np.dtype(np.str_)
     if x in (str, np.str_):
         return np.dtype(np.str_)

arkouda/pandas/extension/_arkouda_array.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence, TypeVar
+from typing import TYPE_CHECKING, Any, Sequence, TypeVar, Union, overload
 from typing import cast as type_cast
 
 import numpy as np
 
 from numpy import ndarray
 from numpy.typing import NDArray
 from pandas.api.extensions import ExtensionArray
-
-from arkouda.numpy.dtypes import dtype as ak_dtype
+from pandas.core.dtypes.dtypes import ExtensionDtype
 
 from ._arkouda_extension_array import ArkoudaExtensionArray
 from ._dtypes import (
@@ -167,27 +166,109 @@ def __setitem__(self, key, value):
 
         self._data[key] = value
 
-    def astype(self, dtype, copy: bool = False):
-        # Always hand back a real object-dtype ndarray when object is requested
-        if dtype in (object, np.object_, "object", np.dtype("O")):
-            return self.to_ndarray().astype(object, copy=copy)
+    @overload
+    def astype(self, dtype: np.dtype[Any], copy: bool = True) -> NDArray[Any]: ...
 
-        if isinstance(dtype, _ArkoudaBaseDtype):
-            dtype = dtype.numpy_dtype
+    @overload
+    def astype(self, dtype: ExtensionDtype, copy: bool = True) -> ExtensionArray: ...
 
-        # Server-side cast for numeric/bool
-        try:
-            npdt = np.dtype(dtype)
-        except Exception:
-            return self.to_ndarray().astype(dtype, copy=copy)
+    @overload
+    def astype(self, dtype: Any, copy: bool = True) -> Union[ExtensionArray, NDArray[Any]]: ...
+
+    def astype(
+        self,
+        dtype: Any,
+        copy: bool = True,
+    ) -> Union[ExtensionArray, NDArray[Any]]:
+        """
+        Cast the array to a specified dtype.
+
+        Casting rules:
+
+        * If ``dtype`` requests ``object``, returns a NumPy ``NDArray[Any]`` of
+          dtype ``object`` containing the array values.
+        * Otherwise, the target dtype is normalized using Arkouda's dtype
+          resolution rules.
+        * If the normalized dtype matches the current dtype and ``copy=False``,
+          returns ``self``.
+        * In all other cases, casts the underlying Arkouda array to the target
+          dtype and returns an Arkouda-backed ``ArkoudaExtensionArray``.
+
+        Parameters
+        ----------
+        dtype : Any
+            Target dtype. May be a NumPy dtype, pandas dtype, Arkouda dtype,
+            or any dtype-like object accepted by Arkouda.
+        copy : bool
+            Whether to force a copy when the target dtype matches the current dtype.
+            Default is True.
+
+        Returns
+        -------
+        Union[ExtensionArray, NDArray[Any]]
+            The cast result. Returns a NumPy array only when casting to ``object``;
+            otherwise returns an Arkouda-backed ExtensionArray.
+
+        Examples
+        --------
+        Basic numeric casting returns an Arkouda-backed array:
+
+        >>> import arkouda as ak
+        >>> from arkouda.pandas.extension import ArkoudaArray
+        >>> a = ArkoudaArray(ak.array([1, 2, 3], dtype="int64"))
+        >>> a.astype("float64").to_ndarray()
+        array([1., 2., 3.])
+
+        Casting to the same dtype with ``copy=False`` returns the original object:
+
+        >>> b = a.astype("int64", copy=False)
+        >>> b is a
+        True
+
+        Forcing a copy when the dtype is unchanged returns a new array:
+
+        >>> c = a.astype("int64", copy=True)
+        >>> c is a
+        False
+        >>> c.to_ndarray()
+        array([1, 2, 3])
+
+        Casting to ``object`` materializes the data to a NumPy array:
+
+        >>> a.astype(object)
+        array([1, 2, 3], dtype=object)
+
+        NumPy and pandas dtype objects are also accepted:
+
+        >>> import numpy as np
+        >>> a.astype(np.dtype("bool")).to_ndarray()
+        array([ True,  True,  True])
+        """
+        from arkouda.numpy.dtypes import dtype as ak_dtype
+
+        # --- 1) ExtensionDtype branch (satisfies overload #2) ---
+        if isinstance(dtype, ExtensionDtype):
+            # pandas extension dtypes typically have .numpy_dtype
+            if hasattr(dtype, "numpy_dtype"):
+                dtype = dtype.numpy_dtype
+
+            if copy is False and self.dtype.numpy_dtype == dtype:
+                return self
+
+            casted = self._data.astype(dtype)
+            return type_cast(ExtensionArray, ArkoudaExtensionArray._from_sequence(casted))
+
+        # --- 2) object -> numpy (satisfies overload #1 / general) ---
+        if dtype in (object, np.object_, "object", np.dtype("O")):
+            return self.to_ndarray().astype(object, copy=False)
 
-        from arkouda.numpy.numeric import cast as ak_cast
+        dtype = ak_dtype(dtype)
 
-        if npdt.kind in {"i", "u", "f", "b"}:
-            return type(self)(ak_cast(self._data, ak_dtype(npdt.name)))
+        if copy is False and self.dtype.numpy_dtype == dtype:
+            return self
 
-        # Fallback: local cast
-        return self.to_ndarray().astype(npdt, copy=copy)
+        casted = self._data.astype(dtype)
+        return ArkoudaExtensionArray._from_sequence(casted)
 
     def isna(self) -> NDArray[np.bool_]:
         from arkouda.numpy import isnan

arkouda/pandas/extension/_arkouda_categorical_array.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence, TypeVar
+from typing import TYPE_CHECKING, Any, Sequence, TypeVar, Union, overload
+from typing import cast as type_cast
 
-import numpy as np  # new
+import numpy as np
 
 from numpy import ndarray
+from numpy.typing import NDArray
+from pandas import CategoricalDtype as pd_CategoricalDtype
+from pandas import StringDtype as pd_StringDtype
 from pandas.api.extensions import ExtensionArray
+from pandas.core.dtypes.dtypes import ExtensionDtype
 
 import arkouda as ak
 
 from ._arkouda_array import ArkoudaArray
 from ._arkouda_extension_array import ArkoudaExtensionArray
+from ._arkouda_string_array import ArkoudaStringArray
 from ._dtypes import ArkoudaCategoricalDtype
 
 
@@ -84,8 +90,125 @@ def __getitem__(self, idx):
             return self._data[idx]
         return ArkoudaCategoricalArray(self._data[idx])
 
-    def astype(self, x, dtype):
-        raise NotImplementedError("array_api.astype is not implemented in Arkouda yet")
+    @overload
+    def astype(self, dtype: np.dtype[Any], copy: bool = True) -> NDArray[Any]: ...
+
+    @overload
+    def astype(self, dtype: ExtensionDtype, copy: bool = True) -> ExtensionArray: ...
+
+    @overload
+    def astype(self, dtype: Any, copy: bool = True) -> Union[ExtensionArray, NDArray[Any]]: ...
+
+    def astype(
+        self,
+        dtype: Any,
+        copy: bool = True,
+    ) -> Union[ExtensionArray, NDArray[Any]]:
+        """
+        Cast to a specified dtype.
+
+        * If ``dtype`` is categorical (pandas ``category`` / ``CategoricalDtype`` /
+          ``ArkoudaCategoricalDtype``), returns an Arkouda-backed
+          ``ArkoudaCategoricalArray`` (optionally copied).
+        * If ``dtype`` requests ``object``, returns a NumPy ``ndarray`` of dtype object
+          containing the category labels (materialized to the client).
+        * If ``dtype`` requests a string dtype, returns an Arkouda-backed
+          ``ArkoudaStringArray`` containing the labels as strings.
+        * Otherwise, casts the labels (as strings) to the requested dtype and returns an
+          Arkouda-backed ExtensionArray.
+
+        Parameters
+        ----------
+        dtype : Any
+            Target dtype.
+        copy : bool
+            Whether to force a copy when possible. If categorical-to-categorical and
+            ``copy=True``, attempts to copy the underlying Arkouda ``Categorical`` (if
+            supported). Default is True.
+
+        Returns
+        -------
+        Union[ExtensionArray, NDArray[Any]]
+            The cast result. Returns a NumPy array only when casting to ``object``;
+            otherwise returns an Arkouda-backed ExtensionArray.
+
+        Examples
+        --------
+        Casting to ``category`` returns an Arkouda-backed categorical array:
+
+        >>> import arkouda as ak
+        >>> from arkouda.pandas.extension import ArkoudaCategoricalArray
+        >>> c = ArkoudaCategoricalArray(ak.Categorical(ak.array(["x", "y", "x"])))
+        >>> out = c.astype("category")
+        >>> out is c
+        False
+
+        Forcing a copy when casting to the same categorical dtype returns a new array:
+
+        >>> out2 = c.astype("category", copy=True)
+        >>> out2 is c
+        False
+        >>> out2.to_ndarray()
+        array(['x', 'y', 'x'], dtype='<U...')
+
+        Casting to ``object`` materializes the category labels to a NumPy object array:
+
+        >>> c.astype(object)
+        array(['x', 'y', 'x'], dtype=object)
+
+        Casting to a string dtype returns an Arkouda-backed string array of labels:
+
+        >>> s = c.astype("string")
+        >>> s.to_ndarray()
+        array(['x', 'y', 'x'], dtype='<U1')
+
+        Casting to another dtype casts the labels-as-strings and returns an Arkouda-backed array:
+
+        >>> c_num = ArkoudaCategoricalArray(ak.Categorical(ak.array(["1", "2", "3"])))
+        >>> a = c_num.astype("int64")
+        >>> a.to_ndarray()
+        array([1, 2, 3])
+        """
+        from arkouda.numpy._typing._typing import is_string_dtype_hint
+
+        # --- 1) ExtensionDtype branch first: proves overload #2 returns ExtensionArray ---
+        if isinstance(dtype, ExtensionDtype):
+            if hasattr(dtype, "numpy_dtype"):
+                dtype = dtype.numpy_dtype
+
+            if isinstance(dtype, (ArkoudaCategoricalDtype, pd_CategoricalDtype)) or dtype in (
+                "category",
+            ):
+                if not copy:
+                    return self
+                data = self._data.copy() if hasattr(self._data, "copy") else self._data
+                return type_cast(ExtensionArray, type(self)(data))
+
+            data = self._data.to_strings()
+
+            if isinstance(dtype, pd_StringDtype) or is_string_dtype_hint(dtype):
+                return type_cast(ExtensionArray, ArkoudaStringArray._from_sequence(data))
+
+            casted = data.astype(dtype)
+            return type_cast(ExtensionArray, ArkoudaExtensionArray._from_sequence(casted))
+
+        # --- 2) object -> numpy ---
+        if dtype in (object, np.object_, "object", np.dtype("O")):
+            return self.to_ndarray().astype(object, copy=False)
+
+        if isinstance(dtype, (ArkoudaCategoricalDtype, pd_CategoricalDtype)) or dtype in ("category",):
+            if not copy:
+                return self
+            data = self._data.copy() if hasattr(self._data, "copy") else self._data
+            return type(self)(data)
+
+        data = self._data.to_strings()
+
+        if isinstance(dtype, pd_StringDtype) or is_string_dtype_hint(dtype):
+            return ArkoudaStringArray._from_sequence(data)
+
+        casted = data.astype(dtype)
+        return ArkoudaExtensionArray._from_sequence(casted)
 
     def isna(self):
         return ak.zeros(self._data.size, dtype=ak.bool)

arkouda/pandas/extension/_arkouda_extension_array.py

@@ -332,9 +332,7 @@ def _from_sequence(
         from arkouda.numpy.pdarraycreation import array as ak_array
         from arkouda.numpy.strings import Strings
         from arkouda.pandas.categorical import Categorical
-        from arkouda.pandas.extension._arkouda_array import ArkoudaArray
-        from arkouda.pandas.extension._arkouda_categorical_array import ArkoudaCategoricalArray
-        from arkouda.pandas.extension._arkouda_string_array import ArkoudaStringArray
+        from arkouda.pandas.extension import ArkoudaArray, ArkoudaCategoricalArray, ArkoudaStringArray
 
         # Fast path: already an Arkouda column. Pick the matching subclass.
         if isinstance(scalars, pdarray):