Merge branch 'master' into feature/ci-win

hpat/hiframes/api.py

@@ -1,33 +1,35 @@
 import itertools
-from hpat.utils import _numba_to_c_type_map, unliteral_all
-from numba.targets.arrayobj import make_array
-import llvmlite.binding as ll
-from llvmlite import ir as lir
-import llvmlite.llvmpy.core as lc
 import operator
 from collections import namedtuple
 import pandas as pd
 import numpy as np
+import llvmlite.binding as ll
+from llvmlite import ir as lir
+import llvmlite.llvmpy.core as lc
 
 import numba
 from numba import ir, ir_utils
-from numba.ir_utils import require, mk_unique_var
 from numba import types, cgutils
+from numba.ir_utils import require, mk_unique_var
 import numba.array_analysis
 from numba.typing import signature
 from numba.typing.templates import infer_global, AbstractTemplate, CallableTemplate
 from numba.typing.arraydecl import _expand_integer
-from numba.extending import overload, intrinsic
-from numba.targets.imputils import (impl_ret_new_ref, impl_ret_borrowed, iternext_impl, RefType)
-from numba.targets.arrayobj import _getitem_array1d
-from numba.extending import register_model, models
+from numba.extending import overload, intrinsic, register_model, models
+from numba.targets.imputils import (
+    lower_builtin,
+    impl_ret_untracked,
+    impl_ret_new_ref,
+    impl_ret_borrowed,
+    iternext_impl,
+    RefType)
+from numba.targets.arrayobj import make_array, _getitem_array1d
 
 import hpat
+from hpat.utils import _numba_to_c_type_map, unliteral_all
 from hpat.str_ext import string_type, list_string_array_type
-from hpat.str_arr_ext import (StringArrayType, string_array_type, is_str_arr_typ)
-
 from hpat.set_ext import build_set
-from numba.targets.imputils import lower_builtin, impl_ret_untracked
+from hpat.str_arr_ext import (StringArrayType, string_array_type, is_str_arr_typ)
 from hpat.hiframes.pd_timestamp_ext import (pandas_timestamp_type, datetime_date_type, set_df_datetime_date_lower)
 from hpat.hiframes.pd_series_ext import (
     SeriesType,
@@ -1599,8 +1601,8 @@ class DataFrameTupleIteratorModel(models.StructModel):
     def __init__(self, dmm, fe_type):
         # We use an unsigned index to avoid the cost of negative index tests.
         # XXX array_types[0] is implicit index
-        members = ([('index', types.EphemeralPointer(types.uintp))] + [('array{}'.format(i), arr)
-                                                                       for i, arr in enumerate(fe_type.array_types[1:])])
+        members = ([('index', types.EphemeralPointer(types.uintp))]
+                   + [('array{}'.format(i), arr) for i, arr in enumerate(fe_type.array_types[1:])])
         super(DataFrameTupleIteratorModel, self).__init__(dmm, fe_type, members)
 
     def from_return(self, builder, value):

hpat/hiframes/pd_dataframe_ext.py

@@ -468,8 +468,7 @@ class GetItemTuple(AbstractTemplate):
 
     def generic(self, args, kws):
         tup, idx = args
-        if (not isinstance(tup, types.BaseTuple) or
-                not isinstance(idx, types.IntegerLiteral)):
+        if (not isinstance(tup, types.BaseTuple) or not isinstance(idx, types.IntegerLiteral)):
             return
         idx_val = idx.literal_value
         if isinstance(idx_val, int):
@@ -497,8 +496,7 @@ def getitem_tuple_lower(context, builder, sig, args):
         items = cgutils.unpack_tuple(builder, tup)[idx]
         res = context.make_tuple(builder, sig.return_type, items)
     else:
-        raise NotImplementedError("unexpected index %r for %s"
-                                  % (idx, sig.args[0]))
+        raise NotImplementedError("unexpected index %r for %s" % (idx, sig.args[0]))
     return impl_ret_borrowed(context, builder, sig.return_type, res)
 
 

hpat/hiframes/pd_series_ext.py

@@ -246,6 +246,9 @@ def resolve_values(self, ary):
     def resolve_shape(self, ary):
         return types.Tuple((types.int64,))
 
+    def resolve_index(self, ary):
+        return ary.index
+
     def resolve_str(self, ary):
         assert ary.dtype in (string_type, types.List(string_type))
         # TODO: add dtype to series_str_methods_type

hpat/hiframes/rolling.py

@@ -156,7 +156,9 @@ def func(a, w, c, p):
 @lower_builtin(rolling_fixed, types.Array, types.Integer, types.Boolean,
                types.Boolean, types.functions.Dispatcher)
 def lower_rolling_fixed_apply(context, builder, sig, args):
-    def func(a, w, c, p, f): return roll_fixed_apply(a, w, c, p, f)
+
+    def func(a, w, c, p, f):
+        return roll_fixed_apply(a, w, c, p, f)
     res = context.compile_internal(builder, func, sig, args)
     return impl_ret_borrowed(context, builder, sig.return_type, res)
 
@@ -200,14 +202,22 @@ def func(a, o, w, c, p,):
     return impl_ret_borrowed(context, builder, sig.return_type, res)
 
 
-@lower_builtin(rolling_variable, types.Array, types.Array, types.Integer, types.Boolean,
-               types.Boolean, types.functions.Dispatcher)
+@lower_builtin(
+    rolling_variable,
+    types.Array,
+    types.Array,
+    types.Integer,
+    types.Boolean,
+    types.Boolean,
+    types.functions.Dispatcher)
 def lower_rolling_variable_apply(context, builder, sig, args):
-    def func(a, o, w, c, p, f): return roll_variable_apply(a, o, w, c, p, f)
+
+    def func(a, o, w, c, p, f):
+        return roll_variable_apply(a, o, w, c, p, f)
     res = context.compile_internal(builder, func, sig, args)
     return impl_ret_borrowed(context, builder, sig.return_type, res)
 
-#### adapted from pandas window.pyx ####
+# ** adapted from pandas window.pyx ****
 
 
 comm_border_tag = 22  # arbitrary, TODO: revisit comm tags

hpat/hiframes/sort.py

@@ -309,8 +309,8 @@ def sort_distributed_run(sort_node, array_dists, typemap, calltypes, typingctx,
     # TODO: use *args
     func_text = "def f({}, {}):\n".format(key_name_args, col_name_args)
     func_text += "  key_arrs = ({},)\n".format(key_name_args)
-    func_text += "  data = ({}{})\n".format(col_name_args, "," if len(in_vars) ==
-                                            1 else "")  # single value needs comma to become tuple
+    # single value needs comma to become tuple
+    func_text += "  data = ({}{})\n".format(col_name_args, "," if len(in_vars) == 1 else "")
     func_text += "  hpat.hiframes.sort.local_sort(key_arrs, data, {})\n".format(sort_node.ascending)
     func_text += "  return key_arrs, data\n"
 
@@ -349,8 +349,7 @@ def sort_distributed_run(sort_node, array_dists, typemap, calltypes, typingctx,
 
     ascending_var = ir.Var(scope, mk_unique_var('ascending'), loc)
     typemap[ascending_var.name] = types.bool_
-    nodes.append(
-        ir.Assign(ir.Const(sort_node.ascending, loc), ascending_var, loc))
+    nodes.append(ir.Assign(ir.Const(sort_node.ascending, loc), ascending_var, loc))
 
     # parallel case
     def par_sort_impl(key_arrs, data, ascending):

hpat/tests/test_series.py

@@ -1429,6 +1429,31 @@ def test_impl(S):
         S = pd.Series([np.nan, 2., 3., 5., np.nan, 6., 7.])
         pd.testing.assert_series_equal(hpat_func(S), test_impl(S))
 
+    def test_series_index1(self):
+        def test_impl():
+            A = pd.Series([1, 2, 3], index=['A', 'C', 'B'])
+            return A.index
+
+        hpat_func = hpat.jit(test_impl)
+        np.testing.assert_array_equal(hpat_func(), test_impl())
+
+    def test_series_index2(self):
+        def test_impl():
+            A = pd.Series([1, 2, 3], index=[0, 1, 2])
+            return A.index
+
+        hpat_func = hpat.jit(test_impl)
+        np.testing.assert_array_equal(hpat_func(), test_impl())
+
+    @unittest.skip("Enabel after fixing distributed for get_series_index")
+    def test_series_index3(self):
+        def test_impl():
+            A = pd.Series([1, 2, 3])
+            return A.index
+
+        hpat_func = hpat.jit(test_impl)
+        np.testing.assert_array_equal(hpat_func(), test_impl())
+
 
 if __name__ == "__main__":
     unittest.main()