support df loc/iloc with boolean array index

Author: Ehsan Totoni

hpat/hiframes.py

@@ -129,6 +129,13 @@ def _run_assign(self, assign):
                 return [hiframes_api.Filter(lhs, rhs.value.name, rhs.index,
                                                         self.df_vars, rhs.loc)]
 
+            # df.loc or df.iloc
+            if rhs.op=='getattr' and rhs.value.name in self.df_vars and rhs.attr in ['loc', 'iloc']:
+                # FIXME: treat iloc and loc as regular df variables so getitem
+                # turns them into filter. Only boolean array is supported
+                self.df_vars[lhs] = self.df_vars[rhs.value.name]
+                return []
+
             # if (rhs.op == 'getitem' and rhs.value.name in self.df_cols):
             #     self.col_filters.add(assign)
 

hpat/tests/test_hiframes.py

@@ -34,7 +34,7 @@ def test_impl(n):
         self.assertEqual(count_parfor_OneDs(), 2)
         self.assertTrue(dist_IR_contains('dist_cumsum'))
 
-    def test_filter(self):
+    def test_filter1(self):
         def test_impl(n):
             df = pd.DataFrame({'A': np.ones(n), 'B': np.ones(n)})
             df1 = df[df.A > .5]
@@ -46,6 +46,18 @@ def test_impl(n):
         self.assertEqual(count_array_REPs(), 0)
         self.assertEqual(count_parfor_REPs(), 0)
 
+    def test_filter2(self):
+        def test_impl(n):
+            df = pd.DataFrame({'A': np.ones(n), 'B': np.ones(n)})
+            df1 = df.loc[df.A > .5]
+            return np.sum(df1.B)
+
+        hpat_func = hpat.jit(test_impl)
+        n = 11
+        self.assertEqual(hpat_func(n), test_impl(n))
+        self.assertEqual(count_array_REPs(), 0)
+        self.assertEqual(count_parfor_REPs(), 0)
+
     def test_rolling1(self):
         def test_impl(n):
             df = pd.DataFrame({'A': np.ones(n), 'B': np.random.ranf(n)})