PERF-#6556: do window block-wise

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

@@ -2232,8 +2232,8 @@ def map(
     def window(
         self,
         axis: Union[int, Axis],
-        reduce_fn: Callable,
         window_size: int,
+        reduce_fn: Callable,
         result_schema: Optional[Dict[Hashable, type]] = None,
     ) -> "PandasDataframe":
         """
@@ -2244,9 +2244,9 @@ def window(
         axis : int or modin.core.dataframe.base.utils.Axis
             The axis to slide over.
         reduce_fn : callable(rowgroup|colgroup) -> row|col
-            The reduce function to apply over the data.
+            The reduce function to apply over the data. This must not change metadata.
         window_size : int
-            The number of row/columns to pass to the function.
+            The number of rows/columns to pass to the reduce function.
             (The size of the sliding window).
         result_schema : dict, optional
             Mapping from column labels to data types that represents the types of the output dataframe.
@@ -2256,13 +2256,136 @@ def window(
         PandasDataframe
             A new PandasDataframe with the reduce function applied over windows of the specified
                 axis.
-
-        Notes
-        -----
-        The user-defined reduce function must reduce each window’s column
-        (row if axis=1) down to a single value.
         """
-        pass
+
+        axis = Axis(axis)
+
+        # applies reduction function over entire virtual partition
+        def window_function_complete(virtual_partition):
+            # have to copy the pandas dataframe on ray because it's immutable
+            virtual_partition_copy = virtual_partition.copy()
+            window_result = reduce_fn(virtual_partition_copy)
+            return window_result
+
+        # applies reduction function over one window of virtual partition
+        def window_function_partition(virtual_partition):
+            virtual_partition_copy = virtual_partition.copy()
+            window_result = reduce_fn(virtual_partition_copy)
+            return (
+                window_result.iloc[:, window_size - 1 :]
+                if axis == Axis.COL_WISE
+                else window_result.iloc[window_size - 1 :, :]
+            )
+
+        num_parts = (
+            len(self._partitions[0]) if axis == Axis.COL_WISE else len(self._partitions)
+        )
+        results = []
+
+        for i in range(num_parts):
+            # get the ith partition
+            starting_part = (
+                self._partitions[:, [i]]
+                if axis == Axis.COL_WISE
+                else self._partitions[i]
+            )
+
+            # partitions to join in virtual partition
+            parts_to_join = (
+                [starting_part]
+                if (axis == Axis.ROW_WISE)
+                else [[partition[0]] for partition in starting_part]
+            )
+
+            # used to determine if window continues into next partition or if we can create virtual partition
+            last_window_span = window_size - 1
+            k = i + 1
+
+            while last_window_span > 0 and k < num_parts:
+                # new partition
+                new_parts = (
+                    self._partitions[:, [k]]
+                    if axis == Axis.COL_WISE
+                    else self._partitions[k]
+                )
+                part_len = (
+                    new_parts[0][0].width()
+                    if axis == Axis.COL_WISE
+                    else new_parts[0].length()
+                )
+
+                if last_window_span <= part_len:
+                    if axis == Axis.COL_WISE:
+                        masked_new_parts = [
+                            [
+                                part[0].mask(
+                                    row_labels=slice(None),
+                                    col_labels=slice(0, last_window_span),
+                                )
+                            ]
+                            for part in new_parts
+                        ]
+                        for x, r in enumerate(parts_to_join):
+                            r.append(masked_new_parts[x][0])
+                    else:
+                        masked_new_parts = np.array(
+                            [
+                                part.mask(
+                                    row_labels=slice(0, last_window_span),
+                                    col_labels=slice(None),
+                                )
+                                for part in new_parts
+                            ]
+                        )
+                        parts_to_join.append(masked_new_parts)
+                    break
+                else:
+                    # window continues into next part, so just add this part to parts_to_join
+                    if axis == Axis.COL_WISE:
+                        for x, r in enumerate(parts_to_join):
+                            r.append(new_parts[x][0])
+                    else:
+                        parts_to_join.append(new_parts)
+                    last_window_span -= part_len
+                    k += 1
+
+            # create virtual partition and perform window operation
+            virtual_partitions = (
+                self._partition_mgr_cls.row_partitions(
+                    np.array(parts_to_join), full_axis=False
+                )
+                if axis == Axis.COL_WISE
+                else self._partition_mgr_cls.column_partitions(
+                    np.array(parts_to_join), full_axis=False
+                )
+            )
+
+            if i == 0:
+                reduce_result = [
+                    virtual_partition.apply(window_function_complete)
+                    for virtual_partition in virtual_partitions
+                ]
+            else:
+                reduce_result = [
+                    virtual_partition.apply(window_function_partition)
+                    for virtual_partition in virtual_partitions
+                ]
+
+            # append reduction result to results array
+            if axis == Axis.ROW_WISE:
+                results.append(reduce_result)
+            else:
+                if results == []:
+                    results = [[x] for x in reduce_result]
+                else:
+                    for x, r in enumerate(results):
+                        r.append(reduce_result[x])
+
+        results = np.array(results)
+
+        return self.__constructor__(
+             np.array(results), self.index, self.columns, None, None, result_schema
+        )
 
     @lazy_metadata_decorator(apply_axis="both")
     def fold(self, axis, func, new_columns=None):

modin/core/dataframe/pandas/partitioning/partition_manager.py

@@ -229,7 +229,7 @@ def column_partitions(cls, partitions, full_axis=True):
         ]
 
     @classmethod
-    def row_partitions(cls, partitions):
+    def row_partitions(cls, partitions, full_axis=True):
         """
         List of `BaseDataframeAxisPartition` objects representing row-wise partitions.
 
@@ -250,7 +250,11 @@ def row_partitions(cls, partitions):
         """
         if not isinstance(partitions, list):
             partitions = [partitions]
-        return [cls._row_partition_class(row) for frame in partitions for row in frame]
+        return [
+            cls._row_partition_class(row, full_axis=full_axis) 
+            for frame in partitions 
+            for row in frame
+        ]
 
     @classmethod
     def axis_partition(cls, partitions, axis, full_axis: bool = True):