Repository: modin-project/modin

Issue:
Add range-partitioning implementation for `df.resample()`
Requires #7117 to be finished first


Discussion:

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
@@ -2438,6 +2438,7 @@ def combine_and_apply(
             dtypes=new_dtypes,
         )
 
+    @lazy_metadata_decorator(apply_axis="both")
     def _apply_func_to_range_partitioning(
         self,
         key_columns,
@@ -2447,6 +2448,7 @@ def _apply_func_to_range_partitioning(
         data=None,
         data_key_columns=None,
         level=None,
+        shuffle_func_cls=ShuffleSortFunctions,
         **kwargs,
     ):
         """
@@ -2470,6 +2472,9 @@ def _apply_func_to_range_partitioning(
             Additional key columns from `data`. Will be combined with `key_columns`.
         level : list of ints or labels, optional
             Index level(s) to build the range partitioning for. Can't be specified along with `key_columns`.
+        shuffle_func_cls : cls, default: ShuffleSortFunctions
+            A class implementing ``modin.core.dataframe.pandas.utils.ShuffleFunctions`` to be used
+            as a shuffle function.
         **kwargs : dict
             Additional arguments to forward to the range builder function.
 
@@ -2573,7 +2578,7 @@ def _apply_func_to_range_partitioning(
         else:
             new_partitions = grouper._partitions
 
-        shuffling_functions = ShuffleSortFunctions(
+        shuffling_functions = shuffle_func_cls(
             grouper,
             key_columns,
             ascending[0] if is_list_like(ascending) else ascending,
diff --git a/modin/core/dataframe/pandas/dataframe/utils.py b/modin/core/dataframe/pandas/dataframe/utils.py
--- a/modin/core/dataframe/pandas/dataframe/utils.py
+++ b/modin/core/dataframe/pandas/dataframe/utils.py
@@ -19,7 +19,9 @@
 
 import numpy as np
 import pandas
+from pandas._libs.tslibs import to_offset
 from pandas.core.dtypes.common import is_list_like, is_numeric_dtype
+from pandas.core.resample import _get_timestamp_range_edges
 
 from modin.error_message import ErrorMessage
 from modin.utils import _inherit_docstrings
@@ -122,6 +124,10 @@ class ShuffleSortFunctions(ShuffleFunctions):
         The ideal number of new partitions.
     level : list of strings or ints, or None
         Index level(s) to use as a key. Can't be specified along with `columns`.
+    closed_on_right : bool, default: False
+        Whether to include the right limit in range-partitioning.
+            True:  bins[i - 1] < x <= bins[i]
+            False: bins[i - 1] <= x < bins[i]
     **kwargs : dict
         Additional keyword arguments.
     """
@@ -133,6 +139,7 @@ def __init__(
         ascending: Union[list, bool],
         ideal_num_new_partitions: int,
         level: Optional[list[Union[str, int]]] = None,
+        closed_on_right: bool = False,
         **kwargs: dict,
     ):
         self.frame_len = len(modin_frame)
@@ -142,6 +149,7 @@ def __init__(
         self.kwargs = kwargs.copy()
         self.level = level
         self.columns_info = None
+        self.closed_on_right = closed_on_right
 
     def sample_fn(self, partition: pandas.DataFrame) -> pandas.DataFrame:
         if self.level is not None:
@@ -159,11 +167,11 @@ def pivot_fn(self, samples: "list[pandas.DataFrame]") -> int:
         columns_info: "list[ColumnInfo]" = []
         number_of_groups = 1
         cols = []
-        for col in samples.columns:
+        for i, col in enumerate(samples.columns):
             num_pivots = int(self.ideal_num_new_partitions / number_of_groups)
             if num_pivots < 2 and len(columns_info):
                 break
-            column_val = samples[col].to_numpy()
+            column_val = samples[col]
             cols.append(col)
             is_numeric = is_numeric_dtype(column_val.dtype)
 
@@ -172,7 +180,13 @@ def pivot_fn(self, samples: "list[pandas.DataFrame]") -> int:
             pivots = self.pick_pivots_from_samples_for_sort(
                 column_val, num_pivots, method, key
             )
-            columns_info.append(ColumnInfo(col, pivots, is_numeric))
+            columns_info.append(
+                ColumnInfo(
+                    self.level[i] if self.level is not None else col,
+                    pivots,
+                    is_numeric,
+                )
+            )
             number_of_groups *= len(pivots) + 1
         self.columns_info = columns_info
         return number_of_groups
@@ -190,6 +204,7 @@ def split_fn(
             self.columns_info,
             self.ascending,
             keys_are_index_levels=self.level is not None,
+            closed_on_right=self.closed_on_right,
             **self.kwargs,
         )
 
@@ -270,10 +285,9 @@ def pick_samples_for_quantiles(
         probability = (1 / m) * np.log(num_partitions * length)
         return df.sample(frac=probability)
 
-    @classmethod
     def pick_pivots_from_samples_for_sort(
-        cls,
-        samples: np.ndarray,
+        self,
+        samples: pandas.Series,
         ideal_num_new_partitions: int,
         method: str = "linear",
         key: Optional[Callable] = None,
@@ -288,7 +302,7 @@ def pick_pivots_from_samples_for_sort(
 
         Parameters
         ----------
-        samples : np.ndarray
+        samples : pandas.Series
             The samples computed by ``get_partition_quantiles_for_sort``.
         ideal_num_new_partitions : int
             The ideal number of new partitions.
@@ -302,6 +316,7 @@ def pick_pivots_from_samples_for_sort(
         np.ndarray
             A list of overall quantiles.
         """
+        samples = samples.to_numpy()
         # We don't call `np.unique` on the samples, since if a quantile shows up in multiple
         # partition's samples, this is probably an indicator of skew in the dataset, and we
         # want our final partitions to take this into account.
@@ -313,7 +328,7 @@ def pick_pivots_from_samples_for_sort(
         # If we only desire 1 partition, we need to ensure that we're not trying to find quantiles
         # from an empty list of pivots.
         if len(quantiles) > 0:
-            return cls._find_quantiles(samples, quantiles, method)
+            return self._find_quantiles(samples, quantiles, method)
         return np.array([])
 
     @staticmethod
@@ -322,6 +337,7 @@ def split_partitions_using_pivots_for_sort(
         columns_info: "list[ColumnInfo]",
         ascending: bool,
         keys_are_index_levels: bool = False,
+        closed_on_right: bool = False,
         **kwargs: dict,
     ) -> "tuple[pandas.DataFrame, ...]":
         """
@@ -342,6 +358,10 @@ def split_partitions_using_pivots_for_sort(
             The ascending flag.
         keys_are_index_levels : bool, default: False
             Whether `columns_info` describes index levels or actual columns from `df`.
+        closed_on_right : bool, default: False
+            Whether to include the right limit in range-partitioning.
+                True:  bins[i - 1] < x <= bins[i]
+                False: bins[i - 1] <= x < bins[i]
         **kwargs : dict
             Additional keyword arguments.
 
@@ -402,7 +422,9 @@ def get_group(grp, key, df):
                 cols_to_digitize = cols_to_digitize.squeeze()
 
             if col_info.is_numeric:
-                groupby_col = np.digitize(cols_to_digitize, pivots)
+                groupby_col = np.digitize(
+                    cols_to_digitize, pivots, right=closed_on_right
+                )
                 # `np.digitize` returns results based off of the sort order of the pivots it is passed.
                 # When we only have one unique value in our pivots, `np.digitize` assumes that the pivots
                 # are sorted in ascending order, and gives us results based off of that assumption - so if
@@ -410,7 +432,11 @@ def get_group(grp, key, df):
                 if not ascending and len(np.unique(pivots)) == 1:
                     groupby_col = len(pivots) - groupby_col
             else:
-                groupby_col = np.searchsorted(pivots, cols_to_digitize, side="right")
+                groupby_col = np.searchsorted(
+                    pivots,
+                    cols_to_digitize,
+                    side="left" if closed_on_right else "right",
+                )
                 # Since np.searchsorted requires the pivots to be in ascending order, if we want to sort
                 # in descending order, we need to swap the new indices.
                 if not ascending:
@@ -479,6 +505,178 @@ def _index_to_df_zero_copy(
         return index_data
 
 
+@_inherit_docstrings(ShuffleSortFunctions)
+class ShuffleResample(ShuffleSortFunctions):
+    def __init__(
+        self,
+        modin_frame: "PandasDataframe",
+        columns: Union[str, list],
+        ascending: Union[list, bool],
+        ideal_num_new_partitions: int,
+        resample_kwargs: dict,
+        **kwargs: dict,
+    ):
+        resample_kwargs = resample_kwargs.copy()
+        rule = resample_kwargs.pop("rule")
+
+        if resample_kwargs["closed"] is None:
+            # this rule regarding the default value of 'closed' is inherited
+            # from pandas documentation for 'pandas.DataFrame.resample'
+            if rule in ("ME", "YE", "QE", "BME", "BA", "BQE", "W"):
+                resample_kwargs["closed"] = "right"
+            else:
+                resample_kwargs["closed"] = "left"
+
+        super().__init__(
+            modin_frame,
+            columns,
+            ascending,
+            ideal_num_new_partitions,
+            closed_on_right=resample_kwargs["closed"] == "right",
+            **kwargs,
+        )
+
+        resample_kwargs["freq"] = to_offset(rule)
+        self.resample_kwargs = resample_kwargs
+
+    @staticmethod
+    def pick_samples_for_quantiles(
+        df: pandas.DataFrame,
+        num_partitions: int,
+        length: int,
+    ) -> pandas.DataFrame:
+        # to build proper bins we need min and max timestamp of the whole DatetimeIndex,
+        # so computing it in each partition
+        return pandas.concat([df.min().to_frame().T, df.max().to_frame().T])
+
+    def pick_pivots_from_samples_for_sort(
+        self,
+        samples: np.ndarray,
+        ideal_num_new_partitions: int,
+        method: str = "linear",
+        key: Optional[Callable] = None,
+    ) -> np.ndarray:
+        if key is not None:
+            raise NotImplementedError(key)
+
+        max_value = samples.max()
+
+        first, last = _get_timestamp_range_edges(
+            samples.min(),
+            max_value,
+            self.resample_kwargs["freq"],
+            unit=samples.dt.unit,
+            closed=self.resample_kwargs["closed"],
+            origin=self.resample_kwargs["origin"],
+            offset=self.resample_kwargs["offset"],
+        )
+
+        all_bins = pandas.date_range(
+            start=first,
+            end=last,
+            freq=self.resample_kwargs["freq"],
+            ambiguous=True,
+            nonexistent="shift_forward",
+            unit=samples.dt.unit,
+        )
+
+        all_bins = self._adjust_bin_edges(
+            all_bins,
+            max_value,
+            freq=self.resample_kwargs["freq"],
+            closed=self.resample_kwargs["closed"],
+        )
+
+        # take pivot values with an even interval
+        step = 1 / ideal_num_new_partitions
+        bins = [
+            all_bins[int(len(all_bins) * i * step)]
+            for i in range(1, ideal_num_new_partitions)
+        ]
+        return bins
+
+    def _adjust_bin_edges(
+        self,
+        binner: pandas.DatetimeIndex,
+        end_timestamp,
+        freq,
+        closed,
+    ) -> pandas.DatetimeIndex:
+        """
+        Adjust bin edges.
+
+        This function was copied & simplified from ``pandas.core.resample.TimeGrouper._adjuct_bin_edges()``.
+
+        Parameters
+        ----------
+        binner : pandas.DatetimeIndex
+        end_timestamp : pandas.Timestamp
+        freq : str
+        closed : bool
+
+        Returns
+        -------
+        pandas.DatetimeIndex
+        """
+        # Some hacks for > daily data, see pandas-dev/pandas#1471, pandas-dev/pandas#1458, pandas-dev/pandas#1483
+
+        if freq.name not in ("BME", "ME", "W") and freq.name.split("-")[0] not in (
+            "BQE",
+            "BYE",
+            "QE",
+            "YE",
+            "W",
+        ):
+            return binner
+
+        # If the right end-point is on the last day of the month, roll forwards
+        # until the last moment of that day. Note that we only do this for offsets
+        # which correspond to the end of a super-daily period - "month start", for
+        # example, is excluded.
+        if closed == "right":
+            # GH 21459, GH 9119: Adjust the bins relative to the wall time
+            edges_dti = binner.tz_localize(None)
+            edges_dti = (
+                edges_dti
+                + pandas.Timedelta(days=1, unit=edges_dti.unit).as_unit(edges_dti.unit)
+                - pandas.Timedelta(1, unit=edges_dti.unit).as_unit(edges_dti.unit)
+            )
+            binner = edges_dti.tz_localize(binner.tz)
+
+        # intraday values on last day
+        if binner[-2] > end_timestamp:
+            binner = binner[:-1]
+        return binner
+
+    @staticmethod
+    def split_partitions_using_pivots_for_sort(
+        df: pandas.DataFrame,
+        columns_info: "list[ColumnInfo]",
+        ascending: bool,
+        closed_on_right: bool = True,
+        **kwargs: dict,
+    ) -> "tuple[pandas.DataFrame, ...]":
+        def add_attr(df, timestamp):
+            if "bin_bounds" in df.attrs:
+                df.attrs["bin_bounds"] = (*df.attrs["bin_bounds"], timestamp)
+            else:
+                df.attrs["bin_bounds"] = (timestamp,)
+            return df
+
+        result = ShuffleSortFunctions.split_partitions_using_pivots_for_sort(
+            df, columns_info, ascending, **kwargs
+        )
+        # it's required for each bin to know its bounds in order for resampling to work
+        # properly when down-sampling occurs. Reach here for an example:
+        # https://github.com/modin-project/modin/pull/7140#discussion_r1549246505
+        # We're writing the bounds as 'attrs' to avoid duplications in the final partition
+        for i, pivot in enumerate(columns_info[0].pivots):
+            add_attr(result[i], pivot - pandas.Timedelta(1, unit="ns"))
+            if i + 1 <= len(result):
+                add_attr(result[i + 1], pivot + pandas.Timedelta(1, unit="ns"))
+        return result
+
+
 def lazy_metadata_decorator(apply_axis=None, axis_arg=-1, transpose=False):
     """
     Lazily propagate metadata for the ``PandasDataframe``.
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
@@ -1023,7 +1023,14 @@ def to_datetime(self, *args, **kwargs):
     # END Reduce operations
 
     def _resample_func(
-        self, resample_kwargs, func_name, new_columns=None, df_op=None, *args, **kwargs
+        self,
+        resample_kwargs,
+        func_name,
+        new_columns=None,
+        df_op=None,
+        allow_range_impl=True,
+        *args,
+        **kwargs,
     ):
         """
         Resample underlying time-series data and apply aggregation on it.
@@ -1039,6 +1046,8 @@ def _resample_func(
             Modin frame. If not specified will be computed automaticly.
         df_op : callable(pandas.DataFrame) -> [pandas.DataFrame, pandas.Series], optional
             Preprocessor function to apply to the passed frame before resampling.
+        allow_range_impl : bool, default: True
+            Whether to use range-partitioning if ``RangePartitioning.get() is True``.
         *args : args
             Arguments to pass to the aggregation function.
         **kwargs : kwargs
@@ -1049,9 +1058,35 @@ def _resample_func(
         PandasQueryCompiler
             New QueryCompiler containing the result of resample aggregation.
         """
+        from modin.core.dataframe.pandas.dataframe.utils import ShuffleResample
 
         def map_func(df, resample_kwargs=resample_kwargs):  # pragma: no cover
             """Resample time-series data of the passed frame and apply aggregation function on it."""
+            if len(df) == 0:
+                if resample_kwargs["on"] is not None:
+                    df = df.set_index(resample_kwargs["on"])
+                return df
+            if "bin_bounds" in df.attrs:
+                timestamps = df.attrs["bin_bounds"]
+                if isinstance(df.index, pandas.MultiIndex):
+                    level_to_keep = resample_kwargs["level"]
+                    if isinstance(level_to_keep, int):
+                        to_drop = [
+                            lvl
+                            for lvl in range(df.index.nlevels)
+                            if lvl != level_to_keep
+                        ]
+                    else:
+                        to_drop = [
+                            lvl for lvl in df.index.names if lvl != level_to_keep
+                        ]
+                    df.index = df.index.droplevel(to_drop)
+                    resample_kwargs = resample_kwargs.copy()
+                    resample_kwargs["level"] = None
+                filler = pandas.DataFrame(
+                    np.NaN, index=pandas.Index(timestamps), columns=df.columns
+                )
+                df = pandas.concat([df, filler], copy=False)
             if df_op is not None:
                 df = df_op(df)
             resampled_val = df.resample(**resample_kwargs)
@@ -1072,13 +1107,37 @@ def map_func(df, resample_kwargs=resample_kwargs):  # pragma: no cover
             else:
                 return val
 
-        new_modin_frame = self._modin_frame.apply_full_axis(
-            axis=0, func=map_func, new_columns=new_columns
-        )
+        if resample_kwargs["on"] is None:
+            level = [
+                0 if resample_kwargs["level"] is None else resample_kwargs["level"]
+            ]
+            key_columns = []
+        else:
+            level = None
+            key_columns = [resample_kwargs["on"]]
+
+        if (
+            not allow_range_impl
+            or resample_kwargs["axis"] not in (0, "index")
+            or not RangePartitioning.get()
+        ):
+            new_modin_frame = self._modin_frame.apply_full_axis(
+                axis=0, func=map_func, new_columns=new_columns
+            )
+        else:
+            new_modin_frame = self._modin_frame._apply_func_to_range_partitioning(
+                key_columns=key_columns,
+                level=level,
+                func=map_func,
+                shuffle_func_cls=ShuffleResample,
+                resample_kwargs=resample_kwargs,
+            )
         return self.__constructor__(new_modin_frame)
 
     def resample_get_group(self, resample_kwargs, name, obj):
-        return self._resample_func(resample_kwargs, "get_group", name=name, obj=obj)
+        return self._resample_func(
+            resample_kwargs, "get_group", name=name, allow_range_impl=False, obj=obj
+        )
 
     def resample_app_ser(self, resample_kwargs, func, *args, **kwargs):
         return self._resample_func(
@@ -1110,24 +1169,39 @@ def resample_agg_df(self, resample_kwargs, func, *args, **kwargs):
 
     def resample_transform(self, resample_kwargs, arg, *args, **kwargs):
         return self._resample_func(
-            resample_kwargs, "transform", arg=arg, *args, **kwargs
+            resample_kwargs,
+            "transform",
+            arg=arg,
+            allow_range_impl=False,
+            *args,
+            **kwargs,
         )
 
     def resample_pipe(self, resample_kwargs, func, *args, **kwargs):
         return self._resample_func(resample_kwargs, "pipe", func=func, *args, **kwargs)
 
     def resample_ffill(self, resample_kwargs, limit):
-        return self._resample_func(resample_kwargs, "ffill", limit=limit)
+        return self._resample_func(
+            resample_kwargs, "ffill", limit=limit, allow_range_impl=False
+        )
 
     def resample_bfill(self, resample_kwargs, limit):
-        return self._resample_func(resample_kwargs, "bfill", limit=limit)
+        return self._resample_func(
+            resample_kwargs, "bfill", limit=limit, allow_range_impl=False
+        )
 
     def resample_nearest(self, resample_kwargs, limit):
-        return self._resample_func(resample_kwargs, "nearest", limit=limit)
+        return self._resample_func(
+            resample_kwargs, "nearest", limit=limit, allow_range_impl=False
+        )
 
     def resample_fillna(self, resample_kwargs, method, limit):
         return self._resample_func(
-            resample_kwargs, "fillna", method=method, limit=limit
+            resample_kwargs,
+            "fillna",
+            method=method,
+            limit=limit,
+            allow_range_impl=method is None,
         )
 
     def resample_asfreq(self, resample_kwargs, fill_value):
@@ -1154,6 +1228,7 @@ def resample_interpolate(
             limit_direction=limit_direction,
             limit_area=limit_area,
             downcast=downcast,
+            allow_range_impl=False,
             **kwargs,
         )
 
@@ -1164,16 +1239,20 @@ def resample_nunique(self, resample_kwargs, *args, **kwargs):
         return self._resample_func(resample_kwargs, "nunique", *args, **kwargs)
 
     def resample_first(self, resample_kwargs, *args, **kwargs):
-        return self._resample_func(resample_kwargs, "first", *args, **kwargs)
+        return self._resample_func(
+            resample_kwargs, "first", allow_range_impl=False, *args, **kwargs
+        )
 
     def resample_last(self, resample_kwargs, *args, **kwargs):
-        return self._resample_func(resample_kwargs, "last", *args, **kwargs)
+        return self._resample_func(
+            resample_kwargs, "last", allow_range_impl=False, *args, **kwargs
+        )
 
     def resample_max(self, resample_kwargs, *args, **kwargs):
         return self._resample_func(resample_kwargs, "max", *args, **kwargs)
 
     def resample_mean(self, resample_kwargs, *args, **kwargs):
-        return self._resample_func(resample_kwargs, "median", *args, **kwargs)
+        return self._resample_func(resample_kwargs, "mean", *args, **kwargs)
 
     def resample_median(self, resample_kwargs, *args, **kwargs):
         return self._resample_func(resample_kwargs, "median", *args, **kwargs)
@@ -1204,7 +1283,10 @@ def resample_prod(self, resample_kwargs, min_count, *args, **kwargs):
 
     def resample_size(self, resample_kwargs):
         return self._resample_func(
-            resample_kwargs, "size", new_columns=[MODIN_UNNAMED_SERIES_LABEL]
+            resample_kwargs,
+            "size",
+            new_columns=[MODIN_UNNAMED_SERIES_LABEL],
+            allow_range_impl=False,
         )
 
     def resample_sem(self, resample_kwargs, *args, **kwargs):


Test patch:
diff --git a/modin/pandas/test/dataframe/test_default.py b/modin/pandas/test/dataframe/test_default.py
--- a/modin/pandas/test/dataframe/test_default.py
+++ b/modin/pandas/test/dataframe/test_default.py
@@ -1003,7 +1003,10 @@ def test_resample_specific(rule, closed, label, on, level):
 
     if on is None and level is not None:
         index = pandas.MultiIndex.from_product(
-            [["a", "b", "c"], pandas.date_range("31/12/2000", periods=4, freq="h")]
+            [
+                ["a", "b", "c", "d"],
+                pandas.date_range("31/12/2000", periods=len(pandas_df) // 4, freq="h"),
+            ]
         )
         pandas_df.index = index
         modin_df.index = index
@@ -1011,8 +1014,12 @@ def test_resample_specific(rule, closed, label, on, level):
         level = None
 
     if on is not None:
-        pandas_df[on] = pandas.date_range("22/06/1941", periods=12, freq="min")
-        modin_df[on] = pandas.date_range("22/06/1941", periods=12, freq="min")
+        pandas_df[on] = pandas.date_range(
+            "22/06/1941", periods=len(pandas_df), freq="min"
+        )
+        modin_df[on] = pandas.date_range(
+            "22/06/1941", periods=len(modin_df), freq="min"
+        )
 
     pandas_resampler = pandas_df.resample(
         rule,
diff --git a/modin/pandas/test/utils.py b/modin/pandas/test/utils.py
--- a/modin/pandas/test/utils.py
+++ b/modin/pandas/test/utils.py
@@ -175,8 +175,11 @@
 test_groupby_data = {f"col{i}": np.arange(NCOLS) % NGROUPS for i in range(NROWS)}
 
 test_data_resample = {
-    "data": {"A": range(12), "B": range(12)},
-    "index": pandas.date_range("31/12/2000", periods=12, freq="h"),
+    "data": {
+        f"col{i}": random_state.randint(RAND_LOW, RAND_HIGH, size=NROWS)
+        for i in range(10)
+    },
+    "index": pandas.date_range("31/12/2000", periods=NROWS, freq="h"),
 }
 
 test_data_with_duplicates = {