Repository: modin-project/modin

Issue:
PERF: compute `lengths` in `rebalance_partitions` when possible



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
@@ -2607,6 +2607,7 @@ def concat(
         """
         axis = Axis(axis)
         new_widths = None
+        new_lengths = None
 
         def _compute_new_widths():
             widths = None
@@ -2653,9 +2654,11 @@ def _compute_new_widths():
                 new_widths = _compute_new_widths()
             else:
                 new_widths = partition_sizes_along_axis
-        new_partitions = self._partition_mgr_cls.concat(
+        new_partitions, new_lengths2 = self._partition_mgr_cls.concat(
             axis.value, left_parts, right_parts
         )
+        if new_lengths is None:
+            new_lengths = new_lengths2
         new_dtypes = None
         if axis == Axis.ROW_WISE:
             new_index = self.index.append([other.index for other in others])
@@ -2674,14 +2677,15 @@ def _compute_new_widths():
             # frame. Typically, if we know the length for any partition in a
             # row, we know the length for the first partition in the row. So
             # just check the lengths of the first column of partitions.
-            new_lengths = []
-            if new_partitions.size > 0:
-                for part in new_partitions.T[0]:
-                    if part._length_cache is not None:
-                        new_lengths.append(part.length())
-                    else:
-                        new_lengths = None
-                        break
+            if not new_lengths:
+                new_lengths = []
+                if new_partitions.size > 0:
+                    for part in new_partitions.T[0]:
+                        if part._length_cache is not None:
+                            new_lengths.append(part.length())
+                        else:
+                            new_lengths = None
+                            break
         else:
             new_columns = self.columns.append([other.columns for other in others])
             new_index = joined_index
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -601,6 +601,8 @@ def concat(cls, axis, left_parts, right_parts):
         -------
         np.ndarray
             A new NumPy array with concatenated partitions.
+        list[int] or None
+            Row lengths if possible to compute it.
 
         Notes
         -----
@@ -625,7 +627,7 @@ def concat(cls, axis, left_parts, right_parts):
         if axis == 0:
             return cls.rebalance_partitions(result)
         else:
-            return result
+            return result, None
 
     @classmethod
     def to_pandas(cls, partitions):
@@ -1341,6 +1343,8 @@ def rebalance_partitions(cls, partitions):
         np.ndarray
             A NumPy array with the same; or new, rebalanced, partitions, depending on the execution
             engine and storage format.
+        list[int] or None
+            Row lengths if possible to compute it.
         """
         if Engine.get() in ["Ray", "Dask"] and StorageFormat.get() == "Pandas":
             # Rebalancing partitions is currently only implemented for PandasOnRay and PandasOnDask.
@@ -1354,7 +1358,7 @@ def rebalance_partitions(cls, partitions):
                 num_existing_partitions
                 <= ideal_num_new_partitions * max_excess_of_num_partitions
             ):
-                return partitions
+                return partitions, None
             # If any partition has an unknown length, give each axis partition
             # roughly the same number of row partitions. We use `_length_cache` here
             # to avoid materializing any unmaterialized lengths.
@@ -1369,7 +1373,7 @@ def rebalance_partitions(cls, partitions):
                 chunk_size = compute_chunksize(
                     num_existing_partitions, ideal_num_new_partitions, min_block_size=1
                 )
-                return np.array(
+                new_partitions = np.array(
                     [
                         cls.column_partitions(
                             partitions[i : i + chunk_size],
@@ -1382,6 +1386,7 @@ def rebalance_partitions(cls, partitions):
                         )
                     ]
                 )
+                return new_partitions, None
 
             # If we know the number of rows in every partition, then we should try
             # instead to give each new partition roughly the same number of rows.
@@ -1413,9 +1418,8 @@ def rebalance_partitions(cls, partitions):
                 # new partition have length `ideal_partition_size`, and the second
                 # partition has the remainder.
                 if partition_size > ideal_partition_size * max_excess_of_num_partitions:
-                    new_last_partition_size = ideal_partition_size - sum(
-                        row[0].length() for row in partitions[start:stop]
-                    )
+                    prev_length = sum(row[0].length() for row in partitions[start:stop])
+                    new_last_partition_size = ideal_partition_size - prev_length
                     partitions = np.insert(
                         partitions,
                         stop + 1,
@@ -1425,10 +1429,20 @@ def rebalance_partitions(cls, partitions):
                         ],
                         0,
                     )
+                    # TODO: explicit `_length_cache` computing may be avoided after #4903 is merged
+                    for obj in partitions[stop + 1]:
+                        obj._length_cache = partition_size - (
+                            prev_length + new_last_partition_size
+                        )
+
                     partitions[stop, :] = [
                         obj.mask(slice(None, new_last_partition_size), slice(None))
                         for obj in partitions[stop]
                     ]
+                    # TODO: explicit `_length_cache` computing may be avoided after #4903 is merged
+                    for obj in partitions[stop]:
+                        obj._length_cache = new_last_partition_size
+
                     partition_size = ideal_partition_size
                 # The new virtual partitions are not `full_axis`, even if they
                 # happen to span all rows in the dataframe, because they are
@@ -1440,5 +1454,7 @@ def rebalance_partitions(cls, partitions):
                     cls.column_partitions(partitions[start : stop + 1], full_axis=False)
                 )
                 start = stop + 1
-            return np.array(new_partitions)
-        return partitions
+            new_partitions = np.array(new_partitions)
+            lengths = [part.length() for part in new_partitions[:, 0]]
+            return new_partitions, lengths
+        return partitions, None


Test patch: