capitalize, casefold, title, swapcase with use prange & work with None (#575)

* parallel and add mask for methods: capitalize, casefold, swapcase, title

* correct master

* change methods: capitalize, casefold, swapcase, title

* correct doc

* correct doc

* add variable len_arr

* change tests

* use subtest

* fix problem with PEP8

Author: Rubtsowa

sdc/datatypes/hpat_pandas_stringmethods_functions.py

@@ -87,6 +87,8 @@ def hpat_pandas_stringmethods_upper_impl(self):
 from sdc.utilities.sdc_typing_utils import TypeChecker
 from sdc.datatypes.hpat_pandas_stringmethods_types import StringMethodsType
 from sdc.utilities.utils import sdc_overload_method
+from sdc.hiframes.api import get_nan_mask
+from sdc.str_arr_ext import str_arr_set_na_by_mask, create_str_arr_from_list
 
 _hpat_pandas_stringmethods_autogen_global_dict = {
     'pandas': pandas,
@@ -1037,10 +1039,13 @@ def hpat_pandas_stringmethods_capitalize(self):
     ty_checker.check(self, StringMethodsType)
 
     def hpat_pandas_stringmethods_capitalize_impl(self):
+        mask = get_nan_mask(self._data._data)
         item_count = len(self._data)
-        result = [''] * item_count
+        res_list = [''] * item_count
         for idx in numba.prange(item_count):
-            result[idx] = self._data._data[idx].capitalize()
+            res_list[idx] = self._data._data[idx].capitalize()
+        str_arr = create_str_arr_from_list(res_list)
+        result = str_arr_set_na_by_mask(str_arr, mask)
 
         return pandas.Series(result, self._data._index, name=self._data._name)
 
@@ -1053,10 +1058,13 @@ def hpat_pandas_stringmethods_title(self):
     ty_checker.check(self, StringMethodsType)
 
     def hpat_pandas_stringmethods_title_impl(self):
+        mask = get_nan_mask(self._data._data)
         item_count = len(self._data)
-        result = [''] * item_count
+        res_list = [''] * item_count
         for idx in numba.prange(item_count):
-            result[idx] = self._data._data[idx].title()
+            res_list[idx] = self._data._data[idx].title()
+        str_arr = create_str_arr_from_list(res_list)
+        result = str_arr_set_na_by_mask(str_arr, mask)
 
         return pandas.Series(result, self._data._index, name=self._data._name)
 
@@ -1069,10 +1077,13 @@ def hpat_pandas_stringmethods_swapcase(self):
     ty_checker.check(self, StringMethodsType)
 
     def hpat_pandas_stringmethods_swapcase_impl(self):
+        mask = get_nan_mask(self._data._data)
         item_count = len(self._data)
-        result = [''] * item_count
+        res_list = [''] * item_count
         for idx in numba.prange(item_count):
-            result[idx] = self._data._data[idx].swapcase()
+            res_list[idx] = self._data._data[idx].swapcase()
+        str_arr = create_str_arr_from_list(res_list)
+        result = str_arr_set_na_by_mask(str_arr, mask)
 
         return pandas.Series(result, self._data._index, name=self._data._name)
 
@@ -1085,10 +1096,13 @@ def hpat_pandas_stringmethods_casefold(self):
     ty_checker.check(self, StringMethodsType)
 
     def hpat_pandas_stringmethods_casefold_impl(self):
+        mask = get_nan_mask(self._data._data)
         item_count = len(self._data)
-        result = [''] * item_count
+        res_list = [''] * item_count
         for idx in numba.prange(item_count):
-            result[idx] = self._data._data[idx].casefold()
+            res_list[idx] = self._data._data[idx].casefold()
+        str_arr = create_str_arr_from_list(res_list)
+        result = str_arr_set_na_by_mask(str_arr, mask)
 
         return pandas.Series(result, self._data._index, name=self._data._name)
 

sdc/hiframes/api.py

@@ -570,7 +570,11 @@ def get_nan_mask(arr):
 def get_nan_mask_overload(arr):
 
     def get_nan_mask_via_isna_impl(arr):
-        return np.array([isna(arr, i) for i in np.arange(len(arr))])
+        len_arr = len(arr)
+        res = np.empty(len_arr, dtype=np.bool_)
+        for i in numba.prange(len_arr):
+            res[i] = isna(arr, i)
+        return res
 
     if isinstance(arr, types.Array):
         dtype = arr.dtype

sdc/str_arr_ext.py

@@ -74,7 +74,6 @@ def typer(string_list=None):
 def iternext_str_array(context, builder, sig, args, result):
     """
     Implementation of iternext() for the StringArrayIterator type
-
     :param context: context descriptor
     :param builder: llvmlite IR Builder
     :param sig: iterator signature

sdc/tests/test_series.py

@@ -1301,7 +1301,6 @@ def test_impl(A, n):
             -3      6
             -3      3
             dtype: int64
-
             >>>S.loc[0:-3]
              0    6
             -3    6
@@ -3196,32 +3195,48 @@ def test_impl(S):
             return S.str.capitalize()
 
         sdc_func = self.jit(test_impl)
-        s = pd.Series(test_global_input_data_unicode_kind4)
-        pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
+        test_data = [test_global_input_data_unicode_kind4,
+                     ['lower', None, 'CAPITALS', None, 'this is a sentence', 'SwApCaSe', None]]
+        for data in test_data:
+            with self.subTest(data=data):
+                s = pd.Series(data)
+                pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
 
     def test_series_title_str(self):
         def test_impl(S):
             return S.str.title()
 
         sdc_func = self.jit(test_impl)
-        s = pd.Series(test_global_input_data_unicode_kind4)
-        pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
+        test_data = [test_global_input_data_unicode_kind4,
+                     ['lower', None, 'CAPITALS', None, 'this is a sentence', 'SwApCaSe', None]]
+        for data in test_data:
+            with self.subTest(data=data):
+                s = pd.Series(data)
+                pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
 
     def test_series_swapcase_str(self):
         def test_impl(S):
             return S.str.swapcase()
 
         sdc_func = self.jit(test_impl)
-        s = pd.Series(test_global_input_data_unicode_kind4)
-        pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
+        test_data = [test_global_input_data_unicode_kind4,
+                     ['lower', None, 'CAPITALS', None, 'this is a sentence', 'SwApCaSe', None]]
+        for data in test_data:
+            with self.subTest(data=data):
+                s = pd.Series(data)
+                pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
 
     def test_series_casefold_str(self):
         def test_impl(S):
             return S.str.casefold()
 
         sdc_func = self.jit(test_impl)
-        s = pd.Series(test_global_input_data_unicode_kind4)
-        pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
+        test_data = [test_global_input_data_unicode_kind4,
+                     ['lower', None, 'CAPITALS', None, 'this is a sentence', 'SwApCaSe', None]]
+        for data in test_data:
+            with self.subTest(data=data):
+                s = pd.Series(data)
+                pd.testing.assert_series_equal(sdc_func(s), test_impl(s))
 
     @sdc_limitation
     def test_series_append_same_names(self):