fix iloc with global column index

Author: Ehsan Totoni

hpat/hiframes/dataframe_pass.py

@@ -264,7 +264,12 @@ def _run_getitem(self, assign, rhs):
             assert isinstance(ind_def, ir.Expr) and ind_def.op == 'build_tuple'
 
             if self._is_df_iloc_var(rhs.value):
-                col_ind = guard(find_const, self.func_ir, ind_def.items[1])
+                # find_const doesn't support globals so use static value
+                # TODO: fix find_const()
+                if rhs.op == 'static_getitem':
+                    col_ind = rhs.index[1]
+                else:
+                    col_ind = guard(find_const, self.func_ir, ind_def.items[1])
                 col_name = df_typ.columns[col_ind]
             else:  # df.loc
                 col_name = guard(find_const, self.func_ir, ind_def.items[1])

hpat/tests/test_dataframe.py

@@ -16,6 +16,8 @@ def inner_get_column(df):
     # df2['D'] = np.ones(3)
     return df.A
 
+COL_IND = 0
+
 class TestDataFrame(unittest.TestCase):
     def test_create1(self):
         def test_impl(n):
@@ -229,6 +231,16 @@ def test_impl(df, n):
         df = pd.DataFrame({'A': np.arange(n), 'B': np.arange(n)**2})
         np.testing.assert_array_equal(hpat_func(df, n), test_impl(df, n))
 
+    def test_iloc5(self):
+        # test iloc with global value
+        def test_impl(df):
+            return df.iloc[:,COL_IND].values
+
+        hpat_func = hpat.jit(test_impl)
+        n = 11
+        df = pd.DataFrame({'A': np.arange(n), 'B': np.arange(n)**2})
+        np.testing.assert_array_equal(hpat_func(df), test_impl(df))
+
     def test_loc1(self):
         def test_impl(df):
             return df.loc[:,'B'].values