Slicing Dask Arrays with Index Arrays

dask_index_arrs.py

import itertools
import numbers
import operator

import numpy as np
import dask.array as da

try:
    from functools import reduce
except ImportError:
    pass

def slice_index_array(a, idx):
    # Validate this is not a single array.
    if not isinstance(idx, tuple):
        raise TypeError("`idx` must be a `tuple` of 1-D arrays")

    # Cast as Dask Arrays.
    a = da.asarray(a)
    idx = tuple(da.asarray(e) for e in idx)

    # Validate the dimensionality and shape of `idx` and `a`.
    if not all([issubclass(e.dtype.type, numbers.Integral) for e in idx]):
        raise TypeError("`idx` must be of integral type")
    if not all(e.ndim == 1 for e in idx):
        raise ValueError("`idx` must contain only 1-D arrays")
    if not np.allclose(idx[0].shape[0], [e.shape[0] for e in idx], equal_nan=True):
        raise ValueError("`idx` must contain the same number of points")
    if a.ndim != len(idx):
        raise ValueError("`idx` and `a` must share the same dimensionality")

    # Construct indices for `a`.
    a_indices = tuple(
        da.arange(s, chunks=c) for s, c in zip(a.shape, a.chunks)
    )

    # Create a series of masks indicating where each coordinate is.
    # Then apply those masks in each chunk to get the values needed.
    idx_vals = da.atop(
        lambda arr, *args: np.fromiter(
            map(
                lambda e: arr[e].item(),
                reduce(
                    np.logical_and,
                    map(
                        np.equal.outer,
                        args[::2],
                        map(
                            lambda i, e: e[i * (None,) + (slice(None),) + (len(args) // 2 - i - 1) * (None,)],
                            range(len(args) // 2),
                            args[1::2]
                        )
                    )
                )
            ),
            dtype=a.dtype
        ),
        (0,),
        a, tuple(range(1, 1 + a.ndim)),
        *itertools.chain(*[
            (e_0, (0,), e_i, (i,))
            for i, (e_0, e_i) in enumerate(zip(idx, a_indices), start=1)
        ]),
        dtype=bool,
        concatenate=True
    )

    return idx_vals

The memory usage is quite substantial unless the chunks are small.

an_arr = da.ones((1024,1024))
slice_index_array(an_arr, an_arr.nonzero()).min().compute()

/home/roat/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/array/blockwise.py:241: UserWarning: The da.atop function has moved to da.blockwise
  warnings.warn("The da.atop function has moved to da.blockwise")
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<timed exec> in <module>

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/base.py in compute(self, **kwargs)
    163         dask.base.compute
    164         """
--> 165         (result,) = compute(self, traverse=False, **kwargs)
    166         return result
    167 

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/base.py in compute(*args, **kwargs)
    434     keys = [x.__dask_keys__() for x in collections]
    435     postcomputes = [x.__dask_postcompute__() for x in collections]
--> 436     results = schedule(dsk, keys, **kwargs)
    437     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    438 

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     79         get_id=_thread_get_id,
     80         pack_exception=pack_exception,
---> 81         **kwargs
     82     )
     83 

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/local.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, **kwargs)
    484                         _execute_task(task, data)  # Re-execute locally
    485                     else:
--> 486                         raise_exception(exc, tb)
    487                 res, worker_id = loads(res_info)
    488                 state["cache"][key] = res

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/local.py in reraise(exc, tb)
    314     if exc.__traceback__ is not tb:
    315         raise exc.with_traceback(tb)
--> 316     raise exc
    317 
    318 

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    220     try:
    221         task, data = loads(task_info)
--> 222         result = _execute_task(task, data)
    223         id = get_id()
    224         result = dumps((result, id))

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in <listcomp>(.0)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in <listcomp>(.0)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    117         func, args = arg[0], arg[1:]
    118         args2 = [_execute_task(a, cache) for a in args]
--> 119         return func(*args2)
    120     elif not ishashable(arg):
    121         return arg

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/optimization.py in __call__(self, *args)
    980         if not len(args) == len(self.inkeys):
    981             raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
--> 982         return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
    983 
    984     def __reduce__(self):

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in get(dsk, out, cache)
    147     for key in toposort(dsk):
    148         task = dsk[key]
--> 149         result = _execute_task(task, cache)
    150         cache[key] = result
    151     result = _execute_task(out, cache)

~/.local/share/virtualenvs/notebooks-xbxHQDT_/lib/python3.6/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    117         func, args = arg[0], arg[1:]
    118         args2 = [_execute_task(a, cache) for a in args]
--> 119         return func(*args2)
    120     elif not ishashable(arg):
    121         return arg

<ipython-input-2-8e9343db9780> in <lambda>(arr, *args)
     49                             lambda i, e: e[i * (None,) + (slice(None),) + (len(args) // 2 - i - 1) * (None,)],
     50                             range(len(args) // 2),
---> 51                             args[1::2]
     52                         )
     53                     )

MemoryError: Unable to allocate array with shape (1048576, 1024, 1024) and data type bool

My guess is there have already been changes in Dask that include support for this in a reasonable way. Would look there instead of this Gist.