ENH: allow non-nano in DatetimeArray, TimedeltaArray._simple_new (#46901)

Author: jbrockmendel

pandas/_libs/tslibs/__init__.py

@@ -25,6 +25,7 @@
  "Tick",
  "BaseOffset",
  "tz_compare",
+ "is_unitless",
 ]
 
 from pandas._libs.tslibs import dtypes
@@ -39,6 +40,7 @@
 from pandas._libs.tslibs.np_datetime import (
  OutOfBoundsDatetime,
  OutOfBoundsTimedelta,
+ is_unitless,
 )
 from pandas._libs.tslibs.offsets import (
  BaseOffset,

pandas/_libs/tslibs/np_datetime.pyi

@@ -9,3 +9,4 @@ def py_td64_to_tdstruct(td64: int, unit: int) -> dict: ...
 def astype_overflowsafe(
  arr: np.ndarray, dtype: np.dtype, copy: bool = ...
 ) -> np.ndarray: ...
+def is_unitless(dtype: np.dtype) -> bool: ...

pandas/_libs/tslibs/np_datetime.pyx

@@ -92,6 +92,18 @@ def py_get_unit_from_dtype(dtype):
  return get_unit_from_dtype(dtype)
 
 
+def is_unitless(dtype: cnp.dtype) -> bool:
+ """
+ Check if a datetime64 or timedelta64 dtype has no attached unit.
+ """
+ if dtype.type_num not in [cnp.NPY_DATETIME, cnp.NPY_TIMEDELTA]:
+ raise ValueError("is_unitless dtype must be datetime64 or timedelta64")
+ cdef:
+ NPY_DATETIMEUNIT unit = get_unit_from_dtype(dtype)
+
+ return unit == NPY_DATETIMEUNIT.NPY_FR_GENERIC
+
+
 # ----------------------------------------------------------------------
 # Comparison
 

pandas/_libs/tslibs/timedeltas.pyi

@@ -88,6 +88,8 @@ class Timedelta(timedelta):
  # GH 46171
  # While Timedelta can return pd.NaT, having the constructor return
  # a Union with NaTType makes things awkward for users of pandas
+ @classmethod
+ def _from_value_and_reso(cls, value: np.int64, reso: int) -> Timedelta: ...
  @property
  def days(self) -> int: ...
  @property

pandas/core/arrays/datetimes.py

@@ -31,6 +31,7 @@
  iNaT,
  ints_to_pydatetime,
  is_date_array_normalized,
+ is_unitless,
  normalize_i8_timestamps,
  timezones,
  to_offset,
@@ -335,7 +336,12 @@ def _simple_new( # type: ignore[override]
  cls, values: np.ndarray, freq: BaseOffset | None = None, dtype=DT64NS_DTYPE
  ) -> DatetimeArray:
  assert isinstance(values, np.ndarray)
- assert values.dtype == DT64NS_DTYPE
+ assert dtype.kind == "M"
+ if isinstance(dtype, np.dtype):
+ # TODO: once non-nano DatetimeTZDtype is implemented, require that
+ # dtype's reso match values's reso
+ assert dtype == values.dtype
+ assert not is_unitless(dtype)
 
  result = super()._simple_new(values, dtype)
  result._freq = freq
@@ -761,7 +767,7 @@ def _add_offset(self, offset) -> DatetimeArray:
  else:
  values = self
  result = offset._apply_array(values).view("M8[ns]")
- result = DatetimeArray._simple_new(result)
+ result = DatetimeArray._simple_new(result, dtype=result.dtype)
  result = result.tz_localize(self.tz)
 
  except NotImplementedError:

pandas/core/arrays/timedeltas.py

@@ -29,6 +29,7 @@
  precision_from_unit,
 )
 from pandas._libs.tslibs.fields import get_timedelta_field
+from pandas._libs.tslibs.np_datetime import py_get_unit_from_dtype
 from pandas._libs.tslibs.timedeltas import (
  array_to_timedelta64,
  ints_to_pytimedelta,
@@ -40,6 +41,7 @@
  npt,
 )
 from pandas.compat.numpy import function as nv
+from pandas.util._decorators import cache_readonly
 from pandas.util._validators import validate_endpoints
 
 from pandas.core.dtypes.astype import astype_td64_unit_conversion
@@ -154,8 +156,15 @@ class TimedeltaArray(dtl.TimelikeOps):
  # Note: ndim must be defined to ensure NaT.__richcmp__(TimedeltaArray)
  # operates pointwise.
 
+ @cache_readonly
+ def _reso(self):
+ return py_get_unit_from_dtype(self.dtype)
+
  def _box_func(self, x: np.timedelta64) -> Timedelta | NaTType:
- return Timedelta(x, unit="ns")
+ y = x.view("i8")
+ if y == NaT.value:
+ return NaT
+ return Timedelta._from_value_and_reso(y, reso=self._reso)
 
  @property
  # error: Return type "dtype" of "dtype" incompatible with return type
@@ -174,7 +183,7 @@ def dtype(self) -> np.dtype: # type: ignore[override]
  -------
  numpy.dtype
  """
- return TD64NS_DTYPE
+ return self._ndarray.dtype
 
  # ----------------------------------------------------------------
  # Constructors
@@ -244,11 +253,13 @@ def __init__(
  def _simple_new( # type: ignore[override]
  cls, values: np.ndarray, freq: BaseOffset | None = None, dtype=TD64NS_DTYPE
  ) -> TimedeltaArray:
- assert dtype == TD64NS_DTYPE, dtype
+ # Require td64 dtype, not unit-less, matching values.dtype
+ assert isinstance(dtype, np.dtype) and dtype.kind == "m"
+ assert not tslibs.is_unitless(dtype)
  assert isinstance(values, np.ndarray), type(values)
- assert values.dtype == TD64NS_DTYPE
+ assert dtype == values.dtype
 
- result = super()._simple_new(values=values, dtype=TD64NS_DTYPE)
+ result = super()._simple_new(values=values, dtype=dtype)
  result._freq = freq
  return result
 
@@ -262,7 +273,7 @@ def _from_sequence(
  data, inferred_freq = sequence_to_td64ns(data, copy=copy, unit=None)
  freq, _ = dtl.validate_inferred_freq(None, inferred_freq, False)
 
- return cls._simple_new(data, freq=freq)
+ return cls._simple_new(data, dtype=data.dtype, freq=freq)
 
  @classmethod
  def _from_sequence_not_strict(
@@ -286,7 +297,7 @@ def _from_sequence_not_strict(
  if explicit_none:
  freq = None
 
- result = cls._simple_new(data, freq=freq)
+ result = cls._simple_new(data, dtype=data.dtype, freq=freq)
 
  if inferred_freq is None and freq is not None:
  # this condition precludes `freq_infer`
@@ -330,7 +341,8 @@ def _generate_range(cls, start, end, periods, freq, closed=None):
  if not right_closed:
  index = index[:-1]
 
- return cls._simple_new(index.view("m8[ns]"), freq=freq)
+ td64values = index.view("m8[ns]")
+ return cls._simple_new(td64values, dtype=td64values.dtype, freq=freq)
 
  # ----------------------------------------------------------------
  # DatetimeLike Interface

pandas/tests/arrays/test_datetimes.py

@@ -11,6 +11,17 @@
 from pandas.core.arrays import DatetimeArray
 
 
+class TestNonNano:
+ @pytest.mark.parametrize("unit,reso", [("s", 7), ("ms", 8), ("us", 9)])
+ @pytest.mark.xfail(reason="_box_func is not yet patched to get reso right")
+ def test_non_nano(self, unit, reso):
+ arr = np.arange(5, dtype=np.int64).view(f"M8[{unit}]")
+ dta = DatetimeArray._simple_new(arr, dtype=arr.dtype)
+
+ assert dta.dtype == arr.dtype
+ assert dta[0]._reso == reso
+
+
 class TestDatetimeArrayComparisons:
  # TODO: merge this into tests/arithmetic/test_datetime64 once it is
  # sufficiently robust

pandas/tests/arrays/test_timedeltas.py

@@ -7,6 +7,16 @@
 from pandas.core.arrays import TimedeltaArray
 
 
+class TestNonNano:
+ @pytest.mark.parametrize("unit,reso", [("s", 7), ("ms", 8), ("us", 9)])
+ def test_non_nano(self, unit, reso):
+ arr = np.arange(5, dtype=np.int64).view(f"m8[{unit}]")
+ tda = TimedeltaArray._simple_new(arr, dtype=arr.dtype)
+
+ assert tda.dtype == arr.dtype
+ assert tda[0]._reso == reso
+
+
 class TestTimedeltaArray:
  @pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])
  def test_astype_int(self, dtype):

pandas/tests/tslibs/test_api.py

@@ -50,6 +50,7 @@ def test_namespace():
  "tz_convert_from_utc_single",
  "to_offset",
  "tz_compare",
+ "is_unitless",
  ]
 
  expected = set(submodules + api)

pandas/tests/tslibs/test_np_datetime.py

@@ -4,13 +4,36 @@
 from pandas._libs.tslibs.np_datetime import (
  OutOfBoundsDatetime,
  astype_overflowsafe,
+ is_unitless,
  py_get_unit_from_dtype,
  py_td64_to_tdstruct,
 )
 
 import pandas._testing as tm
 
 
+def test_is_unitless():
+ dtype = np.dtype("M8[ns]")
+ assert not is_unitless(dtype)
+
+ dtype = np.dtype("datetime64")
+ assert is_unitless(dtype)
+
+ dtype = np.dtype("m8[ns]")
+ assert not is_unitless(dtype)
+
+ dtype = np.dtype("timedelta64")
+ assert is_unitless(dtype)
+
+ msg = "dtype must be datetime64 or timedelta64"
+ with pytest.raises(ValueError, match=msg):
+ is_unitless(np.dtype(np.int64))
+
+ msg = "Argument 'dtype' has incorrect type"
+ with pytest.raises(TypeError, match=msg):
+ is_unitless("foo")
+
+
 def test_get_unit_from_dtype():
  # datetime64
  assert py_get_unit_from_dtype(np.dtype("M8[Y]")) == 0