Skip to content

BUG: rolling with Int64 #43174

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

Merged
merged 5 commits into from
Aug 31, 2021
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
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
2 changes: 1 addition & 1 deletion doc/source/whatsnew/v1.4.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,7 @@ Groupby/resample/rolling
- Bug in :meth:`pandas.DataFrame.ewm`, where non-float64 dtypes were silently failing (:issue:`42452`)
- Bug in :meth:`pandas.DataFrame.rolling` operation along rows (``axis=1``) incorrectly omits columns containing ``float16`` and ``float32`` (:issue:`41779`)
- Bug in :meth:`Resampler.aggregate` did not allow the use of Named Aggregation (:issue:`32803`)
-
- Bug in :meth:`Series.rolling` when the :class:`Series` ``dtype`` was ``Int64`` (:issue:`43016`)

Reshaping
^^^^^^^^^
Expand Down
34 changes: 19 additions & 15 deletions pandas/_libs/algos_common_helper.pxi.in
Original file line number Diff line number Diff line change
Expand Up @@ -37,27 +37,27 @@ def ensure_object(object arr):
{{py:

# name, c_type, dtype
dtypes = [('float64', 'FLOAT64', 'float64'),
('float32', 'FLOAT32', 'float32'),
('int8', 'INT8', 'int8'),
('int16', 'INT16', 'int16'),
('int32', 'INT32', 'int32'),
('int64', 'INT64', 'int64'),
('uint8', 'UINT8', 'uint8'),
('uint16', 'UINT16', 'uint16'),
('uint32', 'UINT32', 'uint32'),
('uint64', 'UINT64', 'uint64'),
# ('platform_int', 'INT', 'int_'),
# ('object', 'OBJECT', 'object_'),
dtypes = [('float64', 'FLOAT64', 'float64', 'nan'),
('float32', 'FLOAT32', 'float32', 'nan'),
('int8', 'INT8', 'int8', None),
('int16', 'INT16', 'int16', None),
('int32', 'INT32', 'int32', None),
('int64', 'INT64', 'int64', None),
('uint8', 'UINT8', 'uint8', None),
('uint16', 'UINT16', 'uint16', None),
('uint32', 'UINT32', 'uint32', None),
('uint64', 'UINT64', 'uint64', None),
# ('platform_int', 'INT', 'int_', None),
# ('object', 'OBJECT', 'object_', None),
]

def get_dispatch(dtypes):

for name, c_type, dtype in dtypes:
yield name, c_type, dtype
for name, c_type, dtype, na_val in dtypes:
yield name, c_type, dtype, na_val
}}

{{for name, c_type, dtype in get_dispatch(dtypes)}}
{{for name, c_type, dtype, na_val in get_dispatch(dtypes)}}


def ensure_{{name}}(object arr, copy=True):
Expand All @@ -66,6 +66,10 @@ def ensure_{{name}}(object arr, copy=True):
return arr
else:
return arr.astype(np.{{dtype}}, copy=copy)
{{if na_val == "nan"}}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

im really not sure about this. ensure_foo is fast ATM, plus this introduces circular dependency

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you think ensure_foo should not be used directly on ExtentionArrays i.e. ensure_foo(arr.to_numpy(...))?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i guess. mainly i think that pd.NA is way more trouble than its worth

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm i did request this, but agree the performance tradeoff might not be worthit.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i think maybe making an ensure_*_with_na is prob better here (and make it the caller responsible if we have an EA)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I went with @jbrockmendel's original isinstance check instead for now. For the ensure_foo_with_na feature, probably should be another issue to discuss because it would be potentially useful anywhere ensure_foo is called.

elif hasattr(arr, "to_numpy"):
return arr.to_numpy(np.{{dtype}}, copy=copy, na_value=np.{{na_val}})
{{endif}}
else:
return np.array(arr, dtype=np.{{dtype}})

Expand Down
9 changes: 9 additions & 0 deletions pandas/tests/window/test_dtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import pytest

from pandas import (
NA,
DataFrame,
Series,
)
Expand Down Expand Up @@ -76,6 +77,14 @@ def test_series_dtypes(method, data, expected_data, coerce_int, dtypes, min_peri
tm.assert_almost_equal(result, expected)


def test_series_nullable_int(any_signed_int_ea_dtype):
# GH 43016
s = Series([0, 1, NA], dtype=any_signed_int_ea_dtype)
result = s.rolling(2).mean()
expected = Series([np.nan, 0.5, np.nan])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't we actually replace and use pd.NA as the output? or is this too big of a change?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally, but currently all rolling/expanding/ewm results always return np.float64 (which is documented as well). So that would be a big change to return the same dtype as the caller.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, do we have an issue about this? we should make this change in 1.4

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm no was speaking about supporting pd.NA specifically (i think rolling_apply is totally fine) if you can add one

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I modified that issue to generally support same input/output dtypes (which include ExtensionDtypes that support pd.NA)

tm.assert_series_equal(result, expected)


@pytest.mark.parametrize(
"method, expected_data, min_periods",
[
Expand Down