Repository: modin-project/modin

Issue:
Implement `cat.codes` using map rather than full-axis approach
After the #6222 is merged, every partition now knows the global categorical codes, meaning that there's no need to use the full-axis fold anymore as we can now simply map over partitions and still receive valid codes
https://github.com/modin-project/modin/blob/7029a700cffe854053c841f156fa2f9fe61dfed6/modin/core/storage_formats/pandas/query_compiler.py#L4043-L4045


Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/base/dataframe/dataframe.py b/modin/core/dataframe/base/dataframe/dataframe.py
--- a/modin/core/dataframe/base/dataframe/dataframe.py
+++ b/modin/core/dataframe/base/dataframe/dataframe.py
@@ -93,6 +93,7 @@ def map(
         function: Callable,
         axis: Optional[Union[int, Axis]] = None,
         dtypes: Optional[str] = None,
+        new_columns: Optional[List[Hashable]] = None,
     ) -> "ModinDataframe":
         """
         Apply a user-defined function row-wise if `axis`=0, column-wise if `axis`=1, and cell-wise if `axis` is None.
@@ -107,6 +108,9 @@ def map(
             The data types for the result. This is an optimization
             because there are functions that always result in a particular data
             type, and this allows us to avoid (re)computing it.
+        new_columns : List[Hashable], optional
+            New column labels of the result, its length has to be identical
+            to the older columns. If not specified, old column labels are preserved.
 
         Returns
         -------
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
@@ -1920,7 +1920,12 @@ def tree_reduce(
         return self._compute_tree_reduce_metadata(axis.value, reduce_parts)
 
     @lazy_metadata_decorator(apply_axis=None)
-    def map(self, func: Callable, dtypes: Optional[str] = None) -> "PandasDataframe":
+    def map(
+        self,
+        func: Callable,
+        dtypes: Optional[str] = None,
+        new_columns: Optional[pandas.Index] = None,
+    ) -> "PandasDataframe":
         """
         Perform a function that maps across the entire dataset.
 
@@ -1932,6 +1937,9 @@ def map(self, func: Callable, dtypes: Optional[str] = None) -> "PandasDataframe"
             The data types for the result. This is an optimization
             because there are functions that always result in a particular data
             type, and this allows us to avoid (re)computing it.
+        new_columns : pandas.Index, optional
+            New column labels of the result, its length has to be identical
+            to the older columns. If not specified, old column labels are preserved.
 
         Returns
         -------
@@ -1939,16 +1947,25 @@ def map(self, func: Callable, dtypes: Optional[str] = None) -> "PandasDataframe"
             A new dataframe.
         """
         new_partitions = self._partition_mgr_cls.map_partitions(self._partitions, func)
+        if new_columns is not None and self.has_materialized_columns:
+            assert len(new_columns) == len(
+                self.columns
+            ), "New column's length must be identical to the previous columns"
+        elif new_columns is None:
+            new_columns = self.copy_columns_cache()
         if isinstance(dtypes, str) and dtypes == "copy":
             dtypes = self.copy_dtypes_cache()
         elif dtypes is not None and not isinstance(dtypes, pandas.Series):
+            if isinstance(new_columns, ModinIndex):
+                # Materializing lazy columns in order to build dtype's index
+                new_columns = new_columns.get(return_lengths=False)
             dtypes = pandas.Series(
-                [np.dtype(dtypes)] * len(self.columns), index=self.columns
+                [np.dtype(dtypes)] * len(new_columns), index=new_columns
             )
         return self.__constructor__(
             new_partitions,
             self.copy_index_cache(),
-            self.copy_columns_cache(),
+            new_columns,
             self._row_lengths_cache,
             self._column_widths_cache,
             dtypes=dtypes,
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -4028,21 +4028,12 @@ def sort_columns_by_row_values(self, rows, ascending=True, **kwargs):
 
     # Cat operations
     def cat_codes(self):
-        def func(df) -> np.ndarray:
-            # `df` is supposed to be consisted of multiple partitions,
-            # which should be concatenated before applying a function.
-            # `pd.concat` doesn't preserve categorical dtype
-            # if the dfs have categorical columns
-            # so we intentionaly restore the right dtype.
-            # TODO: revert the change when https://github.com/pandas-dev/pandas/issues/51362 is fixed.
+        def func(df: pandas.DataFrame) -> pandas.DataFrame:
             ser = df.iloc[:, 0]
-            if ser.dtype != "category":
-                ser = ser.astype("category", copy=False)
+            assert ser.dtype == "category"
             return ser.cat.codes.to_frame(name=MODIN_UNNAMED_SERIES_LABEL)
 
-        res = self._modin_frame.fold(
-            axis=0, func=func, new_columns=[MODIN_UNNAMED_SERIES_LABEL]
-        )
+        res = self._modin_frame.map(func=func, new_columns=[MODIN_UNNAMED_SERIES_LABEL])
         return self.__constructor__(res, shape_hint="column")
 
     # END Cat operations


Test patch: