Repository: modin-project/modin

Issue:
preserve index and columns on `repartition`
`repartition` function does not change indexes, so we can use already materialized indexes.


Discussion:

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
@@ -2187,6 +2187,7 @@ def apply_full_axis(
         new_columns=None,
         dtypes=None,
         keep_partitioning=True,
+        sync_labels=True,
     ):
         """
         Perform a function across an entire axis.
@@ -2210,6 +2211,10 @@ def apply_full_axis(
         keep_partitioning : boolean, default: True
             The flag to keep partition boundaries for Modin Frame.
             Setting it to True disables shuffling data from one partition to another.
+        sync_labels : boolean, default: True
+            Synchronize external indexes (`new_index`, `new_columns`) with internal indexes.
+            This could be used when you're certain that the indices in partitions are equal to
+            the provided hints in order to save time on syncing them.
 
         Returns
         -------
@@ -2228,6 +2233,7 @@ def apply_full_axis(
             dtypes=dtypes,
             other=None,
             keep_partitioning=keep_partitioning,
+            sync_labels=sync_labels,
         )
 
     @lazy_metadata_decorator(apply_axis="both")
@@ -2619,6 +2625,7 @@ def broadcast_apply_full_axis(
         enumerate_partitions=False,
         dtypes=None,
         keep_partitioning=True,
+        sync_labels=True,
     ):
         """
         Broadcast partitions of `other` Modin DataFrame and apply a function along full axis.
@@ -2649,6 +2656,10 @@ def broadcast_apply_full_axis(
         keep_partitioning : boolean, default: True
             The flag to keep partition boundaries for Modin Frame.
             Setting it to True disables shuffling data from one partition to another.
+        sync_labels : boolean, default: True
+            Synchronize external indexes (`new_index`, `new_columns`) with internal indexes.
+            This could be used when you're certain that the indices in partitions are equal to
+            the provided hints in order to save time on syncing them.
 
         Returns
         -------
@@ -2683,10 +2694,35 @@ def broadcast_apply_full_axis(
             kw["dtypes"] = pandas.Series(
                 [np.dtype(dtypes)] * len(kw["columns"]), index=kw["columns"]
             )
+
+        if not keep_partitioning:
+            if kw["row_lengths"] is None and new_index is not None:
+                if axis == 0:
+                    kw["row_lengths"] = get_length_list(
+                        axis_len=len(new_index), num_splits=new_partitions.shape[0]
+                    )
+                elif (
+                    axis == 1
+                    and self._row_lengths_cache is not None
+                    and len(new_index) == sum(self._row_lengths_cache)
+                ):
+                    kw["row_lengths"] = self._row_lengths_cache
+            if kw["column_widths"] is None and new_columns is not None:
+                if axis == 1:
+                    kw["column_widths"] = get_length_list(
+                        axis_len=len(new_columns),
+                        num_splits=new_partitions.shape[1],
+                    )
+                elif (
+                    axis == 0
+                    and self._column_widths_cache is not None
+                    and len(new_columns) == sum(self._column_widths_cache)
+                ):
+                    kw["column_widths"] = self._column_widths_cache
         result = self.__constructor__(new_partitions, **kw)
-        if new_index is not None:
+        if sync_labels and new_index is not None:
             result.synchronize_labels(axis=0)
-        if new_columns is not None:
+        if sync_labels and new_columns is not None:
             result.synchronize_labels(axis=1)
         return result
 
diff --git a/modin/core/execution/python/implementations/pandas_on_python/partitioning/partition.py b/modin/core/execution/python/implementations/pandas_on_python/partitioning/partition.py
--- a/modin/core/execution/python/implementations/pandas_on_python/partitioning/partition.py
+++ b/modin/core/execution/python/implementations/pandas_on_python/partitioning/partition.py
@@ -143,7 +143,7 @@ def put(cls, obj):
         PandasOnPythonDataframePartition
             New ``PandasOnPythonDataframePartition`` object.
         """
-        return cls(obj.copy())
+        return cls(obj.copy(), len(obj.index), len(obj.columns))
 
     @classmethod
     def preprocess_func(cls, func):
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -5090,7 +5090,12 @@ def repartition(self, axis=None):
         for _ax in axes:
             new_query_compiler = new_query_compiler.__constructor__(
                 new_query_compiler._modin_frame.apply_full_axis(
-                    _ax, lambda df: df, keep_partitioning=False
+                    _ax,
+                    lambda df: df,
+                    new_index=self._modin_frame._index_cache,
+                    new_columns=self._modin_frame._columns_cache,
+                    keep_partitioning=False,
+                    sync_labels=False,
                 )
             )
         return new_query_compiler


Test patch: