Repository: modin-project/modin

Issue:
Avoid re-putting a distributed Series to the engine's object store at `.isin()`



Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -1329,7 +1329,7 @@ def conj(df, *args, **kwargs):
     #      we should avoid leaking of the high-level objects to the query compiler level.
     #      (Modin issue #3106)
     #   2. Spread **kwargs into actual arguments (Modin issue #3108).
-    def isin(self, **kwargs):  # noqa: PR02
+    def isin(self, values, ignore_indices=False, **kwargs):  # noqa: PR02
         """
         Check for each element of `self` whether it's contained in passed `values`.
 
@@ -1337,6 +1337,8 @@ def isin(self, **kwargs):  # noqa: PR02
         ----------
         values : list-like, modin.pandas.Series, modin.pandas.DataFrame or dict
             Values to check elements of self in.
+        ignore_indices : bool, default: False
+            Whether to execute ``isin()`` only on an intersection of indices.
         **kwargs : dict
             Serves the compatibility purpose. Does not affect the result.
 
@@ -1346,10 +1348,16 @@ def isin(self, **kwargs):  # noqa: PR02
             Boolean mask for self of whether an element at the corresponding
             position is contained in `values`.
         """
-        # We drop `shape_hint` argument that may be passed from the API layer.
-        # BaseQC doesn't need to know how to handle it.
-        kwargs.pop("shape_hint", None)
-        return DataFrameDefault.register(pandas.DataFrame.isin)(self, **kwargs)
+        shape_hint = kwargs.pop("shape_hint", None)
+        if isinstance(values, type(self)) and ignore_indices:
+            # Pandas logic is that it ignores indexing if 'values' is a 1D object
+            values = values.to_pandas().squeeze(axis=1)
+        if shape_hint == "column":
+            return SeriesDefault.register(pandas.Series.isin)(self, values, **kwargs)
+        else:
+            return DataFrameDefault.register(pandas.DataFrame.isin)(
+                self, values, **kwargs
+            )
 
     def isna(self):
         """
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -1404,12 +1404,37 @@ def stack(self, level, dropna):
 
     # Map partitions operations
     # These operations are operations that apply a function to every partition.
+    def isin(self, values, ignore_indices=False, shape_hint=None):
+        if isinstance(values, type(self)):
+            # HACK: if we don't cast to pandas, then the execution engine will try to
+            # propagate the distributed Series to workers and most likely would have
+            # some performance problems.
+            # TODO: A better way of doing so could be passing this `values` as a query compiler
+            # and broadcast accordingly.
+            values = values.to_pandas()
+            if ignore_indices:
+                # Pandas logic is that it ignores indexing if 'values' is a 1D object
+                values = values.squeeze(axis=1)
+
+        def isin_func(df, values):
+            if shape_hint == "column":
+                df = df.squeeze(axis=1)
+            res = df.isin(values)
+            if isinstance(res, pandas.Series):
+                res = res.to_frame(
+                    MODIN_UNNAMED_SERIES_LABEL if res.name is None else res.name
+                )
+            return res
+
+        return Map.register(isin_func, shape_hint=shape_hint, dtypes=np.bool_)(
+            self, values
+        )
+
     abs = Map.register(pandas.DataFrame.abs, dtypes="copy")
     applymap = Map.register(pandas.DataFrame.applymap)
     conj = Map.register(lambda df, *args, **kwargs: pandas.DataFrame(np.conj(df)))
     convert_dtypes = Map.register(pandas.DataFrame.convert_dtypes)
     invert = Map.register(pandas.DataFrame.__invert__, dtypes="copy")
-    isin = Map.register(pandas.DataFrame.isin, dtypes=np.bool_)
     isna = Map.register(pandas.DataFrame.isna, dtypes=np.bool_)
     _isfinite = Map.register(
         lambda df, *args, **kwargs: pandas.DataFrame(np.isfinite(df))
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -1724,8 +1724,14 @@ def isin(self, values, **kwargs):  # noqa: PR01, RT01, D200
         """
         Whether elements in `BasePandasDataset` are contained in `values`.
         """
+        from .series import Series
+
+        ignore_indices = isinstance(values, Series)
+        values = getattr(values, "_query_compiler", values)
         return self.__constructor__(
-            query_compiler=self._query_compiler.isin(values=values, **kwargs)
+            query_compiler=self._query_compiler.isin(
+                values=values, ignore_indices=ignore_indices, **kwargs
+            )
         )
 
     def isna(self):  # noqa: RT01, D200


Test patch:
diff --git a/modin/pandas/test/dataframe/test_iter.py b/modin/pandas/test/dataframe/test_iter.py
--- a/modin/pandas/test/dataframe/test_iter.py
+++ b/modin/pandas/test/dataframe/test_iter.py
@@ -28,6 +28,8 @@
     test_data_values,
     test_data_keys,
     test_data,
+    create_test_dfs,
+    eval_general,
 )
 from modin.pandas.utils import SET_DATAFRAME_ATTRIBUTE_WARNING
 from modin.config import NPartitions
@@ -332,3 +334,37 @@ def test_isin(data):
     modin_result = modin_df.isin(val)
 
     df_equals(modin_result, pandas_result)
+
+
+def test_isin_with_modin_objects():
+    modin_df1, pandas_df1 = create_test_dfs({"a": [1, 2], "b": [3, 4]})
+    modin_series, pandas_series = pd.Series([1, 4, 5, 6]), pandas.Series([1, 4, 5, 6])
+
+    eval_general(
+        (modin_df1, modin_series),
+        (pandas_df1, pandas_series),
+        lambda srs: srs[0].isin(srs[1]),
+    )
+
+    modin_df2 = modin_series.to_frame("a")
+    pandas_df2 = pandas_series.to_frame("a")
+
+    eval_general(
+        (modin_df1, modin_df2),
+        (pandas_df1, pandas_df2),
+        lambda srs: srs[0].isin(srs[1]),
+    )
+
+    # Check case when indices are not matching
+    modin_df1, pandas_df1 = create_test_dfs({"a": [1, 2], "b": [3, 4]}, index=[10, 11])
+
+    eval_general(
+        (modin_df1, modin_series),
+        (pandas_df1, pandas_series),
+        lambda srs: srs[0].isin(srs[1]),
+    )
+    eval_general(
+        (modin_df1, modin_df2),
+        (pandas_df1, pandas_df2),
+        lambda srs: srs[0].isin(srs[1]),
+    )
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -2070,6 +2070,26 @@ def test_isin(data):
     df_equals(modin_result, pandas_result)
 
 
+def test_isin_with_series():
+    modin_series1, pandas_series1 = create_test_series([1, 2, 3])
+    modin_series2, pandas_series2 = create_test_series([1, 2, 3, 4, 5])
+
+    eval_general(
+        (modin_series1, modin_series2),
+        (pandas_series1, pandas_series2),
+        lambda srs: srs[0].isin(srs[1]),
+    )
+
+    # Verify that Series actualy behaves like Series and ignores unmatched indices on '.isin'
+    modin_series1, pandas_series1 = create_test_series([1, 2, 3], index=[10, 11, 12])
+
+    eval_general(
+        (modin_series1, modin_series2),
+        (pandas_series1, pandas_series2),
+        lambda srs: srs[0].isin(srs[1]),
+    )
+
+
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
 def test_isnull(data):
     modin_series, pandas_series = create_test_series(data)