Repository: mars-project/mars

Issue:
[BUG] `dask.persist` cannot work on dask-on-mars
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

`dask.persist` cannot work on dask-on-mars.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
In [1]: import mars

In [2]: mars.new_session()
Web service started at http://0.0.0.0:64573
Out[2]: <mars.deploy.oscar.session.SyncSession at 0x7fa0f8366d90>

In [3]: import dask
   ...: from mars.contrib.dask import mars_scheduler

In [4]: def inc(x):
   ...:     return x + 1
   ...: 

In [5]: dask_task = dask.delayed(inc)(1)

In [6]: dask_task.persist(scheduler=mars_scheduler)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-095d61a2acc0> in <module>
----> 1 dask_task.persist(scheduler=mars_scheduler)

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(self, **kwargs)
    259         dask.base.persist
    260         """
--> 261         (result,) = persist(self, traverse=False, **kwargs)
    262         return result
    263 

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(traverse, optimize_graph, scheduler, *args, **kwargs)
    831         postpersists.append((rebuild, a_keys, state))
    832 
--> 833     results = schedule(dsk, keys, **kwargs)
    834     d = dict(zip(keys, results))
    835     results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]

~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_scheduler(dsk, keys)
     39         Computed values corresponding to the provided keys.
     40     """
---> 41     res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
     42     if not isinstance(res, List):
     43         return [[res]]

~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_dask_get(dsk, keys)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/highlevelgraph.py in __getitem__(self, key)
    737                 pass
    738 
--> 739         raise KeyError(key)
    740 
    741     def __len__(self) -> int:

KeyError: 'i'
```





Discussion:
Hi, @loopyme , do you have time to look at this issue?
Work in progress.
@loopyme Any progress to update?
Message ID: ***@***.***&gt;Already locate the problem, all we have to do is support not only 2D-keys but also 1D-keys.

My development PC was left somewhere, will raise PR by the end of this week.

Resolving patch:
diff --git a/mars/contrib/dask/scheduler.py b/mars/contrib/dask/scheduler.py
--- a/mars/contrib/dask/scheduler.py
+++ b/mars/contrib/dask/scheduler.py
@@ -14,12 +14,13 @@
 
 from dask.core import istask, ishashable
 
-from typing import List, Tuple
+from typing import List, Tuple, Union
 from .utils import reduce
 from ...remote import spawn
+from ...deploy.oscar.session import execute
 
 
-def mars_scheduler(dsk: dict, keys: List[List[str]]):
+def mars_scheduler(dsk: dict, keys: Union[List[List[str]], List[str]]):
     """
     A Dask-Mars scheduler
 
@@ -30,22 +31,29 @@ def mars_scheduler(dsk: dict, keys: List[List[str]]):
     ----------
     dsk: Dict
         Dask graph, represented as a task DAG dictionary.
-    keys: List[List[str]]
-        2d-list of Dask graph keys whose values we wish to compute and return.
+    keys: Union[List[List[str]], List[str]]
+        1d or 2d list of Dask graph keys whose values we wish to compute and return.
 
     Returns
     -------
     Object
-        Computed values corresponding to the provided keys.
+        Computed values corresponding to the provided keys with same dimension.
     """
-    res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
-    if not isinstance(res, List):
-        return [[res]]
-    else:
-        return res
 
+    if isinstance(keys, List) and not isinstance(keys[0], List):  # 1d keys
+        task = execute(mars_dask_get(dsk, keys))
+        if not isinstance(task, List):
+            task = [task]
+        return map(lambda x: x.fetch(), task)
+    else:  # 2d keys
+        res = execute(reduce(mars_dask_get(dsk, keys))).fetch()
+        if not isinstance(res, List):
+            return [[res]]
+        else:
+            return res
 
-def mars_dask_get(dsk: dict, keys: List[List]):
+
+def mars_dask_get(dsk: dict, keys: Union[List[List[str]], List[str]]):
     """
     A Dask-Mars convert function. This function will send the dask graph layers
         to Mars Remote API, generating mars objects correspond to the provided keys.
@@ -54,13 +62,13 @@ def mars_dask_get(dsk: dict, keys: List[List]):
     ----------
     dsk: Dict
         Dask graph, represented as a task DAG dictionary.
-    keys: List[List[str]]
-        2d-list of Dask graph keys whose values we wish to compute and return.
+    keys: Union[List[List[str]], List[str]]
+        1d or 2d list of Dask graph keys whose values we wish to compute and return.
 
     Returns
     -------
     Object
-        Spawned mars objects corresponding to the provided keys.
+        Spawned mars objects corresponding to the provided keys with same dimension.
     """
 
     def _get_arg(a):
@@ -68,7 +76,7 @@ def _get_arg(a):
         if ishashable(a) and a in dsk.keys():
             while ishashable(a) and a in dsk.keys():
                 a = dsk[a]
-            return _execute_task(a)
+            return _spawn_task(a)
         elif not isinstance(a, str) and hasattr(a, "__getitem__"):
             if istask(
                 a
@@ -80,9 +88,14 @@ def _get_arg(a):
                 return type(a)(_get_arg(i) for i in a)
         return a
 
-    def _execute_task(task: tuple):
+    def _spawn_task(task: tuple):
         if not istask(task):
             return _get_arg(task)
         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
 
-    return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]
+    return [
+        [_spawn_task(dsk[k]) for k in keys_d]
+        if isinstance(keys_d, List)
+        else _spawn_task(dsk[keys_d])
+        for keys_d in keys
+    ]


Test patch:
diff --git a/mars/contrib/dask/tests/test_dask.py b/mars/contrib/dask/tests/test_dask.py
--- a/mars/contrib/dask/tests/test_dask.py
+++ b/mars/contrib/dask/tests/test_dask.py
@@ -143,3 +143,41 @@ def inc(x: int):
         assert dask.compute(test_obj) == dask.compute(
             test_obj, scheduler=mars_scheduler
         )
+
+
+@pytest.mark.skipif(not dask_installed, reason="dask not installed")
+def test_persist(setup_cluster):
+    import dask
+
+    def inc(x):
+        return x + 1
+
+    a = dask.delayed(inc)(1)
+    task_mars_persist = dask.delayed(inc)(a.persist(scheduler=mars_scheduler))
+    task_dask_persist = dask.delayed(inc)(a.persist())
+
+    assert task_dask_persist.compute() == task_mars_persist.compute(
+        scheduler=mars_scheduler
+    )
+
+
+@pytest.mark.skipif(not dask_installed, reason="dask not installed")
+def test_partitioned_dataframe_persist(setup_cluster):
+    import numpy as np
+    import pandas as pd
+    from dask import dataframe as dd
+    from pandas._testing import assert_frame_equal
+
+    data = np.random.randn(10000, 100)
+    df = dd.from_pandas(
+        pd.DataFrame(data, columns=[f"col{i}" for i in range(100)]), npartitions=4
+    )
+    df["col0"] = df["col0"] + df["col1"] / 2
+    col2_mean = df["col2"].mean()
+
+    df_mars_persist = df[df["col2"] > col2_mean.persist(scheduler=mars_scheduler)]
+    df_dask_persist = df[df["col2"] > col2_mean.persist()]
+
+    assert_frame_equal(
+        df_dask_persist.compute(), df_mars_persist.compute(scheduler=mars_scheduler)
+    )