FIX-#1775: Fix support for callable in loc/iloc (#1776)

Signed-off-by: Devin Petersohn <devin.petersohn@gmail.com>

Author: devin-petersohn

modin/pandas/indexing.py

@@ -236,6 +236,8 @@ class _LocIndexer(_LocationIndexerBase):
     """An indexer for modin_df.loc[] functionality"""
 
     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)
         if isinstance(row_loc, slice) and row_loc == slice(None):
             # If we're only slicing columns, handle the case with `__getitem__`
@@ -361,6 +363,8 @@ class _iLocIndexer(_LocationIndexerBase):
     """An indexer for modin_df.iloc[] functionality"""
 
     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)
         self._check_dtypes(row_loc)
         self._check_dtypes(col_loc)

modin/pandas/test/test_dataframe.py

@@ -4389,6 +4389,12 @@ def test_iloc(self, request, data):
             modin_df.iloc[:, 0] = modin_df.iloc[:, 1]
             pandas_df.iloc[:, 0] = pandas_df.iloc[:, 1]
             df_equals(modin_df, pandas_df)
+
+            # From issue #1775
+            df_equals(
+                modin_df.iloc[lambda df: df.index.get_indexer_for(df.index[:5])],
+                pandas_df.iloc[lambda df: df.index.get_indexer_for(df.index[:5])],
+            )
         else:
             with pytest.raises(IndexError):
                 modin_df.iloc[0, 1]
@@ -4484,6 +4490,12 @@ def test_loc(self, request, data):
             pandas_df_copy.loc[[1, 2]] = 42
             df_equals(modin_df_copy, pandas_df_copy)
 
+            # From issue #1775
+            df_equals(
+                modin_df.loc[lambda df: df.iloc[:, 0].isin(list(range(1000)))],
+                pandas_df.loc[lambda df: df.iloc[:, 0].isin(list(range(1000)))],
+            )
+
         # From issue #1374
         with pytest.raises(KeyError):
             modin_df.loc["NO_EXIST"]