Repository: numba/numba

Issue:
np.mean() and np.sum() don't support timedelta arrays
Numpy supports it, but Numba fails on them.



Discussion:

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -719,9 +719,10 @@ def array_view(context, builder, sig, args):
 @implement(numpy.sum, types.Kind(types.Array))
 @implement("array.sum", types.Kind(types.Array))
 def array_sum(context, builder, sig, args):
+    zero = sig.return_type(0)
 
     def array_sum_impl(arr):
-        c = 0
+        c = zero
         for v in arr.flat:
             c += v
         return c
@@ -751,13 +752,14 @@ def array_prod_impl(arr):
 def array_cumsum(context, builder, sig, args):
     scalar_dtype = sig.return_type.dtype
     dtype = as_dtype(scalar_dtype)
+    zero = scalar_dtype(0)
 
     def array_cumsum_impl(arr):
         size = 1
         for i in arr.shape:
             size = size * i
         out = numpy.empty(size, dtype)
-        c = 0
+        c = zero
         for idx, v in enumerate(arr.flat):
             c += v
             out[idx] = c
@@ -793,11 +795,12 @@ def array_cumprod_impl(arr):
 @implement(numpy.mean, types.Kind(types.Array))
 @implement("array.mean", types.Kind(types.Array))
 def array_mean(context, builder, sig, args):
+    zero = sig.return_type(0)
 
     def array_mean_impl(arr):
         # Can't use the naive `arr.sum() / arr.size`, as it would return
         # a wrong result on integer sum overflow.
-        c = 0
+        c = zero
         for v in arr.flat:
             c += v
         return c / arr.size
@@ -836,15 +839,35 @@ def array_std_impl(arry):
 @implement(numpy.min, types.Kind(types.Array))
 @implement("array.min", types.Kind(types.Array))
 def array_min(context, builder, sig, args):
-    def array_min_impl(arry):
-        for v in arry.flat:
-            min_value = v
-            break
+    ty = sig.args[0].dtype
+    if isinstance(ty, (types.NPDatetime, types.NPTimedelta)):
+        # NaT is smaller than every other value, but it is
+        # ignored as far as min() is concerned.
+        nat = ty('NaT')
+
+        def array_min_impl(arry):
+            min_value = nat
+            it = arry.flat
+            for v in it:
+                if v != nat:
+                    min_value = v
+                    break
+
+            for v in it:
+                if v != nat and v < min_value:
+                    min_value = v
+            return min_value
 
-        for v in arry.flat:
-            if v < min_value:
+    else:
+        def array_min_impl(arry):
+            for v in arry.flat:
                 min_value = v
-        return min_value
+                break
+
+            for v in arry.flat:
+                if v < min_value:
+                    min_value = v
+            return min_value
     return context.compile_internal(builder, array_min_impl, sig, args)
 
 
@@ -868,6 +891,8 @@ def array_max_impl(arry):
 @implement(numpy.argmin, types.Kind(types.Array))
 @implement("array.argmin", types.Kind(types.Array))
 def array_argmin(context, builder, sig, args):
+    # XXX argmin() is inconsistent with min() on NaT values:
+    # https://github.com/numpy/numpy/issues/6030
     def array_argmin_impl(arry):
         for v in arry.flat:
             min_value = v
diff --git a/numba/types.py b/numba/types.py
--- a/numba/types.py
+++ b/numba/types.py
@@ -99,6 +99,13 @@ def __lt__(self, other):
         # also greater...).
         return self.unit_code < other.unit_code
 
+    def cast_python_value(self, value):
+        cls = getattr(numpy, self.type_name)
+        if self.unit:
+            return cls(value, self.unit)
+        else:
+            return cls(value)
+
 
 @utils.total_ordering
 class NPTimedelta(_NPDatetimeBase):
diff --git a/numba/typing/npdatetime.py b/numba/typing/npdatetime.py
--- a/numba/typing/npdatetime.py
+++ b/numba/typing/npdatetime.py
@@ -74,15 +74,19 @@ def generic(self, args, kws):
         left, right = args
         if isinstance(right, types.NPTimedelta):
             td, other = right, left
+            sig_factory = lambda other: signature(td, other, td)
         elif isinstance(left, types.NPTimedelta):
             td, other = left, right
+            sig_factory = lambda other: signature(td, td, other)
         else:
             return
+        if not isinstance(other, (types.Float, types.Integer)):
+            return
         # Force integer types to convert to signed because it matches
         # timedelta64 semantics better.
-        if other not in types.signed_domain and other not in types.real_domain:
-            return
-        return signature(td, left, right)
+        if isinstance(other, types.Integer):
+            other = types.int64
+        return sig_factory(other)
 
 
 class TimedeltaDivOp(AbstractTemplate):
@@ -99,10 +103,12 @@ def generic(self, args, kws):
             if (npdatetime.can_cast_timedelta_units(left.unit, right.unit)
                 or npdatetime.can_cast_timedelta_units(right.unit, left.unit)):
                 return signature(types.float64, left, right)
-        # Force integer types to convert to signed because it matches
-        # timedelta64 semantics better.
-        elif right in types.signed_domain or right in types.real_domain:
+        elif isinstance(right, (types.Float)):
             return signature(left, left, right)
+        elif isinstance(right, (types.Integer)):
+            # Force integer types to convert to signed because it matches
+            # timedelta64 semantics better.
+            return signature(left, left, types.int64)
 
 
 @builtin


Test patch:
diff --git a/numba/tests/test_array_reductions.py b/numba/tests/test_array_reductions.py
--- a/numba/tests/test_array_reductions.py
+++ b/numba/tests/test_array_reductions.py
@@ -8,7 +8,7 @@
 from numba import typeof, types
 from numba.compiler import compile_isolated
 from numba.numpy_support import as_dtype
-from .support import TestCase, CompilationCache
+from .support import TestCase, CompilationCache, skip_on_numpy_16
 
 
 def array_cumprod(arr):
@@ -292,6 +292,65 @@ def test_std_magnitude(self):
         self.check_aggregation_magnitude(array_std)
         self.check_aggregation_magnitude(array_std_global)
 
+    def _do_check_nptimedelta(self, pyfunc, arr):
+        arrty = typeof(arr)
+
+        cres = compile_isolated(pyfunc, [arrty])
+        cfunc = cres.entry_point
+
+        self.assertPreciseEqual(cfunc(arr), pyfunc(arr))
+        # Even vs. odd size, for np.median
+        self.assertPreciseEqual(cfunc(arr[:-1]), pyfunc(arr[:-1]))
+        # Test with different orders, for np.median
+        arr = arr[::-1].copy()  # Keep 'C' layout
+        self.assertPreciseEqual(cfunc(arr), pyfunc(arr))
+        np.random.shuffle(arr)
+        self.assertPreciseEqual(cfunc(arr), pyfunc(arr))
+        # Test with a NaT
+        arr[arr.size // 2] = 'NaT'
+        self.assertPreciseEqual(cfunc(arr), pyfunc(arr))
+        # Test with all NaTs
+        arr.fill(arrty.dtype('NaT'))
+        self.assertPreciseEqual(cfunc(arr), pyfunc(arr))
+
+    @skip_on_numpy_16
+    def check_npdatetime(self, pyfunc):
+        arr = np.arange(10).astype(dtype='M8[Y]')
+        self._do_check_nptimedelta(pyfunc, arr)
+
+    @skip_on_numpy_16
+    def check_nptimedelta(self, pyfunc):
+        arr = np.arange(10).astype(dtype='m8[s]')
+        self._do_check_nptimedelta(pyfunc, arr)
+
+    def test_min_npdatetime(self):
+        self.check_npdatetime(array_min)
+        self.check_nptimedelta(array_min)
+
+    def test_max_npdatetime(self):
+        self.check_npdatetime(array_max)
+        self.check_nptimedelta(array_max)
+
+    def test_argmin_npdatetime(self):
+        self.check_npdatetime(array_argmin)
+        self.check_nptimedelta(array_argmin)
+
+    def test_argmax_npdatetime(self):
+        self.check_npdatetime(array_argmax)
+        self.check_nptimedelta(array_argmax)
+
+    def test_median_npdatetime(self):
+        self.check_nptimedelta(array_median_global)
+
+    def test_sum_npdatetime(self):
+        self.check_nptimedelta(array_sum)
+
+    def test_cumsum_npdatetime(self):
+        self.check_nptimedelta(array_cumsum)
+
+    def test_mean_npdatetime(self):
+        self.check_nptimedelta(array_mean)
+
     @classmethod
     def install_generated_tests(cls):
         # These form a testing product where each of the combinations are tested
diff --git a/numba/tests/test_npdatetime.py b/numba/tests/test_npdatetime.py
--- a/numba/tests/test_npdatetime.py
+++ b/numba/tests/test_npdatetime.py
@@ -293,6 +293,8 @@ def check(a, b, expected):
             self.assertPreciseEqual(f(a, b), expected)
             self.assertPreciseEqual(f(b, a), expected)
 
+        # non-int64 int * timedelta64
+        check(TD(3), np.uint32(2), TD(6))
         # int * timedelta64
         check(TD(3), 2, TD(6))
         check(TD(3, 'ps'), 2, TD(6, 'ps'))
@@ -314,6 +316,8 @@ def check(a, b, expected):
             self.assertPreciseEqual(div(a, b), expected)
             self.assertPreciseEqual(floordiv(a, b), expected)
 
+        # timedelta64 / non-int64 int
+        check(TD(-3, 'ps'), np.uint32(2), TD(-1, 'ps'))
         # timedelta64 / int
         check(TD(3), 2, TD(1))
         check(TD(-3, 'ps'), 2, TD(-1, 'ps'))
diff --git a/numba/tests/test_support.py b/numba/tests/test_support.py
--- a/numba/tests/test_support.py
+++ b/numba/tests/test_support.py
@@ -6,7 +6,7 @@
 
 from numba import utils
 from numba import unittest_support as unittest
-from .support import TestCase
+from .support import TestCase, skip_on_numpy_16
 
 
 DBL_EPSILON = 2**-52
@@ -254,6 +254,30 @@ def test_arrays(self):
         self.eq(a, b, prec='double')
         self.ne(a, c, prec='double')
 
+    @skip_on_numpy_16
+    def test_npdatetime(self):
+        a = np.datetime64('1900', 'Y')
+        b = np.datetime64('1900', 'Y')
+        c = np.datetime64('1900-01-01', 'D')
+        d = np.datetime64('1901', 'Y')
+        self.eq(a, b)
+        # Different unit
+        self.ne(a, c)
+        # Different value
+        self.ne(a, d)
+
+    @skip_on_numpy_16
+    def test_nptimedelta(self):
+        a = np.timedelta64(1, 'h')
+        b = np.timedelta64(1, 'h')
+        c = np.timedelta64(60, 'm')
+        d = np.timedelta64(2, 'h')
+        self.eq(a, b)
+        # Different unit
+        self.ne(a, c)
+        # Different value
+        self.ne(a, d)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/numba/tests/test_types.py b/numba/tests/test_types.py
--- a/numba/tests/test_types.py
+++ b/numba/tests/test_types.py
@@ -7,9 +7,12 @@
 import gc
 import weakref
 
+import numpy as np
+
 from numba.utils import IS_PY3
 from numba import abstracttypes, types, typing
 from numba import unittest_support as unittest
+from numba.npdatetime import NPDATETIME_SUPPORTED
 from .support import TestCase
 
 
@@ -193,6 +196,16 @@ def test_call_notation(self):
         # Value cast
         self.assertPreciseEqual(i(42.5), 42)
         self.assertPreciseEqual(d(-5), -5.0)
+        if NPDATETIME_SUPPORTED:
+            ty = types.NPDatetime('Y')
+            self.assertPreciseEqual(ty('1900'), np.datetime64('1900', 'Y'))
+            self.assertPreciseEqual(ty('NaT'), np.datetime64('NaT', 'Y'))
+            ty = types.NPTimedelta('s')
+            self.assertPreciseEqual(ty(5), np.timedelta64(5, 's'))
+            self.assertPreciseEqual(ty('NaT'), np.timedelta64('NaT', 's'))
+            ty = types.NPTimedelta('')
+            self.assertPreciseEqual(ty(5), np.timedelta64(5))
+            self.assertPreciseEqual(ty('NaT'), np.timedelta64('NaT'))
 
     def test_bitwidth_number_types(self):
         """