Sum/nansum numpy-like impl (#570)

* Sum/nansum impl

* Use numba isnan dtype

Co-authored-by: Alexander Kalistratov <alexander.kalistratov@intel.com>

Author: 1e-to

sdc/functions/numpy_like.py

@@ -36,6 +36,7 @@
 
 from numba import types, jit, prange, numpy_support, literally
 from numba.errors import TypingError
+from numba.targets.arraymath import get_isnan
 
 import sdc
 from sdc.utilities.sdc_typing_utils import TypeChecker
@@ -47,6 +48,14 @@ def astype(self, dtype):
     pass
 
 
+def sum(self):
+    pass
+
+
+def nansum(self):
+    pass
+
+
 @sdc_overload(astype)
 def sdc_astype_overload(self, dtype):
     """
@@ -106,3 +115,61 @@ def sdc_astype_number_impl(self, dtype):
         return sdc_astype_number_impl
 
     ty_checker.raise_exc(self.dtype, 'str or type', 'self.dtype')
+
+
+@sdc_overload(sum)
+def sdc_sum_overload(self):
+    """
+    Intel Scalable Dataframe Compiler Developer Guide
+    *************************************************
+    Parallel replacement of numpy.sum.
+    .. only:: developer
+       Test: python -m sdc.runtests sdc.tests.test_sdc_numpy -k sum
+    """
+
+    dtype = self.dtype
+    isnan = get_isnan(dtype)
+    if not isinstance(self, types.Array):
+        return None
+
+    if isinstance(dtype, types.Number):
+        def sdc_sum_number_impl(self):
+            length = len(self)
+            result = 0
+            for i in prange(length):
+                if not isnan(self[i]):
+                    result += self[i]
+                else:
+                    return numpy.nan
+
+            return result
+
+        return sdc_sum_number_impl
+
+
+@sdc_overload(nansum)
+def sdc_sum_overload(self):
+    """
+    Intel Scalable Dataframe Compiler Developer Guide
+    *************************************************
+    Parallel replacement of numpy.nansum.
+    .. only:: developer
+       Test: python -m sdc.runtests sdc.tests.test_sdc_numpy -k nansum
+    """
+
+    dtype = self.dtype
+    isnan = get_isnan(dtype)
+    if not isinstance(self, types.Array):
+        return None
+
+    if isinstance(dtype, types.Number):
+        def sdc_nansum_number_impl(self):
+            length = len(self)
+            result = 0
+            for i in prange(length):
+                if not numpy.isnan(self[i]):
+                    result += self[i]
+
+            return result
+
+        return sdc_nansum_number_impl

sdc/tests/test_sdc_numpy.py

@@ -142,5 +142,35 @@ def sdc_impl(a, t):
                 with self.subTest(data=case, type=type_):
                     np.testing.assert_array_equal(sdc_func(a, type_), ref_impl(a, type_))
 
+    def test_nansum(self):
+        def ref_impl(a):
+            return np.nansum(a)
+
+        def sdc_impl(a):
+            return numpy_like.nansum(a)
+
+        sdc_func = self.jit(sdc_impl)
+
+        cases = [[5, 2, 0, 333, -4], [3.3, 5.4, np.nan, 7.9, np.nan]]
+        for case in cases:
+            a = np.array(case)
+            with self.subTest(data=case):
+                np.testing.assert_array_equal(sdc_func(a), ref_impl(a))
+
+    def test_sum(self):
+        def ref_impl(a):
+            return np.sum(a)
+
+        def sdc_impl(a):
+            return numpy_like.sum(a)
+
+        sdc_func = self.jit(sdc_impl)
+
+        cases = [[5, 2, 0, 333, -4], [3.3, 5.4, np.nan, 7.9, np.nan]]
+        for case in cases:
+            a = np.array(case)
+            with self.subTest(data=case):
+                np.testing.assert_array_equal(sdc_func(a), ref_impl(a))
+
 if __name__ == "__main__":
     unittest.main()

sdc/tests/tests_perf/test_perf_numpy.py

@@ -99,6 +99,15 @@ def _test_case(self, cases, name, total_data_length, data_num=1, input_data=test
         CE(type_='Numba', code='data.astype(np.int64)', jitted=True),
         CE(type_='SDC', code='sdc.functions.numpy_like.astype(data, np.int64)', jitted=True),
     ], usecase_params='data'),
+    TC(name='nansum', size=[10 ** 7], call_expr=[
+        CE(type_='Python', code='np.nansum(data)', jitted=False),
+        CE(type_='SDC', code='sdc.functions.numpy_like.nansum(data)', jitted=True),
+    ], usecase_params='data'),
+    TC(name='sum', size=[10 ** 7], call_expr=[
+        CE(type_='Python', code='np.sum(data)', jitted=False),
+        CE(type_='Numba', code='np.sum(data)', jitted=True),
+        CE(type_='SDC', code='sdc.functions.numpy_like.sum(data)', jitted=True),
+    ], usecase_params='data'),
 ]
 
 generate_test_cases(cases, TestFunctions, 'function')