Skip to content

[ArrayManager] Ensure to store datetimelike data as DatetimeArray/TimedeltaArray (and not ndarray) #40147

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
Show file tree
Hide file tree
Changes from 1 commit
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
14 changes: 13 additions & 1 deletion pandas/core/internals/array_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
)
from pandas.core.dtypes.common import (
is_bool_dtype,
is_datetime64_ns_dtype,
is_dtype_equal,
is_extension_array_dtype,
is_numeric_dtype,
Expand All @@ -53,7 +54,11 @@
)

import pandas.core.algorithms as algos
from pandas.core.arrays import ExtensionArray
from pandas.core.arrays import (
DatetimeArray,
ExtensionArray,
TimedeltaArray,
)
from pandas.core.arrays.sparse import SparseDtype
from pandas.core.construction import (
ensure_wrapped_if_datetimelike,
Expand Down Expand Up @@ -113,6 +118,7 @@ def __init__(

if verify_integrity:
self._axes = [ensure_index(ax) for ax in axes]
self.arrays = [ensure_wrapped_if_datetimelike(arr) for arr in arrays]
self._verify_integrity()

def make_empty(self: T, axes=None) -> T:
Expand Down Expand Up @@ -721,6 +727,8 @@ def fast_xs(self, loc: int) -> ArrayLike:
temp_dtype = dtype.numpy_dtype
elif is_extension_array_dtype(dtype):
temp_dtype = "object"
elif is_datetime64_ns_dtype(dtype) or is_timedelta64_ns_dtype(dtype):
temp_dtype = "object"
elif is_dtype_equal(dtype, str):
temp_dtype = "object"
else:
Expand All @@ -729,6 +737,10 @@ def fast_xs(self, loc: int) -> ArrayLike:
result = np.array([arr[loc] for arr in self.arrays], dtype=temp_dtype)
if isinstance(dtype, ExtensionDtype):
result = dtype.construct_array_type()._from_sequence(result, dtype=dtype)
elif is_datetime64_ns_dtype(dtype):
result = DatetimeArray._from_sequence(result, dtype=dtype)._data
elif is_timedelta64_ns_dtype(dtype):
Copy link
Member

Choose a reason for hiding this comment

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

we have a little-used is_ea_or_datetimelike_dtype, could use an analogous helper to get DatetimeArray/TimedeltaArray in these cases (not for this PR)

result = TimedeltaArray._from_sequence(result, dtype=dtype)._data
Copy link
Member Author

Choose a reason for hiding this comment

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

The reason I am doing this here is because this fails:

In [4]: np.array([pd.NaT], dtype="M8[ns]")
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-99d42e913a1c> in <module>
----> 1 np.array([pd.NaT], dtype="M8[ns]")

ValueError: cannot convert float NaN to integer

(which is what the np.array([arr[loc] for arr in self.arrays], dtype=temp_dtype) above can be doing if the resulting dtype is M/m8)

@jbrockmendel do you know if that's something we do elsewhere as well / there is some existing code for this?

Copy link
Member

Choose a reason for hiding this comment

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

IIRC we can get BM to mess up if we get here with non-consolidated all-td64 blocks.

I think the thing to do (also similar change in the BM method) is to define result on L737 as a list, only wrap with np.array if none of these conditions hold

Copy link
Member Author

Choose a reason for hiding this comment

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

IIRC we can get BM to mess up if we get here with non-consolidated all-td64 blocks.

You can't get here in that case, though, since this is an ArrayManager method. And with the current BM you don't have this problem since it doesn't store DatetimeArray/TimedeltaArray with numpy dtypes.

I think the thing to do (also similar change in the BM method) is to define result on L737 as a list, only wrap with np.array if none of these conditions hold

Thanks, that's a good idea

Copy link
Member Author

Choose a reason for hiding this comment

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

Updated with this simplification.

return result

def iget(self, i: int) -> SingleBlockManager:
Expand Down
1 change: 1 addition & 0 deletions pandas/tests/frame/methods/test_rename.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,7 @@ def test_rename_multiindex(self):
renamed = df.rename(index={"foo1": "foo3", "bar2": "bar3"}, level=0)
tm.assert_index_equal(renamed.index, new_index)

@td.skip_array_manager_not_yet_implemented # TODO(ArrayManager) setitem copy/view
def test_rename_nocopy(self, float_frame):
renamed = float_frame.rename(columns={"C": "foo"}, copy=False)
renamed["foo"] = 1.0
Expand Down