Re-implement df.reset_index based on new structure (#849)

Author: densmirn

sdc/datatypes/hpat_pandas_dataframe_functions.py

@@ -2804,90 +2804,43 @@ def _df_set_column_unicode_key_impl(self, key, value):
     ty_checker.raise_exc(key, 'str', 'key')
 
 
-def sdc_pandas_dataframe_reset_index_codegen(drop, all_params, columns):
+def sdc_pandas_dataframe_reset_index_codegen(drop, all_params, columns, column_loc):
     """
     Example of generated implementation:
         def _df_reset_index_impl(self, level=None, drop=False, inplace=False, col_level=0, col_fill=""):
-          old_index = self.index
-          result_0 = get_dataframe_data(self, 0)
-          result_1 = get_dataframe_data(self, 1)
-          result_2 = get_dataframe_data(self, 2)
-          return pandas.DataFrame({"index": old_index, "A": result_0, "B": result_1, "C": result_2})
+          result_0 = self._data[0][0]
+          result_1 = self._data[0][1]
+          return pandas.DataFrame({"A": result_0, "B": result_1})
     """
     result_name = []
     all_params_str = ', '.join(all_params)
     func_lines = [f'def _df_reset_index_impl({all_params_str}):']
-    df = all_params[0]
-    if not drop.literal_value:
-        old_index = 'old_index'
-        func_lines += [f'  {old_index} = {df}.index']
-        result_name.append((old_index, 'index'))
-    for i, c in enumerate(columns):
-        result_c = f'result_{i}'
-        func_lines += [
-            f'  result_{i} = get_dataframe_data({df}, {i})'
-        ]
-        result_name.append((result_c, c))
-    data = ', '.join(f'"{column_name}": {column}' for column, column_name in result_name)
-    func_lines += [f'  return pandas.DataFrame({{{data}}})']
-    func_text = '\n'.join(func_lines)
-
-    global_vars = {'pandas': pandas,
-                   'numpy': numpy,
-                   'get_dataframe_data': get_dataframe_data}
-
-    return func_text, global_vars
-
-
-def sdc_pandas_dataframe_reset_index_impl(self, drop=False):
-    all_params = ['self', 'level=None', 'drop=False', 'inplace=False', 'col_level=0', 'col_fill=""']
-
-    func_text, global_vars = sdc_pandas_dataframe_reset_index_codegen(drop, all_params, self.columns)
-    loc_vars = {}
-    exec(func_text, global_vars, loc_vars)
-    _apply_impl = loc_vars[f'_df_reset_index_impl']
-
-    return _apply_impl
-
-
-def sdc_pandas_dataframe_reset_index_default_codegen(drop, all_params, columns):
-    """
-    Example of generated implementation:
-        def _df_reset_index_impl(self, level=None, drop=False, inplace=False, col_level=0, col_fill=""):
-          old_index = self.index
-          result_0 = get_dataframe_data(self, 0)
-          result_1 = get_dataframe_data(self, 1)
-          return pandas.DataFrame({"index": old_index, "A": result_0, "B": result_1})
-    """
-    result_name = []
-    all_params_str = ', '.join(all_params)
-    func_lines = [f'def _df_reset_index_impl({all_params_str}):']
-    df = all_params[0]
     if not drop:
         old_index = 'old_index'
-        func_lines += [f'  {old_index} = {df}.index']
+        func_lines += [f'  {old_index} = self.index']
         result_name.append((old_index, 'index'))
     for i, c in enumerate(columns):
+        col_loc = column_loc[c]
+        type_id, col_id = col_loc.type_id, col_loc.col_id
         result_c = f'result_{i}'
         func_lines += [
-            f'  result_{i} = get_dataframe_data({df}, {i})'
+            f'  result_{i} = self._data[{type_id}][{col_id}]'
         ]
         result_name.append((result_c, c))
     data = ', '.join(f'"{column_name}": {column}' for column, column_name in result_name)
     func_lines += [f'  return pandas.DataFrame({{{data}}})']
     func_text = '\n'.join(func_lines)
 
     global_vars = {'pandas': pandas,
-                   'numpy': numpy,
-                   'get_dataframe_data': get_dataframe_data}
+                   'numpy': numpy}
 
     return func_text, global_vars
 
 
-def sdc_pandas_dataframe_reset_index_impl_default(self, drop=False):
+def sdc_pandas_dataframe_reset_index_impl(self, drop=False):
     all_params = ['self', 'level=None', 'drop=False', 'inplace=False', 'col_level=0', 'col_fill=""']
-
-    func_text, global_vars = sdc_pandas_dataframe_reset_index_default_codegen(drop, all_params, self.columns)
+    func_text, global_vars = sdc_pandas_dataframe_reset_index_codegen(drop, all_params,
+                                                                      self.columns, self.column_loc)
     loc_vars = {}
     exec(func_text, global_vars, loc_vars)
     _apply_impl = loc_vars[f'_df_reset_index_impl']
@@ -2946,12 +2899,9 @@ def sdc_pandas_dataframe_reset_index(self, level=None, drop=False, inplace=False
     if not (level is None or isinstance(level, types.Omitted)):
         raise TypingError('{} Unsupported parameter level. Given: {}'.format(func_name, level))
 
-    if not (isinstance(drop, (types.Omitted, types.Boolean)) or drop is False):
+    if not isinstance(drop, (types.Omitted, types.Boolean, bool)):
         ty_checker.raise_exc(drop, 'bool', 'drop')
 
-    if isinstance(drop, types.Omitted):
-        drop = False
-
     if not (inplace is False or isinstance(inplace, types.Omitted)):
         raise TypingError('{} Unsupported parameter inplace. Given: {}'.format(func_name, inplace))
 
@@ -2961,10 +2911,13 @@ def sdc_pandas_dataframe_reset_index(self, level=None, drop=False, inplace=False
     if not (col_fill == '' or isinstance(col_fill, types.Omitted)):
         raise TypingError('{} Unsupported parameter col_fill. Given: {}'.format(func_name, col_fill))
 
-    if not isinstance(drop, types.Literal):
-        if isinstance(drop, bool):
-            return sdc_pandas_dataframe_reset_index_impl_default(self, drop=drop)
-        else:
-            raise SDCLimitation('{} only work with Boolean literals drop.'.format(func_name))
+    if isinstance(drop, types.Literal):
+        literal_drop = drop.literal_value
+        return sdc_pandas_dataframe_reset_index_impl(self, drop=literal_drop)
+    elif isinstance(drop, types.Omitted):
+        return sdc_pandas_dataframe_reset_index_impl(self, drop=drop.value)
+    elif isinstance(drop, bool):
+        return sdc_pandas_dataframe_reset_index_impl(self, drop=drop)
+
+    raise SDCLimitation('Method {}(). Parameter drop is only supported as a literal.'.format(func_name))
 
-    return sdc_pandas_dataframe_reset_index_impl(self, drop=drop)

sdc/tests/test_dataframe.py

@@ -1495,41 +1495,71 @@ def test_impl(df, drop):
             with self.subTest(drop=drop):
                 with self.assertRaises(Exception) as raises:
                     hpat_func(df, drop)
-                msg = 'only work with Boolean literals drop'
-                self.assertIn(msg.format(types.bool_), str(raises.exception))
+                msg = 'drop is only supported as a literal'
+                self.assertIn(msg, str(raises.exception))
 
-    @dfRefactoringNotImplemented
-    def test_df_reset_index_drop_false_index_int(self):
-        def test_impl(df):
-            return df.reset_index(drop=False)
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
+    def test_df_reset_index_drop_literal_index_int(self):
+        def gen_test_impl(drop):
+            def test_impl(df):
+                if drop == False:  # noqa
+                    return df.reset_index(drop=False)
+                else:
+                    return df.reset_index(drop=True)
+            return test_impl
 
-        df = pd.DataFrame({'A': [1.0, 2.0, np.nan, 1.0],
-                           'B': np.arange(4.0)}, index=[5, 8, 4, 6])
-        hpat_func = self.jit(test_impl)
+        df = pd.DataFrame({
+            'A': [1.0, 2.0, np.nan, 1.0],
+            'B': np.arange(4.0)
+        }, index=[5, 8, 4, 6])
+        for drop in [True, False]:
+            with self.subTest(drop=drop):
+                test_impl = gen_test_impl(drop)
+                hpat_func = self.jit(test_impl)
+                pd.testing.assert_frame_equal(hpat_func(df), test_impl(df))
 
-        pd.testing.assert_frame_equal(hpat_func(df), test_impl(df))
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame boxing
+    def test_df_reset_index_drop_literal_index_int_no_unboxing(self):
+        def gen_test_impl(drop):
+            def test_impl():
+                df = pd.DataFrame({
+                    'A': [1.0, 2.0, np.nan, 1.0],
+                    'B': np.arange(4.0)
+                }, index=[5, 8, 4, 6])
+                if drop == False:  # noqa
+                    return df.reset_index(drop=False)
+                else:
+                    return df.reset_index(drop=True)
+            return test_impl
 
-    @dfRefactoringNotImplemented
-    def test_df_reset_index_drop_true_index_int(self):
+        for drop in [True, False]:
+            with self.subTest(drop=drop):
+                test_impl = gen_test_impl(drop)
+                hpat_func = self.jit(test_impl)
+                pd.testing.assert_frame_equal(hpat_func(), test_impl())
+
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
+    def test_df_reset_index_drop_default_index_int(self):
         def test_impl(df):
-            return df.reset_index(drop=True)
+            return df.reset_index()
 
         df = pd.DataFrame({'A': [1.0, 2.0, np.nan, 1.0],
                            'B': np.arange(4.0)}, index=[5, 8, 4, 6])
         hpat_func = self.jit(test_impl)
 
         pd.testing.assert_frame_equal(hpat_func(df), test_impl(df))
 
-    @dfRefactoringNotImplemented
-    def test_df_reset_index_drop_default_index_int(self):
-        def test_impl(df):
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame boxing
+    def test_df_reset_index_drop_default_index_int_no_unboxing(self):
+        def test_impl():
+            df = pd.DataFrame({
+                'A': [1.0, 2.0, np.nan, 1.0],
+                'B': np.arange(4.0)
+            }, index=[5, 8, 4, 6])
             return df.reset_index()
 
-        df = pd.DataFrame({'A': [1.0, 2.0, np.nan, 1.0],
-                           'B': np.arange(4.0)}, index=[5, 8, 4, 6])
         hpat_func = self.jit(test_impl)
-
-        pd.testing.assert_frame_equal(hpat_func(df), test_impl(df))
+        pd.testing.assert_frame_equal(hpat_func(), test_impl())
 
     @skip_numba_jit
     def test_df_reset_index_empty_df(self):