dask-contrib / dask-awkward

Native Dask collection for awkward arrays, and the library to use it.
https://dask-awkward.readthedocs.io
BSD 3-Clause "New" or "Revised" License
61 stars 19 forks source link

Implement `dak.backend` as an overload of `ak.backend` #497

Closed jpivarski closed 7 months ago

jpivarski commented 7 months ago

The implementation would be really easy:

# in dask-awkward

def backend(*arrays):
    ak_arrays = [x._meta if isinstance(x, dak.Array) else x for x in arrays]
    return ak.backend(*ak_arrays)

but it would fix scikit-hep/awkward#3077. See the discussion there, particularly the end of https://github.com/scikit-hep/awkward/issues/3077#issuecomment-2049843773.

martindurant commented 7 months ago

Do you intend to make this a PR? I don't see any problem with it.

jpivarski commented 7 months ago

In which module should I put it?

martindurant commented 7 months ago

Hmm... I have no better idea than utils, which catches everything else.

jpivarski commented 7 months ago

utils.py doesn't have user-facing functions. The describe.py/reducers.py/structure.py is reminiscent of the way Awkward 1 organized its functions (same module names), before we switched to having one file per ak.* function, which I like much better, even if some of the files are really small.

In the describe.py/reducers.py/structure.py taxonomy, I guess backend would go into describe.py, so I'll put it there.

I'll open PR.

agoose77 commented 7 months ago

I'm about to head out the door, but this would always return typetracer for meta arrays, which is probably not what we want.

jpivarski commented 7 months ago

In fact, I wrote the test as assert dak.backend(daa) == "typetracer". Coffea code is using ak.backend(something) to determine if something is a typetracer-backed array. Right now, that's equivalent to being a dask-awkward array and they could have checked isinstance(something, (dak.Array, dak.Record, dak.Scalar)) instead, but checking the backend may be preferable.