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/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -23,6 +23,7 @@
 from .datetimes import SeriesDatetimeMethod
 from .string_ import SeriesStringMethod
 from .transform import df_transform, series_transform
+from .isin import isin
 
 
 def _install():
@@ -51,6 +52,7 @@ def _install():
         setattr(t, 'fillna', fillna)
         setattr(t, 'ffill', ffill)
         setattr(t, 'bfill', bfill)
+        setattr(t, 'isin', isin)
     for t in INDEX_TYPE:
         setattr(t, 'rechunk', rechunk)
 
diff --git a/mars/dataframe/base/isin.py b/mars/dataframe/base/isin.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/base/isin.py
@@ -0,0 +1,114 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from pandas.api.types import is_list_like
+
+from ... import opcodes as OperandDef
+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
+from ..operands import DataFrameOperand, DataFrameOperandMixin, ObjectType
+
+
+class DataFrameIsin(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = OperandDef.ISIN
+
+    _input = KeyField('input')
+    _values = AnyField('values')
+
+    def __init__(self, values=None, object_type=None, **kw):
+        if object_type is None:
+            object_type = ObjectType.series
+        super().__init__(_values=values, _object_type=object_type, **kw)
+
+    @property
+    def input(self):
+        return self._input
+
+    @property
+    def values(self):
+        return self._values
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        self._input = self._inputs[0]
+        if len(inputs) == 2:
+            self._values = self._inputs[1]
+
+    def __call__(self, elements):
+        inputs = [elements]
+        if isinstance(self._values, (SERIES_TYPE, TENSOR_TYPE)):
+            inputs.append(self._values)
+        return self.new_series(inputs, shape=elements.shape, dtype=np.dtype('bool'),
+                               index_value=elements.index_value, name=elements.name)
+
+    @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()
+
+        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)
+            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)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        elements = ctx[op.inputs[0].key]
+        if len(op.inputs) == 2:
+            values = ctx[op.inputs[1].key]
+        else:
+            values = op.values
+        try:
+            ctx[op.outputs[0].key] = elements.isin(values)
+        except ValueError:
+            # buffer read-only
+            ctx[op.outputs[0].key] = elements.copy().isin(values)
+
+
+def isin(elements, values):
+    # TODO(hetao): support more type combinations, for example, DataFrame.isin.
+    if is_list_like(values):
+        values = list(values)
+    elif not isinstance(values, (SERIES_TYPE, TENSOR_TYPE)):
+        raise TypeError('only list-like objects are allowed to be passed to isin(), '
+                        'you passed a [{}]'.format(type(values)))
+    if not isinstance(elements, SERIES_TYPE):  # pragma: no cover
+        raise NotImplementedError('Unsupported parameter types: %s and %s' %
+                                  (type(elements), type(values)))
+    op = DataFrameIsin(values)
+    return op(elements)


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
@@ -668,3 +668,55 @@ def testDatetimeMethod(self):
             _ = series.dt.non_exist
 
         self.assertIn('ceil', dir(series.dt))
+
+    def testSeriesIsin(self):
+        # one chunk in multiple chunks
+        a = from_pandas_series(pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), chunk_size=10)
+        b = from_pandas_series(pd.Series([2, 1, 9, 3]), chunk_size=2)
+
+        r = a.isin(b).tiles()
+        for i, c in enumerate(r.chunks):
+            self.assertEqual(c.index, (i,))
+            self.assertEqual(c.dtype, np.dtype('bool'))
+            self.assertEqual(c.shape, (10,))
+            self.assertEqual(len(c.op.inputs), 2)
+            self.assertEqual(c.op.object_type, ObjectType.series)
+            self.assertEqual(c.op.inputs[0].index, (i,))
+            self.assertEqual(c.op.inputs[0].shape, (10,))
+            self.assertEqual(c.op.inputs[1].index, (0,))
+            self.assertEqual(c.op.inputs[1].shape, (4,))  # has been rechunked
+
+        # multiple chunk in one chunks
+        a = from_pandas_series(pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), chunk_size=2)
+        b = from_pandas_series(pd.Series([2, 1, 9, 3]), chunk_size=4)
+
+        r = a.isin(b).tiles()
+        for i, c in enumerate(r.chunks):
+            self.assertEqual(c.index, (i,))
+            self.assertEqual(c.dtype, np.dtype('bool'))
+            self.assertEqual(c.shape, (2,))
+            self.assertEqual(len(c.op.inputs), 2)
+            self.assertEqual(c.op.object_type, ObjectType.series)
+            self.assertEqual(c.op.inputs[0].index, (i,))
+            self.assertEqual(c.op.inputs[0].shape, (2,))
+            self.assertEqual(c.op.inputs[1].index, (0,))
+            self.assertEqual(c.op.inputs[1].shape, (4,))
+
+        # multiple chunk in multiple chunks
+        a = from_pandas_series(pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), chunk_size=2)
+        b = from_pandas_series(pd.Series([2, 1, 9, 3]), chunk_size=2)
+
+        r = a.isin(b).tiles()
+        for i, c in enumerate(r.chunks):
+            self.assertEqual(c.index, (i,))
+            self.assertEqual(c.dtype, np.dtype('bool'))
+            self.assertEqual(c.shape, (2,))
+            self.assertEqual(len(c.op.inputs), 2)
+            self.assertEqual(c.op.object_type, ObjectType.series)
+            self.assertEqual(c.op.inputs[0].index, (i,))
+            self.assertEqual(c.op.inputs[0].shape, (2,))
+            self.assertEqual(c.op.inputs[1].index, (0,))
+            self.assertEqual(c.op.inputs[1].shape, (4,))  # has been rechunked
+
+        with self.assertRaises(TypeError):
+            _ = a.isin('sth')
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
@@ -24,6 +24,7 @@
 from mars.dataframe.datasource.series import from_pandas as from_pandas_series
 from mars.dataframe.datasource.index import from_pandas as from_pandas_index
 from mars.session import new_session
+from mars.tensor import tensor
 from mars.tests.core import TestBase, require_cudf, ExecutorForTest
 from mars.utils import lazy_import
 
@@ -760,3 +761,59 @@ def testDatetimeMethodExecution(self):
         result = self.executor.execute_dataframe(r, concat=True)[0]
         expected = s.dt.days
         pd.testing.assert_series_equal(result, expected)
+
+    def testSeriesIsin(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])
+        sa = from_pandas_series(a, chunk_size=10)
+        sb = from_pandas_series(b, chunk_size=2)
+
+        result = self.executor.execute_dataframe(sa.isin(sb), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)
+
+        # multiple chunk in one chunks
+        a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        b = pd.Series([2, 1, 9, 3])
+        sa = from_pandas_series(a, chunk_size=2)
+        sb = from_pandas_series(b, chunk_size=4)
+
+        result = self.executor.execute_dataframe(sa.isin(sb), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)
+
+        # multiple chunk in multiple chunks
+        a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        b = pd.Series([2, 1, 9, 3])
+        sa = from_pandas_series(a, chunk_size=2)
+        sb = from_pandas_series(b, chunk_size=2)
+
+        result = self.executor.execute_dataframe(sa.isin(sb), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)
+
+        a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        b = pd.Series([2, 1, 9, 3])
+        sa = from_pandas_series(a, chunk_size=2)
+
+        result = self.executor.execute_dataframe(sa.isin(b), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)
+
+        a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        b = np.array([2, 1, 9, 3])
+        sa = from_pandas_series(a, chunk_size=2)
+        sb = tensor(b, chunk_size=3)
+
+        result = self.executor.execute_dataframe(sa.isin(sb), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)
+
+        a = pd.Series([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        b = {2, 1, 9, 3}  # set
+        sa = from_pandas_series(a, chunk_size=2)
+
+        result = self.executor.execute_dataframe(sa.isin(b), concat=True)[0]
+        expected = a.isin(b)
+        pd.testing.assert_series_equal(result, expected)