Repository: dask/dask

Issue:
da.map_blocks with drop_axis stopped working in dask 1.1.4 
Hi!

First and foremost a big Thank You for the amazing work that you're doing with dask! It's a fantastic piece of software! 

We're using dask as parallelization engine for a large-scale data analysis package (still in early development). Our computational routines work on independent chunks of one big source data-set and then write to a single file (accessible to all workers). We used `map_blocks` for this in dask 1.0.0 and things worked great but unfortunately the same code does not work any more in dask 1.1.4. Here's a MCVE:

```python
import dask
import dask.array as da
import numpy as np

# in reality used to write `chk` to a (pre-formatted) file
def da_dummy(chk, block_info=None):
    idx = block_info[0]["chunk-location"][-1]
    #
    # use `idx` to store `chk` at correct location
    # 
    return idx

np_a = np.ones((10,50*20))
a = da.from_array(np_a, chunks=(10,20))

# works in dask 1.0.0, raises an IndexError in dask 1.1.4
result = a.map_blocks(da_dummy, dtype="int", chunks=(1,), drop_axis=[0])
res = result.compute()
```
In dask 1.0.0 this snippet produces an array `res` that looks like this:
```python
array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
       17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
       34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49])
```
In dask 1.1.4 the same code triggers an `IndexError: index 2 is out of bounds for axis 0 with size 2`, here the full traceback:

<details>

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/Documents/job/dask_fail.py in <module>
     15 
     16 # works in dask 1.0.0, raises an IndexError in dask 1.1.4
---> 17 result = a.map_blocks(da_dummy, dtype="int", chunks=(1,), drop_axis=[0])
     18 res = result.compute()
     19 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in map_blocks(self, func, *args, **kwargs)
   1569     @wraps(map_blocks)
   1570     def map_blocks(self, func, *args, **kwargs):
-> 1571         return map_blocks(func, self, *args, **kwargs)
   1572 
   1573     def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in map_blocks(func, *args, **kwargs)
    534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
--> 536                     for i in shapes}
    537 
    538             v = copy.copy(v)  # Need to copy and unpack subgraph callable

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <dictcomp>(.0)
    534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
--> 536                     for i in shapes}
    537 
    538             v = copy.copy(v)  # Need to copy and unpack subgraph callable

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <listcomp>(.0)
    532                         'num-chunks': num_chunks[i],
    533                         'array-location': [(starts[i][ij][j], starts[i][ij][j + 1])
--> 534                                            for ij, j in enumerate(old_k[1:])],
    535                         'chunk-location': old_k[1:]}
    536                     for i in shapes}

IndexError: index 2 is out of bounds for axis 0 with size 2
```

</details>

A (desperate) attempt at somehow working around this was to not use `drop_axis` and instead leave the chunk-size as is and return a bogus tuple, but no luck either:

```python

import dask
import dask.array as da
import numpy as np

np_a = np.ones((10,50*20))
a = da.from_array(np_a, chunks=(10,20))

# desperate workaround for dask 1.1.4 - still not working
def da_dummy2(chk, block_info=None):
    print(block_info) # show what's going on in here...
    idx = block_info[0]["chunk-location"][-1]
    return (idx, 1)
result = a.map_blocks(da_dummy2, dtype="int", chunks=(1, 1))
res = result.compute()
```

This gives an `IndexError: tuple index out of range`, full traceback here:

<details>

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
~/Documents/job/dask_fail.py in <module>
     20     return (idx, 1)
     21 result = a.map_blocks(da_dummy2, dtype="int", chunks=(1, 1))
---> 22 res = result.compute()
     23 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in compute(self, **kwargs)
    154         dask.base.compute
    155         """
--> 156         (result,) = compute(self, traverse=False, **kwargs)
    157         return result
    158 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in compute(*args, **kwargs)
    397     postcomputes = [x.__dask_postcompute__() for x in collections]
    398     results = schedule(dsk, keys, **kwargs)
--> 399     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    400 
    401 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/base.py in <listcomp>(.0)
    397     postcomputes = [x.__dask_postcompute__() for x in collections]
    398     results = schedule(dsk, keys, **kwargs)
--> 399     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    400 
    401 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in finalize(results)
    777     while isinstance(results2, (tuple, list)):
    778         if len(results2) > 1:
--> 779             return concatenate3(results)
    780         else:
    781             results2 = results2[0]

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in concatenate3(arrays)
   3495     if not ndim:
   3496         return arrays
-> 3497     chunks = chunks_from_arrays(arrays)
   3498     shape = tuple(map(sum, chunks))
   3499 

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in chunks_from_arrays(arrays)
   3325 
   3326     while isinstance(arrays, (list, tuple)):
-> 3327         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   3328         arrays = arrays[0]
   3329         dim += 1

/usr/local/software/anaconda/envs/p3/lib/python3.6/site-packages/dask/array/core.py in <listcomp>(.0)
   3325 
   3326     while isinstance(arrays, (list, tuple)):
-> 3327         result.append(tuple([shape(deepfirst(a))[dim] for a in arrays]))
   3328         arrays = arrays[0]
   3329         dim += 1

IndexError: tuple index out of range
```

</details>

Any help with this is greatly appreciated. Thank you! 


Discussion:
Thanks for the report, confirm locally.

These are a tad tricky to debug. I didn't come up with anything concrete in 15 minutes or so. Will give it another shot later.
Hi Tom! 

I was wondering if you already had a chance to further investigate the root cause of the problem? Thank you!
@jcrist @jakirkham if either of you have time can I ask you to look into
this?

On Wed, Apr 3, 2019 at 5:20 AM Stefan Fuertinger <notifications@github.com>
wrote:

> Hi Tom!
>
> I was wondering if you already had a chance to further investigate the
> root cause of the problem? Thank you!
>
> —
> 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/4584#issuecomment-479429833>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszIXmiglgSVy4GLWIN7V0tj_2tXsEks5vdIB3gaJpZM4brWsE>
> .
>

Sure, the next step would be to run `git bisect` and see what commit that identifies. If you are able to try that @pantaray, that would be very helpful. If not, I can give it a go using what you have provided.
I took a brief look at this.  This is one of the unfortuante collisions in complex logic of all of the keyword arguments within `map_blocks`.  In particular here we're using both `drop_axis` and `block_info`.  It would be good to fix, but this is exotic enough and that code is hairy enough that, for me at least, it's something of a low priority.  

@pantaray I have a couple of suggestions:
1.  
Whoops, misfire.  A couple suggestions:

1.  Find a way to not use map_blocks with both of these keyword arguments at the same time.  Maybe you return an output array that is of the same number of dimensions as your input?
2.  Dive in and fix this in the codebase.  

    The relevant code (I think) is here

https://github.com/dask/dask/blob/d8a4f3be9a6ec685261ed043a45ecc019055ec0d/dask/array/core.py#L588-L589

To be clear, I'm definitely acknowledging that this is a bug and that it should be fixed.  I'm also being somewhat honest here that given the complexity and edge-case-ness of this it's unlikely to be at the top of any of the maintainer's queues any time soon.  Honestly we probably shouldn't be supporting all of the keyword arguments in `map_blocks` and should instead be diverting people towards suggestions like 1 above, to find some other, perhaps less magical way.

Hope this helps somewhat.  Sorry for the lack of support.
@mrocklin Thank you for the explanation! 

Just to clarify: our base problem is that we have a dask array holding about 100GB of data scattered across multiple workers in distributed memory. How can we save the array to disk?

Hence, the only reason we we're using this particular `map_blocks`-construct is to collect and save results computed by multiple workers to a single file. You mentioning the "edge-case-ness" of this approach makes me wonder what are the best practices for achieving this? We're definitely not hell-bent on using `map_blocks` for this and are absolutely open to employing third-party storage formats like HDF5 containers to save our results. However, both `store` and `to_hdf` don't work with multiprocessing schedulers (as discussed in #2488 or #3074). What would you recommend to use instead?

Thank you again for your time and effort!
[Zarr]( https://zarr.readthedocs.io ) works reasonably well for parallel writes. I'm not unbiased though as I do work on it some. That said, I came to it as a Dask user interested in addressing this same problem.
Sorry for the delay, we've spent the last few days testing various options for tackling this. 

@jakirkham: thank you for the suggestion! 

We ran a bunch of performance tests on our cluster using 


1. **Sequential HDF5 writing:** workers pump their chunks of the distributed array `a` into a single container via `map_blocks` followed by `persist()` (not `compute()` to avoid the above errors). Writing is protected (and hence serialized) using a distributed lock. 
2. **Single worker HDF5 writing:** a single worker fetches pieces of the array using `compute()` and stores these pieces in a HDF5 dataset
3. **zarr:** create a `zarr` array `z` and use `da.to_zarr(a, z)` to store the distributed array `a`
4. **Virtual HDF5 datasets:** every worker saves its chunk(s) of `a` into separate HDF5 containers via `map_blocks`/`persist()`. The generated HDF files are subsequently used to create a single virtual dataset in our target HDF5 container. 

In our tests `zarr` arrays and generating multiple HDF5 containers showed the best performance (by a mile) with `zarr` being by far the most convenient solution from a coding point of view (a two liner). We ended up going with virtual HDF5 datasets since most of our source files are HDF5 containers already. I'm happy to provide a MCVE in case anyone's interested. 

Thank you again for your help!
Certainly virtual HDF5 datasets are an interesting solution to explore here. I wonder if we should include an example in the docs. Thoughts from others here?
Have run `git bisect` on this and narrowed the issue down to commit ( https://github.com/dask/dask/commit/9f870d17e6fcae00ab010c86a5b182b2ecf214b8 ). Right before that commit everything works correctly and with that commit this error emerges.
Should add that I don't see this error on `master`. However it no longer gives the original result (on 1.0.0) either. So there is still a bug (albeit a different one). Using `git bisect` narrowed down the behavior change to commit ( https://github.com/dask/dask/commit/a67550745fea122e710e398a34fd7651452c885e ). For clarity, this is what `res` looks like with this change.

```python
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0])
```
Yep, looks like this also came up in issue ( https://github.com/dask/dask/issues/4712 ). Where the last change was added to address this issue. Let's see if @bmerry has thoughts on this particular case.
When working on #4712 I do recall looking at the code and thinking that it didn't look like it would handle `drop_axis`. If I have time I'll take a look later this week to see if I can see how to fix it. It does get fairly hairy once you're combining `new_axis`, `drop_axis` and `block_info`, so I don't want to promise anything.

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -536,7 +536,7 @@ def map_blocks(func, *args, **kwargs):
         chunks2 = []
         for i, (c, nb) in enumerate(zip(chunks, out.numblocks)):
             if isinstance(c, tuple):
-                # We only check cases where numblocks > 1. Becuase of
+                # We only check cases where numblocks > 1. Because of
                 # broadcasting, we can't (easily) validate the chunks
                 # when the number of blocks is 1.
                 # See https://github.com/dask/dask/issues/4299 for more.
@@ -555,16 +555,19 @@ def map_blocks(func, *args, **kwargs):
         num_chunks = {}
         shapes = {}
 
-        for i, arg in enumerate(args):
-            if isinstance(arg, Array):
-                starts[i] = [np.cumsum((0,) + c) for c in arg.chunks]
+        for i, (arg, in_ind) in enumerate(argpairs):
+            if in_ind is not None:
                 shapes[i] = arg.shape
-                num_chunks[i] = arg.numblocks
-        for k, v in kwargs.items():
-            if isinstance(v, Array):
-                starts[k] = [np.cumsum((0,) + c) for c in v.chunks]
-                shapes[k] = arg.shape
-                num_chunks[i] = arg.numblocks
+                if drop_axis:
+                    # We concatenate along dropped axes, so we need to treat them
+                    # as if there is only a single chunk.
+                    starts[i] = [(np.cumsum((0,) + arg.chunks[j])
+                                  if ind in out_ind else np.array([0, arg.shape[j]]))
+                                 for j, ind in enumerate(in_ind)]
+                    num_chunks[i] = tuple(len(s) - 1 for s in starts[i])
+                else:
+                    starts[i] = [np.cumsum((0,) + c) for c in arg.chunks]
+                    num_chunks[i] = arg.numblocks
         out_starts = [np.cumsum((0,) + c) for c in out.chunks]
 
         for k, v in dsk.items():
@@ -572,21 +575,16 @@ def map_blocks(func, *args, **kwargs):
             v = v[0]
             [(key, task)] = v.dsk.items()  # unpack subgraph callable
 
-            if new_axis:
-                # anything using the keys in dsk is incorrect, as the
-                # original array doesn't have values for `new_axis`.
-                old_k = tuple(x for i, x in enumerate(k[1:]) if i not in new_axis)
-            else:
-                old_k = k[1:]
-
+            # Get position of chunk, indexed by axis labels
+            location = {out_ind[i]: loc for i, loc in enumerate(k[1:])}
             info = {}
             for i, shape in shapes.items():
                 # Compute chunk key in the array, taking broadcasting into
                 # account. We don't directly know which dimensions are
                 # broadcast, but any dimension with only one chunk can be
                 # treated as broadcast.
-                arr_k = old_k[-len(shape):]
-                arr_k = tuple(j if num_chunks[i][ij] > 1 else 0 for ij, j in enumerate(arr_k))
+                arr_k = tuple(location.get(ind, 0) if num_chunks[i][j] > 1 else 0
+                              for j, ind in enumerate(argpairs[i][1]))
                 info[i] = {'shape': shape,
                            'num-chunks': num_chunks[i],
                            'array-location': [(starts[i][ij][j], starts[i][ij][j + 1])


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
@@ -1115,7 +1115,7 @@ def func(a, b, c, block_info=None):
     assert_eq(z, x + x + 1 + 100)
 
 
-def test_map_blocks_block_info_with_axis():
+def test_map_blocks_block_info_with_new_axis():
     # https://github.com/dask/dask/issues/4298
     values = da.from_array(np.array(['a', 'a', 'b', 'c']), 2)
 
@@ -1145,6 +1145,40 @@ def func(x, block_info=None):
     assert_eq(z, np.ones((4, 3), dtype='f8'))
 
 
+def test_map_blocks_block_info_with_drop_axis():
+    # https://github.com/dask/dask/issues/4584
+    values = da.from_array(np.array([[1, 2, 4],
+                                     [8, 16, 32],
+                                     [64, 128, 256],
+                                     [1024, 2048, 4096]], dtype='u4'), (2, 1))
+
+    def func(x, block_info=None):
+        assert set(block_info.keys()) == {0, None}
+        assert block_info[0]['shape'] == (4, 3)
+        # drop_axis concatenates along the dropped dimension, hence not (2, 3)
+        assert block_info[0]['num-chunks'] == (2, 1)
+        assert block_info[None]['shape'] == (4,)
+        assert block_info[None]['num-chunks'] == (2,)
+        assert block_info[None]['chunk-shape'] == (2,)
+        assert block_info[None]['dtype'] == np.dtype('u4')
+
+        assert block_info[0]['chunk-location'] in {(0, 0), (1, 0)}
+
+        if block_info[0]['chunk-location'] == (0, 0):
+            assert block_info[0]['array-location'] == [(0, 2), (0, 3)]
+            assert block_info[None]['chunk-location'] == (0,)
+            assert block_info[None]['array-location'] == [(0, 2)]
+        elif block_info[0]['chunk-location'] == (1, 0):
+            assert block_info[0]['array-location'] == [(2, 4), (0, 3)]
+            assert block_info[None]['chunk-location'] == (1,)
+            assert block_info[None]['array-location'] == [(2, 4)]
+
+        return np.sum(x, axis=1, dtype='u4')
+
+    z = values.map_blocks(func, drop_axis=1, dtype='u4')
+    assert_eq(z, np.array([7, 56, 448, 7168], dtype='u4'))
+
+
 def test_map_blocks_block_info_with_broadcast():
     expected0 = [
         {'shape': (3, 4), 'num-chunks': (1, 2), 'array-location': [(0, 3), (0, 2)], 'chunk-location': (0, 0)},