From 2fa5d5d32cf5f726a23fb059b769f4739d8c7deb Mon Sep 17 00:00:00 2001 From: nevrohelios Date: Sun, 24 Aug 2025 00:14:12 +0530 Subject: [PATCH 1/3] REF: use _cast_pointwise_result in map --- pandas/core/arrays/arrow/array.py | 3 ++- pandas/core/arrays/base.py | 3 ++- pandas/core/arrays/masked.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py index 221191773186e..fa0eaa5c05279 100644 --- a/pandas/core/arrays/arrow/array.py +++ b/pandas/core/arrays/arrow/array.py @@ -1594,7 +1594,8 @@ def to_numpy( def map(self, mapper, na_action: Literal["ignore"] | None = None): if is_numeric_dtype(self.dtype): - return map_array(self.to_numpy(), mapper, na_action=na_action) + result = map_array(self.to_numpy(), mapper, na_action=na_action) + return self._cast_pointwise_result(result) else: return super().map(mapper, na_action) diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 1cd10a9eef9d1..1c5e7881c215c 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -2530,7 +2530,8 @@ def map(self, mapper, na_action: Literal["ignore"] | None = None): If the function returns a tuple with more than one element a MultiIndex will be returned. """ - return map_array(self, mapper, na_action=na_action) + results = map_array(self, mapper, na_action=na_action) + return self._cast_pointwise_result(results) # ------------------------------------------------------------------------ # GroupBy Methods diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index 0402452e484ea..e56ca7b0fed39 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -1395,7 +1395,8 @@ 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): - return map_array(self.to_numpy(), mapper, na_action=na_action) + result = map_array(self.to_numpy(), mapper, na_action=na_action) + return self._cast_pointwise_result(result) @overload def any( From f8f579fdd7211ac7b52595fe290ab6dc90e94bac Mon Sep 17 00:00:00 2001 From: nevrohelios Date: Tue, 9 Sep 2025 01:44:12 +0530 Subject: [PATCH 2/3] added preser_root param --- .gitignore | 3 +++ doc/source/reference/extensions.rst | 1 + pandas/core/arrays/arrow/array.py | 10 +++++++--- pandas/core/arrays/base.py | 14 ++++++++++++-- pandas/core/arrays/masked.py | 9 +++++++-- 5 files changed, 30 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index d951f3fb9cbad..84d761ac112ce 100644 --- a/.gitignore +++ b/.gitignore @@ -141,3 +141,6 @@ doc/source/savefig/ # Pyodide/WASM related files # ############################## /.pyodide-xbuildenv-* + +local.py +.venv/ \ No newline at end of file diff --git a/doc/source/reference/extensions.rst b/doc/source/reference/extensions.rst index e412793a328a3..1f5d19d004a4c 100644 --- a/doc/source/reference/extensions.rst +++ b/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 diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py index fa0eaa5c05279..947a5e73a8a71 100644 --- a/pandas/core/arrays/arrow/array.py +++ b/pandas/core/arrays/arrow/array.py @@ -1592,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( diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index 1c5e7881c215c..abd5407c0a82b 100644 --- a/pandas/core/arrays/base.py +++ b/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 diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index e56ca7b0fed39..0f4b191cf283e 100644 --- a/pandas/core/arrays/masked.py +++ b/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( From 494f1b0d469361c0f91d5d532f5488ac9b32cf7f Mon Sep 17 00:00:00 2001 From: nevrohelios Date: Tue, 9 Sep 2025 09:43:26 +0530 Subject: [PATCH 3/3] format --- .gitignore | 2 +- doc/source/reference/extensions.rst | 2 +- pandas/core/arrays/base.py | 9 ++++++--- pandas/core/arrays/masked.py | 9 ++++++--- 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 84d761ac112ce..03f035db527fe 100644 --- a/.gitignore +++ b/.gitignore @@ -143,4 +143,4 @@ doc/source/savefig/ /.pyodide-xbuildenv-* local.py -.venv/ \ No newline at end of file +.venv/ diff --git a/doc/source/reference/extensions.rst b/doc/source/reference/extensions.rst index 1f5d19d004a4c..ff93a3fd25104 100644 --- a/doc/source/reference/extensions.rst +++ b/doc/source/reference/extensions.rst @@ -58,7 +58,7 @@ objects. api.extensions.ExtensionArray.isin api.extensions.ExtensionArray.isna api.extensions.ExtensionArray.ravel - api.extensions.ExtensionArray.map + api.extensions.ExtensionArray.map api.extensions.ExtensionArray.repeat api.extensions.ExtensionArray.searchsorted api.extensions.ExtensionArray.shift diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index a9a8563ca5eeb..c37ee6c98079c 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -2516,9 +2516,12 @@ 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, - preserve_dtype: bool = False): + def map( + self, + mapper, + na_action: Literal["ignore"] | None = None, + preserve_dtype: bool = False, + ): """ Map values using an input mapping or function. diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index ff7d3d41c2052..9d0ee3f833218 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -1395,9 +1395,12 @@ 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, - preserve_dtype: bool = False): + 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) if preserve_dtype: