Repository: modin-project/modin
Pull request: FIX-#4358: MultiIndex `loc` shouldn't drop levels for full-key lookups
<!--
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. -->
Currently, we try to drop levels for full-key MultiIndex lookups which cause errors since there are no levels to be dropped from the intermediate DataFrame. This PR addresses this issue and adds a check before trying to drop levels.

As a side note, we should definitely rework how we handle lookups for `MultiIndex` lookups. This is not very well abstracted in Modin as is and we're starting to see more bugs related to this. 

- [x] commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] 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 #4358 
- [x] tests added and passing
- [x] 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 -->
- [x] 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:
MultiIndex: DataFrame.loc with a tuple on columns defined as MultiIndex returns an error
### System information
- **OS Platform and Distribution**: MacOS Monterey 12.2.1
- **Modin version** 0.14
- **Python version**: 3.8.12
- **Code we can use to reproduce**:

### Using modin: 
```python
import ray
ray.init()
import modin.pandas as pd
import numpy as np

arrays = [
    ["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"],
    ["one", "two", "one", "two", "one", "two", "one", "two"],
]
columns =  pd.MultiIndex.from_tuples(zip(*arrays), names=['a', 'b'])
data = [np.random.randint(10, 50, len(columns)) for _ in range(30)]
df = pd.DataFrame(data, columns=columns)

# bug happens here
df.loc[:, ("bar", "two")]
```

### Same operation in vanilla pandas 
```python
import pandas as pd
import numpy as np

arrays = [
    ["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"],
    ["one", "two", "one", "two", "one", "two", "one", "two"],
]
columns =  pd.MultiIndex.from_tuples(zip(*arrays), names=['a', 'b'])
data = [np.random.randint(10, 50, len(columns)) for _ in range(30)]
df = pd.DataFrame(data, columns=columns)

df.loc[:, ("bar", "two")]
```

### Describe the problem
When using .loc with tuples on a multiIndex, a ValueError appears

### Source code / logs
**traceback:**
```shell
Traceback (most recent call last):
  File "/VIRTUALENV/lib/python3.8/site-packages/ipython-7.32.0-py3.8.egg/IPython/core/interactiveshell.py", line 3524, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-4-59ea95655168>", line 1, in <cell line: 1>
    df.loc[:, ("bar", "two")]
  File "/VIRTUALENV/lib/python3.8/site-packages/modin/pandas/indexing.py", line 674, in __getitem__
    result.columns = result.columns.droplevel(list(range(len(col_loc_as_list))))
  File "/VIRTUALENV/lib/python3.8/site-packages/pandas/core/indexes/base.py", line 2063, in droplevel
    return self._drop_level_numbers(levnums)
  File "/VIRTUALENV/lib/python3.8/site-packages/pandas/core/indexes/base.py", line 2074, in _drop_level_numbers
    raise ValueError(
ValueError: Cannot remove 2 levels from an index with 2 levels: at least one level must be left.
```

Hints from issue discussion:
It will however work when explicitely converting the tuple to a MultiIndex: 

```python
df.loc[:, pd.MultiIndex.from_tuples([ ("bar", "two") ]) ]
```
@neltacigreb thank you for the detailed bug report. I reproduced the issue with your code on Modin commit 6d05c31a24da9776723a179c74b70338b9302a82.

The bug happens when the column selectors in the `loc` select one column at each level of the column MultiIndex. Modin needs to stop trying to drop all levels of the result `DataFrame` [here](https://github.com/modin-project/modin/blob/e5e9634357e60925a5a70e56a1d4882d269f533a/modin/pandas/indexing.py#L674). Instead, Modin should make a `Series` whose name comes from the columns selected, e.g. in the example pandas snippet above, the Series name is `('bar', 'bars')`. I guess we should make `_LocationIndexerBase. __getitem__` return a `Series` with the correct column name in this case.
I tried taking a stab at this bug and was able to trace it down for this specific case which is basically what @mvashishtha was able to triage. The axis we squeeze the resulting DataFrame is incorrect and causes some issues (in `__getitem__` in the base indexer). However, while we can come up with a basic fix for now, it's become clearer that the behavior for MultiIndexes is a bit more complicated than we would like to think. Here's a motivating example of some unclear behavior:
```python
df = pd.DataFrame([1, 2], index=[['a', 'a'], ['b', 'b']])
df.loc[('a', 'b'), 0])
```
The issue here is that we have duplicates in the MultiIndex and this causes `loc` to return the correct Series in the intermediate representation. However, the rest of our code tries to incorrectly drop levels, leading to the same error. We see in this current issue. 

 I'm hesitant to merge in a fix that will add more spaghetti code in our code base and I'd prefer that we have some further discussion about the MultiIndex case. Thoughts? @modin-project/modin-core @modin-project/modin-contributors 
I think it's better to show rather than tell, so folks can take a look at what I tried in this draft PR: https://github.com/modin-project/modin/pull/4608

Commit message: Add comment, refactor tests

Patch to review:
diff --git a/docs/release_notes/release_notes-0.16.0.rst b/docs/release_notes/release_notes-0.16.0.rst
index 8592e5ebe95..769f1861e29 100644
--- a/docs/release_notes/release_notes-0.16.0.rst
+++ b/docs/release_notes/release_notes-0.16.0.rst
@@ -21,6 +21,7 @@ Key Features and Updates
   * FIX-#4641: Reindex pandas partitions in `df.describe()` (#4651)
   * FIX-#4634: Check for FrozenList as `by` in `df.groupby()` (#4667)
   * FIX-#4491: Wait for all partitions in parallel in benchmark mode (#4656)
+  * FIX-#4358: MultiIndex `loc` shouldn't drop levels for full-key lookups (#4608)
 * Performance enhancements
   * PERF-#4182: Add cell-wise execution for binary ops, fix bin ops for empty dataframes (#4391)
   * PERF-#4288: Improve perf of `groupby.mean` for narrow data (#4591)
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index 09619eaf5ad..d4256adf7ae 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -283,6 +283,8 @@ def __init__(self, modin_df):
         self.qc = modin_df._query_compiler
         self.row_scalar = False
         self.col_scalar = False
+        self.row_multiindex_full_lookup = False
+        self.col_multiindex_full_lookup = False
 
     def __getitem__(self, row_lookup, col_lookup, ndim):
         """
@@ -331,11 +333,17 @@ def __getitem__(self, row_lookup, col_lookup, ndim):
         elif ndim == 0:
             axis = None
         else:
+            # We are in the case where ndim == 1
+            # The axis we squeeze on depends on whether we are looking for an exact
+            # value or a subset of rows and columns. Knowing if we have a full MultiIndex
+            # lookup or scalar lookup can help us figure out whether we need to squeeze
+            # on the row or column index.
             axis = (
                 None
-                if self.col_scalar and self.row_scalar
+                if (self.col_scalar and self.row_scalar)
+                or (self.row_multiindex_full_lookup and self.col_multiindex_full_lookup)
                 else 1
-                if self.col_scalar
+                if self.col_scalar or self.col_multiindex_full_lookup
                 else 0
             )
         return self.df.__constructor__(query_compiler=qc_view).squeeze(axis=axis)
@@ -528,6 +536,39 @@ def _handle_boolean_masking(self, row_loc, col_loc):
         # Passing `slice(None)` as a row indexer since we've just applied it
         return type(self)(masked_df)[(slice(None), col_loc)]
 
+    def _multiindex_possibly_contains_key(self, axis, key):
+        """
+        Determine if a MultiIndex row/column possibly contains a key.
+
+        Check to see if the current DataFrame has a MultiIndex row/column and if it does,
+        check to see if the key is potentially a full key-lookup such that the number of
+        levels match up with the length of the tuple key.
+
+        Parameters
+        ----------
+        axis : int
+            0 or 1 for row/column.
+        key : Any
+            Lookup key for MultiIndex row/column.
+
+        Returns
+        -------
+        bool
+            If the MultiIndex possibly contains the given key.
+
+        Notes
+        -----
+        This function only returns False if we have a partial key lookup. It's
+        possibly that this function returns True for a key that does NOT exist
+        since we only check the length of the tuple to match the number of levels
+        in the MultiIndex row/colunmn.
+        """
+        if not self.qc.has_multiindex(axis=axis):
+            return False
+
+        multiindex = self.df.index if axis == 0 else self.df.columns
+        return isinstance(key, tuple) and len(key) == len(multiindex.levels)
+
 
 class _LocIndexer(_LocationIndexerBase):
     """
@@ -563,6 +604,22 @@ def __getitem__(self, key):
         self.row_scalar = is_scalar(row_loc)
         self.col_scalar = is_scalar(col_loc)
 
+        # The thought process here is that we should check to see that we have a full key lookup
+        # for a MultiIndex DataFrame. If that's the case, then we should not drop any levels
+        # since our resulting intermediate dataframe will have dropped these for us already.
+        # Thus, we need to make sure we don't try to drop these levels again. The logic here is
+        # kind of hacked together. Ideally, we should handle this properly in the lower-level
+        # implementations, but this will have to be engineering properly later.
+        self.row_multiindex_full_lookup = self._multiindex_possibly_contains_key(
+            axis=0, key=row_loc
+        )
+        self.col_multiindex_full_lookup = self._multiindex_possibly_contains_key(
+            axis=1, key=col_loc
+        )
+        levels_already_dropped = (
+            self.row_multiindex_full_lookup or self.col_multiindex_full_lookup
+        )
+
         if isinstance(row_loc, Series) and is_boolean_array(row_loc):
             return self._handle_boolean_masking(row_loc, col_loc)
 
@@ -577,6 +634,7 @@ def __getitem__(self, key):
         if (
             isinstance(result, (Series, DataFrame))
             and result._query_compiler.has_multiindex()
+            and not levels_already_dropped
         ):
             if (
                 isinstance(result, Series)
@@ -596,6 +654,7 @@ def __getitem__(self, key):
         if (
             hasattr(result, "columns")
             and not isinstance(col_loc_as_list, slice)
+            and not levels_already_dropped
             and result._query_compiler.has_multiindex(axis=1)
             and all(
                 col_loc_as_list[i] in result.columns.levels[i]
diff --git a/modin/pandas/test/dataframe/test_indexing.py b/modin/pandas/test/dataframe/test_indexing.py
index 407ba5ad14f..74e23efc690 100644
--- a/modin/pandas/test/dataframe/test_indexing.py
+++ b/modin/pandas/test/dataframe/test_indexing.py
@@ -562,6 +562,44 @@ def test_loc_multi_index():
     df_equals(modin_df.loc[modin_df.index[:7]], pandas_df.loc[pandas_df.index[:7]])
 
 
+def test_loc_multi_index_with_tuples():
+    arrays = [
+        ["bar", "bar", "baz", "baz"],
+        ["one", "two", "one", "two"],
+    ]
+    nrows = 5
+    columns = pd.MultiIndex.from_tuples(zip(*arrays), names=["a", "b"])
+    data = np.arange(0, nrows * len(columns)).reshape(nrows, len(columns))
+    modin_df, pandas_df = create_test_dfs(data, columns=columns)
+    eval_general(modin_df, pandas_df, lambda df: df.loc[:, ("bar", "two")])
+
+
+def test_loc_multi_index_duplicate_keys():
+    modin_df, pandas_df = create_test_dfs([1, 2], index=[["a", "a"], ["b", "b"]])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("a", "b"), 0])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("a", "b"), :])
+
+
+def test_loc_multi_index_both_axes():
+    multi_index = pd.MultiIndex.from_tuples(
+        [("r0", "rA"), ("r1", "rB")], names=["Courses", "Fee"]
+    )
+    cols = pd.MultiIndex.from_tuples(
+        [
+            ("Gasoline", "Toyota"),
+            ("Gasoline", "Ford"),
+            ("Electric", "Tesla"),
+            ("Electric", "Nio"),
+        ]
+    )
+    data = [[100, 300, 900, 400], [200, 500, 300, 600]]
+    modin_df, pandas_df = create_test_dfs(data, columns=cols, index=multi_index)
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("r0", "rA"), :])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[:, ("Gasoline", "Toyota")])
+    # Test below bc of issue #4684
+    # df_equals(df.loc[("r0"), ("Gasoline", "Toyota")], pdf.loc[("r0"), ("Gasoline", "Toyota")])
+
+
 def test_loc_empty():
     pandas_df = pandas.DataFrame(index=range(5))
     modin_df = pd.DataFrame(index=range(5))


Code review comments:
diff_hunk: @@ -528,6 +536,39 @@ def _handle_boolean_masking(self, row_loc, col_loc):
         # Passing `slice(None)` as a row indexer since we've just applied it
         return type(self)(masked_df)[(slice(None), col_loc)]
 
+    def _multiindex_possibly_contains_key(self, axis, key):
+        """
+        Determine if a MultiIndex row/column possibly contains a key.
+
+        Check to see if the current DataFrame has a MultiIndex row/column and if it does,
+        check to see if the key is potentially a full key-lookup such that the number of
+        levels match up with the length of the tuple key.
+
+        Parameters
+        ----------
+        axis : int
+            0 or 1 for row/column.
+        key : Any
+            Lookup key for MultiIndex row/column.
+
+        Returns
+        -------
+        bool
+            If the MultiIndex possibly contains the given key.
+
+        Notes
+        -----
+        This function only returns False if we have a partial key lookup. It's
+        possibly that this function returns True for a key that does NOT exist
+        since we only check the length of the tuple to match the number of levels
path: modin/pandas/indexing.py
text: @user1:
```suggestion
        possible that this function returns True for a key that does NOT exist
        since we only check the length of the `key` tuple to match the number of levels
```

diff_hunk: @@ -562,6 +562,44 @@ def test_loc_multi_index():
     df_equals(modin_df.loc[modin_df.index[:7]], pandas_df.loc[pandas_df.index[:7]])
 
 
+def test_loc_multi_index_with_tuples():
+    arrays = [
+        ["bar", "bar", "baz", "baz"],
+        ["one", "two", "one", "two"],
+    ]
+    nrows = 5
+    columns = pd.MultiIndex.from_tuples(zip(*arrays), names=["a", "b"])
+    data = np.arange(0, nrows * len(columns)).reshape(nrows, len(columns))
+    modin_df, pandas_df = create_test_dfs(data, columns=columns)
+    eval_general(modin_df, pandas_df, lambda df: df.loc[:, ("bar", "two")])
+
+
+def test_loc_multi_index_duplicate_keys():
+    modin_df, pandas_df = create_test_dfs([1, 2], index=[["a", "a"], ["b", "b"]])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("a", "b"), 0])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("a", "b"), :])
+
+
+def test_loc_multi_index_both_axes():
+    multi_index = pd.MultiIndex.from_tuples(
+        [("r0", "rA"), ("r1", "rB")], names=["Courses", "Fee"]
+    )
+    cols = pd.MultiIndex.from_tuples(
+        [
+            ("Gasoline", "Toyota"),
+            ("Gasoline", "Ford"),
+            ("Electric", "Tesla"),
+            ("Electric", "Nio"),
+        ]
+    )
+    data = [[100, 300, 900, 400], [200, 500, 300, 600]]
+    modin_df, pandas_df = create_test_dfs(data, columns=cols, index=multi_index)
+    eval_general(modin_df, pandas_df, lambda df: df.loc[("r0", "rA"), :])
+    eval_general(modin_df, pandas_df, lambda df: df.loc[:, ("Gasoline", "Toyota")])
+    # Test below bc of issue #4684
+    # df_equals(df.loc[("r0"), ("Gasoline", "Toyota")], pdf.loc[("r0"), ("Gasoline", "Toyota")])
path: modin/pandas/test/dataframe/test_indexing.py
text: @user2:
```suggestion
    # FIXME: Test below fails because of issue #4684
    # eval_general(modin_df, pandas_df, lambda df: df.loc[("r0"), ("Gasoline", "Toyota")]])
```
or, even better, parametrize the test by `.loc[]` parameter and xfail the test on one of those, like
```python
@user1.mark.parametrize(
    "loc_label",
    [
        (("r0", "rA"), slice(None)),
        (slice(None), ("Gasoline", "Toyota")),
        pytest.param(
            (("r0"), ("Gasoline", "Toyota")),
            marks=pytest.mark.xfail("See Modin issue #4684"),
        ),
    ]
)
def test_loc_multi_index_both_axes(loc_label):
    # setup here the dataframes
    eval_general(modin_df, pandas_df, lambda df: df.loc[loc_label])
```
that way, when that issue is fixed if someone forgets to un-skip the label, CI should turn red on test which we expect to fail but it passed.

@author:
Alternatively, I could just get rid of this test for now because I already mention it in the Github Issue.

@user3:
@user2 I am generally against `xfail`. I think we should only run a test if we hope it will pass. We should track bugs on GitHub and add reproducing test cases when we make fixes. If we do want a note in the code about a bug, I prefer adding a TODO with a link to the issue. I think the extra code and testing time for `xfail` is not worth the assurance that we will catch changes that happen to fix a bug.