Repository: dask/dask
Pull request: Change type of Dask array when meta type changes
Previously da.map_blocks would always return a Dask Array,
even when we mapped a function like pandas.Series.

Now we convert the underlying collection type to a dataframe

Fixes https://github.com/dask/dask/issues/5606

cc @quasiben 


Problem statement:
[FEA] Dask CuPy Array -> Dask CuDF DataFrame
In dask, we can create a dataframe from a dask array like the following:

> dd.from_dask_array(da.random.random(n_rows))

In the above, random numpy data is generated then transformed to a Pandas dataframe.  It would be really cool if I could do something similar with CuPy arrays -> CuDF DataFrames:

```python
import cupy
import dask.array as da
import dask.dataframe as dd

rs = da.random.RandomState(RandomState=cupy.random.RandomState)
dd.from_dask_array(rs.random(n_rows))
```

Maybe in `from_dask_array` we could detect if the obj implemented the `__cuda_array_interface__` and proceed accordingly ?

Hints from issue discussion:
I suspect that `x.map_blocks(cudf.Series)` will work ok and that you don't need the special `from_dask_array` function.
I need to do a bit more than that -- though not much:

```python
import cudf
import cupy
import dask.array as da
import dask.dataframe as dd
from dask.dataframe.utils import make_meta

rs = da.random.RandomState(RandomState=cupy.random.RandomState)
x = rs.random(10, chunks=(5,))
sr_x = dd.from_dask_array(x.map_blocks(cudf.Series, meta=make_meta(cudf.Series())), columns='x')

y = rs.randint(10, size=(10,), chunks=(5,))
sr_y = dd.from_dask_array(y.map_blocks(cudf.Series, meta=make_meta(cudf.Series())), columns='y')

dd.concat([sr_x, sr_y], axis=1)
```

What are your thoughts about usability of the above ?  Could this be simplified ?


> I suspect that x.map_blocks(cudf.Series) will work ok and that you don't need the special from_dask_array function.

What stops this from working?  If it's easy I'd rather fix this than work around it.
Hmm, looking at this again i think above is incorrect -- the `from_dask_array` is still going to cast the data to a pd.Series.  So I'm not sure how the dask array object can get wrapped into a Dask Dataframe object

```python
In [36]: x.map_blocks(cudf.Series, meta=make_meta(cudf.Series()))
Out[36]: dask.array<Series, shape=(100000, 100000), dtype=float64, chunksize=(1000, 1000), chunktype=numpy.ndarray>
```
> What stops this from working? If it's easy I'd rather fix this than work around it.

dtype issues:

```python
In [42]: x.map_blocks(cudf.Series)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-42-9a7f50f7b816> in <module>
----> 1 x.map_blocks(cudf.Series)

/datasets/bzaitlen/GitRepos/dask/dask/array/core.py in map_blocks(self, func, *args, **kwargs)
   2088     @wraps(map_blocks)
   2089     def map_blocks(self, func, *args, **kwargs):
-> 2090         return map_blocks(func, self, *args, **kwargs)
   2091
   2092     def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):

/datasets/bzaitlen/GitRepos/dask/dask/array/core.py in map_blocks(func, name, token, dtype, chunks, drop_axis, new_axis, meta, *args, **kwargs)
    591
    592     if dtype is None and meta is None:
--> 593         dtype = apply_infer_dtype(func, args, original_kwargs, "map_blocks")
    594
    595     if drop_axis:

/datasets/bzaitlen/GitRepos/dask/dask/array/core.py in apply_infer_dtype(func, args, kwargs, funcname, suggest_dtype, nout)
    373         msg = None
    374     if msg is not None:
--> 375         raise ValueError(msg)
    376     return o.dtype if nout is None else tuple(e.dtype for e in o)
    377

ValueError: `dtype` inference failed in `map_blocks`.

Please specify the dtype explicitly using the `dtype` kwarg.

Original error is below:
------------------------
BufferSentryError('ndim mismatch')

Traceback:
---------
  File "/datasets/bzaitlen/GitRepos/dask/dask/array/core.py", line 350, in apply_infer_dtype
    o = func(*args, **kwargs)
  File "/datasets/bzaitlen/GitRepos/cudf/python/cudf/cudf/core/series.py", line 124, in __init__
    data, nan_as_null=nan_as_null, dtype=dtype, name=name
  File "/datasets/bzaitlen/GitRepos/cudf/python/cudf/cudf/core/column/column.py", line 1226, in as_column
    data = as_column(rmm.to_device(arbitrary), nan_as_null=nan_as_null)
  File "/datasets/bzaitlen/GitRepos/cudf/python/cudf/cudf/core/column/column.py", line 1169, in as_column
    data = as_column(Buffer(arbitrary))
  File "/datasets/bzaitlen/GitRepos/cudf/python/cudf/cudf/core/buffer.py", line 44, in __init__
    _BufferSentry(self.mem).ndim(1)
  File "/datasets/bzaitlen/GitRepos/cudf/python/cudf/cudf/core/buffer.py", line 213, in ndim
    raise BufferSentryError("ndim mismatch")
```
OK, for dask dataframe we use this

https://github.com/dask/dask/blob/0805d391718c1efc6f4730b90407fde8e32a2bd8/dask/dataframe/core.py#L5864-L5888

Which allows map_partitions calls that change the underlying meta type to return other kinds of Dask collections.  It looks like we don't currently do this for Dask array.  I can take a look at this starting next week.

@quasiben to unblock you I would recommend making a function that creates a single cudf dataframe and then using that function multiple times with dask.delayed

```python
def make_random_dataframe(n: int) -> cudf.DataFrame:
    return cudf.DataFrame({...: cupy.random.random})
    ...

partitions = [dask.delayed(make_random_dataframe, pure=False)(10000000) for _ in range(npartitions]
df = dd.from_delayed(partitions)
```

Commit message: Explicitly support series objects in da.reduction

Patch to review:
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 240baa35ae9..d58c1b7dfe1 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -139,7 +139,7 @@ def blockwise(
     if new:
         raise ValueError("Unknown dimension", new)
 
-    from .core import Array, unify_chunks, normalize_arg
+    from .core import unify_chunks, normalize_arg
 
     if align_arrays:
         chunkss, arrays = unify_chunks(*args)
@@ -231,12 +231,12 @@ def blockwise(
         from .utils import compute_meta
 
         meta = compute_meta(func, dtype, *args[::2], **kwargs)
-    if meta is not None:
-        return Array(graph, out, chunks, meta=meta)
-    else:
-        return Array(graph, out, chunks, dtype=dtype)
+    return new_da_object(graph, out, chunks, meta=meta, dtype=dtype)
 
 
 def atop(*args, **kwargs):
     warnings.warn("The da.atop function has moved to da.blockwise")
     return blockwise(*args, **kwargs)
+
+
+from .core import new_da_object
diff --git a/dask/array/core.py b/dask/array/core.py
index 8ca6d68f20f..0991fc18f92 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -54,6 +54,9 @@
     ndimlist,
     format_bytes,
     typename,
+    is_dataframe_like,
+    is_series_like,
+    is_index_like,
 )
 from ..core import quote
 from ..delayed import delayed, Delayed
@@ -590,7 +593,13 @@ def map_blocks(
     original_kwargs = kwargs
 
     if dtype is None and meta is None:
-        dtype = apply_infer_dtype(func, args, original_kwargs, "map_blocks")
+        try:
+            meta = func(
+                *[getattr(arg, "_meta", arg) for arg in args],
+                **{k: getattr(v, "_meta", v) for k, v in original_kwargs.items()}
+            )
+        except Exception:
+            dtype = apply_infer_dtype(func, args, original_kwargs, "map_blocks")
 
     if drop_axis:
         out_ind = tuple(x for i, x in enumerate(out_ind) if i not in drop_axis)
@@ -4720,4 +4729,19 @@ def from_npy_stack(dirname, mmap_mode="r"):
     return Array(dsk, name, chunks, dtype)
 
 
+def new_da_object(dsk, name, chunks, meta=None, dtype=None):
+    """Generic constructor for dask.dataframe objects.
+
+    Decides the appropriate output class based on the type of `meta` provided.
+    """
+    if is_dataframe_like(meta) or is_series_like(meta) or is_index_like(meta):
+        from ..dataframe.core import new_dd_object
+
+        assert all(len(c) == 1 for c in chunks[1:])
+        divisions = [None] * (len(chunks[0]) + 1)
+        return new_dd_object(dsk, name, meta, divisions)
+    else:
+        return Array(dsk, name=name, chunks=chunks, meta=meta, dtype=dtype)
+
+
 from .utils import meta_from_array
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 9bbdc11c5d8..d085c25adf5 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -21,7 +21,15 @@
 from .numpy_compat import ma_divide, divide as np_divide
 from ..base import tokenize
 from ..highlevelgraph import HighLevelGraph
-from ..utils import ignoring, funcname, Dispatch, deepmap, getargspec, derived_from
+from ..utils import (
+    ignoring,
+    funcname,
+    Dispatch,
+    deepmap,
+    getargspec,
+    derived_from,
+    is_series_like,
+)
 from .. import config
 
 # Generic functions to support chunks of different types
@@ -149,6 +157,8 @@ def reduction(
         chunk = partial(chunk, dtype=dtype)
     if "dtype" in getargspec(aggregate).args:
         aggregate = partial(aggregate, dtype=dtype)
+    if is_series_like(x):
+        x = x.values
 
     # Map chunk across all blocks
     inds = tuple(range(x.ndim))
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index e477514e55b..b9976eb0150 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -4231,3 +4231,26 @@ def test_rechunk_auto():
     y = x.rechunk()
 
     assert y.npartitions == 1
+
+
+def test_map_blocks_series():
+    pd = pytest.importorskip("pandas")
+    import dask.dataframe as dd
+
+    x = da.ones(10, chunks=(5,))
+    s = x.map_blocks(pd.Series)
+    assert isinstance(s, dd.Series)
+    assert s.npartitions == x.npartitions
+    assert_eq(s, s)
+
+
+@pytest.mark.xfail(reason="need to remove singleton index dimension")
+def test_map_blocks_dataframe():
+    pd = pytest.importorskip("pandas")
+    import dask.dataframe as dd
+
+    x = da.ones((10, 2), chunks=(5, 2))
+    s = x.map_blocks(pd.DataFrame)
+    assert isinstance(s, dd.DataFrame)
+    assert s.npartitions == x.npartitions
+    assert_eq(s, s)


Code review comments:
diff_hunk: @@ -590,7 +593,13 @@ def map_blocks(
     original_kwargs = kwargs
 
     if dtype is None and meta is None:
-        dtype = apply_infer_dtype(func, args, original_kwargs, "map_blocks")
+        try:
+            meta = func(
+                *[getattr(arg, "_meta", arg) for arg in args],
+                **{k: getattr(v, "_meta", v) for k, v in original_kwargs.items()}
+            )
path: dask/array/core.py
text: @author:
@user1 it would be good to get your thoughts on this.  

This seems to work fine, but it seems like something that you would have thought of before, so I suspect that there is some reason why it was not done.

@user1:
It seems very similar to what `compute_meta` is doing, specifically https://github.com/dask/dask/blob/master/dask/array/utils.py#L113-L116. Maybe you wanted to use `compute_meta` instead, or is there something different here that I'm missing?

If the question instead is "why was there no support for `map_blocks` yet?", then probably it's part of the incomplete work for `_meta` support, as tracked in https://github.com/dask/dask/issues/4888.

diff_hunk: @@ -4720,4 +4729,19 @@ def from_npy_stack(dirname, mmap_mode="r"):
     return Array(dsk, name, chunks, dtype)
 
 
+def new_da_object(dsk, name, chunks, meta=None, dtype=None):
+    """Generic constructor for dask.dataframe objects.
path: dask/array/core.py
text: @user1:
```suggestion
    """Generic constructor for dask.array or dask.dataframe objects.
```

Possible copy-paste leftover.

It seems like we can't have a `new_?_object`, since this requires a `chunks` argument, while `new_dd_object` doesn't.