Repository: modin-project/modin

Issue:
BUG: modin can't share categories between multiple frames when pandas can
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import pandas

df = pandas.DataFrame(
    {
        "a": ["a", "a", "b", "c", "c", "d", "b"],
        "b": [0, 0, 0, 1, 1, 1, 1]
    }
)

cat_df = df.astype({"a": "category"})

fi = cat_df.loc[df.b == 0]
se = cat_df.loc[df.b == 1]

print(fi.dtypes["a"].categories.equals(se.dtypes["a"].categories)) # True

print(fi["a"].cat.codes) # [0, 0, 1]    
print(se["a"].cat.codes) # [2, 2, 3, 1] <-- shares categories with 'fi', thus codes start from 1

from modin.test.storage_formats.pandas.test_internals import construct_modin_df_by_scheme

df = construct_modin_df_by_scheme(df, {"row_lengths": [2, 2, 2, 1], "column_widths": [2]})
cat_df = df.astype({"a": "category"})

fi = cat_df.loc[df.b == 0]
se = cat_df.loc[df.b == 1]
print(fi.dtypes["a"].categories.equals(se.dtypes["a"].categories)) # False

print(fi["a"].cat.codes) # [0, 0, 1]
print(se["a"].cat.codes) # [1, 1, 2, 0] <-- have independent categories, thus codes start from 0
```


### Issue Description

Pandas do share categorical values between offspring frames thus making it possible to treat their codes as comparable to each other.

In modin, a categorical frame often doesn't even have defined categorical values at all as the `.dtypes` are often unknown so each partition has its own categorical values. One of the consequences of this is that when building a full-axis column partition we lose categorical types at the result because there's just no full-axis table for categories and was found more beneficial to discard categories at all rather than building categories for the whole column (#2513).

Although we still can ask for a full-axis function to re-build categories on demand (however it still doesn't allows us to share them between multiple frames):
https://github.com/modin-project/modin/blob/bb0950d06ce4bbdd87433ded82d118d803ed47a0/modin/core/storage_formats/pandas/query_compiler.py#L3451-L3453

### Expected Behavior

We apparently would like to have a sharable full-axis categorical object between all partitions and an ability to share it with other frames.

### Error Logs

<details>
There is no error logs, so putting here a plain output of the reproducer

```python-traceback

True
0    0
1    0
2    1
dtype: int8
3    2
4    2
5    3
6    1
dtype: int8
False
0    0
1    0
2    1
dtype: int8
3    1
4    1
5    2
6    0
dtype: int8

```

</details>


### Installed Versions

<details>

Replace this line with the output of pd.show_versions()

</details>



Discussion:
https://github.com/pandas-dev/pandas/issues/51362 might be related to this.
Apparently casting to categories using full-axis function fixes everything:
```diff
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index 30925e40..1e38bde6 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -1201,6 +1201,7 @@ class PandasDataframe(ClassLogger):
         columns = col_dtypes.keys()
         # Create Series for the updated dtypes
         new_dtypes = self.dtypes.copy()
+        full_axis_cast = False
         for i, column in enumerate(columns):
             dtype = col_dtypes[column]
             if (
@@ -1220,6 +1221,7 @@ class PandasDataframe(ClassLogger):
                 # We cannot infer without computing the dtype if
                 elif isinstance(new_dtype, str) and new_dtype == "category":
                     new_dtypes = None
+                    full_axis_cast = True
                     break
                 else:
                     new_dtypes[column] = new_dtype
@@ -1228,9 +1230,15 @@ class PandasDataframe(ClassLogger):
             """Compute new partition frame with dtypes updated."""
             return df.astype({k: v for k, v in col_dtypes.items() if k in df})

-        new_frame = self._partition_mgr_cls.map_partitions(
-            self._partitions, astype_builder
-        )
+        if 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._index_cache,
```
This happens because now each partition will have the whole categorical values, meaning that `pd.concat` will not discard categorical dtypes when building a full-axis column partition.
But the approach above have significant downsides:
1. The cast itself becomes more expensive as it's now performed in a full-axis function
2. Each partition is now obliged to store the whole encoding table, which can cause severe memory issues (the same as we're fixing with multiindex right now #5632)

@dchigarev Is it possible that modin will give incorrect results, while working with category data due to this problem (during join or groupby)?

For instance, when I perform groupby on category columns
not sure about join and groupby (maybe there are cases where it could be broken, however, simple scenarios should work fine, most of the time it just discards categorical dtype) but we already had a case with the `.map()` function that was affected by this malfunction:

<details><summary>example</summary>

```python
import pandas
from modin.test.storage_formats.pandas.test_internals import construct_modin_df_by_scheme

def run_scenario(df):
    print(f"\n== running scenario with {type(df)=} ==")

    cat_df = df.astype({"a": "category"})

    # in pandas 'fi' and 'se' would still share the same categorical values,
    # when with modin they will now completely independent
    fi = cat_df.loc[df.b == 0]
    se = cat_df.loc[df.b == 1]

    fi["a"] = fi["a"].cat.codes
    se["a"] = se["a"].cat.codes

    print(fi["a"]) # pandas: [0, 0, 1] | modin: [0, 0, 1]
    print(se["a"]) # pandas: [2, 2, 3, 1] | modin: [1, 1, 2, 0]

    se_groupby = se.groupby("a").sum().squeeze(axis=1)

    print(fi["a"].map(se_groupby)) # pandas: [nan, nan, 1.0] | modin: [1, 1, 2]

pandas_df = pandas.DataFrame(
    {
        "a": ["a", "a", "b", "c", "c", "d", "b"],
        "b": [0, 0, 0, 1, 1, 1, 1]
    }
)
modin_df = construct_modin_df_by_scheme(pandas_df, {"row_lengths": [2, 2, 2, 1], "column_widths": [2]})

run_scenario(pandas_df)
run_scenario(modin_df)
```

</details>

Resolving patch:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/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(),
diff --git a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py b/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
--- a/modin/experimental/core/execution/native/implementations/hdk_on_native/dataframe/dataframe.py
+++ b/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,


Test patch:
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/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)