Overload df.getitem with boolean series idx (#583)

Author: densmirn

sdc/datatypes/hpat_pandas_dataframe_functions.py

@@ -38,7 +38,10 @@
 from numba import types
 from numba.special import literally
 from sdc.hiframes.pd_dataframe_ext import DataFrameType
-from sdc.utilities.sdc_typing_utils import TypeChecker
+from sdc.hiframes.pd_series_type import SeriesType
+from sdc.utilities.sdc_typing_utils import (TypeChecker, check_index_is_numeric,
+                                            check_types_comparable,
+                                            gen_df_impl_generator)
 from sdc.str_arr_ext import StringArrayType
 
 from sdc.hiframes.pd_dataframe_type import DataFrameType
@@ -929,6 +932,32 @@ def sdc_pandas_dataframe_drop_impl(df, _func_name, args, columns):
     return sdc_pandas_dataframe_drop_impl(df, _func_name, args, columns)
 
 
+def df_getitem_bool_series_idx_main_codelines(self, idx):
+    """Generate main code lines for df.getitem"""
+    func_lines = ['  self_length = len(get_dataframe_data(self, 0))',
+                  '  trimmed_idx_data = idx._data[:self_length]']
+
+    if isinstance(self.index, types.NoneType):
+        func_lines += ['  self_index = numpy.arange(self_length)']
+    else:
+        func_lines += ['  self_index = self._index']
+
+    results = []
+    for i, col in enumerate(self.columns):
+        res_data = f'res_data_{i}'
+        func_lines += [
+            f'  data_{i} = get_dataframe_data(self, {i})',
+            f'  series = pandas.Series(data_{i}, index=self_index, name="{col}")',
+            f'  {res_data} = series[trimmed_idx_data]',
+        ]
+        results.append((col, res_data))
+
+    data = ', '.join(f'"{col}": {data}' for col, data in results)
+    func_lines += [f'  return pandas.DataFrame({{{data}}}, index=self_index[trimmed_idx_data])']
+
+    return func_lines
+
+
 def df_index_codelines(self):
     """Generate code lines to get or create index of DF"""
     if isinstance(self.index, types.NoneType):
@@ -941,6 +970,11 @@ def df_index_codelines(self):
     return func_lines
 
 
+def df_getitem_key_error_codelines():
+    """Generate code lines to raise KeyError"""
+    return ['  raise KeyError("Column is not in the DataFrame")']
+
+
 def df_getitem_slice_idx_main_codelines(self, idx):
     """Generate main code lines for df.getitem with idx of slice"""
     results = []
@@ -978,6 +1012,35 @@ def df_getitem_tuple_idx_main_codelines(self, literal_idx):
     return func_lines
 
 
+def df_getitem_bool_series_codegen(self, idx):
+    """
+    Example of generated implementation with provided index:
+        def _df_getitem_bool_series_idx_impl(self, idx):
+          self_length = len(get_dataframe_data(self, 0))
+          trimmed_idx_data = idx._data[:self_length]
+          self_index = self._index
+          data_0 = get_dataframe_data(self, 0)
+          series = pandas.Series(data_0, index=self_index, name="A")
+          res_data_0 = series[trimmed_idx_data]
+          data_1 = get_dataframe_data(self, 1)
+          series = pandas.Series(data_1, index=self_index, name="B")
+          res_data_1 = series[trimmed_idx_data]
+          return pandas.DataFrame({"A": res_data_0, "B": res_data_1}, index=self_index[trimmed_idx_data])
+    """
+    func_lines = ['def _df_getitem_bool_series_idx_impl(self, idx):']
+    if self.columns:
+        func_lines += df_getitem_bool_series_idx_main_codelines(self, idx)
+    else:
+        # raise KeyError if input DF is empty
+        func_lines += df_getitem_key_error_codelines()
+
+    func_text = '\n'.join(func_lines)
+    global_vars = {'pandas': pandas, 'numpy': numpy,
+                   'get_dataframe_data': get_dataframe_data}
+
+    return func_text, global_vars
+
+
 def df_getitem_slice_idx_codegen(self, idx):
     """
     Example of generated implementation with provided index:
@@ -994,7 +1057,7 @@ def _df_getitem_slice_idx_impl(self, idx)
         func_lines += df_getitem_slice_idx_main_codelines(self, idx)
     else:
         # raise KeyError if input DF is empty
-        func_lines += ['  raise KeyError']
+        func_lines += df_getitem_key_error_codelines()
 
     func_text = '\n'.join(func_lines)
     global_vars = {'pandas': pandas, 'numpy': numpy,
@@ -1022,7 +1085,7 @@ def _df_getitem_tuple_idx_impl(self, idx)
         func_lines += df_getitem_tuple_idx_main_codelines(self, literal_idx)
     else:
         # raise KeyError if input DF is empty or idx is invalid
-        func_lines += ['  raise KeyError']
+        func_lines += df_getitem_key_error_codelines()
 
     func_text = '\n'.join(func_lines)
     global_vars = {'pandas': pandas, 'numpy': numpy,
@@ -1031,28 +1094,17 @@ def _df_getitem_tuple_idx_impl(self, idx)
     return func_text, global_vars
 
 
-def gen_df_getitem_impl_generator(codegen, impl_name):
-    """Generate generator of df.getitem"""
-    def _df_getitem_impl_generator(self, idx):
-        func_text, global_vars = codegen(self, idx)
-
-        loc_vars = {}
-        exec(func_text, global_vars, loc_vars)
-        _impl = loc_vars[impl_name]
-
-        return _impl
-
-    return _df_getitem_impl_generator
-
-
-gen_df_getitem_slice_idx_impl = gen_df_getitem_impl_generator(
+gen_df_getitem_slice_idx_impl = gen_df_impl_generator(
     df_getitem_slice_idx_codegen, '_df_getitem_slice_idx_impl')
-gen_df_getitem_tuple_idx_impl = gen_df_getitem_impl_generator(
+gen_df_getitem_tuple_idx_impl = gen_df_impl_generator(
     df_getitem_tuple_idx_codegen, '_df_getitem_tuple_idx_impl')
+gen_df_getitem_bool_series_idx_impl = gen_df_impl_generator(
+    df_getitem_bool_series_codegen, '_df_getitem_bool_series_idx_impl')
 
 
 @sdc_overload(operator.getitem)
 def sdc_pandas_dataframe_getitem(self, idx):
+    ty_checker = TypeChecker('Operator getitem().')
 
     if not isinstance(self, DataFrameType):
         return None
@@ -1069,7 +1121,7 @@ def _df_getitem_str_literal_idx_impl(self, idx):
                 data = get_dataframe_data(self, col_idx)
                 return pandas.Series(data, index=self._index, name=idx)
             else:
-                raise KeyError
+                raise KeyError('Column is not in the DataFrame')
 
         return _df_getitem_str_literal_idx_impl
 
@@ -1082,12 +1134,30 @@ def _df_getitem_unicode_idx_impl(self, idx):
         return _df_getitem_unicode_idx_impl
 
     if isinstance(idx, types.Tuple):
-        return gen_df_getitem_tuple_idx_impl(self, idx)
+        if all([isinstance(item, types.StringLiteral) for item in idx]):
+            return gen_df_getitem_tuple_idx_impl(self, idx)
 
     if isinstance(idx, types.SliceType):
         return gen_df_getitem_slice_idx_impl(self, idx)
 
-    ty_checker = TypeChecker('Operator getitem().')
+    if isinstance(idx, SeriesType) and isinstance(idx.dtype, types.Boolean):
+        self_index_is_none = isinstance(self.index, types.NoneType)
+        idx_index_is_none = isinstance(idx.index, types.NoneType)
+
+        if self_index_is_none and not idx_index_is_none:
+            if not check_index_is_numeric(idx):
+                ty_checker.raise_exc(idx.index.dtype, 'number', 'idx.index.dtype')
+
+        if not self_index_is_none and idx_index_is_none:
+            if not check_index_is_numeric(self):
+                ty_checker.raise_exc(idx.index.dtype, self.index.dtype, 'idx.index.dtype')
+
+        if not self_index_is_none and not idx_index_is_none:
+            if not check_types_comparable(self.index, idx.index):
+                ty_checker.raise_exc(idx.index.dtype, self.index.dtype, 'idx.index.dtype')
+
+        return gen_df_getitem_bool_series_idx_impl(self, idx)
+
     ty_checker.raise_exc(idx, 'str', 'idx')
 
 

sdc/hiframes/pd_dataframe_ext.py

@@ -349,18 +349,17 @@ def df_getitem_overload(df, ind):
             index = df.columns.index(ind.literal_value)
             return lambda df, ind: sdc.hiframes.api.init_series(df._data[index])
 
-
-@infer_global(operator.getitem)
-class GetItemDataFrame(AbstractTemplate):
-    key = operator.getitem
-
-    def generic(self, args, kws):
-        df, idx = args
-        # df1 = df[df.A > .5]
-        if (isinstance(df, DataFrameType)
-                and isinstance(idx, (SeriesType, types.Array))
-                and idx.dtype == types.bool_):
-            return signature(df, *args)
+    @infer_global(operator.getitem)
+    class GetItemDataFrame(AbstractTemplate):
+        key = operator.getitem
+
+        def generic(self, args, kws):
+            df, idx = args
+            # df1 = df[df.A > .5]
+            if (isinstance(df, DataFrameType)
+                    and isinstance(idx, (SeriesType, types.Array))
+                    and idx.dtype == types.bool_):
+                return signature(df, *args)
 
 
 @infer

sdc/tests/test_dataframe.py

@@ -1236,17 +1236,37 @@ def test_impl(df, start, end):
         ref_result = test_impl(df, start, end)
         pd.testing.assert_frame_equal(jit_result, ref_result)
 
-    @skip_sdc_jit('DF.getitem unsupported Series name')
     def _test_df_getitem_tuple_idx(self, df):
-        def test_impl(df):
-            # pd.df.getitem does not support idx as a tuple
-            return df[['A', 'C']]
+        def gen_test_impl(do_jit=False):
+            def test_impl(df):
+                if do_jit == True:  # noqa
+                    return df[('A', 'C')]
+                else:
+                    return df[['A', 'C']]
 
-        # SDC pd.df.getitem does not support idx as a list
-        sdc_func = self.jit(lambda df: df[('A', 'C')])
+            return test_impl
 
+        test_impl = gen_test_impl()
+        sdc_func = self.jit(gen_test_impl(do_jit=True))
+
+        pd.testing.assert_frame_equal(sdc_func(df), test_impl(df))
+
+    def _test_df_getitem_bool_series_idx(self, df):
+        def test_impl(df):
+            return df[df['A'] == -1.]
+
+        sdc_func = self.jit(test_impl)
         pd.testing.assert_frame_equal(sdc_func(df), test_impl(df))
 
+    def _test_df_getitem_bool_series_even_idx(self, df):
+        def test_impl(df, series):
+            return df[series]
+
+        s = pd.Series([False, True] * 5)
+
+        sdc_func = self.jit(test_impl)
+        pd.testing.assert_frame_equal(sdc_func(df, s), test_impl(df, s))
+
     @skip_sdc_jit('DF.getitem unsupported exceptions')
     def test_df_getitem_str_literal_idx_exception_key_error(self):
         def test_impl(df):
@@ -1292,6 +1312,14 @@ def test_df_getitem_idx(self):
                 self._test_df_getitem_slice_idx(df)
                 self._test_df_getitem_unbox_slice_idx(df, 1, 3)
                 self._test_df_getitem_tuple_idx(df)
+                self._test_df_getitem_bool_series_idx(df)
+
+    @skip_sdc_jit('DF.getitem unsupported Series name')
+    def test_df_getitem_idx_no_index(self):
+        dfs = [gen_df(test_global_input_data_float64), pd.DataFrame({'A': []})]
+        for df in dfs:
+            with self.subTest(df=df):
+                self._test_df_getitem_bool_series_even_idx(df)
 
     @skip_sdc_jit('DF.getitem unsupported Series name')
     def test_df_getitem_idx_multiple_types(self):
@@ -1306,6 +1334,12 @@ def test_df_getitem_idx_multiple_types(self):
                 self._test_df_getitem_slice_idx(df)
                 self._test_df_getitem_unbox_slice_idx(df, 1, 3)
                 self._test_df_getitem_tuple_idx(df)
+                self._test_df_getitem_bool_series_even_idx(df)
+
+    @unittest.skip('DF.getitem df[bool_series] unsupported index')
+    def test_df_getitem_bool_series_even_idx_with_index(self):
+        df = gen_df(test_global_input_data_float64, with_index=True)
+        self._test_df_getitem_bool_series_even_idx(df)
 
     @unittest.skip('DF.getitem unsupported integer columns')
     def test_df_getitem_int_literal_idx(self):

sdc/utilities/sdc_typing_utils.py

@@ -168,3 +168,17 @@ def find_common_dtype_from_numpy_dtypes(array_types, scalar_types):
     numba_common_dtype = numpy_support.from_dtype(np_common_dtype)
 
     return numba_common_dtype
+
+
+def gen_df_impl_generator(codegen, impl_name):
+    """Generate generator of df methods"""
+    def _df_impl_generator(self, idx):
+        func_text, global_vars = codegen(self, idx)
+
+        loc_vars = {}
+        exec(func_text, global_vars, loc_vars)
+        _impl = loc_vars[impl_name]
+
+        return _impl
+
+    return _df_impl_generator