|
34 | 34 | )
|
35 | 35 | from pandas.core.dtypes.common import (
|
36 | 36 | is_bool_dtype,
|
| 37 | + is_datetime64_ns_dtype, |
37 | 38 | is_dtype_equal,
|
38 | 39 | is_extension_array_dtype,
|
39 | 40 | is_numeric_dtype,
|
|
54 | 55 | )
|
55 | 56 |
|
56 | 57 | import pandas.core.algorithms as algos
|
57 |
| -from pandas.core.arrays import ExtensionArray |
| 58 | +from pandas.core.arrays import ( |
| 59 | + DatetimeArray, |
| 60 | + ExtensionArray, |
| 61 | + TimedeltaArray, |
| 62 | +) |
58 | 63 | from pandas.core.arrays.sparse import SparseDtype
|
59 | 64 | from pandas.core.construction import (
|
60 | 65 | ensure_wrapped_if_datetimelike,
|
@@ -114,6 +119,7 @@ def __init__(
|
114 | 119 |
|
115 | 120 | if verify_integrity:
|
116 | 121 | self._axes = [ensure_index(ax) for ax in axes]
|
| 122 | + self.arrays = [ensure_wrapped_if_datetimelike(arr) for arr in arrays] |
117 | 123 | self._verify_integrity()
|
118 | 124 |
|
119 | 125 | def make_empty(self: T, axes=None) -> T:
|
@@ -716,20 +722,16 @@ def fast_xs(self, loc: int) -> ArrayLike:
|
716 | 722 | """
|
717 | 723 | dtype = _interleaved_dtype(self.arrays)
|
718 | 724 |
|
719 |
| - if isinstance(dtype, SparseDtype): |
720 |
| - temp_dtype = dtype.subtype |
721 |
| - elif isinstance(dtype, PandasDtype): |
722 |
| - temp_dtype = dtype.numpy_dtype |
723 |
| - elif is_extension_array_dtype(dtype): |
724 |
| - temp_dtype = "object" |
725 |
| - elif is_dtype_equal(dtype, str): |
726 |
| - temp_dtype = "object" |
727 |
| - else: |
728 |
| - temp_dtype = dtype |
729 |
| - |
730 |
| - result = np.array([arr[loc] for arr in self.arrays], dtype=temp_dtype) |
| 725 | + values = [arr[loc] for arr in self.arrays] |
731 | 726 | if isinstance(dtype, ExtensionDtype):
|
732 |
| - result = dtype.construct_array_type()._from_sequence(result, dtype=dtype) |
| 727 | + result = dtype.construct_array_type()._from_sequence(values, dtype=dtype) |
| 728 | + # for datetime64/timedelta64, the np.ndarray constructor cannot handle pd.NaT |
| 729 | + elif is_datetime64_ns_dtype(dtype): |
| 730 | + result = DatetimeArray._from_sequence(values, dtype=dtype)._data |
| 731 | + elif is_timedelta64_ns_dtype(dtype): |
| 732 | + result = TimedeltaArray._from_sequence(values, dtype=dtype)._data |
| 733 | + else: |
| 734 | + result = np.array(values, dtype=dtype) |
733 | 735 | return result
|
734 | 736 |
|
735 | 737 | def iget(self, i: int) -> SingleBlockManager:
|
|
0 commit comments