Skip to content

REF: Dispatch ArrowExtensionArray.fillna methods to pad_or_backfill #54820

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 28, 2023
Merged

REF: Dispatch ArrowExtensionArray.fillna methods to pad_or_backfill #54820

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e00630a
REF: Dispatch ArrowExtensionArray.fillna methods to pad_or_backfill
mroeschke Aug 28, 2023
df8a3ea
Merge remote-tracking branch 'upstream/main' into ref/arrow/fill_method
mroeschke Aug 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 23 additions & 8 deletions pandas/core/arrays/arrow/array.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,10 @@
from pandas.core.dtypes.dtypes import DatetimeTZDtype
from pandas.core.dtypes.missing import isna

from pandas.core import roperator
from pandas.core import (
missing,
roperator,
)
from pandas.core.arraylike import OpsMixin
from pandas.core.arrays._arrow_string_mixins import ArrowStringArrayMixin
from pandas.core.arrays.base import (
Expand Down Expand Up @@ -933,6 +936,20 @@ def pad_or_backfill(
# TODO(CoW): Not necessary anymore when CoW is the default
return self.copy()

if limit is not None and limit_area is not None:
method = missing.clean_fill_method(method)
try:
if method == "pad":
return type(self)(pc.fill_null_forward(self._pa_array))
elif method == "backfill":
return type(self)(pc.fill_null_backward(self._pa_array))
except pa.ArrowNotImplementedError:
# ArrowNotImplementedError: Function 'coalesce' has no kernel
# matching input types (duration[ns], duration[ns])
# TODO: remove try/except wrapper if/when pyarrow implements
# a kernel for duration types.
pass

# TODO(3.0): after EA.fillna 'method' deprecation is enforced, we can remove
# this method entirely.
return super().pad_or_backfill(
Expand All @@ -953,9 +970,12 @@ def fillna(
# TODO(CoW): Not necessary anymore when CoW is the default
return self.copy()

if limit is not None or method is not None:
if limit is not None:
return super().fillna(value=value, method=method, limit=limit, copy=copy)

if method is not None:
return super().pad_or_backfill(method=method, limit=limit, copy=copy)

if isinstance(value, (np.ndarray, ExtensionArray)):
# Similar to check_value_size, but we do not mask here since we may
# end up passing it to the super() method.
Expand All @@ -972,12 +992,7 @@ def fillna(
raise TypeError(msg) from err

try:
if method is None:
return type(self)(pc.fill_null(self._pa_array, fill_value=fill_value))
elif method == "pad":
return type(self)(pc.fill_null_forward(self._pa_array))
elif method == "backfill":
return type(self)(pc.fill_null_backward(self._pa_array))
return type(self)(pc.fill_null(self._pa_array, fill_value=fill_value))
except pa.ArrowNotImplementedError:
# ArrowNotImplementedError: Function 'coalesce' has no kernel
# matching input types (duration[ns], duration[ns])
Expand Down