-
-
Notifications
You must be signed in to change notification settings - Fork 18.1k
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.
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
BUG: rolling with datetime ArrowDtype #56370
Changes from 3 commits
c232efd
15e0869
ecc8f4d
baa6b15
c844b18
ca57b50
c57585e
45d7a7c
056268a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -14,7 +14,6 @@ | |
Any, | ||
Callable, | ||
Literal, | ||
cast, | ||
) | ||
|
||
import numpy as np | ||
|
@@ -39,6 +38,7 @@ | |
is_numeric_dtype, | ||
needs_i8_conversion, | ||
) | ||
from pandas.core.dtypes.dtypes import ArrowDtype | ||
from pandas.core.dtypes.generic import ( | ||
ABCDataFrame, | ||
ABCSeries, | ||
|
@@ -104,6 +104,7 @@ | |
NDFrameT, | ||
QuantileInterpolation, | ||
WindowingRankType, | ||
npt, | ||
) | ||
|
||
from pandas import ( | ||
|
@@ -404,11 +405,12 @@ def _insert_on_column(self, result: DataFrame, obj: DataFrame) -> None: | |
result[name] = extra_col | ||
|
||
@property | ||
def _index_array(self): | ||
def _index_array(self) -> npt.NDArray[np.int64] | None: | ||
# TODO: why do we get here with e.g. MultiIndex? | ||
if needs_i8_conversion(self._on.dtype): | ||
idx = cast("PeriodIndex | DatetimeIndex | TimedeltaIndex", self._on) | ||
return idx.asi8 | ||
if isinstance(self._on, (PeriodIndex, DatetimeIndex, TimedeltaIndex)): | ||
return self._on.asi8 | ||
elif isinstance(self._on.dtype, ArrowDtype): | ||
return self._on.to_numpy(dtype=np.int64) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. comment/assertion about self._on.dtype.kind? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yup good point. Added a kind check |
||
return None | ||
|
||
def _resolve_output(self, out: DataFrame, obj: DataFrame) -> DataFrame: | ||
|
@@ -439,7 +441,7 @@ def _apply_series( | |
self, homogeneous_func: Callable[..., ArrayLike], name: str | None = None | ||
) -> Series: | ||
""" | ||
Series version of _apply_blockwise | ||
Series version of _apply_columnwise | ||
""" | ||
obj = self._create_data(self._selected_obj) | ||
|
||
|
@@ -455,7 +457,7 @@ def _apply_series( | |
index = self._slice_axis_for_step(obj.index, result) | ||
return obj._constructor(result, index=index, name=obj.name) | ||
|
||
def _apply_blockwise( | ||
def _apply_columnwise( | ||
self, | ||
homogeneous_func: Callable[..., ArrayLike], | ||
name: str, | ||
|
@@ -614,7 +616,7 @@ def calc(x): | |
return result | ||
|
||
if self.method == "single": | ||
return self._apply_blockwise(homogeneous_func, name, numeric_only) | ||
return self._apply_columnwise(homogeneous_func, name, numeric_only) | ||
else: | ||
return self._apply_tablewise(homogeneous_func, name, numeric_only) | ||
|
||
|
@@ -1236,7 +1238,9 @@ def calc(x): | |
|
||
return result | ||
|
||
return self._apply_blockwise(homogeneous_func, name, numeric_only)[:: self.step] | ||
return self._apply_columnwise(homogeneous_func, name, numeric_only)[ | ||
:: self.step | ||
] | ||
|
||
@doc( | ||
_shared_docs["aggregate"], | ||
|
@@ -1871,7 +1875,7 @@ def _validate(self): | |
# we allow rolling on a datetimelike index | ||
if ( | ||
self.obj.empty | ||
or isinstance(self._on, (DatetimeIndex, TimedeltaIndex, PeriodIndex)) | ||
or (isinstance(self._on, PeriodIndex) or self._on.dtype.kind in "Mm") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. id err on the side of being more explicit about the supported dtypes being ArrowDtype |
||
) and isinstance(self.window, (str, BaseOffset, timedelta)): | ||
self._validate_datetimelike_monotonic() | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
typo here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think so? (If you're referring to
dtype=
, it was intentional to render e.g."dtype=int64[pyarrow] does not have a resolution"
)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe im misunderstanding how f-strings work. id expect this to render "int64[pyarrow]= does not have a resolution"
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
neat!