Repository: modin-project/modin
Pull request: FEAT-#3535: Implement partition shuffling mechanism and algebra sort_by
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

- [x] commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
- [] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [ ] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #3535 ? <!-- issue must be created for each patch -->
- [ ] tests added and passing
- [ ] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->
- [ ] added (Issue Number: PR title (PR Number)) and github username to release notes for next major release <!-- e.g. DOCS-#4077: Add release notes template to docs folder (#4078) -->


Problem statement:
Improved shuffling mechanism for Ray
Currently, we do not have a good general purpose shuffle. Ideally, we would have a mechanism that could split the data based on some split function or based on some features of the data (e.g. the quantiles of the column for numeric data).

Hints from issue discussion:
I have been testing an implementation of a numerical sort (not a general purpose shuffle) using Ray actors on an 800 MB test dataset [here](https://drive.google.com/file/d/1pMm8As4IdR_-PfAeI9ZdIuDpCllO8pF7/view?usp=sharing). You can view my latest script [here](https://github.com/mvashishtha/modin/blob/0a3acc150b0856b858fd90f0afe7b0bcd146579b/sort_non_modin.py). On my Mac (16-inch MacBook Pro, 2019, with MacOS Big Sur 11.5.2, Ray version 1.8.0, Python 3.8.8), I compared three alternatives:

1) The sort script with Ray actors linked above. The sort takes about 253 seconds. 

2) Sorting with Modin by defaulting to Pandas as follows:
```python
import modin.pandas as pd
df = pd.read_csv("test_1mx256.csv")
%time df.sort_values(["col2"])
```

This took 401 seconds.

3) Sorting with Pandas:
```python
import pandas as pd
df = pd.read_csv("test_1mx256.csv")
%time df.sort_values(["col2"])
```
This took 2.58 seconds.

On @devin-petersohn's Windows Machine (Windows 10, Ray 1.7.1, Python 3.8.5), I get:

1) sort script with ray actors: 6.42 seconds
2) Sorting with Modin by defaulting to Pandas: 3.58 seconds
3) Sorting with Pandas: 2.24 seconds

The difference in Modin performance between Mac and Windows is [Ray Issue 20388](https://github.com/ray-project/ray/issues/20388). On Windows, the Actor approach takes about 3x as long as Pandas. and about 2x as long as Modin with default to Pandas. I don't think my implementation performs well enough to add to Modin. Does anyone have any suggestions on how to improve it? Should we use Ray actors at all?

What would happen if we try to simulate a case of not fitting in a single node? I mean, no defaulting is crucial if you have so much data that it won't fit in any single node, then some distributed logic is needed regardless of performance.

Also how does the scaling work here, what if we take a 10m rows input? A 100m rows?

**UPD**: That said, I think we also need some sort of shuffling to better perform groupby aggregations - to be able to parallelize over different groups along with parallelizing over columns (what we do now).
@vnlitvinov, @devin-petersohn is looking into setting up a distributed testing environment. Meanwhile, are you able to test the script with 10m or 100m rows of input? Given how slow it was on my computer on the 1m row file from my previous comment, I don't think I'll be able to test those locally.
I wonder how did you make your input data, could you share a script or other means of generating that?
@vnlitvinov Sure, this is how I got the input data:

```python
import numpy as np
import modin.pandas as pd

df = pd.DataFrame(np.random.randint(0, 100, size=(2**20, 2**8))).add_prefix("col")
df.to_csv("test_1mx256.csv", index=False)
```
These results can be misleading, since the current Modin implementation (in big datasets) is faster than pure pandas, although message about defaulting to pandas appears. However, internally only sorting of the `by` columns is defaulting to pandas, the `reindex` operation is performed by Modin itself, usually in parallel. Shouldn't we stop showing the message?

Machine spec: `112 cores, ~800GB Ram, Ubuntu 21.10`

### Results
<details>
<summary>Comparison of `sort_values` implementations (all times in seconds)</summary>

![image](https://user-images.githubusercontent.com/45976948/145248357-fed429eb-6e12-4c8d-b30d-1904885aee89.png)
</details>

<details>
<summary>Prototype scalability</summary>

![image](https://user-images.githubusercontent.com/45976948/145248603-dc2ccfa5-66f5-4b4e-b5ec-8fc6b6f9f155.png)
</details>

<details>
<summary>Current Modin implementation scalability</summary>

![image](https://user-images.githubusercontent.com/45976948/145248738-88f1e778-9c79-4b60-adc9-e517ad3468bd.png)
</details>


I couldn't get results for dataset with 100M rows due to out of memory:
<details> <summary>prototype error</summary>

```bash
	2021-12-08 15:57:29,080 WARNING worker.py:1215 -- The actor or task with ID e788aad55c0377b5f52597f2084589fe866a0a2115908deb cannot be scheduled right now. You can ignore this message if this Ray cluster is expected to auto-scale or if you specified a runtime_env for this actor or task, which may take time to install.  Otherwise, this is likely due to all cluster resources being claimed by actors. To resolve the issue, consider creating fewer actors or increasing the resources available to this Ray cluster.
	Required resources for this actor or task: {CPU: 1.000000}
	Available resources on this node: {0.000000/8.000000 CPU, 551.207606 GiB/551.207606 GiB memory, 186.264515 GiB/186.264515 GiB object_store_memory, 1.000000/1.000000 node:10.125.130.30}
	In total there are 47 pending tasks and 0 pending actors on this node.
	Traceback (most recent call last):
	  File "sort_values_non_modin.py", line 152, in <module>
	    # sorted_dataframe, defaulting_modin, sort_by_pandas
	  File "sort_values_non_modin.py", line 124, in sorted_dataframe
	    ]
	  File "miniconda3/envs/modin/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 82, in wrapper
	    return func(*args, **kwargs)
	  File "miniconda3/envs/modin/lib/python3.8/site-packages/ray/worker.py", line 1621, in get
	    raise value.as_instanceof_cause()
	ray.exceptions.RayTaskError: ray::ShuffleActor.combine_and_apply_dfs() (pid=2409094, ip=10.125.130.30, repr=<sort_values_non_modin.ShuffleActor object at 0x7f363d90daf0>)
	ray.exceptions.ObjectStoreFullError: Failed to put object 2f41f345ddf68417a8782fa603827f3700652efb0100000001000000 in object store because it is full. Object size is 28017495444 bytes.
	The local object store is full of objects that are still in scope and cannot be evicted. Tip: Use the `ray memory` command to list active objects in the cluster.
```
  </details>
<details><summary>current Modin implementation error</summary>

```bash
	2021-12-08 17:09:11,680 WARNING worker.py:1215 -- Traceback (most recent call last):	
	  File "python/ray/_raylet.pyx", line 723, in ray._raylet.spill_objects_handler
	  File "python/ray/_raylet.pyx", line 726, in ray._raylet.spill_objects_handler
	  File "/miniconda3/envs/modin/lib/python3.8/site-packages/ray/external_storage.py", line 522, in spill_objects
	    return _external_storage.spill_objects(object_refs, owner_addresses)
	  File "miniconda3/envs/modin/lib/python3.8/site-packages/ray/external_storage.py", line 277, in spill_objects
	    return self._write_multiple_objects(f, object_refs,
	  File "miniconda3/envs/modin/lib/python3.8/site-packages/ray/external_storage.py", line 136, in _write_multiple_objects
	    written_bytes = f.write(payload)
	OSError: [Errno 28] No space left on device
	An unexpected internal error occurred while the IO worker was spilling objects: [Errno 28] No space left on device
	SORT_VALUES time: 1492.286622285843
```
</details>
<details> <summary>frequent warnings</summary>

```bash
	Warning
		2021-12-07 23:19:23,395 WARNING worker.py:1215 -- The actor or task with ID ffffffffffffffff66ecd7e592db6fa5e6436cfa01000000 cannot be scheduled right now. You can ignore this message if this Ray cluster is expected to auto-scale or if you specified a runtime_env for this actor or task, which may take time to install.  Otherwise, this is likely due to all cluster resources being claimed by actors. To resolve the issue, consider creating fewer actors or increasing the resources available to this Ray cluster.
		Required resources for this actor or task: {CPU: 1.000000}
		Available resources on this node: {26.000000/32.000000 CPU, 552.116908 GiB/552.116908 GiB memory, 186.264515 GiB/186.264515 GiB object_store_memory, 1.000000/1.000000 node:10.125.130.30}
		In total there are 256 pending tasks and 26 pending actors on this node.
	Warning
		2021-12-08 14:09:14,445 WARNING worker.py:1215 -- WARNING: 334 PYTHON worker processes have been started on node: a63fead710fca53f131514508b9451e3325860f9f368a4d972114fd3 with address: 10.125.130.30. This could be a result of using a large number of actors, or due to tasks blocked in ray.get() calls (see https://github.com/ray-project/ray/issues/3644 for some discussion of workarounds).
```
</details>

What is the current state here? Do we have any particular plans on how to continue from here?
@vnlitvinov I don't have any updates since @anmyachev commented [here](https://github.com/modin-project/modin/issues/3535#issuecomment-989001543). @RehanSD is exploring this problem, but I don't think we've committed to a particular solution.
Hi @anmyachev! Thank you so much for the awesome data + graphs! I have a new implementation of the sort implemented at the ModinDataframe layer without using ray Actors. I was wondering if it would be possible for you to bench this new implementation as well?

New Implementation without profiling: https://github.com/RehanSD/modin/tree/rehan/sort_by
New implementation with profiling: https://github.com/RehanSD/modin/tree/rehan/sort_by_time

Code to time:
```python
import modin.pandas as pd
# import pandas as pd
import numpy as np
from modin.core.dataframe.base.dataframe.utils import Axis
df = pd.DataFrame(np.random.standard_cauchy((1000000, 256)))
df._to_pandas() # Force materialization
def modin():
    import ray
    d = df._query_compiler._modin_frame.sort_by(Axis.ROW_WISE, 2)
    ray.wait([obj.oid for row in d._partitions for obj in row])
    return d
%time modin()
```

I included the code with and without profiling since the print statements + variable fetch across layers of the stack might add extra overhead.
@RehanSD yes, let me do it

Data size: 1M x 256 cols

UPD:
Results from `sort_by` branch (min from 10 launches):
- 8 CPUS:   2.85 sec
- 16 CPUS: 2.79 sec
- 32 CPUS: 3.51 sec

Some changes in bench code:
- move `ray` import outside `modin` function
- time measurement with `time.time` function

UPD:
Results from `sort_by_time` branch (min from 10 launches):
- [sort_values_8_cpus.txt](https://github.com/modin-project/modin/files/8062892/sort_values_8_cpus.txt)
- [sort_values_16_cpus.txt](https://github.com/modin-project/modin/files/8062894/sort_values_16_cpus.txt)
- [sort_values_32_cpus.txt](https://github.com/modin-project/modin/files/8062895/sort_values_32_cpus.txt)


UPD:
Results for different data sizes:
![image](https://user-images.githubusercontent.com/45976948/154987088-1d53495b-e137-4fde-b8d1-69a7f97d5ba9.png)

Note: Column with "Prototype" label is actually Mahesh' implementation (I forgot to change the column name, sorry)

<details> <summary>Scalability:</summary>

![image](https://user-images.githubusercontent.com/45976948/154987346-4ba869e8-679a-44f5-b2c5-bb0ae3326e57.png)


</details>

Commit message: Address review comments

Signed-off-by: Rehan Durrani <rehan@ponder.io>

Patch to review:
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index 3f31942162f..30d6a17fe46 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -38,6 +38,7 @@
     Axis,
     JoinType,
 )
+from modin.core.dataframe.pandas.dataframe.utils import build_sort_functions
 
 if TYPE_CHECKING:
     from modin.core.dataframe.base.interchange.dataframe_protocol.dataframe import (
@@ -1961,11 +1962,13 @@ def map_fn(df):
             new_dtypes,
         )
 
+    @lazy_metadata_decorator(apply_axis="both")
     def sort_by(
         self,
         axis: Union[int, Axis],
         columns: Union[str, List[str]],
         ascending: bool = True,
+        **kwargs,
     ) -> "PandasDataframe":
         """
         Logically reorder rows (columns if axis=1) lexicographically by the data in a column or set of columns.
@@ -1978,13 +1981,110 @@ def sort_by(
             Column label(s) to use to determine lexicographical ordering.
         ascending : boolean, default: True
             Whether to sort in ascending or descending order.
+        **kwargs : dict
+            Keyword arguments to pass when sorting partitions.
 
         Returns
         -------
         PandasDataframe
             A new PandasDataframe sorted into lexicographical order by the specified column(s).
         """
-        pass
+        if not isinstance(columns, list):
+            columns = [columns]
+        # When we do a sort on the result of Series.value_counts, we don't rename the index until
+        # after everything is done, which causes an error when sorting the partitions, since the
+        # index and the column share the same name, when in actuality, the index's name should be
+        # None. This fixes the indexes name beforehand in that case, so that the sort works.
+
+        def sort_function(df):
+            index_renaming = None
+            if any(name in df.columns for name in df.index.names):
+                index_renaming = df.index.names
+                df.index = df.index.set_names([None] * len(df.index.names))
+            df = df.sort_values(by=columns, ascending=ascending, **kwargs)
+            if index_renaming is not None:
+                df.index = df.index.set_names(index_renaming)
+            return df
+
+        axis = Axis(axis)
+        if axis != Axis.ROW_WISE:
+            raise NotImplementedError(
+                f"Algebra sort only implemented row-wise. {axis.name} sort not implemented yet!"
+            )
+
+        # If this df is empty, we don't want to try and shuffle or sort.
+        if len(self.axes[0]) == 0 or len(self.axes[1]) == 0:
+            return self.copy()
+        # If this df only has one row partition, we don't want to do a shuffle and sort - we can
+        # just do a full-axis sort.
+        if len(self._partitions) == 1:
+            return self.apply_full_axis(
+                1,
+                sort_function,
+            )
+        if self.dtypes[columns[0]] == object:
+            # This means we are not sorting numbers, so we need our quantiles to not try
+            # arithmetic on the values.
+            method = "inverted_cdf"
+        else:
+            method = "linear"
+
+        shuffling_functions = build_sort_functions(
+            self,
+            columns[0],
+            method,
+            ascending[0] if is_list_like(ascending) else ascending,
+            **kwargs,
+        )
+        major_col_partition_index = self.columns.get_loc(columns[0])
+        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(
+            self._partitions,
+            index,
+            shuffling_functions,
+            sort_function,
+        )
+        new_axes = self.axes
+        if kwargs.get("ignore_index", False):
+            new_axes[axis.value] = RangeIndex(len(new_axes[axis.value]))
+        else:
+            new_axes[axis.value] = self._compute_axis_labels_and_lengths(
+                axis.value, new_partitions
+            )[0]
+        # If we have a MultiIndex, but the first partition is empty, which may happen when
+        # the dataframe is small, `_compute_axis_labels_and_lengths` will return us a flattened
+        # MultiIndex - i.e. an Index consisting of tuples. This is because the MultiIndex from
+        # the remaining partitions is appended to an empty flat Index, which results in a
+        # flattened index. To work around this, we need to convert this flattened Index back
+        # into a MultiIndex.
+        if isinstance(self.axes[axis.value], pandas.MultiIndex) and isinstance(
+            new_axes[axis.value][0], tuple
+        ):
+            new_axes[axis.value] = pandas.MultiIndex.from_tuples(
+                new_axes[axis.value].values
+            )
+        new_axes[axis.value] = new_axes[axis.value].set_names(
+            self.axes[axis.value].names
+        )
+        new_lengths = [None, None]
+        (
+            new_partitions,
+            new_lengths[axis.value],
+        ) = self._partition_mgr_cls.rebalance_partitions(new_partitions)
+        col_parts = new_partitions[0]
+        new_lengths[axis.value ^ 1] = [part.width() for part in col_parts]
+        new_modin_frame = self.__constructor__(
+            new_partitions, *new_axes, *new_lengths, self.dtypes
+        )
+        if kwargs.get("ignore_index", False):
+            new_modin_frame._propagate_index_objs(axis=0)
+        return new_modin_frame
 
     @lazy_metadata_decorator(apply_axis="both")
     def filter(self, axis: Union[Axis, int], condition: Callable) -> "PandasDataframe":
diff --git a/modin/core/dataframe/pandas/dataframe/utils.py b/modin/core/dataframe/pandas/dataframe/utils.py
new file mode 100644
index 00000000000..1e0f15eb59f
--- /dev/null
+++ b/modin/core/dataframe/pandas/dataframe/utils.py
@@ -0,0 +1,278 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team licenses this file to you under the
+# Apache License, Version 2.0 (the "License"); you may not use this file except in
+# compliance with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+"""Collection of algebra utility functions, used to shuffle data across partitions."""
+
+import numpy as np
+import pandas
+from typing import Callable, Union, Optional, TYPE_CHECKING
+from collections import namedtuple
+
+if TYPE_CHECKING:
+    from modin.core.dataframe.pandas.dataframe.dataframe import PandasDataframe
+
+
+ShuffleFunctions = namedtuple(
+    "ShuffleFunctions", ["sample_function", "pivot_function", "split_function"]
+)
+
+
+def build_sort_functions(
+    modin_frame: "PandasDataframe",
+    column: str,
+    method: str,
+    ascending: Union[list, bool],
+    **kwargs: dict,
+) -> ShuffleFunctions:
+    """
+    Return a named tuple containing the functions necessary to perform a sort.
+
+    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.
+    ascending : bool
+        The ascending flag.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    ShuffleFunctions :
+        A named tuple containing the functions to pick quantiles, choose pivot points, and split
+        partitions for sorting.
+    """
+
+    def sample_fn(partition):
+        return pick_samples_for_quantiles(
+            partition, column, len(modin_frame._partitions), len(modin_frame.index)
+        )
+
+    def pivot_fn(samples):
+        key = kwargs.get("key", None)
+        return pick_pivots_from_samples_for_sort(modin_frame, samples, method, key)
+
+    def split_fn(partition, pivots):
+        return split_partitions_using_pivots_for_sort(
+            modin_frame, partition, column, 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,
+    column: str,
+    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.
+    column : str
+        The columns to pick quantiles from. Only the first column in the list will be used.
+    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[column].sample(frac=probability).to_numpy()
+
+
+def pick_pivots_from_samples_for_sort(
+    df: "PandasDataframe",
+    samples: "list[np.ndarray]",
+    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
+    ----------
+    df : PandasDataframe
+        The Modin Dataframe calling this function.
+    samples : list[np.ndarray]
+        The samples computed by ``get_partition_quantiles_for_sort``.
+    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 = len(df._partitions)
+    quantiles = [i / num_quantiles for i in range(1, num_quantiles)]
+    overall_quantiles = _find_quantiles(all_pivots, quantiles, method)
+    return overall_quantiles
+
+
+def split_partitions_using_pivots_for_sort(
+    modin_frame: "PandasDataframe",
+    df: pandas.DataFrame,
+    column: str,
+    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
+    ----------
+    modin_frame : PandasDataframe
+        The Modin Dataframe calling this function.
+    df : pandas.Dataframe
+        The partition to split.
+    column : str
+        The major column to sort by.
+    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 `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 modin_frame.dtypes[column] != object:
+        # `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 modin_frame.dtypes[column] != object:
+        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 and len(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
+    grouped = non_na_rows.groupby(groupby_col)
+    groups = [
+        grouped.get_group(i)
+        if i in grouped.keys
+        else pandas.DataFrame(columns=df.columns).astype(df.dtypes)
+        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)
diff --git a/modin/core/dataframe/pandas/partitioning/partition.py b/modin/core/dataframe/pandas/partitioning/partition.py
index b0a2a381f38..3ae118af5ac 100644
--- a/modin/core/dataframe/pandas/partitioning/partition.py
+++ b/modin/core/dataframe/pandas/partitioning/partition.py
@@ -314,6 +314,27 @@ def width(self):
             self._width_cache = self.apply(self._width_extraction_fn()).get()
         return self._width_cache
 
+    def split(self, split_func, num_splits, *args):
+        """
+        Split the object wrapped by the partition into multiple partitions.
+
+        Parameters
+        ----------
+        split_func : Callable[pandas.DataFrame, List[Any]] -> List[pandas.DataFrame]
+            The function that will split this partition into multiple partitions. The list contains
+            pivots to split by, and will have the same dtype as the major column we are shuffling on.
+        num_splits : int
+            The number of resulting partitions (may be empty).
+        *args : List[Any]
+            Arguments to pass to ``split_func``.
+
+        Returns
+        -------
+        list
+            A list of partitions.
+        """
+        pass
+
     @classmethod
     def empty(cls):
         """
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
index c0f41740030..58d765723c0 100644
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -1495,3 +1495,63 @@ def rebalance_partitions(cls, partitions):
             lengths = [part.length() for part in new_partitions[:, 0]]
             return new_partitions, lengths
         return partitions, None
+
+    @classmethod
+    def shuffle_partitions(
+        cls, partitions, index, shuffle_functions, final_shuffle_func
+    ):
+        """
+        Return shuffled partitions.
+
+        Parameters
+        ----------
+        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.
+        final_shuffle_func : Callable(pandas.DataFrame) -> pandas.DataFrame
+            Function that shuffles the data within each new partition.
+
+        Returns
+        -------
+        np.ndarray
+            A list of row-partitions that have been shuffled.
+        """
+        # 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]
+        # Get each sample to pass in to the pivot function
+        samples = cls.get_objects_from_partitions(samples)
+        pivots = shuffle_functions.pivot_function(samples)
+        # Convert our list of block partitions to row partitions
+        row_partitions = [
+            partition.force_materialization().list_of_block_partitions[0]
+            for partition in cls.row_partitions(partitions)
+        ]
+        # Gather together all of the sub-partitions
+        split_row_partitions = np.array(
+            [
+                partition.split(
+                    shuffle_functions.split_function, len(pivots) + 1, pivots
+                )
+                for partition in row_partitions
+            ]
+        ).T
+        # We need to convert every partition that came from the splits into a full-axis column partition.
+        col_partitioning_func = np.vectorize(
+            lambda partition: cls._row_partition_class(partition)
+        )
+        split_row_partitions = col_partitioning_func(split_row_partitions)
+        new_partitions = [
+            [
+                cls._column_partitions_class(row_partition, full_axis=False).apply(
+                    final_shuffle_func
+                )
+            ]
+            for row_partition in split_row_partitions
+        ]
+        return np.array(new_partitions)
diff --git a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
index 07bc9342ce9..56d26d283cb 100644
--- a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
+++ b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition.py
@@ -301,6 +301,30 @@ def ip(self):
             self._ip_cache = DaskWrapper.materialize(self._ip_cache)
         return self._ip_cache
 
+    def split(self, split_func, num_splits, *args):
+        """
+        Split the object wrapped by the partition into multiple partitions.
+
+        Parameters
+        ----------
+        split_func : Callable[pandas.DataFrame, List[Any]] -> List[pandas.DataFrame]
+            The function that will split this partition into multiple partitions. The list contains
+            pivots to split by, and will have the same dtype as the major column we are shuffling on.
+        num_splits : int
+            The number of resulting partitions (may be empty).
+        *args : List[Any]
+            Arguments to pass to ``split_func``.
+
+        Returns
+        -------
+        list
+            A list of partitions.
+        """
+        outputs = DaskWrapper.deploy(
+            split_func, [self._data] + list(args), num_returns=num_splits, pure=True
+        )
+        return [PandasOnDaskDataframePartition(output) for output in outputs]
+
 
 def apply_func(partition, func, *args, **kwargs):
     """
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
index 6561fcda963..98fab81037d 100644
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
@@ -354,6 +354,37 @@ def ip(self):
             self._ip_cache = RayWrapper.materialize(self._ip_cache)
         return self._ip_cache
 
+    def split(self, split_func, num_splits, *args):
+        """
+        Split the object wrapped by the partition into multiple partitions.
+
+        Parameters
+        ----------
+        split_func : Callable[pandas.DataFrame, List[Any]] -> List[pandas.DataFrame]
+            The function that will split this partition into multiple partitions. The list contains
+            pivots to split by, and will have the same dtype as the major column we are shuffling on.
+        num_splits : int
+            The number of resulting partitions, which may be empty.
+        *args : List[Any]
+            Arguments to pass to ``split_func``.
+
+        Returns
+        -------
+        list
+            A list of partitions.
+        """
+        logger = get_logger()
+        logger.debug(f"ENTER::Partition.split::{self._identity}")
+
+        @ray.remote(num_returns=num_splits)
+        def _split_df(df, split_func, *args):  # pragma: no cover
+            return split_func(df, *args)
+
+        logger.debug(f"SUBMIT::_split_df::{self._identity}")
+        outputs = _split_df.remote(self._data, split_func, *args)
+        logger.debug(f"EXIT::Partition.split::{self._identity}")
+        return [PandasOnRayDataframePartition(output) for output in outputs]
+
 
 @ray.remote(num_returns=2)
 def _get_index_and_columns(df):
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
index 4f3a5ce387b..ca34b504136 100644
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -37,6 +37,7 @@
 import warnings
 
 from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
+from modin.config import Engine
 from modin.error_message import ErrorMessage
 from modin.utils import (
     try_cast_to_pandas,
@@ -3252,6 +3253,12 @@ def iloc_mut(partition, row_internal_indices, col_internal_indices, item):
         return self.__constructor__(new_modin_frame)
 
     def sort_rows_by_column_values(self, columns, ascending=True, **kwargs):
+        # Our algebra sort is only implemented for Engines that support virtual partitioning.
+        if Engine.get() in ["Ray", "Dask", "Unidist"]:
+            new_modin_frame = self._modin_frame.sort_by(
+                0, columns, ascending=ascending, **kwargs
+            )
+            return self.__constructor__(new_modin_frame)
         ignore_index = kwargs.get("ignore_index", False)
         kwargs["ignore_index"] = False
         if not is_list_like(columns):
diff --git a/modin/pandas/test/test_general.py b/modin/pandas/test/test_general.py
index 4b6f3a4acdf..62b6d3993d6 100644
--- a/modin/pandas/test/test_general.py
+++ b/modin/pandas/test/test_general.py
@@ -19,6 +19,7 @@
 from numpy.testing import assert_array_equal
 from modin.utils import get_current_execution, to_pandas
 from modin.test.test_utils import warns_that_defaulting_to_pandas
+from modin.config import Engine
 from pandas.testing import assert_frame_equal
 
 from .utils import (
@@ -636,7 +637,11 @@ def test_unique():
 def test_value_counts(normalize, bins, dropna):
     # We sort indices for Modin and pandas result because of issue #1650
     values = np.array([3, 1, 2, 3, 4, np.nan])
-    with warns_that_defaulting_to_pandas():
+    with (
+        _nullcontext()
+        if Engine.get() in ["Ray", "Dask", "Unidist"]
+        else warns_that_defaulting_to_pandas()
+    ):
         modin_result = sort_index_for_equal_values(
             pd.value_counts(values, normalize=normalize, ascending=False), False
         )
@@ -654,7 +659,11 @@ def test_value_counts(normalize, bins, dropna):
     )
     df_equals(modin_result, pandas_result)
 
-    with warns_that_defaulting_to_pandas():
+    with (
+        _nullcontext()
+        if Engine.get() in ["Ray", "Dask", "Unidist"]
+        else warns_that_defaulting_to_pandas()
+    ):
         modin_result = sort_index_for_equal_values(
             pd.value_counts(values, dropna=dropna, ascending=True), True
         )
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
index f14f77f06b5..a9b077fe0ca 100644
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -3179,7 +3179,7 @@ def test_sort_values(data, ascending, na_position):
     # between `pandas.Series.sort_values`. For this reason, we check that the values are
     # identical instead of the index as well.
     if ascending:
-        df_equals(modin_result, pandas_result)
+        df_equals_with_non_stable_indices(modin_result, pandas_result)
     else:
         np.testing.assert_equal(modin_result.values, pandas_result.values)
 
@@ -3193,7 +3193,7 @@ def test_sort_values(data, ascending, na_position):
     )
     # See above about `ascending=False`
     if ascending:
-        df_equals(modin_series_cp, pandas_series_cp)
+        df_equals_with_non_stable_indices(modin_result, pandas_result)
     else:
         np.testing.assert_equal(modin_series_cp.values, pandas_series_cp.values)
 


Code review comments:
diff_hunk: @@ -2006,83 +2007,84 @@ def sort_function(df):
             return df
 
         axis = Axis(axis)
-        if axis == Axis.ROW_WISE:
-            # If this df is empty, we don't want to try and shuffle or sort.
-            if len(self.axes[0]) == 0 or len(self.axes[1]) == 0:
-                return self.copy()
-            # If this df only has one row partition, we don't want to do a shuffle and sort - we can
-            # just do a full-axis sort.
-            if len(self._partitions) == 1:
-                return self.apply_full_axis(
-                    1,
-                    sort_function,
-                )
-            if self.dtypes[columns[0]] == object:
-                # This means we are not sorting numbers, so we need our quantiles to not try
-                # arithmetic on the values.
-                method = "inverted_cdf"
-            else:
-                method = "linear"
-            from .utils import build_sort_functions
-
-            shuffling_functions = build_sort_functions(
-                self,
-                columns[0],
-                method,
-                ascending if not isinstance(ascending, list) else ascending[0],
-                **kwargs,
+        if axis != Axis.ROW_WISE:
+            raise NotImplementedError(
+                f"Algebra sort only implemented row-wise. {axis.name} sort not implemented yet!"
             )
-            major_col_partition_index = self.columns.get_loc(columns[0])
-            cols_seen = 0
-            index = -1
-            for i, length in enumerate(self._column_widths_cache):
-                cols_seen += length
-                if major_col_partition_index <= cols_seen:
-                    index = i
-                    break
-            new_partitions = self._partition_mgr_cls.shuffle_partitions(
-                self._partitions,
-                index,
-                shuffling_functions.sample_function,
-                shuffling_functions.pivot_function,
-                shuffling_functions.split_function,
+
+        # If this df is empty, we don't want to try and shuffle or sort.
+        if len(self.axes[0]) == 0 or len(self.axes[1]) == 0:
+            return self.copy()
+        # If this df only has one row partition, we don't want to do a shuffle and sort - we can
+        # just do a full-axis sort.
+        if len(self._partitions) == 1:
+            return self.apply_full_axis(
+                1,
                 sort_function,
             )
-            new_axes = self.axes
-            if kwargs.get("ignore_index", False):
-                new_axes[axis.value] = RangeIndex(len(new_axes[axis.value]))
-            else:
-                new_axes[axis.value] = self._compute_axis_labels_and_lengths(
-                    axis.value, new_partitions
-                )[0]
-            # Sometimes when performing a sort for `value_counts` the partitions have a MultiIndex
-            # but the _compute_axis_labels_and_lengths function gives us an Index of tuples instead
-            # of a MultiIndex.
-            if isinstance(self.axes[axis.value], pandas.MultiIndex) and isinstance(
-                new_axes[axis.value][0], tuple
-            ):
-                new_axes[axis.value] = pandas.MultiIndex.from_tuples(
-                    new_axes[axis.value].values
-                )
-            new_axes[axis.value] = new_axes[axis.value].set_names(
-                self.axes[axis.value].names
-            )
-            new_axes[axis.value ^ 1] = self.axes[axis.value ^ 1]
-            new_lengths = [None, None]
-            (
-                new_partitions,
-                new_lengths[axis.value],
-            ) = self._partition_mgr_cls.rebalance_partitions(new_partitions)
-            col_parts = new_partitions[0]
-            new_lengths[axis.value ^ 1] = [part.width() for part in col_parts]
-            new_modin_frame = self.__constructor__(
-                new_partitions, *new_axes, *new_lengths, self.dtypes
+        if self.dtypes[columns[0]] == object:
+            # This means we are not sorting numbers, so we need our quantiles to not try
+            # arithmetic on the values.
+            method = "inverted_cdf"
+        else:
+            method = "linear"
+
+        shuffling_functions = build_sort_functions(
+            self,
+            columns[0],
+            method,
+            ascending[0] if is_list_like(ascending) else ascending,
+            **kwargs,
+        )
+        major_col_partition_index = self.columns.get_loc(columns[0])
+        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(
+            self._partitions,
+            index,
+            shuffling_functions,
+            sort_function,
+        )
+        new_axes = self.axes
+        if kwargs.get("ignore_index", False):
+            new_axes[axis.value] = RangeIndex(len(new_axes[axis.value]))
+        else:
+            new_axes[axis.value] = self._compute_axis_labels_and_lengths(
+                axis.value, new_partitions
+            )[0]
path: modin/core/dataframe/pandas/dataframe/dataframe.py
text: @user1:
I wonder if we should save the second output of `._compute_axis_labels_and_lengths()` (partitions sizes I presume) and pass it in the `__constructor__` down below?

@author:
We call `rebalance_partitions` a few lines lower, and save the lengths there, so I don't think it makes sense to store this value, since we'll get it when we call `rebalance_partitions` anyways.