Skip to content

Commit 75f86bf

Browse files
benwilliamgrahambenwilliamgraham
committed
Fixed array max/min
1 parent 9def383 commit 75f86bf

4 files changed

Lines changed: 164 additions & 45 deletions

File tree

‎numba/cpython/builtins.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -500,7 +500,7 @@ def lower_get_type_min_value(context, builder, sig, args):
500500
elif isinstance(typ, (types.NPDatetime, types.NPTimedelta)):
501501
bw = 64
502502
lty = ir.IntType(bw)
503-
val = types.uint64.minval
503+
val = types.int64.minval + 1 # minval is NaT, so minval + 1 is the smallest value
504504
res = ir.Constant(lty, val)
505505
return impl_ret_untracked(context, builder, lty, res)
506506

‎numba/np/numpy_support.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,8 @@
1212
# re-export
1313
from numba.core.cgutils import is_nonelike # noqa: F401
1414

15-
1615
numpy_version = tuple(map(int, np.__version__.split('.')[:2]))
1716

18-
1917
FROM_DTYPE = {
2018
np.dtype('bool'): types.boolean,
2119
np.dtype('int8'): types.int8,

‎numba/parfors/parfor.py‎

Lines changed: 59 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import numba.core.ir
2828
from numba.core import types, typing, utils, errors, ir, analysis, postproc, rewrites, typeinfer, config, ir_utils
2929
from numba import prange, pndindex
30-
from numba.np.numpy_support import as_dtype
30+
from numba.np.numpy_support import as_dtype, numpy_version
3131
from numba.core.typing.templates import infer_global, AbstractTemplate
3232
from numba.stencils.stencilparfor import StencilPass
3333
from numba.core.extending import register_jitable
@@ -128,13 +128,35 @@ def min_parallel_impl(return_type, arg):
128128
def min_1(in_arr):
129129
return in_arr[()]
130130
elif arg.ndim == 1:
131-
def min_1(in_arr):
132-
numba.parfors.parfor.init_prange()
133-
min_checker(len(in_arr))
134-
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
135-
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
136-
val = min(val, in_arr[i])
137-
return val
131+
if isinstance(arg.dtype, (types.NPDatetime, types.NPTimedelta)):
132+
if numpy_version >= (1, 18):
133+
def min_1(in_arr):
134+
numba.parfors.parfor.init_prange()
135+
min_checker(len(in_arr))
136+
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
137+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
138+
if np.isnat(in_arr[i]):
139+
return in_arr[i]
140+
val = min(val, in_arr[i])
141+
return val
142+
else:
143+
def min_1(in_arr):
144+
numba.parfors.parfor.init_prange()
145+
min_checker(len(in_arr))
146+
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
147+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
148+
if np.isnat(in_arr[i]):
149+
continue
150+
val = min(val, in_arr[i])
151+
return val
152+
else:
153+
def min_1(in_arr):
154+
numba.parfors.parfor.init_prange()
155+
min_checker(len(in_arr))
156+
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
157+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
158+
val = min(val, in_arr[i])
159+
return val
138160
else:
139161
def min_1(in_arr):
140162
numba.parfors.parfor.init_prange()
@@ -150,13 +172,35 @@ def max_parallel_impl(return_type, arg):
150172
def max_1(in_arr):
151173
return in_arr[()]
152174
elif arg.ndim == 1:
153-
def max_1(in_arr):
154-
numba.parfors.parfor.init_prange()
155-
max_checker(len(in_arr))
156-
val = numba.cpython.builtins.get_type_min_value(in_arr.dtype)
157-
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
158-
val = max(val, in_arr[i])
159-
return val
175+
if isinstance(arg.dtype, (types.NPDatetime, types.NPTimedelta)):
176+
if numpy_version >= (1, 18):
177+
def min_1(in_arr):
178+
numba.parfors.parfor.init_prange()
179+
min_checker(len(in_arr))
180+
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
181+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
182+
if np.isnat(in_arr[i]):
183+
return in_arr[i]
184+
val = max(val, in_arr[i])
185+
return val
186+
else:
187+
def min_1(in_arr):
188+
numba.parfors.parfor.init_prange()
189+
min_checker(len(in_arr))
190+
val = numba.cpython.builtins.get_type_max_value(in_arr.dtype)
191+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
192+
if np.isnat(in_arr[i]):
193+
continue
194+
val = max(val, in_arr[i])
195+
return val
196+
else:
197+
def max_1(in_arr):
198+
numba.parfors.parfor.init_prange()
199+
max_checker(len(in_arr))
200+
val = numba.cpython.builtins.get_type_min_value(in_arr.dtype)
201+
for i in numba.parfors.parfor.internal_prange(len(in_arr)):
202+
val = max(val, in_arr[i])
203+
return val
160204
else:
161205
def max_1(in_arr):
162206
numba.parfors.parfor.init_prange()

‎numba/tests/test_npdatetime.py‎

Lines changed: 104 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -550,20 +550,23 @@ def check(a):
550550
(TD,) * len(TD_CASES) + (DT,) * len(TD_CASES)):
551551
check(typ(*case))
552552

553-
def test_min_max(self):
554-
min_f = self.jit(min_usecase)
555-
max_f = self.jit(max_usecase)
553+
def _test_min_max(self, usecase):
554+
f = self.jit(usecase)
556555
def check(a, b):
557-
self.assertPreciseEqual(min_f(a, b), min(a, b))
558-
self.assertPreciseEqual(max_f(a, b), max(a, b))
556+
self.assertPreciseEqual(f(a, b), usecase(a, b))
559557

560-
check(TD(1), TD(2))
561-
check(TD(1), TD(1))
562-
check(TD(2), TD(1))
563-
check(TD(1, 's'), TD(2, 's'))
564-
check(TD(1, 's'), TD(1, 's'))
565-
check(TD(2, 's'), TD(1, 's'))
566-
check(TD('Nat'), TD('Nat'))
558+
for cases in (
559+
(TD(0), TD(1), TD(2), TD('NaT')),
560+
(TD(0, 's'), TD(1, 's'), TD(2, 's'), TD('NaT', 's')),
561+
):
562+
for a, b in itertools.product(cases, cases):
563+
check(a, b)
564+
565+
def test_min(self):
566+
self._test_min_max(min_usecase)
567+
568+
def test_max(self):
569+
self._test_min_max(max_usecase)
567570

568571

569572
class TestTimedeltaArithmeticNoPython(TestTimedeltaArithmetic):
@@ -793,21 +796,23 @@ def check_lt(a, b, expected):
793796
check_lt(a, b + np.timedelta64(1, unit), True)
794797
check_lt(b - np.timedelta64(1, unit), a, True)
795798

796-
def test_min_max(self):
797-
min_f = self.jit(min_usecase)
798-
max_f = self.jit(max_usecase)
799+
def _test_min_max(self, usecase):
800+
f = self.jit(usecase)
799801
def check(a, b):
800-
self.assertPreciseEqual(min_f(a, b), min(a, b))
801-
self.assertPreciseEqual(max_f(a, b), max(a, b))
802+
self.assertPreciseEqual(f(a, b), usecase(a, b))
803+
804+
for cases in (
805+
(DT(0, 'ns'), DT(1, 'ns'), DT(2, 'ns'), DT('NaT', 'ns')),
806+
(DT(0, 's'), DT(1, 's'), DT(2, 's'), DT('NaT', 's')),
807+
):
808+
for a, b in itertools.product(cases, cases):
809+
check(a, b)
802810

803-
check(DT(1, 'ns'), DT(2, 'ns'))
804-
check(DT(1, 'ns'), DT(1, 'ns'))
805-
check(DT(2, 'ns'), DT(1, 'ns'))
806-
check(DT(1, 's'), DT(2, 's'))
807-
check(DT(1, 's'), DT(1, 's'))
808-
check(DT(2, 's'), DT(1, 's'))
809-
check(DT('Nat', 'ns'), DT('Nat', 'ns'))
811+
def test_min(self):
812+
self._test_min_max(min_usecase)
810813

814+
def test_max(self):
815+
self._test_min_max(max_usecase)
811816

812817
class TestDatetimeArithmeticNoPython(TestDatetimeArithmetic):
813818

@@ -955,18 +960,50 @@ def impl(arr):
955960
cfunc = njit(parallel=parallel)(impl)
956961

957962
test_cases = [
963+
np.array([
964+
DT(0, "ns"),
965+
DT(1, "ns"),
966+
DT(2, "ns"),
967+
DT(3, "ns"),
968+
]),
958969
np.array([
959970
DT("2011-01-01", "ns"),
960971
DT("1971-02-02", "ns"),
972+
DT("1900-01-01", "ns"),
961973
DT("2021-03-03", "ns"),
962974
DT("2004-12-07", "ns"),
963975
]),
964976
np.array([
965977
DT("2011-01-01", "D"),
966978
DT("1971-02-02", "D"),
979+
DT("1900-01-01", "D"),
967980
DT("2021-03-03", "D"),
968981
DT("2004-12-07", "D"),
969982
]),
983+
np.array([
984+
DT("2011-01-01", "ns"),
985+
DT("1971-02-02", "ns"),
986+
DT("1900-01-01", "ns"),
987+
DT("2021-03-03", "ns"),
988+
DT("2004-12-07", "ns"),
989+
DT("NaT", "ns"),
990+
]),
991+
np.array([
992+
DT("NaT", "ns"),
993+
DT("2011-01-01", "ns"),
994+
DT("1971-02-02", "ns"),
995+
DT("1900-01-01", "ns"),
996+
DT("2021-03-03", "ns"),
997+
DT("2004-12-07", "ns"),
998+
]),
999+
np.array([
1000+
DT("NaT", "ns"),
1001+
]),
1002+
np.array([
1003+
DT("NaT", "ns"),
1004+
DT("NaT", "ns"),
1005+
DT("NaT", "ns"),
1006+
]),
9701007
np.array([
9711008
TD(1, "ns"),
9721009
TD(2, "ns"),
@@ -979,25 +1016,65 @@ def impl(arr):
9791016
TD(3, "D"),
9801017
TD(4, "D"),
9811018
]),
1019+
np.array([
1020+
TD("NaT", "ns"),
1021+
TD(1, "ns"),
1022+
TD(2, "ns"),
1023+
TD(3, "ns"),
1024+
TD(4, "ns"),
1025+
]),
1026+
np.array([
1027+
TD(1, "ns"),
1028+
TD(2, "ns"),
1029+
TD(3, "ns"),
1030+
TD(4, "ns"),
1031+
TD("NaT", "ns"),
1032+
]),
1033+
np.array([
1034+
TD("NaT", "ns"),
1035+
]),
1036+
np.array([
1037+
TD("NaT", "ns"),
1038+
TD("NaT", "ns"),
1039+
TD("NaT", "ns"),
1040+
]),
9821041
]
9831042

9841043
for arr in test_cases:
985-
self.assertEqual(py_func(arr), cfunc(arr))
986-
987-
1044+
py_res = py_func(arr)
1045+
c_res = cfunc(arr)
1046+
if np.isnat(py_res) or np.isnat(c_res):
1047+
self.assertTrue(np.isnat(py_res))
1048+
self.assertTrue(np.isnat(c_res))
1049+
else:
1050+
self.assertEqual(py_res, c_res)
9881051

9891052
def test_min_func(self):
1053+
self._test_min_max(min, False, False)
1054+
1055+
def test_np_min_func(self):
9901056
self._test_min_max(np.min, False, False)
9911057

9921058
@skip_parfors_unsupported
9931059
def test_min_func_parallel(self):
1060+
self._test_min_max(min, True, False)
1061+
1062+
@skip_parfors_unsupported
1063+
def test_np_min_func_parallel(self):
9941064
self._test_min_max(np.min, True, False)
9951065

9961066
def test_max_func(self):
1067+
self._test_min_max(max, False, False)
1068+
1069+
def test_np_max_func(self):
9971070
self._test_min_max(np.max, False, False)
9981071

9991072
@skip_parfors_unsupported
10001073
def test_max_func_parallel(self):
1074+
self._test_min_max(max, True, False)
1075+
1076+
@skip_parfors_unsupported
1077+
def test_np_max_func_parallel(self):
10011078
self._test_min_max(np.max, True, False)
10021079

10031080
def test_min_method(self):

0 commit comments

Comments
 (0)