Repository: dask/dask
File: dask/array/core.py (Python)

Code before:
        # Determine the start and end of each unique key. We will loop over this
        flag = np.concatenate([[True], sorted_keys[1:] != sorted_keys[:-1], [True]])
        (key_bounds,) = flag.nonzero()

        name = "vindex-slice-" + token
        vindex_merge_name = "vindex-merge-" + token
        dsk = {}
        for okey in other_blocks:
            merge_inputs = defaultdict(list)
            merge_indexer = defaultdict(list)
            for i, (start, stop) in enumerate(
                zip(key_bounds[:-1], key_bounds[1:], strict=True)
            ):
                slicer = slice(start, stop)
                key = sorted_keys[start]
                outblock, *input_blocks = np.unravel_index(key, ravel_shape)
                inblock = [_[slicer] for _ in sorted_inblock_idxs]
                dsk[keyname(name, i, okey)] = (
                    _vindex_transpose,
                    (
                        _vindex_slice,
                        (x.name,) + interleave_none(okey, input_blocks),
                        interleave_none(full_slices, inblock),
                    ),
                    axis,
                )
                merge_inputs[outblock].append(keyname(name, i, okey))
                merge_indexer[outblock].append(sorted_outblock_idx[slicer])

            for i in merge_inputs.keys():
                dsk[keyname(vindex_merge_name, i, okey)] = (
                    _vindex_merge,
                    merge_indexer[i],
                    merge_inputs[i],
                )

        result_1d = Array(
            HighLevelGraph.from_collections(out_name, dsk, dependencies=[x]),
            out_name,
            chunks,
            x.dtype,
            meta=x._meta,
        )
        return result_1d.reshape(broadcast_shape + result_1d.shape[1:])

    # output has a zero dimension, just create a new zero-shape array with the
    # same dtype
    from dask.array.wrap import empty

    result_1d = empty(
        tuple(map(sum, chunks)), chunks=chunks, dtype=x.dtype, name=out_name

Review comment (suggestion):
perhaps we should just scatter out `sorted_inblock_idxs` as dask arrays and embed a reference to a dask array sliced by these slicers in the task instead of the full array? The graph is quite big for the example in https://github.com/dask/dask/issues/11018

Code after:
        ravel_shape = (n_chunks + 1, *_subset_to_indexed_axes(x.numblocks))
        keys = np.ravel_multi_index([outblocks, *block_idxs], ravel_shape)
        # by sorting the data that needs to be inserted in the graph here,
        # we can slice in the hot loop below instead of using fancy indexing which will
        # always copy inside the hot loop.
        sortidx = np.argsort(keys, axis=None)
        sorted_keys = keys.flat[sortidx]  # flattens
        sorted_inblock_idxs = [_.flat[sortidx] for _ in inblock_idxs]
        sorted_outblock_idx = outblock_idx.flat[sortidx]
        # Determine the start and end of each unique key. We will loop over this
        flag = np.concatenate([[True], sorted_keys[1:] != sorted_keys[:-1], [True]])
        (key_bounds,) = flag.nonzero()

        name = "vindex-slice-" + token
        vindex_merge_name = "vindex-merge-" + token
        dsk = {}
        for okey in other_blocks:
            merge_inputs = defaultdict(list)
            merge_indexer = defaultdict(list)
            for i, (start, stop) in enumerate(
                zip(key_bounds[:-1], key_bounds[1:], strict=True)
            ):
                slicer = slice(start, stop)
                key = sorted_keys[start]
                outblock, *input_blocks = np.unravel_index(key, ravel_shape)
                inblock = [_[slicer] for _ in sorted_inblock_idxs]
                dsk[keyname(name, i, okey)] = (
                    _vindex_transpose,
                    (
                        _vindex_slice,
                        (x.name,) + interleave_none(okey, input_blocks),
                        interleave_none(full_slices, inblock),
                    ),
                    axis,
                )
                merge_inputs[outblock].append(keyname(name, i, okey))
                merge_indexer[outblock].append(sorted_outblock_idx[slicer])

            for i in merge_inputs.keys():
                dsk[keyname(vindex_merge_name, i, okey)] = (
                    _vindex_merge,
                    merge_indexer[i],
                    merge_inputs[i],
                )

        result_1d = Array(
            HighLevelGraph.from_collections(out_name, dsk, dependencies=[x]),
            out_name,
            chunks,
            x.dtype,
            meta=x._meta,