Re-implement df.copy based on new structure (#854)

densmirn · web-flow · commit 1f00705c5ace · 2020-05-22T17:29:14.000+03:00
diff --git a/sdc/datatypes/hpat_pandas_dataframe_functions.py b/sdc/datatypes/hpat_pandas_dataframe_functions.py
@@ -548,22 +548,27 @@ def head_overload(df, n=5):
 
 def _dataframe_codegen_copy(func_params, series_params, df):
     """
-    Example func_text for func_name='copy' columns=('A', 'B', 'C'):
+    Example func_text for func_name='copy' columns=('A', 'B'):
         def _df_copy_impl(df, deep=True):
-            series_0 = pandas.Series(df._data[0])
-            result_0 = series_0.copy(deep=deep)
-            series_1 = pandas.Series(df._data[1])
-            result_1 = series_1.copy(deep=deep)
-            return pandas.DataFrame({"A": result_0, "B": result_1}, index=df._index)
+          data_0 = df._data[0][0]
+          series_0 = pandas.Series(data_0, name='A')
+          result_0 = series_0.copy(deep=deep)
+          data_1 = df._data[1][0]
+          series_1 = pandas.Series(data_1, name='B')
+          result_1 = series_1.copy(deep=deep)
+          return pandas.DataFrame({"A": result_0, "B": result_1}, index=df._index)
     """
     results = []
     series_params_str = ', '.join(kwsparams2list(series_params))
     func_params_str = ', '.join(kwsparams2list(func_params))
     func_lines = [f"def _df_copy_impl(df, {func_params_str}):"]
     index = df_index_codegen_all(df)
     for i, c in enumerate(df.columns):
+        col_loc = df.column_loc[c]
+        type_id, col_id = col_loc.type_id, col_loc.col_id
         result_c = f"result_{i}"
-        func_lines += [f"  series_{i} = pandas.Series(df._data[{i}], name='{c}')",
+        func_lines += [f"  data_{i} = df._data[{type_id}][{col_id}]",
+                       f"  series_{i} = pandas.Series(data_{i}, name='{c}')",
                        f"  {result_c} = series_{i}.copy({series_params_str})"]
         results.append((c, result_c))
 
diff --git a/sdc/tests/test_dataframe.py b/sdc/tests/test_dataframe.py
@@ -1362,7 +1362,7 @@ def impl(a):
                     )
                     pd.testing.assert_frame_equal(sdc_func(df), ref_impl(df))
 
-    @dfRefactoringNotImplemented
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame unboxing
     def test_df_copy(self):
         def test_impl(df, deep):
             return df.copy(deep=deep)
@@ -1380,6 +1380,26 @@ def test_impl(df, deep):
                 with self.subTest(index=idx, deep=deep):
                     pd.testing.assert_frame_equal(sdc_func(df, deep), test_impl(df, deep))
 
+    @dfRefactoringNotImplemented  # required re-implementing DataFrame boxing
+    def test_df_copy_no_unboxing(self):
+        def test_impl(idx, deep):
+            df = pd.DataFrame({
+                'A': [3.2, np.nan, 7.0, 3.3, np.nan],
+                'B': [3, 4, 1, 0, 222],
+                'C': [True, True, False, False, True],
+                'D': ['a', 'dd', 'c', '12', None]
+            }, index=idx)
+            return df.copy(deep=deep)
+
+        sdc_impl = sdc.jit(test_impl)
+        indexes = [[3, 4, 2, 6, 1], ['a', 'b', 'c', 'd', 'e'], None]
+        cases_deep = [None, True, False]
+        for idx, deep in product(indexes, cases_deep):
+            with self.subTest(index=idx, deep=deep):
+                jit_result = sdc_impl(idx, deep)
+                ref_result = test_impl(idx, deep)
+                pd.testing.assert_frame_equal(jit_result, ref_result)
+
     @dfRefactoringNotImplemented  # required re-implementing DataFrame boxing
     def test_pct_change1(self):
         def test_impl(n):
