Repository: modin-project/modin

Issue:
Automatically pick a key column with the most unique values when building range-partitioning
When making a range-partitioning, modin currently chooses a key column to be the first in a list that was specified to group/sort on. The amount of parallelism in range-partitioning directly depends on the amount of unique values in a key column, that said, picking a random column for the matter of a key column may not always be a good idea.

It's proposed to sample multiple columns at the [sampling stage](https://github.com/modin-project/modin/blob/f994acacb7cc0f375c2fb22f1cb879713b9e0541/modin/core/dataframe/pandas/dataframe/utils.py#L127) and then pick one that has more unique values.

Here's an example of how picking the right column as a key can help to improve performance:

```python
import modin.pandas as pd
import numpy as np
from timeit import default_timer as timer


import modin.config as cfg
cfg.CpuCount.put(16)
cfg.ExperimentalGroupbyImpl.put(True)
cfg.BenchmarkMode.put(True)

pd.DataFrame([np.arange(cfg.MinPartitionSize.get()) for _ in range(cfg.NPartitions.get() ** 2)]).to_numpy()

NROWS = 10_000_000

df = pd.DataFrame({
    "by_col1": np.tile(np.arange(2), NROWS // 2),
    "by_col2": np.tile(np.arange(10_000), NROWS // 10_000),
    "col3": np.random.randint(0, NROWS, NROWS),
    "col4": np.random.randint(0, NROWS, NROWS),
    "col5": np.random.randint(0, NROWS, NROWS),
    "col6": np.random.randint(0, NROWS, NROWS),
    "col7": np.random.randint(0, NROWS, NROWS),
    "col8": np.random.randint(0, NROWS, NROWS),
    "col9": np.random.randint(0, NROWS, NROWS),
})

t1 = timer()
# the MORE unique column comes first and will be used for building the range-partitioning
df.groupby(["by_col2", "by_col1"]).median()
print(timer() - t1) # ~1.14s

t1 = timer()
# the LESS unique column comes first and will be used for building the range-partitioning
df.groupby(["by_col1", "by_col2"]).median()
print(timer() - t1) # ~2.2s

```


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
@@ -31,7 +31,7 @@
 from pandas._libs.lib import no_default
 from typing import List, Hashable, Optional, Callable, Union, Dict, TYPE_CHECKING
 
-from modin.config import Engine, IsRayCluster
+from modin.config import Engine, IsRayCluster, NPartitions
 from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
 from modin.core.storage_formats.pandas.utils import get_length_list
 from modin.error_message import ErrorMessage
@@ -44,7 +44,7 @@
     JoinType,
 )
 from modin.core.dataframe.pandas.dataframe.utils import (
-    build_sort_functions,
+    ShuffleSortFunctions,
     lazy_metadata_decorator,
 )
 from modin.core.dataframe.pandas.metadata import (
@@ -2386,15 +2386,15 @@ def combine_and_apply(
         )
 
     def _apply_func_to_range_partitioning(
-        self, key_column, func, ascending=True, **kwargs
+        self, key_columns, func, ascending=True, **kwargs
     ):
         """
         Reshuffle data so it would be range partitioned and then apply the passed function row-wise.
 
         Parameters
         ----------
-        key_column : hashable
-            Column name to build the range partitioning for.
+        key_columns : list of hashables
+            Columns to build the range partitioning for.
         func : callable(pandas.DataFrame) -> pandas.DataFrame
             Function to apply against partitions.
         ascending : bool, default: True
@@ -2411,7 +2411,8 @@ def _apply_func_to_range_partitioning(
         if self._partitions.shape[0] == 1:
             return self.apply_full_axis(axis=1, func=func)
 
-        ideal_num_new_partitions = len(self._partitions)
+        # don't want to inherit over-partitioning so doing this 'min' check
+        ideal_num_new_partitions = min(len(self._partitions), NPartitions.get())
         m = len(self.index) / ideal_num_new_partitions
         sampling_probability = (1 / m) * np.log(
             ideal_num_new_partitions * len(self.index)
@@ -2432,56 +2433,48 @@ def _apply_func_to_range_partitioning(
                 # simply combine all partitions and apply the sorting to the whole dataframe
                 return self.combine_and_apply(func=func)
 
-        if is_numeric_dtype(self.dtypes[key_column]):
-            method = "linear"
+            if ideal_num_new_partitions < len(self._partitions):
+                if len(self._partitions) % ideal_num_new_partitions == 0:
+                    joining_partitions = np.split(
+                        self._partitions, ideal_num_new_partitions
+                    )
+                else:
+                    step = round(len(self._partitions) / ideal_num_new_partitions)
+                    joining_partitions = np.split(
+                        self._partitions,
+                        range(step, len(self._partitions), step),
+                    )
+
+                new_partitions = np.array(
+                    [
+                        self._partition_mgr_cls.column_partitions(
+                            ptn_grp, full_axis=False
+                        )
+                        for ptn_grp in joining_partitions
+                    ]
+                )
+            else:
+                new_partitions = self._partitions
         else:
-            # This means we are not sorting numbers, so we need our quantiles to not try
-            # arithmetic on the values.
-            method = "inverted_cdf"
+            new_partitions = self._partitions
 
-        shuffling_functions = build_sort_functions(
+        shuffling_functions = ShuffleSortFunctions(
             self,
-            key_column,
-            method,
+            key_columns,
             ascending[0] if is_list_like(ascending) else ascending,
             ideal_num_new_partitions,
             **kwargs,
         )
-        if ideal_num_new_partitions < len(self._partitions):
-            if len(self._partitions) % ideal_num_new_partitions == 0:
-                joining_partitions = np.split(
-                    self._partitions, ideal_num_new_partitions
-                )
-            else:
-                joining_partitions = np.split(
-                    self._partitions,
-                    range(
-                        0,
-                        len(self._partitions),
-                        round(len(self._partitions) / ideal_num_new_partitions),
-                    )[1:],
-                )
 
-            new_partitions = np.array(
-                [
-                    self._partition_mgr_cls.column_partitions(ptn_grp, full_axis=False)
-                    for ptn_grp in joining_partitions
-                ]
-            )
-        else:
-            new_partitions = self._partitions
+        # here we want to get indices of those partitions that hold the key columns
+        key_indices = self.columns.get_indexer_for(key_columns)
+        partition_indices = np.unique(
+            np.digitize(key_indices, np.cumsum(self.column_widths))
+        )
 
-        major_col_partition_index = self.columns.get_loc(key_column)
-        cols_seen = 0
-        index = -1
-        for i, length in enumerate(self.column_widths):
-            cols_seen += length
-            if major_col_partition_index < cols_seen:
-                index = i
-                break
         new_partitions = self._partition_mgr_cls.shuffle_partitions(
             new_partitions,
-            index,
+            partition_indices,
             shuffling_functions,
             func,
         )
@@ -2543,7 +2536,7 @@ def sort_function(df):  # pragma: no cover
             )
 
         result = self._apply_func_to_range_partitioning(
-            key_column=columns[0], func=sort_function, ascending=ascending, **kwargs
+            key_columns=[columns[0]], func=sort_function, ascending=ascending, **kwargs
         )
 
         result.set_axis_cache(self.copy_axis_cache(axis.value ^ 1), axis=axis.value ^ 1)
@@ -3751,7 +3744,7 @@ def apply_func(df):  # pragma: no cover
             return result
 
         result = self._apply_func_to_range_partitioning(
-            key_column=by[0],
+            key_columns=by,
             func=apply_func,
         )
 
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
@@ -17,287 +17,411 @@
 import pandas
 from typing import Callable, Union, Optional, TYPE_CHECKING
 from collections import namedtuple
+from pandas.core.dtypes.common import is_numeric_dtype, is_list_like
+import abc
+
+from modin.error_message import ErrorMessage
+from modin.utils import _inherit_docstrings
 
 if TYPE_CHECKING:
     from modin.core.dataframe.pandas.dataframe.dataframe import PandasDataframe
 
-
-ShuffleFunctions = namedtuple(
-    "ShuffleFunctions", ["sample_function", "pivot_function", "split_function"]
-)
+ColumnInfo = namedtuple("ColumnInfo", ["name", "pivots", "is_numeric"])
 
 
-def build_sort_functions(
-    modin_frame: "PandasDataframe",
-    column: str,
-    method: str,
-    ascending: Union[list, bool],
-    ideal_num_new_partitions: int,
-    **kwargs: dict,
-) -> ShuffleFunctions:
+class ShuffleFunctions:
     """
-    Return a named tuple containing the functions necessary to perform a sort.
+    Defines an interface to perform the sampling, quantiles picking, and the splitting stages for the range-partitioning building.
 
     Parameters
     ----------
     modin_frame : PandasDataframe
-        The frame calling these sort functions.
-    column : str
-        The major column name to sort by.
-    method : str
-        The method to use for picking quantiles.
+        The frame to build the range-partitioning for.
+    columns : str or list of strings
+        The column/columns to use as a key.
     ascending : bool
-        The ascending flag.
+        Whether the ranges should be in ascending or descending order.
     ideal_num_new_partitions : int
         The ideal number of new partitions.
     **kwargs : dict
         Additional keyword arguments.
-
-    Returns
-    -------
-    ShuffleFunctions :
-        A named tuple containing the functions to pick quantiles, choose pivot points, and split
-        partitions for sorting.
-    """
-    frame_len = len(modin_frame.index)
-    is_column_numeric = pandas.api.types.is_numeric_dtype(modin_frame.dtypes[column])
-
-    def sample_fn(partition):
-        return pick_samples_for_quantiles(
-            partition[column], ideal_num_new_partitions, frame_len
-        )
-
-    def pivot_fn(samples):
-        key = kwargs.get("key", None)
-        return pick_pivots_from_samples_for_sort(
-            samples, ideal_num_new_partitions, method, key
-        )
-
-    def split_fn(partition, pivots):
-        return split_partitions_using_pivots_for_sort(
-            partition, column, is_column_numeric, pivots, ascending, **kwargs
-        )
-
-    return ShuffleFunctions(
-        sample_function=sample_fn, pivot_function=pivot_fn, split_function=split_fn
-    )
-
-
-def _find_quantiles(
-    df: Union[pandas.DataFrame, pandas.Series], quantiles: list, method: str
-) -> np.ndarray:
-    """
-    Find quantiles of a given dataframe using the specified method.
-
-    We use this method to provide backwards compatibility with NumPy versions < 1.23 (e.g. when
-    the user is using Modin in compat mode). This is basically a wrapper around `np.quantile` that
-    ensures we provide the correct `method` argument - i.e. if we are dealing with objects (which
-    may or may not support algebra), we do not want to use a method to find quantiles that will
-    involve algebra operations (e.g. mean) between the objects, since that may fail.
-
-    Parameters
-    ----------
-    df : pandas.DataFrame or pandas.Series
-        The data to pick quantiles from.
-    quantiles : list[float]
-        The quantiles to compute.
-    method : str
-        The method to use. `linear` if dealing with numeric types, otherwise `inverted_cdf`.
-
-    Returns
-    -------
-    np.ndarray
-        A NumPy array with the quantiles of the data.
-    """
-    if method == "linear":
-        # This is the default method for finding quantiles, so it does not need to be specified,
-        # which keeps backwards compatibility with older versions of NumPy that do not have a
-        # `method` keyword argument in np.quantile.
-        return np.quantile(df, quantiles)
-    else:
-        try:
-            return np.quantile(df, quantiles, method=method)
-        except Exception:
-            # In this case, we're dealing with an array of objects, but the current version of
-            # NumPy does not have a `method` kwarg. We need to use the older kwarg, `interpolation`
-            # instead.
-            return np.quantile(df, quantiles, interpolation="lower")
-
-
-def pick_samples_for_quantiles(
-    df: pandas.DataFrame,
-    num_partitions: int,
-    length: int,
-) -> np.ndarray:
-    """
-    Pick samples over the given partition.
-
-    This function picks samples from the given partition using the TeraSort algorithm - each
-    value is sampled with probability 1 / m * ln(n * t) where m = total_length / num_partitions,
-    t = num_partitions, and n = total_length.
-
-    Parameters
-    ----------
-    df : pandas.Dataframe
-        The masked dataframe to pick samples from.
-    num_partitions : int
-        The number of partitions.
-    length : int
-        The total length.
-
-    Returns
-    -------
-    np.ndarray:
-        The samples for the partition.
-
-    Notes
-    -----
-    This sampling algorithm is inspired by TeraSort. You can find more information about TeraSort
-    and the sampling algorithm at https://www.cse.cuhk.edu.hk/~taoyf/paper/sigmod13-mr.pdf.
     """
-    m = length / num_partitions
-    probability = (1 / m) * np.log(num_partitions * length)
-    return df.sample(frac=probability).to_numpy()
 
-
-def pick_pivots_from_samples_for_sort(
-    samples: "list[np.ndarray]",
-    ideal_num_new_partitions: int,
-    method: str = "linear",
-    key: Optional[Callable] = None,
-) -> np.ndarray:
+    def __init__(
+        self, modin_frame, columns, ascending, ideal_num_new_partitions, **kwargs
+    ):
+        pass
+
+    @abc.abstractmethod
+    def sample_fn(self, partition: pandas.DataFrame) -> pandas.DataFrame:
+        """
+        Pick samples over the given partition.
+
+        Parameters
+        ----------
+        partition : pandas.DataFrame
+
+        Returns
+        -------
+        pandas.DataFrame:
+            The samples for the partition.
+        """
+        pass
+
+    @abc.abstractmethod
+    def pivot_fn(self, samples: "list[pandas.DataFrame]") -> int:
+        """
+        Determine quantiles from the given samples and save it for the future ``.split_fn()`` calls.
+
+        Parameters
+        ----------
+        samples : list of pandas.DataFrames
+
+        Returns
+        -------
+        int
+            The number of bins the ``.split_fn()`` will return.
+        """
+        pass
+
+    @abc.abstractmethod
+    def split_fn(self, partition: pandas.DataFrame) -> "tuple[pandas.DataFrame, ...]":
+        """
+        Split the given dataframe into the range-partitions defined by the preceding call of the ``.pivot_fn()``.
+
+        Parameters
+        ----------
+        partition : pandas.DataFrame
+
+        Returns
+        -------
+        tuple of pandas.DataFrames
+
+        Notes
+        -----
+        In order to call this method you must call the ``.pivot_fn()`` first.
+        """
+        pass
+
+
+@_inherit_docstrings(ShuffleFunctions)
+class ShuffleSortFunctions(ShuffleFunctions):
     """
-    Determine quantiles from the given samples.
-
-    This function takes as input the quantiles calculated over all partitions from
-    `sample_func` defined above, and determines a final NPartitions.get() quantiles
-    to use to roughly sort the entire dataframe. It does so by collating all the samples
-    and computing NPartitions.get() quantiles for the overall set.
+    Perform the sampling, quantiles picking, and the splitting stages for the range-partitioning building.
 
     Parameters
     ----------
-    samples : list[np.ndarray]
-        The samples computed by ``get_partition_quantiles_for_sort``.
+    modin_frame : PandasDataframe
+        The frame to build the range-partitioning for.
+    columns : str or list of strings
+        The column/columns to use as a key.
+    ascending : bool
+        Whether the ranges should be in ascending or descending order.
     ideal_num_new_partitions : int
         The ideal number of new partitions.
-    method : str, default: linear
-        The method to use when picking quantiles.
-    key : Callable, default: None
-        The key to use on the samples when picking pivots.
-
-    Returns
-    -------
-    np.ndarray
-        A list of overall quantiles.
-    """
-    # 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.
-    # We need to use numpy to concatenate the samples since the sample from each partition is
-    # a NumPy array, and we want one flattened array of samples.
-    all_pivots = np.concatenate(samples).flatten()
-    if key is not None:
-        all_pivots = key(all_pivots)
-    # We don't want to pick very many quantiles if we have a very small dataframe.
-    num_quantiles = ideal_num_new_partitions
-    quantiles = [i / num_quantiles for i in range(1, num_quantiles)]
-    # 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 _find_quantiles(all_pivots, quantiles, method)
-    return np.array([])
-
-
-def split_partitions_using_pivots_for_sort(
-    df: pandas.DataFrame,
-    column: str,
-    is_numeric_column: bool,
-    pivots: np.ndarray,
-    ascending: bool,
-    **kwargs: dict,
-) -> "tuple[pandas.DataFrame, ...]":
-    """
-    Split the given dataframe into the partitions specified by `pivots`.
-
-    This function takes as input a row-axis partition, as well as the quantiles determined
-    by the `pivot_func` defined above. It then splits the input dataframe into NPartitions.get()
-    dataframes, with the elements in the i-th split belonging to the i-th partition, as determined
-    by the quantiles we're using.
-
-    Parameters
-    ----------
-    df : pandas.Dataframe
-        The partition to split.
-    column : str
-        The major column to sort by.
-    is_numeric_column : bool
-        Whether the passed `column` has numeric type (int, float).
-    pivots : np.ndarray
-        The quantiles to use to split the data.
-    ascending : bool
-        The ascending flag.
     **kwargs : dict
         Additional keyword arguments.
-
-    Returns
-    -------
-    tuple[pandas.DataFrame]
-        A tuple of the splits from this partition.
     """
-    if len(pivots) == 0:
-        # We can return the dataframe with zero changes if there were no pivots passed
-        return (df,)
-    # If `ascending=False` and we are dealing with a numeric dtype, we can pass in a reversed list
-    # of pivots, and `np.digitize` will work correctly. For object dtypes, we use `np.searchsorted`
-    # which breaks when we reverse the pivots.
-    if not ascending and is_numeric_column:
-        # `key` is already applied to `pivots` in the `pick_pivots_from_samples_for_sort` function.
-        pivots = pivots[::-1]
-    key = kwargs.pop("key", None)
-    na_index = df[column].isna()
-    na_rows = df[na_index]
-    non_na_rows = df[~na_index]
-    cols_to_digitize = non_na_rows[column]
-    if key is not None:
-        cols_to_digitize = key(cols_to_digitize)
-
-    if is_numeric_column:
-        groupby_col = np.digitize(cols_to_digitize.squeeze(), pivots)
-        # `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
-        # we actually want to sort in descending order, we need to swap the new indices.
-        if not ascending and len(np.unique(pivots)) == 1:
-            groupby_col = len(pivots) - groupby_col
-    else:
-        groupby_col = np.searchsorted(pivots, cols_to_digitize.squeeze(), side="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:
-            groupby_col = len(pivots) - groupby_col
-    if len(non_na_rows) == 1:
-        groups = [
-            # taking an empty slice for an index's metadata
-            pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(df.dtypes)
-            if i != groupby_col
-            else non_na_rows
-            for i in range(len(pivots) + 1)
-        ]
-    else:
-        grouped = non_na_rows.groupby(groupby_col)
-        groups = [
-            grouped.get_group(i) if i in grouped.keys
-            # taking an empty slice for an index's metadata
-            else pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(
-                df.dtypes
+
+    def __init__(
+        self,
+        modin_frame: "PandasDataframe",
+        columns: Union[str, list],
+        ascending: Union[list, bool],
+        ideal_num_new_partitions: int,
+        **kwargs: dict,
+    ):
+        self.frame_len = len(modin_frame.index)
+        self.ideal_num_new_partitions = ideal_num_new_partitions
+        self.columns = columns if is_list_like(columns) else [columns]
+        self.ascending = ascending
+        self.kwargs = kwargs.copy()
+        self.columns_info = None
+
+    def sample_fn(self, partition: pandas.DataFrame) -> pandas.DataFrame:
+        return self.pick_samples_for_quantiles(
+            partition[self.columns], self.ideal_num_new_partitions, self.frame_len
+        )
+
+    def pivot_fn(self, samples: "list[pandas.DataFrame]") -> int:
+        key = self.kwargs.get("key", None)
+        samples = pandas.concat(samples, axis=0, copy=False)
+
+        columns_info: "list[ColumnInfo]" = []
+        number_of_groups = 1
+        cols = []
+        for col in 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()
+            cols.append(col)
+            is_numeric = is_numeric_dtype(column_val.dtype)
+
+            # When we are not sorting numbers, we need our quantiles to not do arithmetic on the values
+            method = "linear" if is_numeric else "inverted_cdf"
+            pivots = self.pick_pivots_from_samples_for_sort(
+                column_val, num_pivots, method, key
             )
-            for i in range(len(pivots) + 1)
-        ]
-    index_to_insert_na_vals = -1 if kwargs.get("na_position", "last") == "last" else 0
-    groups[index_to_insert_na_vals] = pandas.concat(
-        [groups[index_to_insert_na_vals], na_rows]
-    ).astype(df.dtypes)
-    return tuple(groups)
+            columns_info.append(ColumnInfo(col, pivots, is_numeric))
+            number_of_groups *= len(pivots) + 1
+        self.columns_info = columns_info
+        return number_of_groups
+
+    def split_fn(
+        self,
+        partition: pandas.DataFrame,
+    ) -> "tuple[pandas.DataFrame, ...]":
+        ErrorMessage.catch_bugs_and_request_email(
+            failure_condition=self.columns_info is None,
+            extra_log="The 'split_fn' doesn't have proper metadata, the probable reason is that it was called before 'pivot_fn'",
+        )
+        return self.split_partitions_using_pivots_for_sort(
+            partition,
+            self.columns_info,
+            self.ascending,
+            **self.kwargs,
+        )
+
+    @staticmethod
+    def _find_quantiles(
+        df: Union[pandas.DataFrame, pandas.Series], quantiles: list, method: str
+    ) -> np.ndarray:
+        """
+        Find quantiles of a given dataframe using the specified method.
+
+        We use this method to provide backwards compatibility with NumPy versions < 1.23 (e.g. when
+        the user is using Modin in compat mode). This is basically a wrapper around `np.quantile` that
+        ensures we provide the correct `method` argument - i.e. if we are dealing with objects (which
+        may or may not support algebra), we do not want to use a method to find quantiles that will
+        involve algebra operations (e.g. mean) between the objects, since that may fail.
+
+        Parameters
+        ----------
+        df : pandas.DataFrame or pandas.Series
+            The data to pick quantiles from.
+        quantiles : list[float]
+            The quantiles to compute.
+        method : str
+            The method to use. `linear` if dealing with numeric types, otherwise `inverted_cdf`.
+
+        Returns
+        -------
+        np.ndarray
+            A NumPy array with the quantiles of the data.
+        """
+        if method == "linear":
+            # This is the default method for finding quantiles, so it does not need to be specified,
+            # which keeps backwards compatibility with older versions of NumPy that do not have a
+            # `method` keyword argument in np.quantile.
+            return np.unique(np.quantile(df, quantiles))
+        else:
+            try:
+                return np.unique(np.quantile(df, quantiles, method=method))
+            except Exception:
+                # In this case, we're dealing with an array of objects, but the current version of
+                # NumPy does not have a `method` kwarg. We need to use the older kwarg, `interpolation`
+                # instead.
+                return np.unique(np.quantile(df, quantiles, interpolation="lower"))
+
+    @staticmethod
+    def pick_samples_for_quantiles(
+        df: pandas.DataFrame,
+        num_partitions: int,
+        length: int,
+    ) -> pandas.DataFrame:
+        """
+        Pick samples over the given partition.
+
+        This function picks samples from the given partition using the TeraSort algorithm - each
+        value is sampled with probability 1 / m * ln(n * t) where m = total_length / num_partitions,
+        t = num_partitions, and n = total_length.
+
+        Parameters
+        ----------
+        df : pandas.Dataframe
+            The masked dataframe to pick samples from.
+        num_partitions : int
+            The number of partitions.
+        length : int
+            The total length.
+
+        Returns
+        -------
+        pandas.DataFrame:
+            The samples for the partition.
+
+        Notes
+        -----
+        This sampling algorithm is inspired by TeraSort. You can find more information about TeraSort
+        and the sampling algorithm at https://www.cse.cuhk.edu.hk/~taoyf/paper/sigmod13-mr.pdf.
+        """
+        m = length / num_partitions
+        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,
+        ideal_num_new_partitions: int,
+        method: str = "linear",
+        key: Optional[Callable] = None,
+    ) -> np.ndarray:
+        """
+        Determine quantiles from the given samples.
+
+        This function takes as input the quantiles calculated over all partitions from
+        `sample_func` defined above, and determines a final NPartitions.get() quantiles
+        to use to roughly sort the entire dataframe. It does so by collating all the samples
+        and computing NPartitions.get() quantiles for the overall set.
+
+        Parameters
+        ----------
+        samples : np.ndarray
+            The samples computed by ``get_partition_quantiles_for_sort``.
+        ideal_num_new_partitions : int
+            The ideal number of new partitions.
+        method : str, default: linear
+            The method to use when picking quantiles.
+        key : Callable, default: None
+            The key to use on the samples when picking pivots.
+
+        Returns
+        -------
+        np.ndarray
+            A list of overall quantiles.
+        """
+        # 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.
+        if key is not None:
+            samples = key(samples)
+        # We don't want to pick very many quantiles if we have a very small dataframe.
+        num_quantiles = ideal_num_new_partitions
+        quantiles = [i / num_quantiles for i in range(1, num_quantiles)]
+        # 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 np.array([])
+
+    @staticmethod
+    def split_partitions_using_pivots_for_sort(
+        df: pandas.DataFrame,
+        columns_info: "list[ColumnInfo]",
+        ascending: bool,
+        **kwargs: dict,
+    ) -> "tuple[pandas.DataFrame, ...]":
+        """
+        Split the given dataframe into the partitions specified by `pivots` in `columns_info`.
+
+        This function takes as input a row-axis partition, as well as the quantiles determined
+        by the `pivot_func` defined above. It then splits the input dataframe into NPartitions.get()
+        dataframes, with the elements in the i-th split belonging to the i-th partition, as determined
+        by the quantiles we're using.
+
+        Parameters
+        ----------
+        df : pandas.Dataframe
+            The partition to split.
+        columns_info : list of ColumnInfo
+            Information regarding keys and pivots for range partitioning.
+        ascending : bool
+            The ascending flag.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        tuple[pandas.DataFrame]
+            A tuple of the splits from this partition.
+        """
+        if len(columns_info) == 0:
+            # We can return the dataframe with zero changes if there were no pivots passed
+            return (df,)
+
+        na_index = (
+            df[[col_info.name for col_info in columns_info]].isna().squeeze(axis=1)
+        )
+        if na_index.ndim == 2:
+            na_index = na_index.any(axis=1)
+        na_rows = df[na_index]
+        non_na_rows = df[~na_index]
+
+        def get_group(grp, key, df):
+            """Get a group with the `key` from the `grp`, if it doesn't exist return an empty slice of `df`."""
+            try:
+                return grp.get_group(key)
+            except KeyError:
+                return pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(
+                    df.dtypes
+                )
+
+        groupby_codes = []
+        group_keys = []
+        for col_info in columns_info:
+            pivots = col_info.pivots
+            if len(pivots) == 0:
+                continue
+            # If `ascending=False` and we are dealing with a numeric dtype, we can pass in a reversed list
+            # of pivots, and `np.digitize` will work correctly. For object dtypes, we use `np.searchsorted`
+            # which breaks when we reverse the pivots.
+            if not ascending and col_info.is_numeric:
+                # `key` is already applied to `pivots` in the `pick_pivots_from_samples_for_sort` function.
+                pivots = pivots[::-1]
+            group_keys.append(range(len(pivots) + 1))
+            key = kwargs.pop("key", None)
+            cols_to_digitize = non_na_rows[col_info.name]
+            if key is not None:
+                cols_to_digitize = key(cols_to_digitize)
+
+            if col_info.is_numeric:
+                groupby_col = np.digitize(cols_to_digitize.squeeze(), pivots)
+                # `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
+                # we actually want to sort in descending order, we need to swap the new indices.
+                if not ascending and len(np.unique(pivots)) == 1:
+                    groupby_col = len(pivots) - groupby_col
+            else:
+                groupby_col = np.searchsorted(
+                    pivots, cols_to_digitize.squeeze(), side="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:
+                    groupby_col = len(pivots) - groupby_col
+            groupby_codes.append(groupby_col)
+
+        if len(group_keys) == 0:
+            # We can return the dataframe with zero changes if there were no pivots passed
+            return (df,)
+        elif len(group_keys) == 1:
+            group_keys = group_keys[0]
+        else:
+            group_keys = pandas.MultiIndex.from_product(group_keys)
+
+        if len(non_na_rows) == 1:
+            groups = [
+                # taking an empty slice for an index's metadata
+                pandas.DataFrame(index=df.index[:0], columns=df.columns).astype(
+                    df.dtypes
+                )
+                if key != groupby_codes[0]
+                else non_na_rows
+                for key in group_keys
+            ]
+        else:
+            grouped = non_na_rows.groupby(groupby_codes)
+            groups = [get_group(grouped, key, df) for key in group_keys]
+        index_to_insert_na_vals = (
+            -1 if kwargs.get("na_position", "last") == "last" else 0
+        )
+        groups[index_to_insert_na_vals] = pandas.concat(
+            [groups[index_to_insert_na_vals], na_rows]
+        ).astype(df.dtypes)
+        return tuple(groups)
 
 
 def lazy_metadata_decorator(apply_axis=None, axis_arg=-1, transpose=False):
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
@@ -23,6 +23,7 @@
 import pandas
 from pandas._libs.lib import no_default
 import warnings
+from typing import TYPE_CHECKING
 
 from modin.error_message import ErrorMessage
 from modin.core.storage_formats.pandas.utils import compute_chunksize
@@ -32,6 +33,9 @@
 
 import os
 
+if TYPE_CHECKING:
+    from modin.core.dataframe.pandas.dataframe.utils import ShuffleFunctions
+
 
 def wait_computations_if_benchmark_mode(func):
     """
@@ -1563,7 +1567,11 @@ def rebalance_partitions(cls, partitions):
     @classmethod
     @wait_computations_if_benchmark_mode
     def shuffle_partitions(
-        cls, partitions, index, shuffle_functions, final_shuffle_func
+        cls,
+        partitions,
+        index,
+        shuffle_functions: "ShuffleFunctions",
+        final_shuffle_func,
     ):
         """
         Return shuffled partitions.
@@ -1572,10 +1580,10 @@ def shuffle_partitions(
         ----------
         partitions : np.ndarray
             The 2-d array of partitions to shuffle.
-        index : int
-            The index of partitions corresponding to the partitions that contain the column to sample.
-        shuffle_functions : NamedTuple
-            A named tuple containing the functions that we will be using to perform this shuffle.
+        index : int or list of ints
+            The index(es) of the column partitions corresponding to the partitions that contain the column to sample.
+        shuffle_functions : ShuffleFunctions
+            An object implementing the functions that we will be using to perform this shuffle.
         final_shuffle_func : Callable(pandas.DataFrame) -> pandas.DataFrame
             Function that shuffles the data within each new partition.
 
@@ -1587,23 +1595,28 @@ def shuffle_partitions(
         # Mask the partition that contains the column that will be sampled.
         masked_partitions = partitions[:, index]
         # Sample each partition
-        sample_func = cls.preprocess_func(shuffle_functions.sample_function)
-        samples = [partition.apply(sample_func) for partition in masked_partitions]
+        sample_func = cls.preprocess_func(shuffle_functions.sample_fn)
+        if masked_partitions.ndim == 1:
+            samples = [partition.apply(sample_func) for partition in masked_partitions]
+        else:
+            samples = [
+                cls._row_partition_class(row_part, full_axis=False).apply(sample_func)
+                for row_part in masked_partitions
+            ]
         # Get each sample to pass in to the pivot function
         samples = cls.get_objects_from_partitions(samples)
-        pivots = np.unique(shuffle_functions.pivot_function(samples))
+        num_bins = shuffle_functions.pivot_fn(samples)
         # Convert our list of block partitions to row partitions. We need to create full-axis
         # row partitions since we need to send the whole partition to the split step as otherwise
         # we wouldn't know how to split the block partitions that don't contain the shuffling key.
         row_partitions = cls.row_partitions(partitions)
-        if len(pivots):
+        if num_bins > 1:
             # Gather together all of the sub-partitions
             split_row_partitions = np.array(
                 [
                     partition.split(
-                        shuffle_functions.split_function,
-                        num_splits=len(pivots) + 1,
-                        f_args=(pivots,),
+                        shuffle_functions.split_fn,
+                        num_splits=num_bins,
                         # The partition's metadata will never be accessed for the split partitions,
                         # thus no need to compute it.
                         extract_metadata=False,


Test patch:
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -21,6 +21,10 @@
 from modin.distributed.dataframe.pandas import from_partitions
 from modin.core.storage_formats.pandas.utils import split_result_of_axis_func_pandas
 from modin.utils import try_cast_to_pandas
+from modin.core.dataframe.pandas.dataframe.utils import (
+    ShuffleSortFunctions,
+    ColumnInfo,
+)
 
 import numpy as np
 import pandas
@@ -841,10 +845,6 @@ def test_split_partitions_kernel(
         Duplicate pivot values cause empty partitions to be produced. This parameter helps
         to verify that the function still behaves correctly in such cases.
     """
-    from modin.core.dataframe.pandas.dataframe.utils import (
-        split_partitions_using_pivots_for_sort,
-    )
-
     random_state = np.random.RandomState(42)
 
     df = pandas.DataFrame(
@@ -865,11 +865,15 @@ def test_split_partitions_kernel(
 
     # Randomly reordering rows in the dataframe
     df = df.reindex(random_state.permutation(df.index))
-    bins = split_partitions_using_pivots_for_sort(
+    bins = ShuffleSortFunctions.split_partitions_using_pivots_for_sort(
         df,
-        col_name,
-        is_numeric_column=pandas.api.types.is_numeric_dtype(df.dtypes[col_name]),
-        pivots=pivots,
+        [
+            ColumnInfo(
+                name=col_name,
+                is_numeric=pandas.api.types.is_numeric_dtype(df.dtypes[col_name]),
+                pivots=pivots,
+            )
+        ],
         ascending=ascending,
     )
 
@@ -904,10 +908,6 @@ def test_split_partitions_with_empty_pivots(col_name, ascending):
     This test verifies that the splitting function performs correctly when an empty pivots list is passed.
     The expected behavior is to return a single split consisting of the exact copy of the input dataframe.
     """
-    from modin.core.dataframe.pandas.dataframe.utils import (
-        split_partitions_using_pivots_for_sort,
-    )
-
     df = pandas.DataFrame(
         {
             "numeric_col": range(9),
@@ -915,11 +915,15 @@ def test_split_partitions_with_empty_pivots(col_name, ascending):
         }
     )
 
-    result = split_partitions_using_pivots_for_sort(
+    result = ShuffleSortFunctions.split_partitions_using_pivots_for_sort(
         df,
-        col_name,
-        is_numeric_column=pandas.api.types.is_numeric_dtype(df.dtypes[col_name]),
-        pivots=[],
+        [
+            ColumnInfo(
+                name=col_name,
+                is_numeric=pandas.api.types.is_numeric_dtype(df.dtypes[col_name]),
+                pivots=[],
+            )
+        ],
         ascending=ascending,
     )
     # We're expecting to recieve a single split here
@@ -939,17 +943,12 @@ def test_shuffle_partitions_with_empty_pivots(ascending):
 
     assert modin_frame._partitions.shape == (1, 1)
 
-    from modin.core.dataframe.pandas.dataframe.utils import (
-        build_sort_functions,
-    )
-
     column_name = modin_frame.columns[1]
 
-    shuffle_functions = build_sort_functions(
+    shuffle_functions = ShuffleSortFunctions(
         # These are the parameters we pass in the `.sort_by()` implementation
         modin_frame,
-        column=column_name,
-        method="inverted_cdf",
+        columns=column_name,
         ascending=ascending,
         ideal_num_new_partitions=1,
     )
@@ -973,10 +972,6 @@ def test_split_partition_preserve_names(ascending):
     This test verifies that the dataframes being split by ``split_partitions_using_pivots_for_sort``
     preserve their index/column names.
     """
-    from modin.core.dataframe.pandas.dataframe.utils import (
-        split_partitions_using_pivots_for_sort,
-    )
-
     df = pandas.DataFrame(
         {
             "numeric_col": range(9),
@@ -989,11 +984,9 @@ def test_split_partition_preserve_names(ascending):
 
     # Pivots that contain empty bins
     pivots = [2, 2, 5, 7]
-    splits = split_partitions_using_pivots_for_sort(
+    splits = ShuffleSortFunctions.split_partitions_using_pivots_for_sort(
         df,
-        column="numeric_col",
-        is_numeric_column=True,
-        pivots=pivots,
+        [ColumnInfo(name="numeric_col", is_numeric=True, pivots=pivots)],
         ascending=ascending,
     )