added preserve_root param for Extension array

Author: NevroHelios

doc/source/reference/extensions.rst

@@ -58,6 +58,7 @@ objects.
       api.extensions.ExtensionArray.isin
       api.extensions.ExtensionArray.isna
       api.extensions.ExtensionArray.ravel
+      api.extensions.ExtensionArray.map
       api.extensions.ExtensionArray.repeat
       api.extensions.ExtensionArray.searchsorted
       api.extensions.ExtensionArray.shift

pandas/core/arrays/arrow/array.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import functools
-import math
 import operator
 import re
 import textwrap
@@ -17,15 +16,13 @@
 import warnings
 
 import numpy as np
-import pandas as pd
 
 from pandas._libs import lib
 from pandas._libs.tslibs import (
     Timedelta,
     Timestamp,
     timezones,
 )
-from sqlalchemy import values
 from pandas.compat import (
     HAS_PYARROW,
     pa_version_under12p1,
@@ -404,26 +401,16 @@ def _cast_pointwise_result(self, values) -> ArrayLike:
         if len(values) == 0:
             # Retain our dtype
             return self[:0].copy()
-        if not isinstance(self.dtype, ArrowDtype):
-            return super()._cast_pointwise_result(values)
-        
+
         try:
-            if (
-            isinstance(values, (np.ndarray, list, tuple))
-            and np.asarray(values).dtype == np.bool_
-            ):
-                arr = pa.array(values, from_pandas=True)
-            else:
-                arr = pa.array(values, type=self._pa_array.type, from_pandas=True)
+            arr = pa.array(values, from_pandas=True)
         except (ValueError, TypeError):
             # e.g. test_by_column_values_with_same_starting_value with nested
             #  values, one entry of which is an ArrowStringArray
             #  or test_agg_lambda_complex128_dtype_conversion for complex values
             return super()._cast_pointwise_result(values)
-        if pa.types.is_duration(self._pa_array.type) and pa.types.is_floating(arr.type):
-            # just return a numpy array / float Series
-            return np.asarray(values, dtype="float64")
-        elif pa.types.is_duration(arr.type):
+
+        if pa.types.is_duration(arr.type):
             # workaround for https://github.com/apache/arrow/issues/40620
             result = ArrowExtensionArray._from_sequence(values)
             if pa.types.is_duration(self._pa_array.type):
@@ -444,79 +431,28 @@ def _cast_pointwise_result(self, values) -> ArrayLike:
                 dtype = ArrowDtype(pa.duration("s"))
                 result = result.astype(dtype)  # type: ignore[assignment]
             return result
-    
-        elif pa.types.is_timestamp(arr.type) and pa.types.is_timestamp(self.dtype.pyarrow_dtype):
-            # Preserve the original array's timestamp unit (i.e. us/ns/...)
-            original_unit = self.dtype.pyarrow_dtype.unit
-            tz = arr.type.tz
-        
-            # Only convert if units don't match
-            if arr.type.unit != original_unit:
-                target_pa_dtype = pa.timestamp(original_unit, tz=tz)
-                arr = arr.cast(target_pa_dtype)
-        
-            # Create ArrowExtensionArray with the processed array
-            return self._from_pyarrow_array(arr)
-        elif pa.types.is_floating(self._pa_array.type):
-            try:
-                if self._pa_array.type == pa.float32():
-                    coerced = [
-                        None if (v is None or v is pd.NA or (isinstance(v, float) and np.isnan(v)))
-                        else np.float32(v)
-                        for v in values
-                    ]
-                    arr = pa.array(coerced, type=pa.float32(), from_pandas=True)
-                else:
-                    arr = pa.array(values, type=self._pa_array.type, from_pandas=True)
-            except pa.ArrowInvalid:
-                arr = pa.array(values, from_pandas=True)
+
         elif pa.types.is_date(arr.type) and pa.types.is_date(self._pa_array.type):
             arr = arr.cast(self._pa_array.type)
         elif pa.types.is_time(arr.type) and pa.types.is_time(self._pa_array.type):
             arr = arr.cast(self._pa_array.type)
         elif pa.types.is_decimal(arr.type) and pa.types.is_decimal(self._pa_array.type):
             arr = arr.cast(self._pa_array.type)
-        elif is_numeric_dtype(self.dtype):
-            if pa.types.is_integer(self._pa_array.type):
-                try:
-                    # Handle the case where Python map gives floats (e.g., 1 → 1.0)
-                    floats = [v for v in values if isinstance(v, float) and v is not None]
-                    if floats and all(math.isfinite(v) and v.is_integer() for v in floats):
-                        arr = pa.array([int(v) if isinstance(v, float) else v for v in values],
-                                    type=self._pa_array.type,
-                                    from_pandas=True)
-                        return self._from_pyarrow_array(arr)
-
-                    # Special handling for unsigned integers
-                    if pa.types.is_unsigned_integer(self._pa_array.type):
-                        if any((isinstance(v, (int, np.integer, float)) and v is not None and v < 0)
-                            for v in values):
-                            # Promote to signed int (wider type to hold negatives)
-                            signed_type = pa.int16() if self._pa_array.type == pa.uint8() else pa.int64()
-                            arr = pa.array(values, type=signed_type, from_pandas=True)
-                        else:
-                            arr = pa.array(values, type=self._pa_array.type, from_pandas=True)
-                    else:
-                        arr = pa.array(values, type=self._pa_array.type, from_pandas=True)
-                except (pa.ArrowInvalid, OverflowError):
-                    arr = pa.array(values, from_pandas=True)
-
-        elif (
-            (pa.types.is_integer(arr.type) and pa.types.is_integer(self._pa_array.type))
-            or (pa.types.is_floating(arr.type) and pa.types.is_integer(self._pa_array.type))
-        ):
+        elif pa.types.is_integer(arr.type) and pa.types.is_integer(self._pa_array.type):
             try:
                 arr = arr.cast(self._pa_array.type)
             except pa.lib.ArrowInvalid:
                 # e.g. test_combine_add if we can't cast
                 pass
-        elif pa.types.is_floating(arr.type) and pa.types.is_floating(self._pa_array.type):
+        elif pa.types.is_floating(arr.type) and pa.types.is_floating(
+            self._pa_array.type
+        ):
             try:
                 arr = arr.cast(self._pa_array.type)
             except pa.lib.ArrowInvalid:
                 # e.g. test_combine_add if we can't cast
                 pass
-        
+
         if isinstance(self.dtype, StringDtype):
             if pa.types.is_string(arr.type) or pa.types.is_large_string(arr.type):
                 # ArrowStringArray preserves dtype.na_value
@@ -526,7 +462,6 @@ def _cast_pointwise_result(self, values) -> ArrayLike:
                 #  result instead
                 return super()._cast_pointwise_result(values)
             return ArrowExtensionArray(arr)
-        
         return self._from_pyarrow_array(arr)
 
     @classmethod
@@ -794,28 +729,15 @@ def __getitem__(self, item: PositionalIndexer):
             return self._from_pyarrow_array(value)
         else:
             pa_type = self._pa_array.type
-            # Special case: timestamp : avoid overflow
-            if pa.types.is_timestamp(pa_type):
-                return pd.Timestamp(value.as_py())  
-            # Special case: duration
-            if pa.types.is_duration(pa_type):
-                return pd.Timedelta(value.as_py())
             scalar = value.as_py()
-
             if scalar is None:
                 return self._dtype.na_value
-            elif pa.types.is_timestamp(pa_type):
+            elif pa.types.is_timestamp(pa_type) and pa_type.unit != "ns":
                 # GH 53326
-                ts = pd.Timestamp(scalar)
-                if pa_type.unit != "ns":
-                    return ts.as_unit(pa_type.unit)
-                return ts
-            elif pa.types.is_duration(pa_type):
+                return Timestamp(scalar).as_unit(pa_type.unit)
+            elif pa.types.is_duration(pa_type) and pa_type.unit != "ns":
                 # GH 53326
-                td = pd.Timedelta(scalar)
-                if pa_type.unit != "ns":
-                    return td.as_unit(pa_type.unit)
-                return td
+                return Timedelta(scalar).as_unit(pa_type.unit)
             else:
                 return scalar
 
@@ -1670,12 +1592,16 @@ def to_numpy(
             result[~mask] = data[~mask]._pa_array.to_numpy()
         return result
 
-    def map(self, mapper, na_action: Literal["ignore"] | None = None):
+    def map(self, mapper,
+            na_action: Literal["ignore"] | None = None,
+            preserve_dtype: bool = False):
         if is_numeric_dtype(self.dtype):
             result = map_array(self.to_numpy(), mapper, na_action=na_action)
-            return self._cast_pointwise_result(result)
+            if preserve_dtype:
+                result = self._cast_pointwise_result(result)
+            return result
         else:
-            return super().map(mapper, na_action)
+            return super().map(mapper, na_action, preserve_dtype=preserve_dtype)
 
     @doc(ExtensionArray.duplicated)
     def duplicated(

pandas/core/arrays/base.py

@@ -2510,7 +2510,9 @@ def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
 
         return arraylike.default_array_ufunc(self, ufunc, method, *inputs, **kwargs)
 
-    def map(self, mapper, na_action: Literal["ignore"] | None = None):
+    def map(self, mapper, 
+            na_action: Literal["ignore"] | None = None,
+            preserve_dtype: bool = False):
         """
         Map values using an input mapping or function.
 
@@ -2522,6 +2524,12 @@ def map(self, mapper, na_action: Literal["ignore"] | None = None):
             If 'ignore', propagate NA values, without passing them to the
             mapping correspondence. If 'ignore' is not supported, a
             ``NotImplementedError`` should be raised.
+        preserve_dtype : bool, default False
+            If True, attempt to cast the elementwise result back to the
+            original ExtensionArray type (and dtype) when possible. This is
+            primarily intended for identity or dtype-preserving mappings.
+            If False, the result of the mapping is returned as produced by
+            the underlying implementation (typically a NumPy ndarray).
 
         Returns
         -------
@@ -2531,7 +2539,9 @@ def map(self, mapper, na_action: Literal["ignore"] | None = None):
             a MultiIndex will be returned.
         """
         results = map_array(self, mapper, na_action=na_action)
-        return self._cast_pointwise_result(results)
+        if preserve_dtype:
+            results = self._cast_pointwise_result(results)
+        return results
 
     # ------------------------------------------------------------------------
     # GroupBy Methods

pandas/core/arrays/masked.py

@@ -1394,9 +1394,14 @@ def max(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
         )
         return self._wrap_reduction_result("max", result, skipna=skipna, axis=axis)
 
-    def map(self, mapper, na_action: Literal["ignore"] | None = None):
+    def map(self, mapper, 
+            na_action: Literal["ignore"] | None = None,
+            preserve_dtype: bool = False):
+        """See ExtensionArray.map."""
         result = map_array(self.to_numpy(), mapper, na_action=na_action)
-        return self._cast_pointwise_result(result)
+        if preserve_dtype:
+            result = self._cast_pointwise_result(result)
+        return result
 
     @overload
     def any(