fix str arr dropna() kernel

Author: Ehsan Totoni

hpat/distributed_analysis.py

@@ -418,6 +418,18 @@ def _analyze_call(self, lhs, rhs, func_var, args, array_dists):
         if fdef == ('num_total_chars', 'hpat.str_arr_ext'):
             return
 
+        if fdef == ('_series_dropna_str_alloc_impl_inner', 'hpat.hiframes.series_kernels'):
+            if lhs not in array_dists:
+                array_dists[lhs] = Distribution.OneD_Var
+            in_dist = array_dists[rhs.args[0].name]
+            out_dist = array_dists[lhs]
+            out_dist = Distribution(min(out_dist.value, in_dist.value))
+            array_dists[lhs] = out_dist
+            # output can cause input REP
+            if out_dist != Distribution.OneD_Var:
+                array_dists[rhs.args[0].name] = out_dist
+            return
+
         if (fdef == ('copy_non_null_offsets', 'hpat.str_arr_ext')
                 or fdef == ('copy_data', 'hpat.str_arr_ext')):
             out_arrname = rhs.args[0].name

hpat/hiframes/series_kernels.py

@@ -70,19 +70,31 @@ def _series_dropna_float_impl(S, name):  # pragma: no cover
 
     return hpat.hiframes.api.init_series(A, None, name)
 
-def _series_dropna_str_alloc_impl(B, name):  # pragma: no cover
-    # local_len to enable 1D_Var dist
+
+# using njit since 1D_var is broken for alloc when there is calculation of len
+@numba.njit(no_cpython_wrapper=True)
+def _series_dropna_str_alloc_impl_inner(B):  # pragma: no cover
     # TODO: test
     # TODO: generalize
-    old_len = hpat.distributed_api.local_len(B)
+    old_len = len(B)
+    na_count = 0
+    for i in range(len(B)):
+        if hpat.str_arr_ext.str_arr_is_na(B, i):
+            na_count += 1
     # TODO: more efficient null counting
-    new_len = old_len - hpat.hiframes.api.init_series(B).isna().sum()
+    new_len = old_len - na_count
     num_chars = hpat.str_arr_ext.num_total_chars(B)
     A = hpat.str_arr_ext.pre_alloc_string_array(new_len, num_chars)
     hpat.str_arr_ext.copy_non_null_offsets(A, B)
     hpat.str_arr_ext.copy_data(A, B)
+    return A
+
+
+def _series_dropna_str_alloc_impl(B, name):  # pragma: no cover
+    A = hpat.hiframes.series_kernels._series_dropna_str_alloc_impl_inner(B)
     return hpat.hiframes.api.init_series(A, None, name)
 
+
 # return the nan value for the type (handle dt64)
 def _get_nan(val):
     return np.nan