FIX-#5722: Use full axis function when casting to "category"

modin/core/dataframe/pandas/dataframe/dataframe.py

@@ -1369,6 +1369,11 @@ def astype(self, col_dtypes, errors: str = "raise"):
         columns = col_dtypes.keys()
         # Create Series for the updated dtypes
         new_dtypes = self.dtypes.copy()
+        # When casting to "category" we have to make up the whole axis partition
+        # to get the properly encoded table of categories. Every block partition
+        # will store the encoded table. That can lead to higher memory footprint.
+        # TODO: Revisit if this hurts users.
+        use_full_axis_cast = False
         for i, column in enumerate(columns):
             dtype = col_dtypes[column]
             if (
@@ -1395,6 +1400,7 @@ def astype(self, col_dtypes, errors: str = "raise"):
                             columns=[column]
                         )[column],
                     )
+                    use_full_axis_cast = True
                 else:
                     new_dtypes[column] = new_dtype
 
@@ -1404,9 +1410,14 @@ def astype_builder(df):
                 {k: v for k, v in col_dtypes.items() if k in df}, errors=errors
             )
 
-        new_frame = self._partition_mgr_cls.map_partitions(
-            self._partitions, astype_builder
-        )
+        if use_full_axis_cast:
+            new_frame = self._partition_mgr_cls.map_axis_partitions(
+                0, self._partitions, astype_builder, keep_partitioning=True
+            )
+        else:
+            new_frame = self._partition_mgr_cls.map_partitions(
+                self._partitions, astype_builder
+            )
         return self.__constructor__(
             new_frame,
             self.copy_index_cache(),

modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py

@@ -1647,10 +1647,11 @@ def cat_codes(self):
         exprs[col_name] = build_if_then_else(
             col_expr.is_null(), null_val, code_expr, get_dtype("int32")
         )
+        dtypes = [expr._dtype for expr in exprs.values()]
 
         return self.__constructor__(
             columns=Index([col_name]),
-            dtypes=pd.Series(self._dtypes[0], index=Index([col_name])),
+            dtypes=pd.Series(dtypes, index=Index(exprs.keys())),
             op=TransformNode(self, exprs),
             index=self._index_cache,
             index_cols=self._index_cols,

modin/pandas/test/test_series.py

@@ -1154,6 +1154,35 @@ def test_astype_categorical(data):
     assert modin_result.dtype == pandas_result.dtype
 
 
+@pytest.mark.parametrize("data", [["a", "a", "b", "c", "c", "d", "b", "d"]])
+@pytest.mark.parametrize(
+    "set_min_partition_size",
+    [2, 4],
+    ids=["four_partitions", "two_partitions"],
+    indirect=True,
+)
+def test_astype_categorical_issue5722(data, set_min_partition_size):
+    modin_series, pandas_series = create_test_series(data)
+
+    modin_result = modin_series.astype("category")
+    pandas_result = pandas_series.astype("category")
+    df_equals(modin_result, pandas_result)
+    assert modin_result.dtype == pandas_result.dtype
+
+    pandas_result1, pandas_result2 = pandas_result.iloc[:4], pandas_result.iloc[4:]
+    modin_result1, modin_result2 = modin_result.iloc[:4], modin_result.iloc[4:]
+
+    # check categories
+    assert pandas_result1.cat.categories.equals(pandas_result2.cat.categories)
+    assert modin_result1.cat.categories.equals(modin_result2.cat.categories)
+    assert pandas_result1.cat.categories.equals(modin_result1.cat.categories)
+    assert pandas_result2.cat.categories.equals(modin_result2.cat.categories)
+
+    # check codes
+    assert_array_equal(pandas_result1.cat.codes.values, modin_result1.cat.codes.values)
+    assert_array_equal(pandas_result2.cat.codes.values, modin_result2.cat.codes.values)
+
+
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
 def test_at(data):
     modin_series, pandas_series = create_test_series(data)