Fix for DataFrame str column created as List of Optional values (#582)

* Fix for DataFrame str column created as List of Optional values

* Applying review comments

* Fixing more remarks

Author: kozlov-alexey

sdc/hiframes/api.py

@@ -1185,7 +1185,9 @@ def generic(self, args, kws):
             and (isinstance(column.dtype, types.Number)
                  or column.dtype == types.boolean)):
             ret_typ = types.Array(column.dtype, 1, 'C')
-        if isinstance(column, types.List) and column.dtype == string_type:
+        if (isinstance(column, types.List)
+            and (column.dtype == string_type
+                 or isinstance(column.dtype, types.Optional) and column.dtype.type == string_type)):
             ret_typ = string_array_type
         if isinstance(column, DatetimeIndexType):
             ret_typ = sdc.hiframes.pd_index_ext._dt_index_data_typ
@@ -1214,7 +1216,10 @@ def fix_df_array_list_impl(column):  # pragma: no cover
         return fix_df_array_list_impl
 
     # convert list of strings to string array
-    if isinstance(column, types.List) and column.dtype == string_type:
+    if (isinstance(column, types.List)
+        and (column.dtype == string_type
+             or isinstance(column.dtype, types.Optional) and column.dtype.type == string_type)):
+
         def fix_df_array_str_impl(column):  # pragma: no cover
             return sdc.str_arr_ext.StringArray(column)
         return fix_df_array_str_impl

sdc/str_arr_ext.py

@@ -670,33 +670,43 @@ def construct_string_array(context, builder):
 @lower_builtin(StringArray, types.UniTuple)
 @lower_builtin(StringArray, types.Tuple)
 def impl_string_array_single(context, builder, sig, args):
-    if isinstance(args[0], types.UniTuple):
-        assert args[0].dtype == string_type
 
-    if isinstance(args[0], types.Tuple):
-        for i in args[0]:
+    arg = args[0]
+    if isinstance(arg, (types.UniTuple, types.List)):
+        assert (arg.dtype == string_type
+                or (isinstance(arg.dtype, types.Optional) and arg.dtype.type == string_type))
+
+    # FIXME: doesn't work for Tuple with None values
+    if isinstance(arg, types.Tuple):
+        for i in arg:
             assert i.dtype == string_type or i.dtype == types.StringLiteral
 
     if not sig.args:  # return empty string array if no args
         res = context.compile_internal(
             builder, lambda: pre_alloc_string_array(0, 0), sig, args)
         return res
 
-    def str_arr_from_list(in_list):
+    def str_arr_from_sequence(in_list):
         n_strs = len(in_list)
         total_chars = 0
         # TODO: use vector to avoid two passes?
         # get total number of chars
-        for s in in_list:
-            total_chars += get_utf8_size(s)
+        nan_mask = np.zeros(n_strs, dtype=np.bool_)
+        for i in numba.prange(n_strs):
+            s = in_list[i]
+            if s is None:
+                nan_mask[i] = True
+            else:
+                total_chars += get_utf8_size(s)
 
         A = pre_alloc_string_array(n_strs, total_chars)
-        for i in range(n_strs):
-            A[i] = in_list[i]
+        for i in np.arange(n_strs):
+            A[i] = '' if nan_mask[i] else in_list[i]
+        str_arr_set_na_by_mask(A, nan_mask)
 
         return A
 
-    res = context.compile_internal(builder, str_arr_from_list, sig, args)
+    res = context.compile_internal(builder, str_arr_from_sequence, sig, args)
     return res
 
 # @lower_builtin(StringArray)

sdc/tests/test_dataframe.py

@@ -33,6 +33,7 @@
 import unittest
 from itertools import permutations, product
 from numba.config import IS_32BITS
+from numba.special import literal_unroll
 
 import sdc
 from sdc.tests.gen_test_data import ParquetGenerator
@@ -1726,6 +1727,47 @@ def test_impl():
         self.assertTrue(isinstance(two, np.ndarray))
         self.assertTrue(isinstance(three, np.ndarray))
 
+    def test_df_iterate_over_columns1(self):
+        """ Verifies iteration over df columns using literal tuple of column indices. """
+        from sdc.hiframes.pd_dataframe_ext import get_dataframe_data
+        from sdc.hiframes.api import get_nan_mask
+
+        @self.jit
+        def jitted_func():
+            df = pd.DataFrame({
+                        'A': ['a', 'b', None, 'a', '', None, 'b'],
+                        'B': ['a', 'b', 'd', 'a', '', 'c', 'b'],
+                        'C': [np.nan, 1, 2, 1, np.nan, 2, 1],
+                        'D': [1, 2, 9, 5, 2, 1, 0]
+            })
+
+            # tuple of literals has to be created in a jitted function, otherwise
+            # col_id won't be literal and unboxing in get_dataframe_data won't compile
+            column_ids = (0, 1, 2, 3)
+            res_nan_mask = np.zeros(len(df), dtype=np.bool_)
+            for col_id in literal_unroll(column_ids):
+                res_nan_mask += get_nan_mask(get_dataframe_data(df, col_id))
+            return res_nan_mask
+
+        # expected is a boolean mask of df rows that have None values
+        expected = np.asarray([True, False, True, False, True, True, False])
+        result = jitted_func()
+        np.testing.assert_array_equal(result, expected)
+
+    def test_df_create_str_with_none(self):
+        """ Verifies creation of a dataframe with a string column from a list of Optional values. """
+        def test_impl():
+            df = pd.DataFrame({
+                        'A': ['a', 'b', None, 'a', '', None, 'b'],
+                        'B': ['a', 'b', 'd', 'a', '', 'c', 'b'],
+                        'C': [np.nan, 1, 2, 1, np.nan, 2, 1]
+            })
+
+            return df['A'].isna()
+        hpat_func = self.jit(test_impl)
+
+        pd.testing.assert_series_equal(hpat_func(), test_impl())
+
 
 if __name__ == "__main__":
     unittest.main()

sdc/tests/test_series.py

@@ -6291,14 +6291,13 @@ def test_impl(A, i, value):
         for series_data in all_data:
             for series_index in indexes:
                 S = pd.Series(series_data, series_index, dtype=dtype)
-                for idx in idxs:
-                    for value in values:
-                        with self.subTest(series=S, idx=idx, value=value):
-                            S1 = S.copy(deep=True)
-                            S2 = S.copy(deep=True)
-                            hpat_func(S1, idx, value)
-                            test_impl(S2, idx, value)
-                            pd.testing.assert_series_equal(S1, S2)
+                for idx, value in product(idxs, values):
+                    with self.subTest(series=S, idx=idx, value=value):
+                        S1 = S.copy(deep=True)
+                        S2 = S.copy(deep=True)
+                        hpat_func(S1, idx, value)
+                        test_impl(S2, idx, value)
+                        pd.testing.assert_series_equal(S1, S2)
 
     @skip_sdc_jit('Not implemented in old-pipeline')
     @skip_numba_jit('Requires StringArray support of operator.eq')
@@ -6365,11 +6364,10 @@ def test_series_setitem_idx_str_series(self):
             integer Series with index of matching dtype and scalar and non scalar assigned values """
 
         n, k = 11, 4
-        np.random.seed(0)
         series_data = np.arange(n)
         series_index = gen_strlist(n, 2, 'abcd123 ')
 
-        idx = create_series_from_values(k, series_index)
+        idx = create_series_from_values(k, series_index, seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [-100,
                           np.array(assigned_values),
@@ -6382,11 +6380,10 @@ def test_series_setitem_idx_float_series(self):
             integer Series with index of matching dtype and scalar and non scalar assigned values """
 
         n, k = 11, 4
-        np.random.seed(0)
         series_data = np.arange(n)
         series_index = np.arange(n, dtype=np.float)
 
-        idx = create_series_from_values(k, series_index)
+        idx = create_series_from_values(k, series_index, seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [
                             -100,
@@ -6404,11 +6401,10 @@ def test_impl(A, i, value):
         hpat_func = self.jit(test_impl)
 
         n, k = 11, 4
-        np.random.seed(0)
         series_data = np.arange(n)
         series_index = np.arange(n)
 
-        idx = create_series_from_values(k, series_index)
+        idx = create_series_from_values(k, series_index, seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [-100,
                           np.array(assigned_values),
@@ -6421,11 +6417,10 @@ def test_series_setitem_idx_int_series2(self):
             integer Series with index of non-matching dtype and scalar and non scalar assigned values """
 
         n, k = 11, 4
-        np.random.seed(0)
         series_data = np.arange(n)
         series_index = gen_strlist(n, 2, 'abcd123 ')
 
-        idx = create_series_from_values(k, np.arange(n))
+        idx = create_series_from_values(k, np.arange(n), seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [-100,
                           np.array(assigned_values),
@@ -6476,12 +6471,11 @@ def test_series_setitem_idx_int_array1(self):
             integer Series with integer index and scalar and non scalar assigned values """
 
         n, k = 11, 4
-        np.random.seed(0)
-
         series_data = np.arange(n)
         series_index = np.arange(n)
 
-        idx = take_k_elements(k, series_index)
+        np.random.seed(0)
+        idx = take_k_elements(k, series_index, seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [
                             -100,
@@ -6496,11 +6490,10 @@ def test_series_setitem_idx_int_array2(self):
             integer Series with string index and scalar and non scalar assigned values """
 
         n, k = 11, 4
-        np.random.seed(0)
         series_data = np.arange(n)
         series_index = gen_strlist(n, 2, 'abcd123 ')
 
-        idx = take_k_elements(k, np.arange(n))
+        idx = take_k_elements(k, np.arange(n), seed=0)
         assigned_values = -10 + np.arange(k) * (-1)
         values_to_test = [
                             -100,

sdc/tests/test_utils.py

@@ -212,28 +212,22 @@ def skip_inline(msg_or_func):
     return wrapper(func) if func else wrapper
 
 
-def take_k_elements(k, data):
-    random_idx = np.arange(len(data))
-    np.random.shuffle(random_idx)
-    return np.asarray(data).take(random_idx[:k])
+def take_k_elements(k, data, repeat=False, seed=None):
+    if seed is not None:
+        np.random.seed(seed)
+    return np.random.choice(np.asarray(data), k, replace=repeat)
 
 
-def create_series_from_values(size, data_values, index_values=None, name=None, unique=True):
+def create_series_from_values(size, data_values, index_values=None, name=None, unique=True, seed=None):
+    if seed is not None:
+        np.random.seed(seed)
 
     min_size = min(size, len(data_values))
     if index_values:
         min_size = min(min_size, len(index_values))
+    repeat = False if unique and min_size == size else True
 
-    if unique and min_size == size:
-        series_data = take_k_elements(size, data_values)
-        series_index = take_k_elements(size, index_values) if index_values else None
-    else:
-        data_values_pos = np.random.randint(0, len(data_values), size)
-        series_data = np.asarray(data_values).take(data_values_pos)
-        if index_values:
-            index_values_pos = np.random.randint(0, len(index_values), size)
-            series_index = np.asarray(index_values).take(index_values_pos)
-        else:
-            series_index = None
+    series_data = take_k_elements(size, data_values, repeat)
+    series_index = take_k_elements(size, index_values, repeat) if index_values else None
 
     return pandas.Series(series_data, series_index, name)