Repository: dask/dask

Issue:
topk vs argtopk Performance
I've run some tests comparing the performance of `topk` vs `argtopk` and it seems like there's a significant difference despite a similar underlying codebase.

Sample script -
```
a = da.random.normal(size=(100000,100),chunks=(1000,100))
b = da.random.normal(size=(100,100000),chunks=(100,10000))

c = a.dot(b)

#top = c.topk(5, axis=-1)
top = c.argtopk(5, axis=-1)

with ProgressBar():
    da.compute(top)
```

When the above is executed with `topk` it completes in 1min 40sec.  When it runs with argtopk, it gets to 4% in 3min 40sec.  I interrupted this test early because a similar experiment took nearly 2 hours to complete - if growth is anywhere near linear, this one would complete in over an hour.

Is this a known issue with argtopk or anything we can do to optimize execution?


Discussion:

Resolving patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -11,6 +11,16 @@
 from ..compatibility import getargspec
 from ..utils import ignoring
 
+try:
+    from numpy import broadcast_to
+except ImportError:  # pragma: no cover
+    broadcast_to = npcompat.broadcast_to
+
+try:
+    from numpy import take_along_axis
+except ImportError:  # pragma: no cover
+    take_along_axis = npcompat.take_along_axis
+
 
 def keepdims_wrapper(a_callable):
     """
@@ -169,54 +179,92 @@ def trim(x, axes=None):
     return x[tuple(slice(ax, -ax if ax else None) for ax in axes)]
 
 
-try:
-    from numpy import broadcast_to
-except ImportError:  # pragma: no cover
-    broadcast_to = npcompat.broadcast_to
-
-
 def topk(a, k, axis, keepdims):
-    """Kernel of topk and argtopk.
+    """ Chunk and combine function of topk
+
     Extract the k largest elements from a on the given axis.
     If k is negative, extract the -k smallest elements instead.
     Note that, unlike in the parent function, the returned elements
     are not sorted internally.
     """
+    assert keepdims is True
     axis = axis[0]
     if abs(k) >= a.shape[axis]:
         return a
+
     a = np.partition(a, -k, axis=axis)
-    # return a[-k:] if k>0 else a[:-k], on arbitrary axis
-    return a[[
-        (slice(-k, None) if k > 0 else slice(None, -k))
-        if i == axis else slice(None)
-        for i in range(a.ndim)
-    ]]
+    k_slice = slice(-k, None) if k > 0 else slice(-k)
+    return a[tuple(k_slice if i == axis else slice(None)
+                   for i in range(a.ndim))]
+
 
+def topk_aggregate(a, k, axis, keepdims):
+    """ Final aggregation function of topk
 
-def topk_postprocess(a, k, axis):
-    """Kernel of topk and argtopk.
-    Post-processes the output of topk, sorting the results internally.
+    Invoke topk one final time and then sort the results internally.
     """
+    assert keepdims is True
+    a = topk(a, k, axis, keepdims)
+    axis = axis[0]
     a = np.sort(a, axis=axis)
-    if k > 0:
-        # a = a[::-1] on arbitrary axis
-        a = a[[
-            slice(None, None, -1) if i == axis else slice(None)
-            for i in range(a.ndim)
-        ]]
-    return a
+    if k < 0:
+        return a
+    return a[tuple(slice(None, None, -1) if i == axis else slice(None)
+                   for i in range(a.ndim))]
 
 
 def argtopk_preprocess(a, idx):
-    """Kernel of argtopk.
-    Preprocess data, by putting it together with its indexes in a recarray
+    """ Preparatory step for argtopk
+
+    Put data together with its original indices in a tuple.
+    """
+    return a, idx
+
+
+def argtopk(a_plus_idx, k, axis, keepdims):
+    """ Chunk and combine function of argtopk
+
+    Extract the indices of the k largest elements from a on the given axis.
+    If k is negative, extract the indices of the -k smallest elements instead.
+    Note that, unlike in the parent function, the returned elements
+    are not sorted internally.
     """
-    # np.core.records.fromarrays won't work if a and idx don't have the same shape
-    res = np.recarray(a.shape, dtype=[('values', a.dtype), ('idx', idx.dtype)])
-    res.values = a
-    res.idx = idx
-    return res
+    assert keepdims is True
+    axis = axis[0]
+
+    if isinstance(a_plus_idx, list):
+        a = np.concatenate([ai for ai, _ in a_plus_idx], axis)
+        idx = np.concatenate([broadcast_to(idxi, ai.shape)
+                              for ai, idxi in a_plus_idx], axis)
+    else:
+        a, idx = a_plus_idx
+
+    if abs(k) >= a.shape[axis]:
+        return a_plus_idx
+
+    idx2 = np.argpartition(a, -k, axis=axis)
+    k_slice = slice(-k, None) if k > 0 else slice(-k)
+    idx2 = idx2[tuple(k_slice if i == axis else slice(None)
+                      for i in range(a.ndim))]
+    return take_along_axis(a, idx2, axis), take_along_axis(idx, idx2, axis)
+
+
+def argtopk_aggregate(a_plus_idx, k, axis, keepdims):
+    """ Final aggregation function of argtopk
+
+    Invoke argtopk one final time, sort the results internally, drop the data
+    and return the index only.
+    """
+    assert keepdims is True
+    a, idx = argtopk(a_plus_idx, k, axis, keepdims)
+    axis = axis[0]
+
+    idx2 = np.argsort(a, axis=axis)
+    idx = take_along_axis(idx, idx2, axis)
+    if k < 0:
+        return idx
+    return idx[tuple(slice(None, None, -1) if i == axis else slice(None)
+                     for i in range(idx.ndim))]
 
 
 def arange(start, stop, step, length, dtype):
diff --git a/dask/array/numpy_compat.py b/dask/array/numpy_compat.py
--- a/dask/array/numpy_compat.py
+++ b/dask/array/numpy_compat.py
@@ -125,19 +125,19 @@ def nancumsum(a, axis=None, dtype=None, out=None):
 
         Examples
         --------
-        >>> np.nancumsum(1) #doctest: +SKIP
+        >>> nancumsum(1)
         array([1])
-        >>> np.nancumsum([1]) #doctest: +SKIP
+        >>> nancumsum([1])
         array([1])
-        >>> np.nancumsum([1, np.nan]) #doctest: +SKIP
+        >>> nancumsum([1, np.nan])
         array([ 1.,  1.])
         >>> a = np.array([[1, 2], [3, np.nan]])
-        >>> np.nancumsum(a) #doctest: +SKIP
+        >>> nancumsum(a)
         array([ 1.,  3.,  6.,  6.])
-        >>> np.nancumsum(a, axis=0) #doctest: +SKIP
+        >>> nancumsum(a, axis=0)
         array([[ 1.,  2.],
                [ 4.,  2.]])
-        >>> np.nancumsum(a, axis=1) #doctest: +SKIP
+        >>> nancumsum(a, axis=1)
         array([[ 1.,  3.],
                [ 3.,  3.]])
 
@@ -186,19 +186,19 @@ def nancumprod(a, axis=None, dtype=None, out=None):
 
         Examples
         --------
-        >>> np.nancumprod(1) #doctest: +SKIP
+        >>> nancumprod(1)
         array([1])
-        >>> np.nancumprod([1]) #doctest: +SKIP
+        >>> nancumprod([1])
         array([1])
-        >>> np.nancumprod([1, np.nan]) #doctest: +SKIP
+        >>> nancumprod([1, np.nan])
         array([ 1.,  1.])
         >>> a = np.array([[1, 2], [3, np.nan]])
-        >>> np.nancumprod(a) #doctest: +SKIP
+        >>> nancumprod(a)
         array([ 1.,  2.,  6.,  6.])
-        >>> np.nancumprod(a, axis=0) #doctest: +SKIP
+        >>> nancumprod(a, axis=0)
         array([[ 1.,  2.],
                [ 3.,  2.]])
-        >>> np.nancumprod(a, axis=1) #doctest: +SKIP
+        >>> nancumprod(a, axis=1)
         array([[ 1.,  2.],
                [ 3.,  3.]])
 
@@ -207,6 +207,133 @@ def nancumprod(a, axis=None, dtype=None, out=None):
         return np.cumprod(a, axis=axis, dtype=dtype, out=out)
 
 
+if LooseVersion(np.__version__) < '1.15.0':
+    # These functions were added in numpy 1.15.0. For previous versions they
+    # are duplicated here
+
+    def _make_along_axis_idx(arr_shape, indices, axis):
+        # compute dimensions to iterate over
+        if not np.issubdtype(indices.dtype, np.integer):
+            raise IndexError('`indices` must be an integer array')
+        if len(arr_shape) != indices.ndim:
+            raise ValueError(
+                "`indices` and `arr` must have the same number of dimensions")
+        shape_ones = (1,) * indices.ndim
+        dest_dims = list(range(axis)) + [None] + list(
+            range(axis + 1, indices.ndim))
+
+        # build a fancy index, consisting of orthogonal aranges, with the
+        # requested index inserted at the right location
+        fancy_index = []
+        for dim, n in zip(dest_dims, arr_shape):
+            if dim is None:
+                fancy_index.append(indices)
+            else:
+                ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1:]
+                fancy_index.append(np.arange(n).reshape(ind_shape))
+
+        return tuple(fancy_index)
+
+    def take_along_axis(arr, indices, axis):
+        """
+        Take values from the input array by matching 1d index and data slices.
+        This iterates over matching 1d slices oriented along the specified axis in
+        the index and data arrays, and uses the former to look up values in the
+        latter. These slices can be different lengths.
+        Functions returning an index along an axis, like `argsort` and
+        `argpartition`, produce suitable indices for this function.
+        .. versionadded:: 1.15.0
+        Parameters
+        ----------
+        arr: ndarray (Ni..., M, Nk...)
+            Source array
+        indices: ndarray (Ni..., J, Nk...)
+            Indices to take along each 1d slice of `arr`. This must match the
+            dimension of arr, but dimensions Ni and Nj only need to broadcast
+            against `arr`.
+        axis: int
+            The axis to take 1d slices along. If axis is None, the input array is
+            treated as if it had first been flattened to 1d, for consistency with
+            `sort` and `argsort`.
+        Returns
+        -------
+        out: ndarray (Ni..., J, Nk...)
+            The indexed result.
+        Notes
+        -----
+        This is equivalent to (but faster than) the following use of `ndindex` and
+        `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices::
+            Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
+            J = indices.shape[axis]  # Need not equal M
+            out = np.empty(Nk + (J,) + Nk)
+            for ii in ndindex(Ni):
+                for kk in ndindex(Nk):
+                    a_1d       = a      [ii + s_[:,] + kk]
+                    indices_1d = indices[ii + s_[:,] + kk]
+                    out_1d     = out    [ii + s_[:,] + kk]
+                    for j in range(J):
+                        out_1d[j] = a_1d[indices_1d[j]]
+        Equivalently, eliminating the inner loop, the last two lines would be::
+                    out_1d[:] = a_1d[indices_1d]
+        See Also
+        --------
+        take : Take along an axis, using the same indices for every 1d slice
+        put_along_axis :
+            Put values into the destination array by matching 1d index and data slices
+        Examples
+        --------
+        For this sample array
+        >>> a = np.array([[10, 30, 20], [60, 40, 50]])
+
+        We can sort either by using sort directly, or argsort and this function
+        >>> np.sort(a, axis=1)
+        array([[10, 20, 30],
+               [40, 50, 60]])
+        >>> ai = np.argsort(a, axis=1); ai
+        array([[0, 2, 1],
+               [1, 2, 0]])
+        >>> take_along_axis(a, ai, axis=1)
+        array([[10, 20, 30],
+               [40, 50, 60]])
+
+        The same works for max and min, if you expand the dimensions:
+        >>> np.expand_dims(np.max(a, axis=1), axis=1)
+        array([[30],
+               [60]])
+        >>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1)
+        >>> ai
+        array([[1],
+               [0]])
+        >>> take_along_axis(a, ai, axis=1)
+        array([[30],
+               [60]])
+
+        If we want to get the max and min at the same time,
+        we can stack the indices first:
+        >>> ai_min = np.expand_dims(np.argmin(a, axis=1), axis=1)
+        >>> ai_max = np.expand_dims(np.argmax(a, axis=1), axis=1)
+        >>> ai = np.concatenate([ai_min, ai_max], axis=1)
+        >>> ai
+        array([[0, 1],
+               [1, 0]])
+        >>> take_along_axis(a, ai, axis=1)
+        array([[10, 30],
+               [40, 60]])
+        """
+        # normalize inputs
+        if axis is None:
+            arr = arr.flat
+            arr_shape = (len(arr),)  # flatiter has no .shape
+            axis = 0
+        else:
+            if axis < 0:
+                axis = arr.ndim + axis
+            arr_shape = arr.shape
+
+        # use the fancy index
+        return arr[_make_along_axis_idx(arr_shape, indices, axis)]
+
+
 def _make_sliced_dtype(dtype, index):
     if LooseVersion(np.__version__) == LooseVersion("1.14.0"):
         return _make_sliced_dtype_np_ge_14(dtype, index)
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -37,11 +37,91 @@ def divide(a, b, dtype=None):
     return f(a, b, dtype=dtype)
 
 
-def reduction(x, chunk, aggregate, axis=None, keepdims=None, dtype=None,
-              split_every=None, combine=None, name=None, out=None):
+def reduction(x, chunk, aggregate, axis=None, keepdims=False, dtype=None,
+              split_every=None, combine=None, name=None, out=None,
+              concatenate=True, output_size=1):
     """ General version of reductions
 
-    >>> reduction(my_array, np.sum, np.sum, axis=0, keepdims=False)  # doctest: +SKIP
+    Parameters
+    ----------
+    x: Array
+        Data being reduced along one or more axes
+    chunk: callable(x_chunk, axis, keepdims)
+        First function to be executed when resolving the dask graph.
+        This function is applied in parallel to all original chunks of x.
+        See below for function parameters.
+    combine: callable(x_chunk, axis, keepdims), optional
+        Function used for intermediate recursive aggregation (see
+        split_every below). If omitted, it defaults to aggregate.
+        If the reduction can be performed in less than 3 steps, it will not
+        be invoked at all.
+    aggregate: callable(x_chunk, axis, keepdims)
+        Last function to be executed when resolving the dask graph,
+        producing the final output. It is always invoked, even when the reduced
+        Array counts a single chunk along the reduced axes.
+    axis: int or sequence of ints, optional
+        Axis or axes to aggregate upon. If omitted, aggregate along all axes.
+    keepdims: boolean, optional
+        Whether the reduction function should preserve the reduced axes,
+        leaving them at size ``output_size``, or remove them.
+    dtype: np.dtype, optional
+        Force output dtype. Defaults to x.dtype if omitted.
+    split_every: int >= 2 or dict(axis: int), optional
+        Determines the depth of the recursive aggregation. If set to or more
+        than the number of input chunks, the aggregation will be performed in
+        two steps, one ``chunk`` function per input chunk and a single
+        ``aggregate`` function at the end. If set to less than that, an
+        intermediate ``combine`` function will be used, so that any one
+        ``combine`` or ``aggregate`` function has no more than ``split_every``
+        inputs. The depth of the aggregation graph will be
+        :math:`log_{split_every}(input chunks along reduced axes)`. Setting to
+        a low value can reduce cache size and network transfers, at the cost of
+        more CPU and a larger dask graph.
+
+        Omit to let dask heuristically decide a good default. A default can
+        also be set globally with the ``split_every`` key in
+        :mod:`dask.config`.
+    name: str, optional
+        Prefix of the keys of the intermediate and output nodes. If omitted it
+        defaults to the function names.
+    out: Array, optional
+        Another dask array whose contents will be replaced. Omit to create a
+        new one. Note that, unlike in numpy, this setting gives no performance
+        benefits whatsoever, but can still be useful  if one needs to preserve
+        the references to a previously existing Array.
+    concatenate: bool, optional
+        If True (the default), the outputs of the ``chunk``/``combine``
+        functions are concatenated into a single np.array before being passed
+        to the ``combine``/``aggregate`` functions. If False, the input of
+        ``combine`` and ``aggregate`` will be either a list of the raw outputs
+        of the previous step or a single output, and the function will have to
+        concatenate it itself. It can be useful to set this to False if the
+        chunk and/or combine steps do not produce np.arrays.
+    output_size: int >= 1, optional
+        Size of the output of the ``aggregate`` function along the reduced
+        axes. Ignored if keepdims is False.
+
+    Returns
+    -------
+    dask array
+
+    **Function Parameters**
+
+    x_chunk: numpy.ndarray
+        Individual input chunk. For ``chunk`` functions, it is one of the
+        original chunks of x. For ``combine`` and ``aggregate`` functions, it's
+        the concatenation of the outputs produced by the previous ``chunk`` or
+        ``combine`` functions. If concatenate=False, it's a list of the raw
+        outputs from the previous functions.
+    axis: tuple
+        Normalized list of axes to reduce upon, e.g. ``(0, )``
+        Scalar, negative, and None axes have been normalized away.
+        Note that some numpy reduction functions cannot reduce along multiple
+        axes at once and strictly require an int in input. Such functions have
+        to be wrapped to cope.
+    keepdims: bool
+        Whether the reduction function should preserve the reduced axes or
+        remove them.
     """
     if axis is None:
         axis = tuple(range(x.ndim))
@@ -60,17 +140,19 @@ def reduction(x, chunk, aggregate, axis=None, keepdims=None, dtype=None,
     inds = tuple(range(x.ndim))
     # The dtype of `tmp` doesn't actually matter, and may be incorrect.
     tmp = atop(chunk, inds, x, inds, axis=axis, keepdims=True, dtype=x.dtype)
-    tmp._chunks = tuple((1, ) * len(c) if i in axis else c for (i, c)
-                        in enumerate(tmp.chunks))
-
+    tmp._chunks = tuple((output_size, ) * len(c) if i in axis else c
+                        for i, c in enumerate(tmp.chunks))
     result = _tree_reduce(tmp, aggregate, axis, keepdims, dtype, split_every,
-                          combine, name=name)
+                          combine, name=name, concatenate=concatenate)
+    if keepdims and output_size != 1:
+        result._chunks = tuple((output_size, ) if i in axis else c
+                               for i, c in enumerate(tmp.chunks))
     return handle_out(out, result)
 
 
 def _tree_reduce(x, aggregate, axis, keepdims, dtype, split_every=None,
-                 combine=None, name=None):
-    """Perform the tree reduction step of a reduction.
+                 combine=None, name=None, concatenate=True):
+    """ Perform the tree reduction step of a reduction.
 
     Lower level, users should use ``reduction`` or ``arg_reduction`` directly.
     """
@@ -82,26 +164,28 @@ def _tree_reduce(x, aggregate, axis, keepdims, dtype, split_every=None,
         n = builtins.max(int(split_every ** (1 / (len(axis) or 1))), 2)
         split_every = dict.fromkeys(axis, n)
     else:
-        split_every = dict((k, v) for (k, v) in enumerate(x.numblocks) if k in axis)
+        raise ValueError("split_every must be a int or a dict")
 
     # Reduce across intermediates
     depth = 1
     for i, n in enumerate(x.numblocks):
         if i in split_every and split_every[i] != 1:
             depth = int(builtins.max(depth, ceil(log(n, split_every[i]))))
-    func = compose(partial(combine or aggregate, axis=axis, keepdims=True),
-                   partial(_concatenate2, axes=axis))
+    func = partial(combine or aggregate, axis=axis, keepdims=True)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=axis))
     for i in range(depth - 1):
         x = partial_reduce(func, x, split_every, True, dtype=dtype,
                            name=(name or funcname(combine or aggregate)) + '-partial')
-    func = compose(partial(aggregate, axis=axis, keepdims=keepdims),
-                   partial(_concatenate2, axes=axis))
+    func = partial(aggregate, axis=axis, keepdims=keepdims)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=axis))
     return partial_reduce(func, x, split_every, keepdims=keepdims, dtype=dtype,
                           name=(name or funcname(aggregate)) + '-aggregate')
 
 
 def partial_reduce(func, x, split_every, keepdims=False, dtype=None, name=None):
-    """Partial reduction across multiple axes.
+    """ Partial reduction across multiple axes.
 
     Parameters
     ----------
@@ -465,7 +549,7 @@ def vnorm(a, ord=None, axis=None, dtype=None, keepdims=False, split_every=None,
 
 
 def _arg_combine(data, axis, argfunc, keepdims=False):
-    """Merge intermediate results from ``arg_*`` functions"""
+    """ Merge intermediate results from ``arg_*`` functions"""
     axis = None if len(axis) == data.ndim or data.ndim == 1 else axis[0]
     vals = data['vals']
     arg = data['arg']
@@ -532,7 +616,7 @@ def nanarg_agg(func, argfunc, data, axis=None, **kwargs):
 
 
 def arg_reduction(x, chunk, combine, agg, axis=None, split_every=None, out=None):
-    """Generic function for argreduction.
+    """ Generic function for argreduction.
 
     Parameters
     ----------
@@ -583,7 +667,7 @@ def arg_reduction(x, chunk, combine, agg, axis=None, split_every=None, out=None)
 
 
 def make_arg_reduction(func, argfunc, is_nan_func=False):
-    """Create a argreduction callable.
+    """ Create an argreduction callable
 
     Parameters
     ----------
@@ -725,13 +809,30 @@ def validate_axis(ndim, axis):
 
 
 def topk(a, k, axis=-1, split_every=None):
-    """Extract the k largest elements from a on the given axis,
+    """ Extract the k largest elements from a on the given axis,
     and return them sorted from largest to smallest.
     If k is negative, extract the -k smallest elements instead,
     and return them sorted from smallest to largest.
 
-    This assumes that ``k`` is small.  All results will be returned in a single
-    chunk along the given axis.
+    This performs best when ``k`` is much smaller than the chunk size. All
+    results will be returned in a single chunk along the given axis.
+
+    Parameters
+    ----------
+    x: Array
+        Data being sorted
+    k: int
+    axis: int, optional
+    split_every: int >=2, optional
+        See :func:`reduce`. This parameter becomes very important when k is
+        on the same order of magnitude of the chunk size or more, as it
+        prevents getting the whole or a significant portion of the input array
+        in memory all at once, with a negative impact on network transfer
+        too when running on distributed.
+
+    Returns
+    -------
+    Selection of x with size abs(k) along the given axis.
 
     Examples
     --------
@@ -744,31 +845,48 @@ def topk(a, k, axis=-1, split_every=None):
     array([1, 3])
     """
     if isinstance(a, int) and isinstance(k, Array):
-        warnings.warn("DeprecationWarning: topk(k, a) has been replaced with topk(a, k)")
+        warnings.warn("DeprecationWarning: topk(k, a) has been replaced with "
+                      "topk(a, k)")
         a, k = k, a
 
     axis = validate_axis(a.ndim, axis)
 
-    kernel = partial(chunk.topk, k=k)
-    res = reduction(a, kernel, kernel, axis=axis, keepdims=True,
-                    dtype=a.dtype, split_every=split_every)
-    # reduction(keepdims=True) sets shape[axis] to 1. Fix it.
-    chunks = list(res.chunks)
-    chunks[axis] = (abs(k), )
-    res = Array(res.dask, res.name, chunks, res.dtype)
+    # chunk and combine steps of the reduction, which recursively invoke
+    # np.partition to pick the top/bottom k elements from the previous step.
+    # The selection is not sorted internally.
+    chunk_combine = partial(chunk.topk, k=k)
+    # aggregate step of the reduction. Internally invokes the chunk/combine
+    # function, then sorts the results internally.
+    aggregate = partial(chunk.topk_aggregate, k=k)
 
-    # Sort result internally
-    return res.map_blocks(chunk.topk_postprocess, k=k, axis=axis, dtype=a.dtype)
+    return reduction(
+        a, chunk=chunk_combine, combine=chunk_combine, aggregate=aggregate,
+        axis=axis, keepdims=True, dtype=a.dtype, split_every=split_every,
+        output_size=abs(k))
 
 
 def argtopk(a, k, axis=-1, split_every=None):
-    """Extract the indices of the k largest elements from a on the given axis,
-    and return them sorted from largest to smallest.
-    If k is negative, extract the indices of the -k smallest elements instead,
-    and return them sorted from smallest to largest.
+    """ Extract the indices of the k largest elements from a on the given axis,
+    and return them sorted from largest to smallest. If k is negative, extract
+    the indices of the -k smallest elements instead, and return them sorted
+    from smallest to largest.
+
+    This performs best when ``k`` is much smaller than the chunk size. All
+    results will be returned in a single chunk along the given axis.
+
+    Parameters
+    ----------
+    x: Array
+        Data being sorted
+    k: int
+    axis: int, optional
+    split_every: int >=2, optional
+        See :func:`topk`. The performance considerations for topk also apply
+        here.
 
-    This assumes that ``k`` is small.  All results will be returned in a single
-    chunk along the given axis.
+    Returns
+    -------
+    Selection of int64 indices of x with size abs(k) along the given axis.
 
     Examples
     --------
@@ -782,13 +900,24 @@ def argtopk(a, k, axis=-1, split_every=None):
     """
     axis = validate_axis(a.ndim, axis)
 
-    # Convert a to a recarray that contains its index
+    # Generate nodes where every chunk is a tuple of (a, original index of a)
     idx = arange(a.shape[axis], chunks=a.chunks[axis], dtype=np.int64)
-    idx = idx[tuple(slice(None) if i == axis else np.newaxis for i in range(a.ndim))]
-    a_rec = a.map_blocks(chunk.argtopk_preprocess, idx,
-                         dtype=[('a', a.dtype), ('idx', idx.dtype)])
-
-    res = topk(a_rec, k, axis=axis, split_every=split_every)
-
-    # Discard values
-    return res['idx']
+    idx = idx[tuple(slice(None) if i == axis else np.newaxis
+                    for i in range(a.ndim))]
+    a_plus_idx = a.map_blocks(chunk.argtopk_preprocess, idx,
+                              dtype=object)
+
+    # chunk and combine steps of the reduction. They acquire in input a tuple
+    # of (a, original indices of a) and return another tuple containing the top
+    # k elements of a and the matching original indices. The selection is not
+    # sorted internally, as in np.argpartition.
+    chunk_combine = partial(chunk.argtopk, k=k)
+    # aggregate step of the reduction. Internally invokes the chunk/combine
+    # function, then sorts the results internally, drops a and returns the
+    # index only.
+    aggregate = partial(chunk.argtopk_aggregate, k=k)
+
+    return reduction(
+        a_plus_idx, chunk=chunk_combine, combine=chunk_combine,
+        aggregate=aggregate, axis=axis, keepdims=True, dtype=np.int64,
+        split_every=split_every, concatenate=False, output_size=abs(k))


Test patch:
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -420,14 +420,14 @@ def test_array_cumreduction_out(func):
     (np.sort, da.topk),
     (np.argsort, da.argtopk),
 ])
-@pytest.mark.parametrize('split_every', [None, 2])
+@pytest.mark.parametrize('split_every', [None, 2, 4, 8])
 def test_topk_argtopk1(npfunc, daskfunc, split_every):
     # Test data
     k = 5
-    a = da.random.random(1000, chunks=250)
+    # Test at least 3 levels of aggregation when split_every=2
+    # to stress the different chunk, combine, aggregate kernels
+    a = da.random.random(800, chunks=100)
     b = da.random.random((10, 20, 30), chunks=(4, 8, 8))
-    c = da.from_array(np.array([(1, 'Hello'), (2, 'World')], dtype=[('foo', int), ('bar', '<U5')]),
-                      chunks=1)
 
     # 1-dimensional arrays
     # top 5 elements, sorted descending
@@ -459,12 +459,6 @@ def test_topk_argtopk1(npfunc, daskfunc, split_every):
     with pytest.raises(ValueError):
         daskfunc(b, -k, axis=3, split_every=split_every)
 
-    # structured arrays
-    assert_eq(npfunc(c, axis=0)[-1:][::-1],
-              daskfunc(c, 1, split_every=split_every))
-    assert_eq(npfunc(c, axis=0)[:1],
-              daskfunc(c, -1, split_every=split_every))
-
 
 def test_topk_argtopk2():
     a = da.random.random((10, 20, 30), chunks=(4, 8, 8))