Repository: modin-project/modin

Issue:
Everything after DataFrame.sample() running out of memory
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux (remote machine) 
- **Modin version** (`modin.__version__`): 0.7.0
- **Python version**: 3.7
- **Code we can use to reproduce**:
```python
import modin.pandas as pd

# read data files
df1 = pd.read_parquet("data1.parquet")
df2 = pd.read_parquet("data2.parquet") 

# do some operations
df = pd.merge(df1, df2, right_index=True, left_index=True, sort=False)

# write out the result
df.to_parquet("new_data.parquet")

# but if you sample down
df = df.sample(n=10000, random_state=42)

# then do anything after it, nothing works
df.to_parquet("new_data2.parquet") #<-- doesn't complete
```
<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->

Everything after `DataFrame.sample()` is not working/killing the process/workers running out of memory (similar to this [issue](https://github.com/dask/dask/issues/4256) on Dask)

The same code runs beautifully without `.sample()` on the same machine. Whenever `.sample()` is summoned, the workers start running out of memory and the run gets killed after a couple of retries. 

What's worth noting here that modin is doing a beautiful job on the complete data file (~ 1M records), but then whenever I sub-sample to say 100,000 records, it starts running out of memory!

PS: I'm running modin + dask

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->



Discussion:
Hi @jadhosn thanks for the report!

>What's worth noting here that modin is doing a beautiful job on the complete data file (~ 1M records), but then whenever I sub-sample to say 100,000 records, it starts running out of memory!

This is interesting, looking through the `sample` [code](https://github.com/modin-project/modin/blob/6961c50ead66f6b6def9631459880f212908f727/modin/pandas/base.py#L2345), it should not be creating copies. I spent some time triaging this here is what I can note (with Dask):

- `sample` is taking way too much time and using only one core the whole time. 2 minutes on 1m lines of code is not reasonable, we should fix this
- Python worker memory does appear to increase even though it should not be. It is going through the internal `mask`, which has a memory and runtime optimized implementation so I would not expect this. 
- If I dig into the implementation a bit, the rearranged nature of the `sample` is causing `mask` to take longer than expected. Example:

```python
# From internal implementation
sample_indices = random_num_gen.choice( 
    np.arange(0, len(df)), size=n, replace=replace, p=weights 
)
df.iloc[sample_indices] # Takes >2 min on 1m rows, creates >8000 partitions
df.iloc[np.sort(sample_indices)] # Takes 173ms on 1m rows
```

I will optimistically tag for next release (Monday) but cannot make a promise here. If you are okay with sorted indices for this use case, I can provide a temporary workaround.
@devin-petersohn thank you for the quick response! That seems like exactly what I'm seeing. 

I am okay with sorted indices (assuming this operation is not expensive in modin as I will be executing it right before the `sample()` call). What workaround did you have in mind? 

This will be overfit to your example code above, but you should be able to adjust as you need.

Code you provided:
```python
# read data files
df1 = pd.read_parquet("data1.parquet")
df2 = pd.read_parquet("data2.parquet") 

# do some operations
df = pd.merge(df1, df2, right_index=True, left_index=True, sort=False)
```

To replace `sample`:
```python
random_state = 42
n = 10_000 # 10,000 records
replace = False
weights = None
random_num_gen = np.random.RandomState(random_state)
sample_indices = random_num_gen.choice( 
    np.arange(0, len(df)), size=n, replace=replace, p=weights 
)
sampled_df = df.iloc[np.sort(sample_indices)] # Takes 173ms on 1m rows
```

This should be near instant from an interactive perspective.
We weren't able to get to this in time for this release, will prioritize for the next release.

Resolving patch:
diff --git a/modin/engines/base/frame/data.py b/modin/engines/base/frame/data.py
--- a/modin/engines/base/frame/data.py
+++ b/modin/engines/base/frame/data.py
@@ -11,9 +11,8 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
-from itertools import groupby
+from collections import OrderedDict
 import numpy as np
-from operator import itemgetter
 import pandas
 from pandas.core.indexes.api import ensure_index
 from pandas.core.dtypes.common import is_numeric_dtype
@@ -329,14 +328,21 @@ def mask(
         Note: If both row_indices and row_numeric_idx are set, row_indices will be used.
             The same rule applied to col_indices and col_numeric_idx.
 
-        Args:
-            row_indices: The row labels to extract.
-            row_numeric_idx: The row indices to extract.
-            col_indices: The column labels to extract.
-            col_numeric_idx: The column indices to extract.
+        Parameters
+        ----------
+        row_indices : list of hashable
+            The row labels to extract.
+        row_numeric_idx : list of int
+            The row indices to extract.
+        col_indices : list of hashable
+            The column labels to extract.
+        col_numeric_idx : list of int
+            The column indices to extract.
 
-        Returns:
-             A new dataframe.
+        Returns
+        -------
+        BasePandasFrame
+             A new BasePandasFrame from the mask provided.
         """
         if (
             row_indices is None
@@ -348,34 +354,34 @@ def mask(
         if row_indices is not None:
             row_numeric_idx = self.index.get_indexer_for(row_indices)
         if row_numeric_idx is not None:
-            row_partitions_list = self._get_dict_of_block_index(
-                1, row_numeric_idx, ordered=True
-            )
-            new_row_lengths = [len(indices) for _, indices in row_partitions_list]
-            new_index = self.index[row_numeric_idx]
-        else:
-            row_partitions_list = [
-                (i, slice(None)) for i in range(len(self._row_lengths))
+            row_partitions_list = self._get_dict_of_block_index(1, row_numeric_idx)
+            new_row_lengths = [
+                len(indices) for _, indices in row_partitions_list.items()
             ]
+            new_index = self.index[sorted(row_numeric_idx)]
+        else:
+            row_partitions_list = {
+                i: slice(None) for i in range(len(self._row_lengths))
+            }
             new_row_lengths = self._row_lengths
             new_index = self.index
 
         if col_indices is not None:
             col_numeric_idx = self.columns.get_indexer_for(col_indices)
         if col_numeric_idx is not None:
-            col_partitions_list = self._get_dict_of_block_index(
-                0, col_numeric_idx, ordered=True
-            )
-            new_col_widths = [len(indices) for _, indices in col_partitions_list]
-            new_columns = self.columns[col_numeric_idx]
+            col_partitions_list = self._get_dict_of_block_index(0, col_numeric_idx)
+            new_col_widths = [
+                len(indices) for _, indices in col_partitions_list.items()
+            ]
+            new_columns = self.columns[sorted(col_numeric_idx)]
             if self._dtypes is not None:
-                new_dtypes = self.dtypes[col_numeric_idx]
+                new_dtypes = self.dtypes[sorted(col_numeric_idx)]
             else:
                 new_dtypes = None
         else:
-            col_partitions_list = [
-                (i, slice(None)) for i in range(len(self._column_widths))
-            ]
+            col_partitions_list = {
+                i: slice(None) for i in range(len(self._column_widths))
+            }
             new_col_widths = self._column_widths
             new_columns = self.columns
             if self._dtypes is not None:
@@ -388,16 +394,16 @@ def mask(
                     self._partitions[row_idx][col_idx].mask(
                         row_internal_indices, col_internal_indices
                     )
-                    for col_idx, col_internal_indices in col_partitions_list
+                    for col_idx, col_internal_indices in col_partitions_list.items()
                     if isinstance(col_internal_indices, slice)
                     or len(col_internal_indices) > 0
                 ]
-                for row_idx, row_internal_indices in row_partitions_list
+                for row_idx, row_internal_indices in row_partitions_list.items()
                 if isinstance(row_internal_indices, slice)
                 or len(row_internal_indices) > 0
             ]
         )
-        return self.__constructor__(
+        intermediate = self.__constructor__(
             new_partitions,
             new_index,
             new_columns,
@@ -405,6 +411,76 @@ def mask(
             new_col_widths,
             new_dtypes,
         )
+        # Check if monotonically increasing, return if it is. Fast track code path for
+        # common case to keep it fast.
+        if (
+            row_numeric_idx is None
+            or len(row_numeric_idx) == 1
+            or np.all(row_numeric_idx[1:] >= row_numeric_idx[:-1])
+        ) and (
+            col_numeric_idx is None
+            or len(col_numeric_idx) == 1
+            or np.all(col_numeric_idx[1:] >= col_numeric_idx[:-1])
+        ):
+            return intermediate
+        # The new labels are often smaller than the old labels, so we can't reuse the
+        # original order values because those were mapped to the original data. We have
+        # to reorder here based on the expected order from within the data.
+        # We create a dictionary mapping the position of the numeric index with respect
+        # to all others, then recreate that order by mapping the new order values from
+        # the old. This information is sent to `reorder_labels`.
+        if row_numeric_idx is not None:
+            row_order_mapping = dict(
+                zip(sorted(row_numeric_idx), range(len(row_numeric_idx)))
+            )
+            new_row_order = [row_order_mapping[idx] for idx in row_numeric_idx]
+        else:
+            new_row_order = None
+        if col_numeric_idx is not None:
+            col_order_mapping = dict(
+                zip(sorted(col_numeric_idx), range(len(col_numeric_idx)))
+            )
+            new_col_order = [col_order_mapping[idx] for idx in col_numeric_idx]
+        else:
+            new_col_order = None
+        return intermediate.reorder_labels(
+            row_numeric_idx=new_row_order, col_numeric_idx=new_col_order
+        )
+
+    def reorder_labels(self, row_numeric_idx=None, col_numeric_idx=None):
+        """Reorder the column and or rows in this DataFrame.
+
+        Parameters
+        ----------
+        row_numeric_idx : list of int, optional
+            The ordered list of new row orders such that each position within the list
+            indicates the new position.
+        col_numeric_idx : list of int, optional
+            The ordered list of new column orders such that each position within the
+            list indicates the new position.
+
+        Returns
+        -------
+        BasePandasFrame
+            A new BasePandasFrame with reordered columns and/or rows.
+        """
+        if row_numeric_idx is not None:
+            ordered_rows = self._frame_mgr_cls.map_axis_partitions(
+                0, self._partitions, lambda df: df.iloc[row_numeric_idx]
+            )
+            row_idx = self.index[row_numeric_idx]
+        else:
+            ordered_rows = self._partitions
+            row_idx = self.index
+        if col_numeric_idx is not None:
+            ordered_cols = self._frame_mgr_cls.map_axis_partitions(
+                1, ordered_rows, lambda df: df.iloc[:, col_numeric_idx]
+            )
+            col_idx = self.columns[col_numeric_idx]
+        else:
+            ordered_cols = ordered_rows
+            col_idx = self.columns
+        return self.__constructor__(ordered_cols, row_idx, col_idx)
 
     def copy(self):
         """Copy this object.
@@ -544,26 +620,23 @@ def _numeric_columns(self, include_bool=True):
                 columns.append(col)
         return columns
 
-    def _get_dict_of_block_index(self, axis, indices, ordered=False):
+    def _get_dict_of_block_index(self, axis, indices):
         """Convert indices to a dict of block index to internal index mapping.
 
-        Note: See `_get_blocks_containing_index` for primary usage. This method
-            accepts a list of indices rather than just a single value, and uses
-            `_get_blocks_containing_index`.
-
-        Args:
-            axis: The axis along which to get the indices
-                (0 - columns, 1 - rows)
-            indices: A list of global indices to convert.
+        Parameters
+        ----------
+        axis : (0 - columns, 1 - rows)
+               The axis along which to get the indices
+        indices : list of int
+                A list of global indices to convert.
 
         Returns
-            For unordered: a dictionary of {block index: list of local indices}.
-            For ordered: a list of tuples mapping block index: list of local indices.
+        -------
+        dictionary mapping int to list of int
+            A mapping from partition to list of internal indices to extract from that
+            partition.
         """
-        if not ordered:
-            indices = np.sort(indices)
-        else:
-            indices = np.array(indices)
+        indices = np.sort(indices)
         if not axis:
             bins = np.array(self._column_widths)
         else:
@@ -581,52 +654,36 @@ def internal(block_idx, global_index):
             )
 
         partition_ids = np.digitize(indices, cumulative)
-        # If the output order doesn't matter or if the indices are monotonically
-        # increasing, the computation is significantly simpler and faster than doing
-        # the zip and groupby.
-        if not ordered or np.all(np.diff(indices) > 0):
-            count_for_each_partition = np.array(
-                [(partition_ids == i).sum() for i in range(len(cumulative))]
-            ).cumsum()
-            # Compute the internal indices and pair those with the partition index.
-            # If the first partition has any values we need to return, compute those
-            # first to make the list comprehension easier. Otherwise, just append the
-            # rest of the values to an empty list.
-            if count_for_each_partition[0] > 0:
-                first_partition_indices = [
-                    (0, internal(0, indices[slice(count_for_each_partition[0])]))
-                ]
-            else:
-                first_partition_indices = []
-            partition_ids_with_indices = first_partition_indices + [
-                (
-                    i,
-                    internal(
-                        i,
-                        indices[
-                            slice(
-                                count_for_each_partition[i - 1],
-                                count_for_each_partition[i],
-                            )
-                        ],
-                    ),
-                )
-                for i in range(1, len(count_for_each_partition))
-                if count_for_each_partition[i] > count_for_each_partition[i - 1]
+        count_for_each_partition = np.array(
+            [(partition_ids == i).sum() for i in range(len(cumulative))]
+        ).cumsum()
+        # Compute the internal indices and pair those with the partition index.
+        # If the first partition has any values we need to return, compute those
+        # first to make the list comprehension easier. Otherwise, just append the
+        # rest of the values to an empty list.
+        if count_for_each_partition[0] > 0:
+            first_partition_indices = [
+                (0, internal(0, indices[slice(count_for_each_partition[0])]))
             ]
-            return (
-                dict(partition_ids_with_indices)
-                if not ordered
-                else partition_ids_with_indices
+        else:
+            first_partition_indices = []
+        partition_ids_with_indices = first_partition_indices + [
+            (
+                i,
+                internal(
+                    i,
+                    indices[
+                        slice(
+                            count_for_each_partition[i - 1],
+                            count_for_each_partition[i],
+                        )
+                    ],
+                ),
             )
-
-        all_partitions_and_idx = zip(partition_ids, indices)
-        # In ordered, we have to maintain the order of the list of indices provided.
-        # This means that we need to return a list instead of a dictionary.
-        return [
-            (k, internal(k, [x for _, x in v]))
-            for k, v in groupby(all_partitions_and_idx, itemgetter(0))
+            for i in range(1, len(count_for_each_partition))
+            if count_for_each_partition[i] > count_for_each_partition[i - 1]
         ]
+        return OrderedDict(partition_ids_with_indices)
 
     def _join_index_objects(self, axis, other_index, how, sort):
         """Joins a pair of index objects (columns or rows) by a given strategy.


Test patch: