Repository: dask/dask
Pull request: stack nd array w/unknown chunks
- [x] Closes #7546
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask` / `isort dask`


Problem statement:
numpy/dask difference in stacking multi-dim arrays with unknown chunk sizes
**What happened**:

When using da.stack and stacking a multi-dim list (of arrays), the allow_unknown_chunksizes parameter is not passed into asarray.  This is needed as asarray calls back to stack if it is passed a sequence, and an error is raised if the unknown chunksize is not the first dimension.

**What you expected to happen**:

da.stack and np.stack to produce the same result when stacking a multi-dim array.

**Minimal Complete Verifiable Example**:

```python
import numpy as np
import pandas as pd
import dask.array as da
import dask.dataframe as dd

a_df = pd.DataFrame({"a": np.arange(12)})
b_df = pd.DataFrame({"b": np.arange(12) * 10})
c_df = pd.DataFrame({"c": np.arange(12) * 100})
d_df = pd.DataFrame({"d": np.arange(12) * 1000})

a_ddf = dd.from_pandas(a_df, sort=False, npartitions=3)
b_ddf = dd.from_pandas(b_df, sort=False, npartitions=3)
c_ddf = dd.from_pandas(c_df, sort=False, npartitions=3)
d_ddf = dd.from_pandas(d_df, sort=False, npartitions=3)

a_x = a_ddf.values
b_x = b_ddf.values
c_x = c_ddf.values
d_x = d_ddf.values

# Works
np.stack([
    [a_df.values, b_df.values],
    [c_df.values, d_df.values],
])

# Does not work
da.stack([
    [a_x, b_x],
    [c_x, d_x]
], allow_unknown_chunksizes=True)
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2021.04.0+16.ge83379d5
- Python version: 3.8.8
- Operating System: MacOS 11.2.3
- Install method (conda, pip, source): conda

Hints from issue discussion:

Commit message: Add allow_unknown_chunksizes to asarray

Patch to review:
diff --git a/dask/array/core.py b/dask/array/core.py
index 776050c9aa6..302b1360a2e 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4090,7 +4090,7 @@ def retrieve_from_ooc(keys, dsk_pre, dsk_post=None):
     return load_dsk
 
 
-def asarray(a, **kwargs):
+def asarray(a, allow_unknown_chunksizes=False, **kwargs):
     """Convert the input to a dask array.
 
     Parameters
@@ -4122,7 +4122,7 @@ def asarray(a, **kwargs):
     elif type(a).__module__.split(".")[0] == "xarray" and hasattr(a, "data"):
         return asarray(a.data)
     elif isinstance(a, (list, tuple)) and any(isinstance(i, Array) for i in a):
-        return stack(a)
+        return stack(a, allow_unknown_chunksizes=allow_unknown_chunksizes, **kwargs)
     elif not isinstance(getattr(a, "shape", None), Iterable):
         a = np.asarray(a)
     return from_array(a, getitem=getter_inline, **kwargs)
@@ -4657,7 +4657,7 @@ def stack(seq, axis=0, allow_unknown_chunksizes=False):
     """
     from . import wrap
 
-    seq = [asarray(a) for a in seq]
+    seq = [asarray(a, allow_unknown_chunksizes=allow_unknown_chunksizes) for a in seq]
 
     if not seq:
         raise ValueError("Need array(s) to stack")
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 1b5783daea4..ddf842706e8 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -440,6 +440,28 @@ def test_stack_unknown_chunksizes():
 
     assert_eq(c_x, np.stack([a_df.values, b_df.values], axis=1))
 
+    m_df = pd.DataFrame({"m": np.arange(12) * 100})
+    n_df = pd.DataFrame({"n": np.arange(12) * 1000})
+
+    m_ddf = dd.from_pandas(m_df, sort=False, npartitions=3)
+    n_ddf = dd.from_pandas(n_df, sort=False, npartitions=3)
+
+    m_x = m_ddf.values
+    n_x = n_ddf.values
+
+    assert np.isnan(m_x.shape[0])
+    assert np.isnan(n_x.shape[0])
+
+    with pytest.raises(ValueError) as exc_info:
+        da.stack([[a_x, b_x], [m_x, n_x]])
+
+    assert "shape" in str(exc_info.value)
+    assert "nan" in str(exc_info.value)
+
+    c_x = da.stack([[a_x, b_x], [m_x, n_x]], allow_unknown_chunksizes=True)
+
+    assert_eq(c_x, np.stack([[a_df.values, b_df.values], [m_df.values, n_df.values]]))
+
 
 def test_concatenate():
     a, b, c = [


Code review comments:
diff_hunk: @@ -4122,7 +4122,7 @@ def asarray(a, **kwargs):
     elif type(a).__module__.split(".")[0] == "xarray" and hasattr(a, "data"):
         return asarray(a.data)
     elif isinstance(a, (list, tuple)) and any(isinstance(i, Array) for i in a):
-        return stack(a)
+        return stack(a, allow_unknown_chunksizes=allow_unknown_chunksizes, **kwargs)
path: dask/array/core.py
text: @user1:
I don't think we should be passing the `kwargs` into this stack call.

```suggestion
        return stack(a, allow_unknown_chunksizes=allow_unknown_chunksizes)
```

@author:
I agree.  Thanks.