Rewrite Series Typing using Array

hpat/compiler.py

@@ -148,8 +148,11 @@ def define_pipelines(self, pm):
         # e.g. need to handle string array exprs before nopython rewrites
         # converts them to arrayexpr.
         # self.add_optimization_stage(pm)
-        pm.add_stage(self.stage_pre_parfor_pass, "Preprocessing for parfors")
+        # hiframes typed pass should be before pre_parfor since variable types
+        # need updating, and A.call to np.call transformation is invalid for
+        # Series (e.g. S.var is not the same as np.var(S))
         pm.add_stage(self.stage_df_typed_pass, "typed hiframes pass")
+        pm.add_stage(self.stage_pre_parfor_pass, "Preprocessing for parfors")
         if not self.flags.no_rewrites:
             pm.add_stage(self.stage_nopython_rewrites, "nopython rewrites")
         if self.flags.auto_parallel.enabled:

hpat/hiframes.py

@@ -1881,20 +1881,13 @@ def _get_renamed_df(self, df_var):
 
 
 def gen_empty_like(in_arr, out_arr):
-    scope = in_arr.scope
-    loc = in_arr.loc
-    # g_np_var = Global(numpy)
-    g_np_var = ir.Var(scope, mk_unique_var("$np_g_var"), loc)
-    g_np = ir.Global('np', np, loc)
-    g_np_assign = ir.Assign(g_np, g_np_var, loc)
-    # attr call: empty_attr = getattr(g_np_var, empty_like)
-    empty_attr_call = ir.Expr.getattr(g_np_var, "empty_like", loc)
-    attr_var = ir.Var(scope, mk_unique_var("$empty_attr_attr"), loc)
-    attr_assign = ir.Assign(empty_attr_call, attr_var, loc)
-    # alloc call: out_arr = empty_attr(in_arr)
-    alloc_call = ir.Expr.call(attr_var, [in_arr], (), loc)
-    alloc_assign = ir.Assign(alloc_call, out_arr, loc)
-    return [g_np_assign, attr_assign, alloc_assign]
+    def f(A):  # pragma: no cover
+        B = np.empty(A.shape, A.dtype)
+    f_block = compile_to_numba_ir(f, {'hpat': hpat, 'np': np}).blocks.popitem()[1]
+    replace_arg_nodes(f_block, [in_arr])
+    nodes = f_block.body[:-3]  # remove none return
+    nodes[-1].target = out_arr
+    return nodes
 
 
 def gen_stencil_call(in_arr, out_arr, kernel_func, index_offsets, fir_globals,
@@ -1918,8 +1911,10 @@ def gen_stencil_call(in_arr, out_arr, kernel_func, index_offsets, fir_globals,
     stencil_nodes.append(ir.Assign(kernel_func, kernel_var, loc))
 
     def f(A, B, f):  # pragma: no cover
-        numba.stencil(f)(A, out=B)
-    f_block = compile_to_numba_ir(f, {'numba': numba}).blocks.popitem()[1]
+        in_arr = hpat.hiframes_api.to_arr_from_series(A)
+        numba.stencil(f)(in_arr, out=B)
+    f_block = compile_to_numba_ir(f, {'numba': numba,
+        'hpat': hpat}).blocks.popitem()[1]
     replace_arg_nodes(f_block, [in_arr, out_arr, kernel_var])
     stencil_nodes += f_block.body[:-3]  # remove none return
     setup_call = stencil_nodes[-2].value

hpat/hiframes_api.py

@@ -22,7 +22,9 @@
 import numpy as np
 from hpat.pd_timestamp_ext import timestamp_series_type, pandas_timestamp_type
 import hpat
-from hpat.pd_series_ext import SeriesType, BoxedSeriesType, string_series_type, arr_to_series_type, arr_to_boxed_series_type, series_to_array_type
+from hpat.pd_series_ext import (SeriesType, BoxedSeriesType,
+    string_series_type, if_arr_to_series_type, arr_to_boxed_series_type,
+    series_to_array_type, if_series_to_array_type)
 
 # from numba.typing.templates import infer_getattr, AttributeTemplate, bound_function
 # from numba import types
@@ -91,7 +93,8 @@ def generic(self, args, kws):
             ret_typ = string_array_type
         else:
             # use typer of np.concatenate
-            ret_typ = numba.typing.npydecl.NdConcatenate(self.context).generic()(arr_list)
+            arr_list_to_arr = if_series_to_array_type(arr_list)
+            ret_typ = numba.typing.npydecl.NdConcatenate(self.context).generic()(arr_list_to_arr)
 
         return signature(ret_typ, arr_list)
 
@@ -656,27 +659,29 @@ def generic(self, args, kws):
         if isinstance(arr, BoxedSeriesType):
             series_type = SeriesType(arr.dtype, 1, 'C')
         else:
-            series_type = arr_to_series_type(arr)
+            series_type = if_arr_to_series_type(arr)
         assert series_type is not None, "unknown type for pd.Series: {}".format(arr)
         return signature(series_type, arr)
 
 @lower_builtin(to_series_type, types.Any)
 def to_series_dummy_impl(context, builder, sig, args):
     return impl_ret_borrowed(context, builder, sig.return_type, args[0])
 
-def if_series_to_array_type(typ, replace_boxed=False):
-    if isinstance(typ, SeriesType):
-        return series_to_array_type(typ, replace_boxed)
-    # XXX: Boxed series variable types shouldn't be replaced in hiframes_typed
-    # it results in cast error for call dummy_unbox_series
-    if replace_boxed and isinstance(typ, BoxedSeriesType):
-        return series_to_array_type(typ, replace_boxed)
-    if isinstance(typ, (types.Tuple, types.UniTuple)):
-        return types.Tuple(
-            [if_series_to_array_type(t, replace_boxed) for t in typ.types])
-    # TODO: other types than can have Series inside: list, set, etc.
-    return typ
 
+def to_arr_from_series(arr):
+    return arr
+
+@infer_global(to_arr_from_series)
+class ToArrFromSeriesType(AbstractTemplate):
+    def generic(self, args, kws):
+        assert not kws
+        assert len(args) == 1
+        arr = args[0]
+        return signature(if_series_to_array_type(arr), arr)
+
+@lower_builtin(to_arr_from_series, types.Any)
+def to_arr_from_series_dummy_impl(context, builder, sig, args):
+    return impl_ret_borrowed(context, builder, sig.return_type, args[0])
 
 # dummy func to convert input series to array type
 def dummy_unbox_series(arr):

hpat/hiframes_typed.py

@@ -8,15 +8,16 @@
                             find_topo_order, gen_np_call, get_definition, guard,
                             find_callname, mk_alloc, find_const, is_setitem,
                             is_getitem)
-from numba.typing.templates import Signature
+from numba.typing.templates import Signature, bound_function, signature
+from numba.typing.arraydecl import ArrayAttribute
 import hpat
 from hpat.utils import get_definitions, debug_prints
 from hpat.hiframes import include_new_blocks, gen_empty_like
-from hpat.hiframes_api import if_series_to_array_type
 from hpat.str_ext import string_type
 from hpat.str_arr_ext import string_array_type, StringArrayType, is_str_arr_typ
 from hpat.pd_series_ext import (SeriesType, string_series_type,
-    series_to_array_type, BoxedSeriesType, dt_index_series_type)
+    series_to_array_type, BoxedSeriesType, dt_index_series_type,
+    if_series_to_array_type)
 
 
 class HiFramesTyped(object):
@@ -54,13 +55,28 @@ def run(self):
             blocks[label].body = new_body
 
         if debug_prints():  # pragma: no cover
-            print("types before Series replacement:", self.typemap)
+            print("--- types before Series replacement:", self.typemap)
+            print("calltypes: ", self.calltypes)
+
         replace_series = {}
         for vname, typ in self.typemap.items():
             if isinstance(typ, SeriesType):
                 # print("replacing series type", vname)
                 new_typ = series_to_array_type(typ)
                 replace_series[vname] = new_typ
+            # replace array.call() variable types
+            if isinstance(typ, types.BoundFunction) and isinstance(typ.this, SeriesType):
+                this = series_to_array_type(typ.this)
+                # TODO: handle string arrays, etc.
+                assert typ.typing_key.startswith('array.')
+                attr = typ.typing_key[len('array.'):]
+                resolver = getattr(ArrayAttribute, 'resolve_'+attr)
+                # methods are either installed with install_array_method or
+                # using @bound_function in arraydecl.py
+                if hasattr(resolver, '__wrapped__'):
+                    resolver = bound_function(typ.typing_key)(resolver.__wrapped__)
+                new_typ = resolver(ArrayAttribute(self.typingctx), this)
+                replace_series[vname] = new_typ
 
         for vname, typ in replace_series.items():
             self.typemap.pop(vname)
@@ -79,7 +95,22 @@ def run(self):
                 # reusing sig.args since some types become Const in sig
                 argtyps = sig.args[:len(call.args)]
                 kwtyps = {name: self.typemap[v.name] for name, v in call.kws}
-                self.typemap[call.func.name].get_call_type(self.typingctx , argtyps, kwtyps)
+
+                new_sig = self.typemap[call.func.name].get_call_type(
+                    self.typingctx , argtyps, kwtyps)
+                # calltypes of things like BoundFunction (array.call) need to
+                # be update for lowering to work
+                # XXX: new_sig could be None for things like np.int32()
+                if call in self.calltypes and new_sig is not None:
+                    old_sig = self.calltypes.pop(call)
+                    # fix types with undefined dtypes in empty_inferred, etc.
+                    return_type = _fix_typ_undefs(new_sig.return_type, old_sig.return_type)
+                    args = tuple(_fix_typ_undefs(a, b) for a,b  in zip(new_sig.args, old_sig.args))
+                    self.calltypes[call] = Signature(return_type, args, new_sig.recvr, new_sig.pysig)
+
+        if debug_prints():  # pragma: no cover
+            print("--- types after Series replacement:", self.typemap)
+            print("calltypes: ", self.calltypes)
 
         self.func_ir._definitions = get_definitions(self.func_ir.blocks)
         return if_series_to_array_type(self.return_type)
@@ -130,7 +161,7 @@ def _run_assign(self, assign):
         return [assign]
 
     def _run_call_hiframes(self, assign, lhs, rhs, func_name):
-        if func_name == 'to_series_type':
+        if func_name in ('to_series_type', 'to_arr_from_series'):
             assign.value = rhs.args[0]
             return [assign]
 
@@ -155,7 +186,7 @@ def f(column):  # pragma: no cover
                     a = column.astype(np.float64)
                 f_block = compile_to_numba_ir(f,
                                               {'hpat': hpat, 'np': np}, self.typingctx,
-                                              (self.typemap[in_arr.name],),
+                                              (if_series_to_array_type(self.typemap[in_arr.name]),),
                                               self.typemap, self.calltypes).blocks.popitem()[1]
                 replace_arg_nodes(f_block, [in_arr])
                 nodes = f_block.body[:-3]
@@ -204,8 +235,8 @@ def _handle_dt_index_binop(self, lhs, rhs, assign):
         f_blocks = compile_to_numba_ir(f,
                                         {'numba': numba, 'np': np, 'hpat': hpat},
                                         self.typingctx,
-                                        (self.typemap[arg1.name],
-                                        self.typemap[arg2.name]),
+                                        (if_series_to_array_type(self.typemap[arg1.name]),
+                                        if_series_to_array_type(self.typemap[arg2.name])),
                                         self.typemap, self.calltypes).blocks
         replace_arg_nodes(f_blocks[min(f_blocks.keys())], [arg1, arg2])
         # replace == expression with result of parfor (S)
@@ -248,8 +279,8 @@ def _handle_string_array_expr(self, lhs, rhs, assign):
             f = loc_vars['f']
             f_blocks = compile_to_numba_ir(f,
                                            {'numba': numba, 'np': np}, self.typingctx,
-                                           (self.typemap[arg1.name],
-                                            self.typemap[arg2.name]),
+                                           (if_series_to_array_type(self.typemap[arg1.name]),
+                                            if_series_to_array_type(self.typemap[arg2.name])),
                                            self.typemap, self.calltypes).blocks
             replace_arg_nodes(f_blocks[min(f_blocks.keys())], [arg1, arg2])
             # replace == expression with result of parfor (S)
@@ -273,7 +304,7 @@ def f(_in_arr):  # pragma: no cover
                 _alloc_size = _in_arr.shape
                 _out_arr = np.empty(_alloc_size, _in_arr.dtype)
 
-        f_block = compile_to_numba_ir(f, {'np': np}, self.typingctx, (self.typemap[in_arr.name],),
+        f_block = compile_to_numba_ir(f, {'np': np}, self.typingctx, (if_series_to_array_type(self.typemap[in_arr.name]),),
                                         self.typemap, self.calltypes).blocks.popitem()[1]
         replace_arg_nodes(f_block, [in_arr])
         nodes = f_block.body[:-3]  # remove none return
@@ -302,8 +333,8 @@ def _handle_str_contains(self, assign, lhs, rhs, fname):
         f_blocks = compile_to_numba_ir(f,
                                        {'numba': numba, 'np': np,
                                            'hpat': hpat}, self.typingctx,
-                                       (self.typemap[str_arr.name],
-                                        self.typemap[pat.name]),
+                                       (if_series_to_array_type(self.typemap[str_arr.name]),
+                                        if_series_to_array_type(self.typemap[pat.name])),
                                        self.typemap, self.calltypes).blocks
         replace_arg_nodes(f_blocks[min(f_blocks.keys())], [str_arr, pat])
         # replace call with result of parfor (S)
@@ -324,7 +355,7 @@ def _handle_df_col_filter(self, lhs_name, rhs, assign):
             index_var = rhs.index
             f_blocks = compile_to_numba_ir(_column_filter_impl_float,
                                            {'numba': numba, 'np': np}, self.typingctx,
-                                           (self.typemap[lhs.name], self.typemap[in_arr.name],
+                                           (if_series_to_array_type(self.typemap[lhs.name]), if_series_to_array_type(self.typemap[in_arr.name]),
                                                self.typemap[index_var.name]),
                                            self.typemap, self.calltypes).blocks
             first_block = min(f_blocks.keys())
@@ -347,7 +378,7 @@ def f(_in_arr, _ind):
 
             f_block = compile_to_numba_ir(f, {'numba': numba, 'np': np,
                                                'hpat': hpat}, self.typingctx,
-                                           (self.typemap[in_arr.name], types.intp),
+                                           (if_series_to_array_type(self.typemap[in_arr.name]), types.intp),
                                            self.typemap, self.calltypes).blocks.popitem()[1]
             replace_arg_nodes(f_block, [in_arr, ind])
             nodes = f_block.body[:-3]  # remove none return
@@ -359,7 +390,7 @@ def f(_in_arr, _ind):
             f_blocks = compile_to_numba_ir(_column_count_impl,
                                            {'numba': numba, 'np': np,
                                                'hpat': hpat}, self.typingctx,
-                                           (self.typemap[in_arr.name],),
+                                           (if_series_to_array_type(self.typemap[in_arr.name]),),
                                            self.typemap, self.calltypes).blocks
             topo_order = find_topo_order(f_blocks)
             first_block = topo_order[0]
@@ -375,8 +406,8 @@ def f(_in_arr, _ind):
             val = rhs.args[2]
             f_blocks = compile_to_numba_ir(_column_fillna_impl,
                                            {'numba': numba, 'np': np}, self.typingctx,
-                                           (self.typemap[out_arr.name], self.typemap[in_arr.name],
-                                               self.typemap[val.name]),
+                                           (if_series_to_array_type(self.typemap[out_arr.name]), if_series_to_array_type(self.typemap[in_arr.name]),
+                                               if_series_to_array_type(self.typemap[val.name])),
                                            self.typemap, self.calltypes).blocks
             first_block = min(f_blocks.keys())
             replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val])
@@ -387,7 +418,7 @@ def f(_in_arr, _ind):
             f_blocks = compile_to_numba_ir(_column_sum_impl,
                                            {'numba': numba, 'np': np,
                                                'hpat': hpat}, self.typingctx,
-                                           (self.typemap[in_arr.name],),
+                                           (if_series_to_array_type(self.typemap[in_arr.name]),),
                                            self.typemap, self.calltypes).blocks
             topo_order = find_topo_order(f_blocks)
             first_block = topo_order[0]
@@ -402,7 +433,7 @@ def f(_in_arr, _ind):
             f_blocks = compile_to_numba_ir(_column_mean_impl,
                                            {'numba': numba, 'np': np,
                                                'hpat': hpat}, self.typingctx,
-                                           (self.typemap[in_arr.name],),
+                                           (if_series_to_array_type(self.typemap[in_arr.name]),),
                                            self.typemap, self.calltypes).blocks
             topo_order = find_topo_order(f_blocks)
             first_block = topo_order[0]
@@ -417,7 +448,7 @@ def f(_in_arr, _ind):
             f_blocks = compile_to_numba_ir(_column_var_impl,
                                            {'numba': numba, 'np': np,
                                                'hpat': hpat}, self.typingctx,
-                                           (self.typemap[in_arr.name],),
+                                           (if_series_to_array_type(self.typemap[in_arr.name]),),
                                            self.typemap, self.calltypes).blocks
             topo_order = find_topo_order(f_blocks)
             first_block = topo_order[0]
@@ -433,6 +464,18 @@ def is_bool_arr(self, varname):
         typ = self.typemap[varname]
         return isinstance(typ, types.npytypes.Array) and typ.dtype == types.bool_
 
+def _fix_typ_undefs(new_typ, old_typ):
+    if isinstance(old_typ, (types.Array, SeriesType)):
+        assert isinstance(new_typ, (types.Array, SeriesType))
+        if new_typ.dtype == types.undefined:
+            return new_typ.copy(old_typ.dtype)
+    if isinstance(old_typ, (types.Tuple, types.UniTuple)):
+        return types.Tuple([_fix_typ_undefs(t, u)
+                                for t, u in zip(new_typ.types, old_typ.types)])
+    # TODO: fix List, Set
+    return new_typ
+
+
 # float columns can have regular np.nan