fix split view in box df, isna, distributed, etc

Author: Ehsan Totoni

hpat/distributed.py

@@ -584,6 +584,13 @@ def f(arr, bag, start, count):  # pragma: no cover
             self._array_counts[lhs] = self._array_counts[in_arr]
             self._array_sizes[lhs] = self._array_sizes[in_arr]
 
+        if (fdef == ('compute_split_view', 'hpat.hiframes.split_impl')
+                and self._is_1D_arr(rhs.args[0].name)):
+            in_arr = rhs.args[0].name
+            self._array_starts[lhs] = self._array_starts[in_arr]
+            self._array_counts[lhs] = self._array_counts[in_arr]
+            self._array_sizes[lhs] = self._array_sizes[in_arr]
+
         if fdef == ('isna', 'hpat.hiframes.api') and self._is_1D_arr(rhs.args[0].name):
             # fix index in call to isna
             arr = rhs.args[0]

hpat/distributed_analysis.py

@@ -379,6 +379,10 @@ def _analyze_call(self, lhs, rhs, func_var, args, array_dists):
             self._meet_array_dists(lhs, rhs.args[0].name, array_dists)
             return
 
+        if fdef == ('compute_split_view', 'hpat.hiframes.split_impl'):
+            self._meet_array_dists(lhs, rhs.args[0].name, array_dists)
+            return
+
         # np.fromfile()
         if fdef == ('file_read', 'hpat.io.np_io'):
             return

hpat/hiframes/api.py

@@ -36,6 +36,7 @@
     update_shuffle_meta,  alloc_pre_shuffle_metadata,
     )
 from hpat.hiframes.join import write_send_buff
+from hpat.hiframes.split_impl import string_array_split_view_type
 
 # XXX: used in agg func output to avoid mutating filter, agg, join, etc.
 # TODO: fix type inferrer and remove this
@@ -446,6 +447,8 @@ def isna_overload(arr, i):
     # TODO: support NaN in list(list(str))
     if arr == list_string_array_type:
         return lambda arr, i: False
+    if arr == string_array_split_view_type:
+        return lambda arr, i: False
     # TODO: extend to other types
     assert isinstance(arr, types.Array)
     dtype = arr.dtype

hpat/hiframes/boxing.py

@@ -189,6 +189,8 @@ def box_dataframe(typ, val, c):
                 elif isinstance(dtype, PDCategoricalDtype):
                     arr_obj = box_categorical_array(arr_typ, arr, c)
                     # context.nrt.incref(builder, arr_typ, arr)
+                elif arr_typ == string_array_split_view_type:
+                    arr_obj = box_str_arr_split_view(arr_typ, arr, c)
                 elif dtype == types.List(string_type):
                     arr_obj = box_list(list_string_array_type, arr, c)
                     # context.nrt.incref(builder, arr_typ, arr)  # TODO required?

hpat/hiframes/hiframes_typed.py

@@ -36,7 +36,8 @@
 from hpat.hiframes.aggregate import Aggregate
 from hpat.hiframes import series_kernels, split_impl
 from hpat.hiframes.series_kernels import series_replace_funcs
-from hpat.hiframes.split_impl import string_array_split_view_type
+from hpat.hiframes.split_impl import (string_array_split_view_type,
+    StringArraySplitViewType)
 
 
 _dt_index_binops = ('==', '!=', '>=', '>', '<=', '<', '-',
@@ -2399,7 +2400,7 @@ def _get_arg(self, f_name, args, kws, arg_no, arg_name, default=None,
 def _fix_typ_undefs(new_typ, old_typ):
     if isinstance(old_typ, (types.Array, SeriesType)):
         assert isinstance(new_typ, (types.Array, SeriesType, StringArrayType,
-            types.List))
+            types.List, StringArraySplitViewType))
         if new_typ.dtype == types.undefined:
             return new_typ.copy(old_typ.dtype)
     if isinstance(old_typ, (types.Tuple, types.UniTuple)):

hpat/hiframes/pd_series_ext.py

@@ -157,7 +157,8 @@ def __init__(self, dmm, fe_type):
 
 
 def series_to_array_type(typ, replace_boxed=False):
-    return _get_series_array_type(typ.dtype)
+    return typ.data
+    # return _get_series_array_type(typ.dtype)
 
 
 def is_series_type(typ):

hpat/utils.py

@@ -289,7 +289,8 @@ def get_slice_step(typemap, func_ir, var):
 def is_array(typemap, varname):
     return (varname in typemap
         and (is_np_array(typemap, varname)
-        or typemap[varname] in (string_array_type, list_string_array_type)
+        or typemap[varname] in (string_array_type, list_string_array_type,
+            hpat.hiframes.split_impl.string_array_split_view_type)
         or isinstance(typemap[varname], hpat.hiframes.pd_series_ext.SeriesType)))
 
 def is_np_array(typemap, varname):