Repository: modin-project/modin
Pull request: FIX-#3515: Make loc with slice and 1 column return dataframe instead of series
The root cause was that in cases like [this](https://github.com/modin-project/modin/issues/3515#issue-1016462060), _LocIndexer.__getitem__ passed a list of columns instead of a scalar column to df.__get__item [here](https://github.com/modin-project/modin/blob/c2b399a0d69baa519dab67acd715061e93986b22/modin/pandas/indexing.py#L539).

The new test case fails without this change.

Signed-off-by: mvashishtha <mahesh@ponder.io>

## What do these changes do?

Fix issue 3515.

- ✅   commit message follows format outlined [here](https://modin.readthedocs.io/en/latest/contributing.html)
- ✅    passes `flake8 modin`
- ✅  passes `black --check modin`
- ✅  signed commit with `git commit -s`
- ✅   Resolves #3515
- ✅  tests added and passing
- ✅  module layout described at `docs/developer/architecture.rst` is up-to-date 


Problem statement:
`loc` with slice for rows and single column returns dataframe instead of series
### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Independent
- **Modin version** (`modin.__version__`): latest
- **Python version**: independent
- **Code we can use to reproduce**:

```python
import numpy as np
import modin.pandas as pd

df = pd.DataFrame(np.random.randint(0, 100, size=(2**10, 2**4))).add_prefix("col")
df.loc[:, "col1"]
```

<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

Hints from issue discussion:

Commit message: FIX-#3515: Make loc with slice and 1 column return dataframe instead of series.

The root cause was that in cases like [this](https://github.com/modin-project/modin/issues/3515#issue-1016462060), _LocIndexer.__getitem__ passed a list of columns instead of a scalar column to df.__get__item [here](https://github.com/modin-project/modin/blob/c2b399a0d69baa519dab67acd715061e93986b22/modin/pandas/indexing.py#L539).

Signed-off-by: mvashishtha <mahesh@ponder.io>

Patch to review:
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index 1d414869409..7ccd4ac3325 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -204,15 +204,11 @@ def _parse_tuple(tup):
     Returns
     -------
     row_loc : list
-        List of row locators.
+        Row locator(s) as a scalar or List.
     col_list : list
-        List of column locators.
+        Column locator(s) as a scalar or List.
     ndim : {0, 1, 2}
         Number of dimensions of located dataset.
-    row_scaler : bool
-        True if `row_loc` is a scalar, False otherwise.
-    col_scaler : bool
-        True if `col_loc` is a scalar, False otherwise.
     """
     row_loc, col_loc = slice(None), slice(None)
 
@@ -225,13 +221,7 @@ def _parse_tuple(tup):
     else:
         row_loc = tup
 
-    ndim = _compute_ndim(row_loc, col_loc)
-    row_scaler = is_scalar(row_loc)
-    col_scaler = is_scalar(col_loc)
-    row_loc = [row_loc] if row_scaler else row_loc
-    col_loc = [col_loc] if col_scaler else col_loc
-
-    return row_loc, col_loc, ndim, row_scaler, col_scaler
+    return row_loc, col_loc, _compute_ndim(row_loc, col_loc)
 
 
 def _compute_ndim(row_loc, col_loc):
@@ -528,7 +518,9 @@ def __getitem__(self, key):
         """
         if callable(key):
             return self.__getitem__(key(self.df))
-        row_loc, col_loc, ndim, self.row_scaler, self.col_scaler = _parse_tuple(key)
+        row_loc, col_loc, ndim = _parse_tuple(key)
+        self.row_scaler = is_scalar(row_loc)
+        self.col_scaler = is_scalar(col_loc)
         if isinstance(row_loc, slice) and row_loc == slice(None):
             # If we're only slicing columns, handle the case with `__getitem__`
             if not isinstance(col_loc, slice):
@@ -551,29 +543,35 @@ def __getitem__(self, key):
         if isinstance(result, Series):
             result._parent = self.df
             result._parent_axis = 0
+        col_loc_as_list = [col_loc] if self.col_scaler else col_loc
+        row_loc_as_list = [row_loc] if self.row_scaler else row_loc
         # Pandas drops the levels that are in the `loc`, so we have to as well.
         if hasattr(result, "index") and isinstance(result.index, pandas.MultiIndex):
             if (
                 isinstance(result, Series)
-                and not isinstance(col_loc, slice)
+                and not isinstance(col_loc_as_list, slice)
                 and all(
-                    col_loc[i] in result.index.levels[i] for i in range(len(col_loc))
+                    col_loc_as_list[i] in result.index.levels[i]
+                    for i in range(len(col_loc_as_list))
                 )
             ):
-                result.index = result.index.droplevel(list(range(len(col_loc))))
-            elif not isinstance(row_loc, slice) and all(
-                not isinstance(row_loc[i], slice)
-                and row_loc[i] in result.index.levels[i]
-                for i in range(len(row_loc))
+                result.index = result.index.droplevel(list(range(len(col_loc_as_list))))
+            elif not isinstance(row_loc_as_list, slice) and all(
+                not isinstance(row_loc_as_list[i], slice)
+                and row_loc_as_list[i] in result.index.levels[i]
+                for i in range(len(row_loc_as_list))
             ):
-                result.index = result.index.droplevel(list(range(len(row_loc))))
+                result.index = result.index.droplevel(list(range(len(row_loc_as_list))))
         if (
             hasattr(result, "columns")
-            and not isinstance(col_loc, slice)
+            and not isinstance(col_loc_as_list, slice)
             and isinstance(result.columns, pandas.MultiIndex)
-            and all(col_loc[i] in result.columns.levels[i] for i in range(len(col_loc)))
+            and all(
+                col_loc_as_list[i] in result.columns.levels[i]
+                for i in range(len(col_loc_as_list))
+            )
         ):
-            result.columns = result.columns.droplevel(list(range(len(col_loc))))
+            result.columns = result.columns.droplevel(list(range(len(col_loc_as_list))))
         # This is done for cases where the index passed in has other state, like a
         # frequency in the case of DateTimeIndex.
         if (
@@ -600,7 +598,7 @@ def __setitem__(self, key, item):
         --------
         pandas.DataFrame.loc
         """
-        row_loc, col_loc, _, row_scaler, col_scaler = _parse_tuple(key)
+        row_loc, col_loc, _ = _parse_tuple(key)
         if isinstance(row_loc, list) and len(row_loc) == 1:
             if row_loc[0] not in self.qc.index:
                 index = self.qc.index.insert(len(self.qc.index), row_loc[0])
@@ -623,7 +621,7 @@ def __setitem__(self, key, item):
                 col_lookup,
                 item,
                 axis=self._determine_setitem_axis(
-                    row_lookup, col_lookup, row_scaler, col_scaler
+                    row_lookup, col_lookup, is_scalar(row_loc), is_scalar(col_loc)
                 ),
             )
 
@@ -666,9 +664,9 @@ def _compute_lookup(self, row_loc, col_loc):
 
         Parameters
         ----------
-        row_loc : slice, list, array or tuple
+        row_loc : scalar, slice, list, array or tuple
             Row locator.
-        col_loc : slice, list, array or tuple
+        col_loc : scalar, slice, list, array or tuple
             Columns locator.
 
         Returns
@@ -678,6 +676,8 @@ def _compute_lookup(self, row_loc, col_loc):
         col_lookup : numpy.ndarray
             List of columns labels.
         """
+        row_loc = [row_loc] if is_scalar(row_loc) else row_loc
+        col_loc = [col_loc] if is_scalar(col_loc) else col_loc
         if is_list_like(row_loc) and len(row_loc) == 1:
             if (
                 isinstance(self.qc.index.values[0], np.datetime64)
@@ -746,7 +746,9 @@ def __getitem__(self, key):
         """
         if callable(key):
             return self.__getitem__(key(self.df))
-        row_loc, col_loc, ndim, self.row_scaler, self.col_scaler = _parse_tuple(key)
+        row_loc, col_loc, ndim = _parse_tuple(key)
+        self.row_scaler = is_scalar(row_loc)
+        self.col_scaler = is_scalar(col_loc)
         self._check_dtypes(row_loc)
         self._check_dtypes(col_loc)
 
@@ -772,7 +774,9 @@ def __setitem__(self, key, item):
         --------
         pandas.DataFrame.iloc
         """
-        row_loc, col_loc, _, row_scaler, col_scaler = _parse_tuple(key)
+        row_loc, col_loc, _ = _parse_tuple(key)
+        row_scaler = is_scalar(row_loc)
+        col_scaler = is_scalar(col_loc)
         self._check_dtypes(row_loc)
         self._check_dtypes(col_loc)
 
@@ -804,6 +808,8 @@ def _compute_lookup(self, row_loc, col_loc):
         col_lookup : numpy.ndarray
             List of columns labels.
         """
+        row_loc = [row_loc] if is_scalar(row_loc) else row_loc
+        col_loc = [col_loc] if is_scalar(col_loc) else col_loc
         if (
             not isinstance(row_loc, slice)
             or isinstance(row_loc, slice)
diff --git a/modin/pandas/test/dataframe/test_indexing.py b/modin/pandas/test/dataframe/test_indexing.py
index fa20c868f6c..4869c2ee3e4 100644
--- a/modin/pandas/test/dataframe/test_indexing.py
+++ b/modin/pandas/test/dataframe/test_indexing.py
@@ -284,6 +284,7 @@ def test_loc(data):
     df_equals(modin_df.loc[0], pandas_df.loc[0])
     df_equals(modin_df.loc[1:, key1], pandas_df.loc[1:, key1])
     df_equals(modin_df.loc[1:2, key1], pandas_df.loc[1:2, key1])
+    df_equals(modin_df.loc[:, key1], pandas_df.loc[:, key1])
 
     # DataFrame
     df_equals(modin_df.loc[[1, 2]], pandas_df.loc[[1, 2]])


Code review comments:
diff_hunk: @@ -204,15 +204,11 @@ def _parse_tuple(tup):
     Returns
     -------
     row_loc : list
path: modin/pandas/indexing.py
text: @user1:
```suggestion
    row_loc : scalar or list
```

diff_hunk: @@ -204,15 +204,11 @@ def _parse_tuple(tup):
     Returns
     -------
     row_loc : list
-        List of row locators.
+        Row locator(s) as a scalar or List.
     col_list : list
path: modin/pandas/indexing.py
text: @user1:
```suggestion
    col_list : scalar or list
```