Repository: dask/dask

Issue:
Allow da.stack to operate on unknown chunk sizes
Currently `da.stack` requires that shape of the passed arrays matches. This can fail when one of more of the chunk sizes is unknown:

```python

In [1]: import dask.array as da

In [2]: import numpy as np

In [3]: a = da.ones((10, 15))

In [4]: b = da.ones((10, 15))

In [5]: a._chunks = ((np.nan,), (15,))

In [6]: b._chunks = ((np.nan,), (15,))

In [7]: da.stack([a, b])
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-7-87775e3d62ee> in <module>
----> 1 da.stack([a, b])

~/sandbox/dask/dask/array/core.py in stack(seq, axis)
   4257             "The first {0} had shape {1}, while array "
   4258             "{2} has shape {3}".format(
-> 4259                 idx[0], seq[0].shape, idx[0] + 1, seq[idx[0]].shape
   4260             )
   4261         )

ValueError: Stacked arrays must have the same shape. The first 0 had shape (nan, 15), while array 1 has shape (nan, 15)

```

We could accept an `allow_unknown_chunksizes` similar to `vstack` to treat missing chunk sizes as equal.


Discussion:
+1 from me.  I suspect that we'll still at least get a runtime error.

On Mon, May 4, 2020 at 1:23 PM Tom Augspurger <notifications@github.com>
wrote:

> Currently da.stack requires that shape of the passed arrays matches. This
> can fail when one of more of the chunk sizes is unknown:
>
> In [1]: import dask.array as da
>
> In [2]: import numpy as np
>
> In [3]: a = da.ones((10, 15))
>
> In [4]: b = da.ones((10, 15))
>
> In [5]: a._chunks = ((np.nan,), (15,))
>
> In [6]: b._chunks = ((np.nan,), (15,))
>
> In [7]: da.stack([a, b])---------------------------------------------------------------------------ValueError                                Traceback (most recent call last)<ipython-input-7-87775e3d62ee> in <module>----> 1 da.stack([a, b])
> ~/sandbox/dask/dask/array/core.py in stack(seq, axis)
>    4257             "The first {0} had shape {1}, while array "
>    4258             "{2} has shape {3}".format(-> 4259                 idx[0], seq[0].shape, idx[0] + 1, seq[idx[0]].shape
>    4260             )
>    4261         )
> ValueError: Stacked arrays must have the same shape. The first 0 had shape (nan, 15), while array 1 has shape (nan, 15)
>
> We could accept an allow_unknown_chunksizes similar to vstack to treat
> missing chunk sizes as equal.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6177>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTBP5RAGIDXYRB3MDGLRP4PZ7ANCNFSM4MZBZNCQ>
> .
>

Oh... perhaps I'll remove the "good first issue" label then. I didn't look closely at the current implementation.
Oh I misunderstood. I think by "still get a runtime error" you mean we'll get a runtime error if the chunk sizes are in fact different.
Right, if the chunk sizes are different then they probably won't get an
error message when they construct the dask array, but they will get an
error when they call compute.  This is less awesome, but still nice to know.

On Mon, May 4, 2020 at 2:21 PM Tom Augspurger <notifications@github.com>
wrote:

> Oh I misunderstood. I think by "still get a runtime error" you mean we'll
> get a runtime error if the chunk sizes are in fact different.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6177#issuecomment-623712714>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTGZZNVQS2IOM64GTATRP4WVRANCNFSM4MZBZNCQ>
> .
>

Would like to work on this.
Should we also change this exception message?
https://github.com/dask/dask/blob/592ab6cea0ea60da87b6e1e8a9a85a10d6af22de/dask/array/core.py#L3793-L3799
We'll want to keep that exception when any of the shapes are *known* to be different.

But we won't raise an exception there when we don't know if the shapes differ because one or more of the arrays has unknown chunk sizes.

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -4216,13 +4216,24 @@ def transposelist(arrays, axes, extradims=0):
     return reshapelist(newshape, result)
 
 
-def stack(seq, axis=0):
+def stack(seq, axis=0, allow_unknown_chunksizes=False):
     """
     Stack arrays along a new axis
 
     Given a sequence of dask arrays, form a new dask array by stacking them
     along a new dimension (axis=0 by default)
 
+     Parameters
+    ----------
+    seq: list of dask.arrays
+    axis: int
+        Dimension along which to align all of the arrays
+    allow_unknown_chunksizes: bool
+        Allow unknown chunksizes, such as come from converting from dask
+        dataframes.  Dask.array is unable to verify that chunks line up.  If
+        data comes from differently aligned sources then this can cause
+        unexpected results.
+
     Examples
     --------
 
@@ -4256,7 +4267,7 @@ def stack(seq, axis=0):
 
     if not seq:
         raise ValueError("Need array(s) to stack")
-    if not all(x.shape == seq[0].shape for x in seq):
+    if not allow_unknown_chunksizes and not all(x.shape == seq[0].shape for x in seq):
         idx = np.where(np.asanyarray([x.shape for x in seq]) != seq[0].shape)[0]
         raise ValueError(
             "Stacked arrays must have the same shape. "


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -362,6 +362,43 @@ def test_stack_rechunk():
     assert_eq(z, np.stack([x.compute(), y.compute()], axis=0))
 
 
+def test_stack_unknown_chunksizes():
+    dd = pytest.importorskip("dask.dataframe")
+    pd = pytest.importorskip("pandas")
+
+    a_df = pd.DataFrame({"x": np.arange(12)})
+    b_df = pd.DataFrame({"y": np.arange(12) * 10})
+
+    a_ddf = dd.from_pandas(a_df, sort=False, npartitions=3)
+    b_ddf = dd.from_pandas(b_df, sort=False, npartitions=3)
+
+    a_x = a_ddf.values
+    b_x = b_ddf.values
+
+    assert np.isnan(a_x.shape[0])
+    assert np.isnan(b_x.shape[0])
+
+    with pytest.raises(ValueError) as exc_info:
+        da.stack([a_x, b_x], axis=0)
+
+    assert "shape" in str(exc_info.value)
+    assert "nan" in str(exc_info.value)
+
+    c_x = da.stack([a_x, b_x], axis=0, allow_unknown_chunksizes=True)
+
+    assert_eq(c_x, np.stack([a_df.values, b_df.values], axis=0))
+
+    with pytest.raises(ValueError) as exc_info:
+        da.stack([a_x, b_x], axis=1)
+
+    assert "shape" in str(exc_info.value)
+    assert "nan" in str(exc_info.value)
+
+    c_x = da.stack([a_x, b_x], axis=1, allow_unknown_chunksizes=True)
+
+    assert_eq(c_x, np.stack([a_df.values, b_df.values], axis=1))
+
+
 def test_concatenate():
     a, b, c = [
         Array(