Repository: mars-project/mars

Issue:
Support {DataFrame,Series}.{mask,where}
Support DataFrame.where. pandas docs see https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.where.html.


Discussion:

Resolving patch:
diff --git a/mars/dataframe/core.py b/mars/dataframe/core.py
--- a/mars/dataframe/core.py
+++ b/mars/dataframe/core.py
@@ -744,6 +744,10 @@ def index(self):
 
         return from_tileable(self)
 
+    @property
+    def axes(self):
+        return [self.index]
+
     def to_tensor(self, dtype=None):
         from ..tensor.datasource.from_dataframe import from_series
         return from_series(self, dtype=dtype)
@@ -1035,6 +1039,10 @@ def columns(self):
 
         return from_pandas_index(self.dtypes.index)
 
+    @property
+    def axes(self):
+        return [self.index, self.columns]
+
 
 class DataFrameData(BaseDataFrameData):
     _type_name = 'DataFrame'
diff --git a/mars/dataframe/indexing/__init__.py b/mars/dataframe/indexing/__init__.py
--- a/mars/dataframe/indexing/__init__.py
+++ b/mars/dataframe/indexing/__init__.py
@@ -15,6 +15,7 @@
 
 def _install():
     from pandas.util import cache_readonly
+    from ..operands import DATAFRAME_TYPE, SERIES_TYPE, INDEX_TYPE
     from .at import at
     from .getitem import dataframe_getitem, series_getitem
     from .iat import iat
@@ -26,7 +27,7 @@ def _install():
     from .set_index import set_index
     from .setitem import dataframe_setitem
     from .reindex import reindex
-    from ..operands import DATAFRAME_TYPE, SERIES_TYPE, INDEX_TYPE
+    from .where import mask, where
 
     for cls in DATAFRAME_TYPE + SERIES_TYPE:
         setattr(cls, 'iloc', cache_readonly(iloc))
@@ -36,6 +37,8 @@ def _install():
         setattr(cls, 'reindex', reindex)
         setattr(cls, 'head', head)
         setattr(cls, 'tail', tail)
+        setattr(cls, 'mask', mask)
+        setattr(cls, 'where', where)
 
     for cls in DATAFRAME_TYPE:
         setattr(cls, 'set_index', set_index)
diff --git a/mars/dataframe/indexing/where.py b/mars/dataframe/indexing/where.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/indexing/where.py
@@ -0,0 +1,332 @@
+# 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
+import pandas as pd
+
+from mars import opcodes
+from mars.core import Base, Entity
+from mars.serialize import AnyField, Int32Field, BoolField, StringField
+from mars.tensor.utils import filter_inputs
+from mars.dataframe.core import DATAFRAME_TYPE, SERIES_TYPE
+from mars.dataframe.operands import DataFrameOperandMixin, DataFrameOperand
+from mars.dataframe.utils import build_df, build_series, validate_axis
+
+
+class DataFrameWhere(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = opcodes.WHERE
+
+    _input = AnyField('input')
+    _cond = AnyField('cond')
+    _other = AnyField('other')
+    _axis = Int32Field('axis')
+    _level = AnyField('level')
+    _errors = StringField('errors')
+    _try_cast = BoolField('try_cast')
+    _replace_true = BoolField('replace_true')
+
+    def __init__(self, input=None, cond=None, other=None,  # pylint: disable=redefined-builtin
+                 axis=None, level=None, errors=None, try_cast=None, replace_true=None, **kw):
+        super().__init__(_input=input, _cond=cond, _other=other, _axis=axis, _level=level,
+                         _errors=errors, _try_cast=try_cast, _replace_true=replace_true, **kw)
+
+    @property
+    def input(self):
+        return self._input
+
+    @property
+    def cond(self):
+        return self._cond
+
+    @property
+    def other(self):
+        return self._other
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @property
+    def level(self):
+        return self._level
+
+    @property
+    def errors(self):
+        return self._errors
+
+    @property
+    def try_cast(self):
+        return self._try_cast
+
+    @property
+    def replace_true(self):
+        return self._replace_true
+
+    def __call__(self, df_or_series):
+        def _check_input_index(obj, axis=None):
+            axis = axis if axis is not None else self.axis
+            if isinstance(obj, DATAFRAME_TYPE) \
+                    and (
+                        df_or_series.columns_value.key != obj.columns_value.key
+                        or df_or_series.index_value.key != obj.index_value.key
+                    ):
+                raise NotImplementedError('Aligning different indices not supported')
+            elif isinstance(obj, SERIES_TYPE) \
+                    and df_or_series.axes[axis].index_value.key != obj.index_value.key:
+                raise NotImplementedError('Aligning different indices not supported')
+
+        _check_input_index(self.cond, axis=0)
+        _check_input_index(self.other)
+
+        if isinstance(df_or_series, DATAFRAME_TYPE):
+            mock_obj = build_df(df_or_series)
+        else:
+            mock_obj = build_series(df_or_series)
+
+        if isinstance(self.other, (pd.DataFrame, DATAFRAME_TYPE)):
+            mock_other = build_df(self.other)
+        elif isinstance(self.other, (pd.Series, SERIES_TYPE)):
+            mock_other = build_series(self.other)
+        else:
+            mock_other = self.other
+
+        result_df = mock_obj.where(np.zeros(mock_obj.shape).astype(bool), other=mock_other,
+                                   axis=self.axis, level=self.level, errors=self.errors,
+                                   try_cast=self.try_cast)
+
+        inputs = filter_inputs([df_or_series, self.cond, self.other])
+        if isinstance(df_or_series, DATAFRAME_TYPE):
+            return self.new_dataframe(inputs, shape=df_or_series.shape,
+                                      dtypes=result_df.dtypes, index_value=df_or_series.index_value,
+                                      columns_value=df_or_series.columns_value)
+        else:
+            return self.new_series(inputs, shape=df_or_series.shape, name=df_or_series.name,
+                                   dtype=result_df.dtype, index_value=df_or_series.index_value)
+
+    def _set_inputs(self, inputs):
+        super()._set_inputs(inputs)
+        inputs_iter = iter(self._inputs)
+        self._input = next(inputs_iter)
+        if isinstance(self._cond, (Base, Entity)):
+            self._cond = next(inputs_iter)
+        if isinstance(self._other, (Base, Entity)):
+            self._other = next(inputs_iter)
+
+    @classmethod
+    def tile(cls, op: "DataFrameWhere"):
+        def rechunk_input(inp, axis=None):
+            axis = axis if axis is not None else op.axis
+            if isinstance(inp, DATAFRAME_TYPE):
+                inp = inp.rechunk(op.input.nsplits)._inplace_tile()
+            elif isinstance(inp, SERIES_TYPE):
+                inp = inp.rechunk({0: op.input.nsplits[axis]})._inplace_tile()
+            return inp
+
+        def get_tiled_chunk(obj, index, axis=None):
+            if isinstance(obj, DATAFRAME_TYPE):
+                return obj.cix[index[0], index[1]]
+            elif isinstance(obj, SERIES_TYPE):
+                axis = axis if axis is not None else op.axis
+                return obj.cix[index[axis], ]
+            else:
+                return obj
+
+        # TODO support axis alignment for three objects
+        cond = rechunk_input(op.cond, axis=0)
+        other = rechunk_input(op.other)
+
+        chunks = []
+        for c in op.input.chunks:
+            cond_chunk = get_tiled_chunk(cond, c.index, axis=0)
+            other_chunk = get_tiled_chunk(other, c.index)
+
+            new_op = op.copy().reset_key()
+            new_op._cond = cond_chunk
+            new_op._other = other_chunk
+
+            inputs = filter_inputs([c, cond_chunk, other_chunk])
+            chunks.append(new_op.new_chunk(inputs, **c.params))
+
+        new_op = op.copy().reset_key()
+        return new_op.new_tileables(op.inputs, chunks=chunks, nsplits=op.input.nsplits,
+                                    **op.input.params)
+
+    @classmethod
+    def execute(cls, ctx, op: "DataFrameWhere"):
+        out_obj = op.outputs[0]
+
+        input_data = ctx[op.input.key]
+        cond = op.cond
+        if isinstance(cond, (Base, Entity)):
+            cond = ctx[cond.key]
+
+        other = op.other
+        if isinstance(other, (Base, Entity)):
+            other = ctx[other.key]
+
+        if op.replace_true:
+            ctx[out_obj.key] = input_data.mask(cond, other, axis=op.axis, level=op.level,
+                                               errors=op.errors, try_cast=op.try_cast)
+        else:
+            ctx[out_obj.key] = input_data.where(cond, other, axis=op.axis, level=op.level,
+                                                errors=op.errors, try_cast=op.try_cast)
+
+
+_doc_template = """
+Replace values where the condition is {replace_true}.
+
+Parameters
+----------
+cond : bool Series/DataFrame, array-like, or callable
+    Where `cond` is False, keep the original value. Where
+    True, replace with corresponding value from `other`.
+    If `cond` is callable, it is computed on the Series/DataFrame and
+    should return boolean Series/DataFrame or array. The callable must
+    not change input Series/DataFrame (though pandas doesn't check it).
+other : scalar, Series/DataFrame, or callable
+    Entries where `cond` is True are replaced with
+    corresponding value from `other`.
+    If other is callable, it is computed on the Series/DataFrame and
+    should return scalar or Series/DataFrame. The callable must not
+    change input Series/DataFrame (though pandas doesn't check it).
+inplace : bool, default False
+    Whether to perform the operation in place on the data.
+axis : int, default None
+    Alignment axis if needed.
+level : int, default None
+    Alignment level if needed.
+errors : str, {{'raise', 'ignore'}}, default 'raise'
+    Note that currently this parameter won't affect
+    the results and will always coerce to a suitable dtype.
+
+    - 'raise' : allow exceptions to be raised.
+    - 'ignore' : suppress exceptions. On error return original object.
+
+try_cast : bool, default False
+    Try to cast the result back to the input type (if possible).
+
+Returns
+-------
+Same type as caller
+
+See Also
+--------
+:func:`DataFrame.{opposite}` : Return an object of same shape as
+    self.
+
+Notes
+-----
+The mask method is an application of the if-then idiom. For each
+element in the calling DataFrame, if ``cond`` is ``False`` the
+element is used; otherwise the corresponding element from the DataFrame
+``other`` is used.
+
+The signature for :func:`DataFrame.where` differs from
+:func:`numpy.where`. Roughly ``df1.where(m, df2)`` is equivalent to
+``np.where(m, df1, df2)``.
+
+For further details and examples see the ``mask`` documentation in
+:ref:`indexing <indexing.where_mask>`.
+
+Examples
+--------
+>>> import mars.tensor as mt
+>>> import mars.dataframe as md
+>>> s = md.Series(range(5))
+>>> s.where(s > 0).execute()
+0    NaN
+1    1.0
+2    2.0
+3    3.0
+4    4.0
+dtype: float64
+
+>>> s.mask(s > 0).execute()
+0    0.0
+1    NaN
+2    NaN
+3    NaN
+4    NaN
+dtype: float64
+
+>>> s.where(s > 1, 10).execute()
+0    10
+1    10
+2    2
+3    3
+4    4
+dtype: int64
+
+>>> df = md.DataFrame(mt.arange(10).reshape(-1, 2), columns=['A', 'B'])
+>>> df.execute()
+   A  B
+0  0  1
+1  2  3
+2  4  5
+3  6  7
+4  8  9
+>>> m = df % 3 == 0
+>>> df.where(m, -df).execute()
+   A  B
+0  0 -1
+1 -2  3
+2 -4 -5
+3  6 -7
+4 -8  9
+>>> df.where(m, -df) == mt.where(m, df, -df).execute()
+      A     B
+0  True  True
+1  True  True
+2  True  True
+3  True  True
+4  True  True
+>>> df.where(m, -df) == df.mask(~m, -df).execute()
+      A     B
+0  True  True
+1  True  True
+2  True  True
+3  True  True
+4  True  True
+"""
+
+
+def _where(df_or_series, cond, other=np.nan, inplace=False, axis=None, level=None, errors='raise',
+           try_cast=False, replace_true=False):
+    if df_or_series.ndim == 2 and getattr(other, 'ndim', 2) == 1 and axis is None:
+        raise ValueError('Must specify axis=0 or 1')
+
+    axis = validate_axis(axis or 0, df_or_series)
+    op = DataFrameWhere(cond=cond, other=other, axis=axis, level=level, errors=errors,
+                        try_cast=try_cast, replace_true=replace_true)
+    result = op(df_or_series)
+    if inplace:
+        df_or_series.data = result.data
+    else:
+        return result
+
+
+def where(df_or_series, cond, other=np.nan, inplace=False, axis=None, level=None, errors='raise',
+          try_cast=False):
+    return _where(df_or_series, cond, other=other, inplace=inplace, axis=axis, level=level,
+                  errors=errors, try_cast=try_cast, replace_true=False)
+
+
+def mask(df_or_series, cond, other=np.nan, inplace=False, axis=None, level=None, errors='raise',
+         try_cast=False):
+    return _where(df_or_series, cond, other=other, inplace=inplace, axis=axis, level=level,
+                  errors=errors, try_cast=try_cast, replace_true=True)
+
+
+mask.__doc__ = _doc_template.format(replace_true=True, opposite='where')
+where.__doc__ = _doc_template.format(replace_true=False, opposite='mask')


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
--- a/mars/dataframe/indexing/tests/test_indexing_execution.py
+++ b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -1081,3 +1081,62 @@ def testReindexExecution(self):
             result = self.executor.execute_dataframe(r, concat=True)[0]
             expected = s_data.reindex(['c2', 'c11', 'c4'], copy=False)
             pd.testing.assert_series_equal(result, expected)
+
+    def testWhereExecution(self):
+        dates = pd.date_range('1/1/2000', periods=20)
+
+        raw_df = pd.DataFrame(np.random.randn(20, 10), index=dates, columns=list('ABCDEFGHIJ'))
+        raw_df2 = pd.DataFrame(np.random.randn(20, 10), index=dates, columns=list('ABCDEFGHIJ'))
+        df = md.DataFrame(raw_df, chunk_size=6)
+        df2 = md.DataFrame(raw_df2, chunk_size=7)
+
+        raw_series = pd.Series(np.random.randn(20), index=dates)
+        raw_series2 = pd.Series(np.random.randn(20), index=dates)
+        raw_series3 = pd.Series(np.random.randn(10), index=list('ABCDEFGHIJ'))
+        series = md.Series(raw_series, chunk_size=6)
+        series2 = md.Series(raw_series2, chunk_size=7)
+        series3 = md.Series(raw_series3, chunk_size=7)
+
+        # tests for dataframes
+        with self.assertRaises(NotImplementedError):
+            df.mask(df < 0, md.DataFrame(np.random.randn(5, 5)))
+        with self.assertRaises(NotImplementedError):
+            df.mask(series < 0, md.Series(np.random.randn(5)), axis=0)
+
+        r = df.mask(df < 0)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw_df.mask(raw_df < 0))
+        r = df.mask(raw_df < 0, df2)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw_df.mask(raw_df < 0, raw_df2))
+
+        # tests for series
+        with self.assertRaises(NotImplementedError):
+            series.mask(series < 0, md.Series(np.random.randn(5)))
+
+        r = series.where(series < 0, 0)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw_series.where(raw_series < 0, 0))
+        r = series.where(series < 0, series2)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw_series.where(raw_series < 0, raw_series2))
+
+        # test for dataframe with series
+        with self.assertRaises(ValueError):
+            df.mask(df < 0, series)
+
+        r = df.mask(df < 0, series, axis=0)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw_df.mask(raw_df < 0, raw_series, axis=0))
+        r = df.mask(series < 0, df2)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw_df.mask(raw_series < 0, raw_df2))
+        r = df.mask(series < 0, series3, axis=1)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                      raw_df.mask(raw_series < 0, raw_series3, axis=1))
+
+        # test inplace
+        new_df = df.copy()
+        new_df.mask(new_df < 0, inplace=True)
+        pd.testing.assert_frame_equal(self.executor.execute_dataframe(new_df, concat=True)[0],
+                                      raw_df.mask(raw_df < 0))
diff --git a/mars/deploy/kubernetes/tests/test_kubernetes.py b/mars/deploy/kubernetes/tests/test_kubernetes.py
--- a/mars/deploy/kubernetes/tests/test_kubernetes.py
+++ b/mars/deploy/kubernetes/tests/test_kubernetes.py
@@ -18,6 +18,7 @@
 import shutil
 import subprocess
 import tempfile
+import time
 import unittest
 import uuid
 from contextlib import contextmanager
@@ -166,12 +167,20 @@ def testRescale(self):
 
             cluster_client.rescale_workers(1)
 
-            worker_items = kube_api.list_namespaced_pod(
-                cluster_client.namespace,
-                label_selector=f'mars/service-type={MarsWorkersConfig.rc_name}').to_dict()
-            self.assertEqual(len([w for w in worker_items['items'] if w['status']['phase'] == 'Running']), 1)
             self.assertEqual(cluster_client._session._sess.count_workers(), 1)
 
+            start_check_time = time.time()
+            n_workers = 2
+            while time.time() - start_check_time < 60:
+                worker_items = kube_api.list_namespaced_pod(
+                    cluster_client.namespace,
+                    label_selector=f'mars/service-type={MarsWorkersConfig.rc_name}').to_dict()
+                n_workers = len([w for w in worker_items['items'] if w['status']['phase'] == 'Running'])
+                if n_workers == 1:
+                    break
+                time.sleep(1)
+            self.assertEqual(n_workers, 1)
+
     @mock.patch('kubernetes.client.CoreV1Api.create_namespaced_replication_controller',
                 new=lambda *_, **__: None)
     def testCreateTimeout(self):