Extend time unit support (#132)

Author: cbourjau

CHANGELOG.rst

@@ -7,6 +7,17 @@
 Changelog
 =========
 
+0.13.0 (2025-05-22)
+-------------------
+
+**New features**
+
+- The :class:`ndonnx.TimeDelta64DType` and :class:`ndonnx.DateTime64DType` gained support for milli and microseconds as units.
+- :func:`ndonnx.where` now promotes time units between the two branches.
+- Addition, multiplication, division, and subtraction between arrays with timedelta or datetime data types now support promotion between time units.
+- Comparison operations between arrays with timedelta or datetime data types now support promotion between time units.
+
+
 0.12.0 (2025-05-15)
 -------------------
 

ndonnx/_typed_array/datetime.py

@@ -30,7 +30,7 @@
     from ndonnx.types import NestedSequence, OnnxShape, PyScalar
 
 
-Unit = Literal["ns", "s"]
+Unit = Literal["ns", "us", "ms", "s"]
 
 _NAT_SENTINEL = onnx.const(np.iinfo(np.int64).min).astype(onnx.int64)
 TIMEARRAY_co = TypeVar("TIMEARRAY_co", bound="TimeBaseArray", covariant=True)
@@ -63,6 +63,9 @@ def __ndx_cast_from__(self, arr: TyArrayBase) -> TIMEARRAY_co:
     def __ndx_result_type__(self, other: DType | PyScalar) -> DType:
         if isinstance(other, int):
             return self
+        if isinstance(other, BaseTimeDType):
+            target_unit = _result_unit(self.unit, other.unit)
+            return type(self)(target_unit)
         return NotImplemented
 
     def __ndx_argument__(self, shape: OnnxShape) -> TIMEARRAY_co:
@@ -277,7 +280,7 @@ def _apply_comp(
         if type(self) is not type(other):
             return NotImplemented
 
-        self, other = _coerce_units(self, other_arr)
+        self, other = _promote_unit(self, other_arr)
 
         data = op(self._data, other._data)
         is_nat = self.is_nat | other.is_nat
@@ -326,10 +329,10 @@ def __ndx_where__(
     ) -> TyArrayBase:
         if not isinstance(other, TyArrayBase):
             return NotImplemented
-        if self.dtype != other.dtype or not isinstance(other, type(self)):
-            return NotImplemented
-
-        return self.dtype._build(onnx.where(cond, self._data, other._data))
+        if isinstance(other, type(self)):
+            a, b = _promote_unit(self, other)
+            return a.dtype._build(onnx.where(cond, a._data, b._data))
+        return NotImplemented
 
     def clip(
         self, /, min: TyArrayBase | None = None, max: TyArrayBase | None = None
@@ -394,10 +397,11 @@ def __add__(self, rhs: TyArrayBase | PyScalar) -> TyArrayTimeDelta:
         if isinstance(rhs, int):
             rhs = TyArrayTimeDelta(onnx.const(rhs), self.dtype.unit)
         if isinstance(rhs, TyArrayTimeDelta):
-            if {self.dtype.unit, rhs.dtype.unit} == {"s", "ns"}:
-                self = self.astype(TimeDelta64DType("ns"))
-                rhs = rhs.astype(TimeDelta64DType("ns"))
-            return _apply_op(self, rhs, operator.add, True)
+            allowed_units = set(get_args(Unit))
+            lhs = self
+            if lhs.dtype.unit in allowed_units and rhs.dtype.unit in allowed_units:
+                lhs, rhs = _promote_unit(lhs, rhs)
+            return _apply_op(lhs, rhs, operator.add, True)
         return NotImplemented
 
     def __radd__(self, lhs: TyArrayBase | PyScalar) -> TyArrayTimeDelta:
@@ -419,10 +423,11 @@ def __sub__(self, rhs: TyArrayBase | PyScalar) -> TyArrayTimeDelta:
         if isinstance(rhs, int):
             rhs = TyArrayTimeDelta(onnx.const(rhs), self.dtype.unit)
         if isinstance(rhs, TyArrayTimeDelta):
-            if {self.dtype.unit, rhs.dtype.unit} == {"s", "ns"}:
-                self = self.astype(TimeDelta64DType("ns"))
-                rhs = rhs.astype(TimeDelta64DType("ns"))
-            return _apply_op(self, rhs, operator.sub, True)
+            allowed_units = set(get_args(Unit))
+            lhs = self
+            if lhs.dtype.unit in allowed_units and rhs.dtype.unit in allowed_units:
+                lhs, rhs = _promote_unit(lhs, rhs)
+            return _apply_op(lhs, rhs, operator.sub, True)
         return NotImplemented
 
     def __rsub__(self, lhs: TyArrayBase | PyScalar) -> TyArrayTimeDelta:
@@ -551,7 +556,7 @@ def __add__(self, rhs: TyArrayBase | PyScalar) -> Self:
         if rhs is NotImplemented:
             return NotImplemented
 
-        lhs, rhs = _coerce_units(self, rhs)
+        lhs, rhs = _promote_unit(self, rhs)
 
         data = lhs._data + rhs._data
         is_nat = lhs.is_nat | rhs.is_nat
@@ -571,7 +576,7 @@ def _sub(self, other, forward: bool):
             return self - other_ if forward else other_ - self
 
         if isinstance(other, TyArrayDateTime):
-            a, b = _coerce_units(self, other)
+            a, b = _promote_unit(self, other)
             is_nat = a.is_nat | b.is_nat
             data = safe_cast(
                 onnx.TyArrayInt64, a._data - b._data if forward else b._data - a._data
@@ -582,7 +587,7 @@ def _sub(self, other, forward: bool):
 
         elif isinstance(other, TyArrayTimeDelta) and forward:
             # *_ due to types of various locals set in the previous if statement
-            a_, b_ = _coerce_units(self, other)
+            a_, b_ = _promote_unit(self, other)
             is_nat = a_.is_nat | b_.is_nat
             data = safe_cast(
                 onnx.TyArrayInt64,
@@ -610,13 +615,10 @@ def __ndx_equal__(self, other) -> onnx.TyArrayBool:
 
         if not isinstance(other, TyArrayDateTime):
             return NotImplemented
-        if self.dtype.unit != other.dtype.unit:
-            raise TypeError(
-                "comparison between different units is not implemented, yet"
-            )
 
-        res = self._data == other._data
-        is_nat = self.is_nat | other.is_nat
+        lhs, rhs = _promote_unit(self, other)
+        res = lhs._data == rhs._data
+        is_nat = lhs.is_nat | rhs.is_nat
 
         return safe_cast(onnx.TyArrayBool, res & ~is_nat)
 
@@ -662,17 +664,16 @@ def _coerce_other(
     return NotImplemented
 
 
-def _coerce_units(a: T1, b: T2) -> tuple[T1, T2]:
-    table: dict[tuple[Unit, Unit], Unit] = {
-        ("ns", "s"): "ns",
-        ("s", "ns"): "ns",
-        ("s", "s"): "s",
-        ("ns", "ns"): "ns",
-    }
-    target = table[(a.dtype.unit, b.dtype.unit)]
-    dtype_a = type(a.dtype)(unit=target)
-    dtype_b = type(b.dtype)(unit=target)
-    return (a.astype(dtype_a), b.astype(dtype_b))
+def _promote_unit(a: T1, b: T2) -> tuple[T1, T2]:
+    unit = _result_unit(a.dtype.unit, b.dtype.unit)
+
+    return a.astype(type(a.dtype)(unit=unit)), b.astype(type(b.dtype)(unit=unit))
+
+
+def _result_unit(a: Unit, b: Unit) -> Unit:
+    ordered_units = ["ns", "us", "ms", "s"]
+    res, _ = sorted([a, b], key=lambda el: ordered_units.index(el))
+    return res  # type: ignore
 
 
 def validate_unit(unit: str) -> Unit:

tests/test_datetime.py

@@ -61,6 +61,8 @@ def test_datetime_from_np_array(ty, unit):
     "scalar, dtype, res_dtype",
     [
         (1, ndx.DateTime64DType("s"), ndx.DateTime64DType("s")),
+        (1, ndx.DateTime64DType("ms"), ndx.DateTime64DType("ms")),
+        (1, ndx.DateTime64DType("us"), ndx.DateTime64DType("us")),
         (1, ndx.DateTime64DType("ns"), ndx.DateTime64DType("ns")),
     ],
 )
@@ -196,9 +198,11 @@ def test_comparisons_timedelta(op, x, y, unit):
         (["NaT"], ["NaT"]),
     ],
 )
-def test_comparisons_datetime(op, x, y, unit):
-    np_x = np.array(x, f"datetime64[{unit}]")
-    np_y = np.array(y, f"datetime64[{unit}]")
+@pytest.mark.parametrize("unit1", get_args(Unit))
+@pytest.mark.parametrize("unit2", get_args(Unit))
+def test_comparisons_datetime(op, x, y, unit1, unit2):
+    np_x = np.array(x, f"datetime64[{unit1}]")
+    np_y = np.array(y, f"datetime64[{unit2}]")
 
     desired = op(np_x, np_y)
     actual = op(ndx.asarray(np_x), ndx.asarray(np_y))
@@ -216,9 +220,11 @@ def test_comparisons_datetime(op, x, y, unit):
     ],
 )
 @pytest.mark.parametrize("forward", [True, False])
-def test_subtraction_datetime_arrays(x, y, unit, forward):
-    np_x = np.array(x, f"datetime64[{unit}]")
-    np_y = np.array(y, f"datetime64[{unit}]")
+@pytest.mark.parametrize("unit1", get_args(Unit))
+@pytest.mark.parametrize("unit2", get_args(Unit))
+def test_subtraction_datetime_arrays(x, y, unit1, unit2, forward):
+    np_x = np.array(x, f"datetime64[{unit1}]")
+    np_y = np.array(y, f"datetime64[{unit2}]")
 
     desired = np_x - np_y if forward else np_y - np_x
     actual = (
@@ -265,13 +271,13 @@ def test_isnan(unit):
     )
 
 
-@pytest.mark.parametrize(
-    "dtype", ["datetime64[s]", "timedelta64[s]", "datetime64[ns]", "timedelta64[ns]"]
-)
-def test_where(dtype):
+@pytest.mark.parametrize("unit1", get_args(Unit))
+@pytest.mark.parametrize("unit2", get_args(Unit))
+@pytest.mark.parametrize("dtype_name", ["datetime64", "timedelta64"])
+def test_where(dtype_name, unit1, unit2):
     cond = np.asarray([False, True, False])
-    np_arr1 = np.asarray(["NaT", 1, 2], dtype=dtype)
-    np_arr2 = np.asarray(["NaT", "NaT", "NaT"], dtype=dtype)
+    np_arr1 = np.asarray(["NaT", 1, 2], dtype=f"{dtype_name}[{unit1}]")
+    np_arr2 = np.asarray(["NaT", "NaT", "NaT"], dtype=f"{dtype_name}[{unit2}]")
 
     expected = np.where(cond, np_arr1, np_arr2)
     actual = ndx.where(ndx.asarray(cond), ndx.asarray(np_arr1), ndx.asarray(np_arr2))
@@ -400,3 +406,19 @@ def test_datetime_dtypes_have_numpy_repr(unit):
     dtype = ndx.DateTime64DType(unit)
 
     assert dtype.unwrap_numpy() == np.dtype(f"datetime64[{unit}]")
+
+
+@pytest.mark.parametrize("dtype_cls", [ndx.DateTime64DType, ndx.TimeDelta64DType])
+@pytest.mark.parametrize("unit1", get_args(Unit))
+@pytest.mark.parametrize("unit2", get_args(Unit))
+def test_result_type(dtype_cls, unit1, unit2):
+    def do(npx):
+        dtype1 = dtype_cls(unit1)
+        dtype2 = dtype_cls(unit2)
+        if npx == np:
+            dtype1 = dtype1.unwrap_numpy()
+            dtype2 = dtype2.unwrap_numpy()
+
+        return npx.result_type(dtype1, dtype2)
+
+    assert do(np) == do(ndx).unwrap_numpy()