FIX-#4808: Set dtypes correctly after column rename. (#4809)

Signed-off-by: mvashishtha <mahesh@ponder.io> Co-authored-by: Karthik Velayutham <karthik.velayutham@gmail.com>
modin-project · Aug 12, 2022 · d42ebec · d42ebec
1 parent 6e1849f
commit d42ebec
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/docs/release_notes/release_notes-0.16.0.rst b/docs/release_notes/release_notes-0.16.0.rst
@@ -32,6 +32,7 @@ Key Features and Updates
   * FIX-#4657: Use `fsspec` for handling s3/http-like paths instead of `s3fs` (#4710)
   * FIX-#4676: drain sub-virtual-partition call queues (#4695)
   * FIX-#4782: Exclude certain non-parquet files in `read_parquet` (#4783)    
+  * FIX-#4808: Set dtypes correctly after column rename (#4809)
 * Performance enhancements
   * PERF-#4182: Add cell-wise execution for binary ops, fix bin ops for empty dataframes (#4391)
   * PERF-#4288: Improve perf of `groupby.mean` for narrow data (#4591)

diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -1754,13 +1754,14 @@ def map_fn(df):
             return df.rename(index=new_row_labels, columns=new_col_labels, level=level)
 
         new_parts = self._partition_mgr_cls.map_partitions(self._partitions, map_fn)
+        new_dtypes = None if self._dtypes is None else self._dtypes.set_axis(new_cols)
         return self.__constructor__(
             new_parts,
             new_index,
             new_cols,
             self._row_lengths,
             self._column_widths,
-            self._dtypes,
+            new_dtypes,
         )
 
     def sort_by(

diff --git a/modin/pandas/test/dataframe/test_map_metadata.py b/modin/pandas/test/dataframe/test_map_metadata.py
@@ -205,6 +205,9 @@ def test_add_prefix(data):
     new_modin_df = modin_df.add_prefix(test_prefix)
     new_pandas_df = pandas_df.add_prefix(test_prefix)
     df_equals(new_modin_df.columns, new_pandas_df.columns)
+    # TODO(https://github.com/modin-project/modin/issues/3804):
+    # make df_equals always check dtypes.
+    df_equals(new_modin_df.dtypes, new_pandas_df.dtypes)
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)