Are these short for other and array or something else? If so I recommend adding the extra two characters here. I think that the cost of extra typing here does not outweigh the loss of information.

Yes. Was trying to keep line length short later on. That said, no strong attachment to this or another variable name.

Can I request a docstring or comment here?

What else would you like to see here (meaning content-wise)?

Tried to add some prose here in the hopes of clarifying things. Please let me know what you think.

OK, so we do this in two passes, first with other_indices and then with the array_indices?

Right. Sorry should have been explicit about this before.

We do the first pass to deal with everything that is static (e.g. known single indexes, slices, lists, NumPy arrays). Then we do a second pass with everything that is dynamic (i.e. Dask Arrays of known and unknown length).

Assumed there are use cases (xarray?) where being able to explicitly select portions of the graph to keep are important. If not, we could adjust our atop call to handle these as well.

Yes, there are probably use cases where the ability to identify the right chunks ahead of time would make a big difference for efficiency.

Yep, that would not surprise me. :)

Might do some light refactoring in here to make the static and dynamic passes more obvious.

Added a comment before this to make the first pass obvious. Also explained in the docstring now.

I definitely don't understand what's going on here, but that's also pretty typical for non-trivial uses of atop in the codebase, much of which is my fault. This is a good learning experience for me to be on the other end of this. I was trying to think of ways to help with this situation. Perhaps comments for simple cases as in this snippet in unify_chunks?

    arginds = [(asarray(a) if ind is not None else a, ind)
               for a, ind in partition(2, args)]  # [x, ij, y, jk]
    args = list(concat(arginds))  # [(x, ij), (y, jk)]

Yeah, it's a bit ugly. Have been trying to thing of ways to make it clearer. Open to adding comments if that would help.

Basically we pass in our array (probably obvious), arange along each dimension in x, and then either a 1-D Dask Array that selects points or slice(None) (if there is no selection or if it was handled in the first pass). Also we interleaved arange and the 1-D selection, which makes this ugly here, but easier to work with in _vindex_dask_arrays_chunk. Expect these two last parts makes things confusing.

I'm a little concerned about the algorithm here. If I understand it correctly, atop here is mapping over all possible pairs of chunks on the indexed and indexing array. This will work okay for small numbers of chunks, but could quickly become intractably slow for larger numbers of chunks.

That said, I can't think of any obvious ways to improve this within dask's data model. You really need to compute some of the indices to know which source chunks they index, but dask (at least the non-distributed version) does not have support for that sort of dynamic graph construction.

It may be that it's always better to collapse a chunked index dask-array into a single chunk. This would increase our work-per-chunk but decrease the scheduling overhead. In common the case of single-dimensional indexers I suspect that this is almost always a good idea.

Tried to add some information about what happens here in a docstring for this function and in a comment above. Suggestions welcome.

Certainly aware that there could be performance problems with large numbers of chunks and would be happy to see them improved. Though it is not unique to this case. That said, it feels like this issue is outside the scope of this PR at present.

It would be reassuring to see more cases where the values of the array are

1. out of order
2. chunked differently.
3. negative

Sorry, could you please explain what you mean by "out of order"?

Have added a couple explicit tests like this. Interestingly the nonzero case also captured these issues too. Though explicit test cases seem preferable.

It would be helpful if this function had a docstring with example input/output.

Added a docstring. Missed the example input/output part when writing the docs/comments. Sorry. Happy to add that as well in a follow-up.

In the interim it would be good to know if the docs/comments added help clarify this function and its arguments. Also would be good to hear how to better present that information.

I think np.in1d could be much faster here (O(n log n) vs O(n^2)). This requires allocating an array for the full outer product of the indices.

Sure, that makes sense.

When I previously was playing with this in the gist, it used only Dask Array operations without atop. So when pushing it all into atop, expect that there remains some things from the old approach that are suboptimal.

Happy to fix any that come up.

Looking back at it, I think the current code was chosen to preserve ordering. For example...

>>> import numpy as np
>>> e_0 = np.array([23, 21, 24])
>>> e_i = np.arange(21, 25)
>>> np.equal.outer(e_0, e_i).nonzero()[1]
array([2, 0, 3])
>>> np.in1d(e_i, e_0).nonzero()[0]
array([0, 2, 3])

Not sure if there is an obvious way to get this from in1d or a similarly performant function, but would be happy to try. Expect this would be a good point to improve.

If we're OK with a pandas dependency (which we would probably rather avoid in dask.array), this version does the lookup in linear time with a hash-table:

In [14]: import pandas as pd

In [15]: pd.Index(e_i).get_indexer(e_0)
Out[15]: array([2, 0, 3])

It looks like this could also be done with np.sort()/np.searchsorted(), which means O(n log n + k log n) time and O(n+k) space:

In [24]: np.searchsorted(np.sort(e_i), e_0)
Out[24]: array([2, 0, 3])

(you should add tests/think carefully about how to properly handle repeated indices)

Thanks for the tip. Will give that a look.

So we do have tests for repeated indices in this PR. Please let me know if we are missing something in particular there.

Should add we don't have tests for non-sequential indices, which we need. Happy to add those.

FWICT searchsorted works well. We can skip the sort as e_i is always sorted. Updated the PR accordingly. Thanks again for the suggestion.

I'm a little concerned about the algorithm here. If I understand it correctly, atop here is mapping over all possible pairs of chunks on the indexed and indexing array. This will work okay for small numbers of chunks, but could quickly become intractably slow for larger numbers of chunks.

That said, I can't think of any obvious ways to improve this within dask's data model. You really need to compute some of the indices to know which source chunks they index, but dask (at least the non-distributed version) does not have support for that sort of dynamic graph construction.

Consider writing the more succinct: (e_0, (0,) if isinstance(e_0, Array) else None, e_i, (i,)))

Yeah I went back and forth on these. The current approach seems nice as you can easily contrast what you will get as a tuple.

I don't understand the variable names e_0 and e_i. I think using variable names like ref_ind and new_ind would make all this code easier to understand.

I was puzzled by this at first, and then I realized that this function only supported 1d dask arrays in the indexer currently.

Consider extending this to support arbitrary dimensionality inputs with broadcasting, or at least raise an informative error message.

I think broadcasting would actually be pretty easy, with broadcast_arrays(), out_axes in atop like what I did in isin and ravel/reshape around searchsorted for handling the chunk indices.

So I tried to extend N-D arrays, but ran into some issues with index arrays that have unknown shape.

Something to watch out for: this operation is lazy, but isn't really distributed friendly. If I understand correctly, concatenate=True means that the indexing arrays essentially get concatenated together across chunks. So things will break if your indexing arrays are too big to fit into memory.

This could conceivably be avoided, but it would require some tricky logic to merge across chunks.

Alright, will give this some thought. Thanks.

AFAICT concatenate=True is just a convenience option that gives one a single NumPy array instead of some list of NumPy arrays that atop would otherwise provide to the kernel function. ( #1609 ) The amount of data given in each case (concatenate=False/concatenate=True) is the same. Historically it sounds like that wasn't always true.

There may be some argument to not performing the concatenation as the memory usage doubles in the process (at least temporarily). That said, one normally chooses chunk sizes knowing that copies may be made (maybe even a few). So this doesn't seem unreasonable to me.

concatenate=True hasn't changed, and indeed atop has always processed multiple inputs at once when output indices are not available. But loading all input chunks corresponding to dropped outputs can still be problematic.

Looking at this again, I think you actually consolidate all indexed chunks along these dimensions, e.g., with array.vindex[key], then array could potentially get computed. This would be problematic in the case where array corresponds to a large tiled image, e.g., a 100k x 100k matrix with chunks of size 2000 x 2000.

That said, I don't think dask (without distributed) is capable in handling this in any sane way, since it really requires dynamic task creation. Otherwise you're stuck with all combinations of indexed and indexer chunks, which will quickly get overwhelming for moderately sized numbers of chunks.