Repository: modin-project/modin

Issue:
BUG:  `.sort_values` fails with a KeyError
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [ ] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import os
import modin.pandas as pd

columns_names = [
    "trip_id",
    "vendor_id",
    "pickup_datetime",
    "dropoff_datetime",
    "store_and_fwd_flag",
    "rate_code_id",
    "pickup_longitude",
    "pickup_latitude",
    "dropoff_longitude",
    "dropoff_latitude",
    "passenger_count",
    "trip_distance",
    "fare_amount",
    "extra",
    "mta_tax",
    "tip_amount",
    "tolls_amount",
    "ehail_fee",
    "improvement_surcharge",
    "total_amount",
    "payment_type",
    "trip_type",
    "pickup",
    "dropoff",
    "cab_type",
    "precipitation",
    "snow_depth",
    "snowfall",
    "max_temperature",
    "min_temperature",
    "average_wind_speed",
    "pickup_nyct2010_gid",
    "pickup_ctlabel",
    "pickup_borocode",
    "pickup_boroname",
    "pickup_ct2010",
    "pickup_boroct2010",
    "pickup_cdeligibil",
    "pickup_ntacode",
    "pickup_ntaname",
    "pickup_puma",
    "dropoff_nyct2010_gid",
    "dropoff_ctlabel",
    "dropoff_borocode",
    "dropoff_boroname",
    "dropoff_ct2010",
    "dropoff_boroct2010",
    "dropoff_cdeligibil",
    "dropoff_ntacode",
    "dropoff_ntaname",
    "dropoff_puma",
]

columns_types = [
        "int64",
        "category",
        "datetime64",
        "datetime64",
        "category",
        "int64",
        "float64",
        "float64",
        "float64",
        "float64",
        "int64",
        "float64",
        "float64",
        "float64",
        "float64",
        "float64",
        "float64",
        "float64",
        "float64",
        "float64",
        "category",
        "float64",
        "category",
        "category",
        "category",
        "float64",
        "int64",
        "float64",
        "int64",
        "int64",
        "float64",
        "float64",
        "float64",
        "float64",
        "category",
        "float64",
        "float64",
        "category",
        "category",
        "category",
        "float64",
        "float64",
        "float64",
        "float64",
        "category",
        "float64",
        "float64",
        "category",
        "category",
        "category",
        "float64",
    ]

def q4(df):
    df["pickup_datetime"] = df["pickup_datetime"].dt.year
    df["trip_distance"] = df["trip_distance"].astype("int64")
    q4_output = (
            df.groupby(["passenger_count", "pickup_datetime", "trip_distance"], sort=False)
            .size()
            .reset_index()
            .sort_values(by=["pickup_datetime", 0], ignore_index=True, ascending=[True, False])
    )
    q4_output.shape  # to trigger real execution

    return q4_output
# trips_xaa.csv could be picked up from s3://modin-datasets
df = pd.read_csv("trips_xaa.csv", names=columns_names, dtype=None, header=None, nrows=None, parse_dates=["pickup_datetime", "dropoff_datetime"])
print(q4(df))
```


### Issue Description

Looks like a regression after sorting speed up commit (https://github.com/modin-project/modin/commit/f0d06357450a0bd7d1ed7520ee8cac2de905e58e). Now I got an exception for the operation with sort_values involved


### Expected Behavior

No exception

### Error Logs

<details>

```python-traceback

Traceback (most recent call last):
  File "t2.py", line 126, in <module>
    print(q4(df))
  File "t2.py", line 116, in q4
    df.groupby(["passenger_count", "pickup_datetime", "trip_distance"], sort=False)
  File "/localdisk/dchigare/repos/modin/modin/logging/logger_decorator.py", line 128, in run_and_log
    return obj(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/pandas/base.py", line 2903, in sort_values
    result = self._query_compiler.sort_rows_by_column_values(
  File "/localdisk/dchigare/repos/modin/modin/logging/logger_decorator.py", line 128, in run_and_log
    return obj(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/storage_formats/pandas/query_compiler.py", line 3258, in sort_rows_by_column_values
    new_modin_frame = self._modin_frame.sort_by(
  File "/localdisk/dchigare/repos/modin/modin/logging/logger_decorator.py", line 128, in run_and_log
    return obj(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 126, in run_f_on_minimally_updated_metadata
    result = f(self, *args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 2078, in sort_by
    new_modin_frame._propagate_index_objs(axis=0)
  File "/localdisk/dchigare/repos/modin/modin/logging/logger_decorator.py", line 128, in run_and_log
    return obj(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 530, in _propagate_index_objs
    self._filter_empties()
  File "/localdisk/dchigare/repos/modin/modin/logging/logger_decorator.py", line 128, in run_and_log
    return obj(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 487, in _filter_empties
    [
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 494, in <listcomp>
    if i < len(self.row_lengths) and self.row_lengths[i] != 0
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 253, in row_lengths
    self._row_lengths_cache = [part.length() for part in row_parts]
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/dataframe.py", line 253, in <listcomp>
    self._row_lengths_cache = [part.length() for part in row_parts]
  File "/localdisk/dchigare/repos/modin/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py", line 282, in length
    self._length_cache = RayWrapper.materialize(self._length_cache)
  File "/localdisk/dchigare/repos/modin/modin/core/execution/ray/common/engine_wrapper.py", line 92, in materialize
    return ray.get(obj_id)
  File "/localdisk/dchigare/miniconda3/envs/modin/lib/python3.8/site-packages/ray/_private/client_mode_hook.py", line 105, in wrapper
    return func(*args, **kwargs)
  File "/localdisk/dchigare/miniconda3/envs/modin/lib/python3.8/site-packages/ray/_private/worker.py", line 2280, in get
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(KeyError): ray::deploy_ray_func() (pid=4021249, ip=10.241.129.69)
  At least one of the input arguments for this task could not be computed:
ray.exceptions.RayTaskError: ray::deploy_ray_func() (pid=4021221, ip=10.241.129.69)
  At least one of the input arguments for this task could not be computed:
ray.exceptions.RayTaskError: ray::_deploy_ray_func() (pid=4012398, ip=10.241.129.69)
  File "/localdisk/dchigare/repos/modin/modin/core/execution/ray/common/engine_wrapper.py", line 44, in _deploy_ray_func
    return func(*args, **kwargs)
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/utils.py", line 70, in split_fn
    return split_partitions_using_pivots_for_sort(
  File "/localdisk/dchigare/repos/modin/modin/core/dataframe/pandas/dataframe/utils.py", line 264, in split_partitions_using_pivots_for_sort
    grouped = non_na_rows.groupby(groupby_col)
  File "/localdisk/dchigare/miniconda3/envs/modin/lib/python3.8/site-packages/pandas/core/frame.py", line 8399, in groupby
    return DataFrameGroupBy(
  File "/localdisk/dchigare/miniconda3/envs/modin/lib/python3.8/site-packages/pandas/core/groupby/groupby.py", line 959, in __init__
    grouper, exclusions, obj = get_grouper(
  File "/localdisk/dchigare/miniconda3/envs/modin/lib/python3.8/site-packages/pandas/core/groupby/grouper.py", line 888, in get_grouper
    raise KeyError(gpr)
KeyError: 38

```

</details>


### Installed Versions

commit 0517922f03f4dada335ea3d098b4051880991405



Discussion:
cc @RehanSD   looks like showstopper for 0.18 if confirmed
Tried reproducing locally with the first 500,000; 1,000,000, and 5,328,631 lines of `trips_xaa.csv` and it works - attempting to get access to the machine Intel is using to see if I can get some more insight into why its failing.

Took a look at the specific line that's failing, and it looks like the `groupby` in the stage where we split the dataframe by its pivots is failing at the following lines in pure pandas:
```python
        elif is_in_axis(gpr):  # df.groupby('name')
            if gpr in obj:
                if validate:
                    obj._check_label_or_level_ambiguity(gpr, axis=axis)
                in_axis, name, gpr = True, gpr, obj[gpr]
                if gpr.ndim != 1:
                    # non-unique columns; raise here to get the name in the
                    # exception message
                    raise ValueError(f"Grouper for '{name}' not 1-dimensional")
                exclusions.add(name)
            elif obj._is_level_reference(gpr, axis=axis):
                in_axis, level, gpr = False, gpr, None
            else:
                raise KeyError(gpr)
```

which is odd, since the key we are grouping by is not in axis. If I re-run the code having modified my pandas import, I note that we don't hit this elif case at all (again, with 5328631 lines)

Modified pandas code:
```python
    for gpr, level in zip(keys, levels):
        print("HI - we're here!")
        print("GPR:", gpr)
        print("level:", level)
        if is_in_obj(gpr):  # df.groupby(df['name'])
            in_axis = True
            exclusions.add(gpr.name)

        elif is_in_axis(gpr):  # df.groupby('name')
            print("HI - I'm in this case!")
            print("GPR Inside elif case:", gpr)
            if gpr in obj:
                if validate:
                    obj._check_label_or_level_ambiguity(gpr, axis=axis)
                in_axis, name, gpr = True, gpr, obj[gpr]
                if gpr.ndim != 1:
                    # non-unique columns; raise here to get the name in the
                    # exception message
                    raise ValueError(f"Grouper for '{name}' not 1-dimensional")
                exclusions.add(name)
            elif obj._is_level_reference(gpr, axis=axis):
                in_axis, level, gpr = False, gpr, None
            else:
                raise KeyError(gpr)
        elif isinstance(gpr, Grouper) and gpr.key is not None:
            # Add key to exclusions
            exclusions.add(gpr.key)
            in_axis = False
        else:
            print("Hitting this else, not in axis")
            in_axis = False
```

Output from running test script:
```python
_apply_list_of_funcs pid=54221) HI - we're here!
<Abridged, since this is before we even hit the sort>
(_apply_list_of_funcs pid=54259) GPR Inside elif case: trip_distance
<Hit the QC sort in the next line>
SORT KEYS: ['pickup_datetime', 0]
(deploy_ray_func pid=54222) HI - we're here!
(deploy_ray_func pid=54222) GPR: None
(deploy_ray_func pid=54222) level: 0
(deploy_ray_func pid=54222) Hitting this else, not in axis
(deploy_ray_func pid=54222) HI - we're here!
(deploy_ray_func pid=54222) GPR: None
(deploy_ray_func pid=54222) level: 1
(deploy_ray_func pid=54222) Hitting this else, not in axis
(deploy_ray_func pid=54222) HI - we're here!
(deploy_ray_func pid=54222) GPR: None
(deploy_ray_func pid=54222) level: 2
(deploy_ray_func pid=54222) Hitting this else, not in axis
<Hit the _split_df function in the next line, which is where KeyError is originating from>
(_split_df pid=54222) HI - we're here!
(_split_df pid=54222) GPR: [15 15 15 15 15 15 15 15  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7
(_split_df pid=54222)   7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7]
(_split_df pid=54222) level: None
(_split_df pid=54222) Hitting this else, not in axis
(_split_df pid=54222) HI - we're here!
(_split_df pid=54222) GPR: [ 7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7
(_split_df pid=54222)   7  7  7  7  7  7  7  7  7  7  7 15 15 15 15 15 15 15]
(_split_df pid=54222) level: None
(_split_df pid=54222) Hitting this else, not in axis
(_split_df pid=54224) HI - we're here!
(_split_df pid=54224) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54224)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54224) level: None
(_split_df pid=54224) Hitting this else, not in axis
(_split_df pid=54215) HI - we're here!
(_split_df pid=54215) GPR: [ 7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7
(_split_df pid=54215)   7  7  7  7  7  7  7  7  7  7 15 15 15 15 15 15 15 15]
(_split_df pid=54215) level: None
(_split_df pid=54215) Hitting this else, not in axis
(_split_df pid=54215) HI - we're here!
(_split_df pid=54215) GPR: [15 15  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7
(_split_df pid=54215)   7  7  7  7  7  7  7 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54215) level: None
(_split_df pid=54215) Hitting this else, not in axis
(_split_df pid=54217) HI - we're here!
(_split_df pid=54217) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54217)  15 15 15 15 15  7  7  7  7  7  7 15 15 15 15 15 15 15]
(_split_df pid=54217) level: None
(_split_df pid=54217) Hitting this else, not in axis
(_split_df pid=54259) HI - we're here!
(_split_df pid=54259) GPR: [ 7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7 15 15 15
(_split_df pid=54259)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54259) level: None
(_split_df pid=54259) Hitting this else, not in axis
(_split_df pid=54259) HI - we're here!
(_split_df pid=54259) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54259)  15 15 15 15 15 15 15  7  7  7  7  7  7  7  7  7  7  7]
(_split_df pid=54259) level: None
(_split_df pid=54259) Hitting this else, not in axis
(_split_df pid=54353) HI - we're here!
(_split_df pid=54353) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15  7  7  7  7  7  7  7  7  7
(_split_df pid=54353)   7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7]
(_split_df pid=54353) level: None
(_split_df pid=54353) Hitting this else, not in axis
(_split_df pid=54354) HI - we're here!
(_split_df pid=54354) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54354)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54354) level: None
(_split_df pid=54354) Hitting this else, not in axis
(_split_df pid=54360) HI - we're here!
(_split_df pid=54360) GPR: [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54360)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15  7  7]
(_split_df pid=54360) level: None
(_split_df pid=54360) Hitting this else, not in axis
(_split_df pid=54360) HI - we're here!
(_split_df pid=54360) GPR: [15 15 15 15  7  7  7  7 15 15 15 15 15 15 15 15 15 15  7  7  7  7  7 15
(_split_df pid=54360)  15 15 15 15 15 15]
(_split_df pid=54360) level: None
(_split_df pid=54360) Hitting this else, not in axis
(_split_df pid=54361) HI - we're here!
(_split_df pid=54361) GPR: [ 7  7  7  7  7  7 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54361)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54361) level: None
(_split_df pid=54361) Hitting this else, not in axis
(_split_df pid=54365) HI - we're here!
(_split_df pid=54365) GPR: [15 15 15 15 15 15 15 15  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7  7
(_split_df pid=54365)   7  7  7  7  7  7  7  7  7  7  7  7 15 15 15 15 15 15]
(_split_df pid=54365) level: None
(_split_df pid=54365) Hitting this else, not in axis
(_split_df pid=54366) HI - we're here!
(_split_df pid=54366) GPR: [ 7  7  7  7  7 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15
(_split_df pid=54366)  15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15]
(_split_df pid=54366) level: None
(_split_df pid=54366) Hitting this else, not in axis
(_split_df pid=54369) HI - we're here!
(_split_df pid=54369) GPR: [7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7
(_split_df pid=54369)  7 7 7 7 7]
(_split_df pid=54369) level: None
(_split_df pid=54369) Hitting this else, not in axis
```
As expected, the grouper is recognized as not in the axis - perhaps we're getting a weird partitioning where there's only one row in this partition, so it thinks 38 is the column name instead of the value?
When I took a look at the error, it seemed super odd that we were getting a KeyError on 38 - nothing seems to call 38, except for in the final split stage, where we groupby the index of the new partition for each row. I figure what must be happening is that either there’s only one row in that partition, or there’s multiple rows, but only one has a Non-NA value, so when we do the groupby, instead of grouping by a non-axis object as its supposed to, pandas attempts to group by 38 (the new partition index) as a column name rather than an out of axis object.

I confirmed this by testing when we have a partition with all NA values except for one row - we get a similar error. Interestingly, a partition with all NA values succeeds. Pushing a fix!

Thank you @Garra1980 - this was a great catch! (and a very interesting edge case)

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
@@ -2040,7 +2040,7 @@ def sort_function(df):
         index = -1
         for i, length in enumerate(self.column_widths):
             cols_seen += length
-            if major_col_partition_index <= cols_seen:
+            if major_col_partition_index < cols_seen:
                 index = i
                 break
         new_partitions = self._partition_mgr_cls.shuffle_partitions(
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
@@ -261,13 +261,21 @@ def split_partitions_using_pivots_for_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)
-    ]
+    if len(non_na_rows) == 1:
+        groups = [
+            pandas.DataFrame(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
+            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]


Test patch:
diff --git a/modin/pandas/test/dataframe/test_join_sort.py b/modin/pandas/test/dataframe/test_join_sort.py
--- a/modin/pandas/test/dataframe/test_join_sort.py
+++ b/modin/pandas/test/dataframe/test_join_sort.py
@@ -37,7 +37,7 @@
     extra_test_parameters,
     default_to_pandas_ignore_string,
 )
-from modin.config import NPartitions
+from modin.config import NPartitions, Engine
 from modin.test.test_utils import warns_that_defaulting_to_pandas
 
 NPartitions.put(4)
@@ -527,6 +527,41 @@ def test_sort_values_with_string_index():
     df_equals(modin_df, pandas_df)
 
 
+@pytest.mark.skipif(
+    Engine.get() not in ["Ray", "Dask", "Unidist"],
+    reason="We only need to test this case where sort does not default to pandas.",
+)
+@pytest.mark.parametrize("ascending", [True, False], ids=["True", "False"])
+@pytest.mark.parametrize("na_position", ["first", "last"], ids=["first", "last"])
+def test_sort_values_with_only_one_non_na_row_in_partition(ascending, na_position):
+    pandas_df = pandas.DataFrame(
+        np.random.rand(1000, 100), columns=[f"col {i}" for i in range(100)]
+    )
+    # Need to ensure that one of the partitions has all NA values except for one row
+    pandas_df.iloc[340:] = np.NaN
+    pandas_df.iloc[-1] = -4.0
+    modin_df = pd.DataFrame(pandas_df)
+    eval_general(
+        modin_df,
+        pandas_df,
+        lambda df: df.sort_values(
+            "col 3", ascending=ascending, na_position=na_position
+        ),
+    )
+
+
+@pytest.mark.skipif(
+    Engine.get() not in ["Ray", "Dask", "Unidist"],
+    reason="We only need to test this case where sort does not default to pandas.",
+)
+def test_sort_values_with_sort_key_on_partition_boundary():
+    modin_df = pd.DataFrame(
+        np.random.rand(1000, 100), columns=[f"col {i}" for i in range(100)]
+    )
+    sort_key = modin_df.columns[modin_df._query_compiler._modin_frame.column_widths[0]]
+    eval_general(modin_df, modin_df._to_pandas(), lambda df: df.sort_values(sort_key))
+
+
 def test_where():
     columns = list("abcdefghij")