PERF-modin-project#7132: Preserve partition lengths in apply_full_axi…

…s with keep_partitioning=False (modin-project#7133)

Signed-off-by: Iaroslav Igoshev <iaroslav.igoshev@intel.com>
Co-authored-by: Dmitry Chigarev <dmitry.chigarev@intel.com>

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

@@ -3439,11 +3439,20 @@ def broadcast_apply_full_axis(
                             index=new_columns,
                         )
                     )
-
-        if not keep_partitioning:
-            if kw["row_lengths"] is None and ModinIndex.is_materialized_index(
-                new_index
+        is_index_materialized = ModinIndex.is_materialized_index(new_index)
+        is_columns_materialized = ModinIndex.is_materialized_index(new_columns)
+        if axis == 0:
+            if (
+                is_columns_materialized
+                and len(new_partitions.shape) > 1
+                and new_partitions.shape[1] == 1
             ):
+                kw["column_widths"] = [len(new_columns)]
+        elif axis == 1:
+            if is_index_materialized and new_partitions.shape[0] == 1:
+                kw["row_lengths"] = [len(new_index)]
+        if not keep_partitioning:
+            if kw["row_lengths"] is None and is_index_materialized:
                 if axis == 0:
                     kw["row_lengths"] = get_length_list(
                         axis_len=len(new_index), num_splits=new_partitions.shape[0]
@@ -3453,11 +3462,7 @@ def broadcast_apply_full_axis(
                         self._row_lengths_cache
                     ):
                         kw["row_lengths"] = self._row_lengths_cache
-                    elif len(new_index) == 1 and new_partitions.shape[0] == 1:
-                        kw["row_lengths"] = [1]
-            if kw["column_widths"] is None and ModinIndex.is_materialized_index(
-                new_columns
-            ):
+            if kw["column_widths"] is None and is_columns_materialized:
                 if axis == 1:
                     kw["column_widths"] = get_length_list(
                         axis_len=len(new_columns),
@@ -3468,42 +3473,27 @@ def broadcast_apply_full_axis(
                         new_columns
                     ) == sum(self._column_widths_cache):
                         kw["column_widths"] = self._column_widths_cache
-                    elif len(new_columns) == 1 and new_partitions.shape[1] == 1:
-                        kw["column_widths"] = [1]
         else:
             if axis == 0:
                 if (
                     kw["row_lengths"] is None
                     and self._row_lengths_cache is not None
-                    and ModinIndex.is_materialized_index(new_index)
+                    and is_index_materialized
                     and len(new_index) == sum(self._row_lengths_cache)
                     # to avoid problems that may arise when filtering empty dataframes
                     and all(r != 0 for r in self._row_lengths_cache)
                 ):
                     kw["row_lengths"] = self._row_lengths_cache
-                if (
-                    kw["column_widths"] is None
-                    and ModinIndex.is_materialized_index(new_columns)
-                    and len(new_partitions.shape) > 1
-                    and new_partitions.shape[1] == 1
-                ):
-                    kw["column_widths"] = [len(new_columns)]
-            if axis == 1:
+            elif axis == 1:
                 if (
                     kw["column_widths"] is None
                     and self._column_widths_cache is not None
-                    and ModinIndex.is_materialized_index(new_columns)
+                    and is_columns_materialized
                     and len(new_columns) == sum(self._column_widths_cache)
                     # to avoid problems that may arise when filtering empty dataframes
                     and all(w != 0 for w in self._column_widths_cache)
                 ):
                     kw["column_widths"] = self._column_widths_cache
-                if (
-                    kw["row_lengths"] is None
-                    and ModinIndex.is_materialized_index(new_index)
-                    and new_partitions.shape[0] == 1
-                ):
-                    kw["row_lengths"] = [len(new_index)]
 
         result = self.__constructor__(
             new_partitions, index=new_index, columns=new_columns, **kw