Repository: mars-project/mars

Issue:
Implement DataFrame.isin and Series.isin
Implements `DataFrame.isin(self, values)`:

> _**values**_ : iterable, Series, DataFrame or dict
> The result will only be true at a location if all the labels match. If values is a Series, that’s the index. If values is a dict, the keys must be the column names, which must match. If values is a DataFrame, then both the index and column labels must match.


Discussion:

Resolving patch:
diff --git a/mars/_version.py b/mars/_version.py
--- a/mars/_version.py
+++ b/mars/_version.py
@@ -15,7 +15,7 @@
 import subprocess
 import os
 
-version_info = (0, 6, 0, 'a2')
+version_info = (0, 6, 0, 'a3')
 _num_index = max(idx if isinstance(v, int) else 0
                  for idx, v in enumerate(version_info))
 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \
diff --git a/mars/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -20,7 +20,7 @@
 from .apply import df_apply, series_apply
 from .fillna import fillna, ffill, bfill
 from .transform import df_transform, series_transform
-from .isin import isin
+from .isin import series_isin, df_isin
 from .checkna import isna, notna, isnull, notnull
 from .dropna import df_dropna, series_dropna
 from .cut import cut
@@ -56,6 +56,7 @@ def _install():
         setattr(t, 'fillna', fillna)
         setattr(t, 'ffill', ffill)
         setattr(t, 'bfill', bfill)
+        setattr(t, 'isin', df_isin)
         setattr(t, 'isna', isna)
         setattr(t, 'isnull', isnull)
         setattr(t, 'notna', notna)
@@ -86,7 +87,7 @@ def _install():
         setattr(t, 'fillna', fillna)
         setattr(t, 'ffill', ffill)
         setattr(t, 'bfill', bfill)
-        setattr(t, 'isin', isin)
+        setattr(t, 'isin', series_isin)
         setattr(t, 'isna', isna)
         setattr(t, 'isnull', isnull)
         setattr(t, 'notna', notna)
diff --git a/mars/dataframe/base/isin.py b/mars/dataframe/base/isin.py
--- a/mars/dataframe/base/isin.py
+++ b/mars/dataframe/base/isin.py
@@ -13,15 +13,16 @@
 # limitations under the License.
 
 import numpy as np
+import pandas as pd
 from pandas.api.types import is_list_like
 
 from ... import opcodes as OperandDef
-from ...core import OutputType
+from ...core import Base, Entity
 from ...serialize import KeyField, AnyField
 from ...tensor.core import TENSOR_TYPE
 from ...tiles import TilesError
 from ...utils import check_chunks_unknown_shape
-from ..core import SERIES_TYPE, INDEX_TYPE
+from ..core import DATAFRAME_TYPE, SERIES_TYPE, INDEX_TYPE
 from ..operands import DataFrameOperand, DataFrameOperandMixin
 
 
@@ -32,8 +33,6 @@ class DataFrameIsin(DataFrameOperand, DataFrameOperandMixin):
     _values = AnyField('values')
 
     def __init__(self, values=None, output_types=None, **kw):
-        if output_types is None:
-            output_types = [OutputType.series]
         super().__init__(_values=values, _output_types=output_types, **kw)
 
     @property
@@ -46,54 +45,101 @@ def values(self):
 
     def _set_inputs(self, inputs):
         super()._set_inputs(inputs)
-        self._input = self._inputs[0]
-        if len(inputs) == 2:
-            self._values = self._inputs[1]
+        inputs_iter = iter(self._inputs)
+        self._input = next(inputs_iter)
+        if len(self._inputs) > 1:
+            if isinstance(self._values, dict):
+                new_values = dict()
+                for k, v in self._values.items():
+                    if isinstance(v, (Base, Entity)):
+                        new_values[k] = next(inputs_iter)
+                    else:
+                        new_values[k] = v
+                self._values = new_values
+            else:
+                self._values = self._inputs[1]
 
     def __call__(self, elements):
         inputs = [elements]
-        if isinstance(self._values, (SERIES_TYPE, TENSOR_TYPE, INDEX_TYPE)):
+        if isinstance(self._values, (Base, Entity)):
             inputs.append(self._values)
-        return self.new_series(inputs, shape=elements.shape, dtype=np.dtype('bool'),
-                               index_value=elements.index_value, name=elements.name)
+        elif isinstance(self._values, dict):
+            for v in self._values.values():
+                if isinstance(v, (Base, Entity)):
+                    inputs.append(v)
+
+        if elements.ndim == 1:
+            return self.new_series(inputs, shape=elements.shape,
+                                   dtype=np.dtype('bool'),
+                                   index_value=elements.index_value,
+                                   name=elements.name)
+        else:
+            dtypes = pd.Series([np.dtype(bool) for _ in elements.dtypes],
+                               index=elements.dtypes.index)
+            return self.new_dataframe(inputs, shape=elements.shape,
+                                      index_value=elements.index_value,
+                                      columns_value=elements.columns_value,
+                                      dtypes=dtypes)
 
     @classmethod
     def tile(cls, op):
         in_elements = op.input
         out_elements = op.outputs[0]
 
-        values = op.values
-        if len(op.inputs) == 2:
-            # make sure arg has known shape when it's a md.Series
-            check_chunks_unknown_shape([op.values], TilesError)
-            values = op.values.rechunk(op.values.shape)._inplace_tile()
+        values_inputs = []
+        if len(op.inputs) > 1:
+            for value in op.inputs[1:]:
+                # make sure arg has known shape when it's a md.Series
+                check_chunks_unknown_shape([value], TilesError)
+                value = value.rechunk(value.shape)._inplace_tile()
+                values_inputs.append(value)
 
         out_chunks = []
         for chunk in in_elements.chunks:
             chunk_op = op.copy().reset_key()
             chunk_inputs = [chunk]
-            if len(op.inputs) == 2:
-                chunk_inputs.append(values.chunks[0])
-            out_chunk = chunk_op.new_chunk(chunk_inputs, shape=chunk.shape,
-                                           dtype=out_elements.dtype,
-                                           index_value=chunk.index_value,
-                                           name=out_elements.name,
-                                           index=chunk.index)
+            if len(op.inputs) > 1:
+                chunk_inputs.extend(v.chunks[0] for v in values_inputs)
+            if out_elements.ndim == 1:
+                out_chunk = chunk_op.new_chunk(chunk_inputs, shape=chunk.shape,
+                                               dtype=out_elements.dtype,
+                                               index_value=chunk.index_value,
+                                               name=out_elements.name,
+                                               index=chunk.index)
+            else:
+                chunk_dtypes = pd.Series([np.dtype(bool) for _ in chunk.dtypes],
+                                         index=chunk.dtypes.index)
+                out_chunk = chunk_op.new_chunk(chunk_inputs, shape=chunk.shape,
+                                               index_value=chunk.index_value,
+                                               columns_value=chunk.columns_value,
+                                               dtypes=chunk_dtypes,
+                                               index=chunk.index)
             out_chunks.append(out_chunk)
 
         new_op = op.copy()
-        return new_op.new_seriess(op.inputs, out_elements.shape,
-                                  nsplits=in_elements.nsplits,
-                                  chunks=out_chunks, dtype=out_elements.dtype,
-                                  index_value=out_elements.index_value, name=out_elements.name)
+        params = out_elements.params
+        params['nsplits'] = in_elements.nsplits
+        params['chunks'] = out_chunks
+        return new_op.new_tileables(op.inputs, kws=[params])
 
     @classmethod
     def execute(cls, ctx, op):
-        elements = ctx[op.inputs[0].key]
-        if len(op.inputs) == 2:
-            values = ctx[op.inputs[1].key]
+        inputs_iter = iter(op.inputs)
+        elements = ctx[next(inputs_iter).key]
+
+        if isinstance(op.values, dict):
+            values = dict()
+            for k, v in op.values.items():
+                if isinstance(v, (Base, Entity)):
+                    values[k] = ctx[next(inputs_iter).key]
+                else:
+                    values[k] = v
         else:
-            values = op.values
+            if isinstance(op.values, (Base, Entity)):
+                values = ctx[next(inputs_iter).key]
+            else:
+                values = op.values
+
         try:
             ctx[op.outputs[0].key] = elements.isin(values)
         except ValueError:
@@ -101,14 +147,136 @@ def execute(cls, ctx, op):
             ctx[op.outputs[0].key] = elements.copy().isin(values.copy())
 
 
-def isin(elements, values):
-    # TODO(hetao): support more type combinations, for example, DataFrame.isin.
+def series_isin(elements, values):
+    """
+    Whether elements in Series are contained in `values`.
+
+    Return a boolean Series showing whether each element in the Series
+    matches an element in the passed sequence of `values` exactly.
+
+    Parameters
+    ----------
+    values : set or list-like
+        The sequence of values to test. Passing in a single string will
+        raise a ``TypeError``. Instead, turn a single string into a
+        list of one element.
+
+    Returns
+    -------
+    Series
+        Series of booleans indicating if each element is in values.
+
+    Raises
+    ------
+    TypeError
+      * If `values` is a string
+
+    See Also
+    --------
+    DataFrame.isin : Equivalent method on DataFrame.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> s = md.Series(['lama', 'cow', 'lama', 'beetle', 'lama',
+    ...                'hippo'], name='animal')
+    >>> s.isin(['cow', 'lama']).execute()
+    0     True
+    1     True
+    2     True
+    3    False
+    4     True
+    5    False
+    Name: animal, dtype: bool
+
+    Passing a single string as ``s.isin('lama')`` will raise an error. Use
+    a list of one element instead:
+
+    >>> s.isin(['lama']).execute()
+    0     True
+    1    False
+    2     True
+    3    False
+    4     True
+    5    False
+    Name: animal, dtype: bool
+    """
     if is_list_like(values):
         values = list(values)
     elif not isinstance(values, (SERIES_TYPE, TENSOR_TYPE, INDEX_TYPE)):
         raise TypeError('only list-like objects are allowed to be passed to isin(), '
                         f'you passed a [{type(values)}]')
-    if not isinstance(elements, SERIES_TYPE):  # pragma: no cover
-        raise NotImplementedError(f'Unsupported parameter types: {type(elements)} and {type(values)}')
     op = DataFrameIsin(values)
     return op(elements)
+
+
+def df_isin(df, values):
+    """
+    Whether each element in the DataFrame is contained in values.
+
+    Parameters
+    ----------
+    values : iterable, Series, DataFrame or dict
+        The result will only be true at a location if all the
+        labels match. If `values` is a Series, that's the index. If
+        `values` is a dict, the keys must be the column names,
+        which must match. If `values` is a DataFrame,
+        then both the index and column labels must match.
+
+    Returns
+    -------
+    DataFrame
+        DataFrame of booleans showing whether each element in the DataFrame
+        is contained in values.
+
+    See Also
+    --------
+    DataFrame.eq: Equality test for DataFrame.
+    Series.isin: Equivalent method on Series.
+    Series.str.contains: Test if pattern or regex is contained within a
+        string of a Series or Index.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> df = md.DataFrame({'num_legs': [2, 4], 'num_wings': [2, 0]},
+    ...                   index=['falcon', 'dog'])
+    >>> df.execute()
+            num_legs  num_wings
+    falcon         2          2
+    dog            4          0
+
+    When ``values`` is a list check whether every value in the DataFrame
+    is present in the list (which animals have 0 or 2 legs or wings)
+
+    >>> df.isin([0, 2]).execute()
+            num_legs  num_wings
+    falcon      True       True
+    dog        False       True
+
+    When ``values`` is a dict, we can pass values to check for each
+    column separately:
+
+    >>> df.isin({'num_wings': [0, 3]}).execute()
+            num_legs  num_wings
+    falcon     False      False
+    dog        False       True
+
+    When ``values`` is a Series or DataFrame the index and column must
+    match. Note that 'falcon' does not match based on the number of legs
+    in df2.
+
+    >>> other = md.DataFrame({'num_legs': [8, 2], 'num_wings': [0, 2]},
+    ...                      index=['spider', 'falcon'])
+    >>> df.isin(other).execute()
+            num_legs  num_wings
+    falcon      True       True
+    dog        False      False
+    """
+    if is_list_like(values) and not isinstance(values, dict):
+        values = list(values)
+    elif not isinstance(values, (SERIES_TYPE, DATAFRAME_TYPE, TENSOR_TYPE, INDEX_TYPE, dict)):
+        raise TypeError('only list-like objects or dict are allowed to be passed to isin(), '
+                        f'you passed a [{type(values)}]')
+    op = DataFrameIsin(values)
+    return op(df)


Test patch:
diff --git a/mars/dataframe/base/tests/test_base.py b/mars/dataframe/base/tests/test_base.py
--- a/mars/dataframe/base/tests/test_base.py
+++ b/mars/dataframe/base/tests/test_base.py
@@ -671,6 +671,9 @@ def testSeriesIsin(self):
         with self.assertRaises(TypeError):
             _ = a.isin('sth')
 
+        with self.assertRaises(TypeError):
+            _ = a.to_frame().isin('sth')
+
     def testDropNA(self):
         # dataframe cases
         df_raw = pd.DataFrame(np.nan, index=range(0, 20), columns=list('ABCDEFGHIJ'))
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -745,7 +745,7 @@ def testDatetimeMethodExecution(self):
         expected = s.dt.days
         pd.testing.assert_series_equal(result, expected)
 
-    def testSeriesIsin(self):
+    def testIsinExecution(self):
         # one chunk in multiple chunks
         a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
         b = pd.Series([2, 1, 9, 3])
@@ -801,6 +801,32 @@ def testSeriesIsin(self):
         expected = a.isin(b)
         pd.testing.assert_series_equal(result, expected)
 
+        rs = np.random.RandomState(0)
+        raw = pd.DataFrame(rs.randint(1000, size=(10, 3)))
+        df = from_pandas_df(raw, chunk_size=(5, 2))
+
+        # set
+        b = {2, 1, raw[1][0]}
+        r = df.isin(b)
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.isin(b)
+        pd.testing.assert_frame_equal(result, expected)
+
+        # mars object
+        b = tensor([2, 1, raw[1][0]], chunk_size=2)
+        r = df.isin(b)
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.isin([2, 1, raw[1][0]])
+        pd.testing.assert_frame_equal(result, expected)
+
+        # dict
+        b = {1: tensor([2, 1, raw[1][0]], chunk_size=2),
+             2: [3, 10]}
+        r = df.isin(b)
+        result = self.executor.execute_dataframe(r, concat=True)[0]
+        expected = raw.isin({1: [2, 1, raw[1][0]], 2: [3, 10]})
+        pd.testing.assert_frame_equal(result, expected)
+
     def testCheckNA(self):
         df_raw = pd.DataFrame(np.nan, index=range(0, 20), columns=list('ABCDEFGHIJ'))
         for _ in range(20):