Interoperability of awkward and dask_awkward

[jrueb]

I think it would be very convenient if one could use awkward functions on dask_awkward arrays straight away.
For example, I think it should be very nice to be able to do

ak.unflatten(some_dask_awkward_array, num)

I would like code in general to work regardless of whether somebody decides to use dask or standard awkward (without dask).
I'm saying this because this functionality is already there in case of numpy. For example I can do

np.reshape(some_awkward_array, newshape)

and also

np.reshape(some_dask_array, newshape)

Of course, reshape and unflatten just serve as examples here and this argument is meant to apply to all common functions awkward and dask_awkward have in common.

Is this something that could be considered to be added?

[jpivarski]

I think this is reasonable: to be able to use the same script on eager ak.Array and lazy dak.Array without having to change all of the ak. to dak..

One way to do it would be to check for dak.Array on the first line of the implementation of every function in src/awkward/operations.py. However, all of these functions already start by going through an OperationErrorContext, so the check could be implemented once with minimal code duplication by adding the

NOT_DASK = object()

for arg in arguments.values():
    if type(arg).__module__.split(".")[0] == "dask_awkward":
        import dask_awkward
        import inspect
        function_name = inspect.stack()[1].function
        dask_equivalent = getattr(dask_awkward, function_name, None)
        if dask_equivalent is None:
            raise NotImplementedError(f"dask-awkward function dak.{} does not exist")
        else:
            return dask_equivalent(**arguments)
    else:
        return NOT_DASK

in the OperationErrorContext.__init__.

Hmmm. The calling function would have to deal with continuing to normal operation on NOT_DASK, which means that some code would need to be added to every function in src/awkward/operations.py, after all. Also, this is abusing the notion of an "error context" because what we're doing here—looking up an equivalent dask-awkward function for each Awkward function we're implementing—has nothing to do with error messages. One thing that the OperationErrorContext does for us is that it puts all of the arguments into a dict, arguments, which we'll need to pass to dask_equivalent, so it's tempting to reuse the OperationErrorContext for that reason. Rewriting every canonicalized list of arguments would turn slightly too much boilerplate into way too much boilerplate.

So some thought is needed to do this nicely, but I agree that it should be done.

[agoose77]

@jpivarski there's also a temptation to remove the need for explicitly enumerating those arguments by introducing a decorator instead. This would benefit from the same approach I think.