-
-
Notifications
You must be signed in to change notification settings - Fork 18.4k
[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
[ArrayManager] Ensure to store datetimelike data as DatetimeArray/TimedeltaArray (and not ndarray) #40147
Changes from 1 commit
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 |
---|---|---|
|
@@ -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, | ||
|
@@ -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, | ||
|
@@ -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: | ||
|
@@ -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: | ||
|
@@ -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): | ||
result = TimedeltaArray._from_sequence(result, dtype=dtype)._data | ||
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. The reason I am doing this here is because this fails:
(which is what the @jbrockmendel do you know if that's something we do elsewhere as well / there is some existing code for this? 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. 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 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.
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.
Thanks, that's a good idea 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. Updated with this simplification. |
||
return result | ||
|
||
def iget(self, i: int) -> SingleBlockManager: | ||
|
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.
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)