feat: support explicit bulk preloading (uproot-only)

[pfackeldey]

This PR allows to bulk preload multiple branches with NanoEventsFactory.from_root. This is useful to group multiple reads into a single request. In the virtual mode the preloaded branches are still kept virtual (wrapped essentially by a lambda) to delay evaluating the DSL until the field is actually needed, see:

# explicitly preload "Jet_pt"
preload = lambda b: b.name == "Jet_pt"

events = NanoEventsFactory.from_root(
  {"nanoaod.root": "Events"}, 
  mode="virtual", 
  access_log=(access_log := []), 
  preload=preload,  # new arg!
).events()

# this is expensive as it wasn't preloaded, needs to be loaded on the spot
%timeit -n 1 -r 1 ak.materialize(events.Jet["eta"])
# 468 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

# this is much faster because it's preloaded, but still the DSL has to run 
%timeit -n 1 -r 1 ak.materialize(events.Jet["pt"])
# 2.19 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

# check what has been accessed
print(access_log)
# ['nJet', 'Jet_eta', 'Jet_pt']

# this is much faster because it is preloaded _and_ the DSL has already been run
%timeit -n 1 -r 1 ak.materialize(events.Jet["pt"])
# 610 μs ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In addition you can trigger bulk loading of everything (single read request for all columns) in the eager mode as follows:

preload = lambda _: True  # match anything

events = NanoEventsFactory.from_root(
  {"nanoaod.root": "Events"}, 
  mode="eager", 
  access_log=(access_log := []), 
  preload=preload,  # new arg!
).events()

Currently, coffea reads with a single request per column.

[ikrommyd]

I'd probably move this further up the list. Maybe after entry start and entry stop? We want this to be visible no?

[pfackeldey]

I didn't mean to mess with the order in case people pass positional args here - that would be a breaking change. Maybe coffea may want to force all of those options to be passed as kw-only at some point?

[ikrommyd]

Ohh yeah you're right about that. I guess I've never seen a person not pass kwargs that's why I said this. Indeed forcing kwargs would be best here from the coffea side.

[ikrommyd]

@lgray what do you think? Should we enforce kwargs in a separate PR? Should we just break the order here and not care? I've actually never seen anybody use positional args here because the first argument treepath is something that's never used because no one is actually passing in an uproot.reading.ReadOnlyDirectory to from_root.
`

[lgray]

Ah - sorry for missing this. I would take it to a separate PR. I hope no one is actually using those and positional arguments though.

[ikrommyd]

Same comment as above, I'd not put this in the end.

[ikrommyd]

I like this @pfackeldey and have only left minor comments. I also like the fact that it is a callable like filter_branch is.

Is it possible to test something like this? I'm not sure because it's still gonna be a VirtualArray buffer but it's callable is just getting an array from this preloaded array dict. I don't think we can see it somehow.

One request though, is it possible to fill out the access log during preloading? I think it's best to see what has been loaded (and that's also a unit test). If that involves moving the access_log around, that's fine.

[pfackeldey]

I like this @pfackeldey and have only left minor comments. I also like the fact that it is a callable like filter_branch is.

Is it possible to test something like this? I'm not sure because it's still gonna be a VirtualArray buffer but it's callable is just getting an array from this preloaded array dict. I don't think we can see it somehow.

One request though, is it possible to fill out the access log during preloading? I think it's best to see what has been loaded (and that's also a unit test). If that involves moving the access_log around, that's fine.

I didn't want to add them to the access_log, because it's an access-log and not a loaded-log. In my understanding, this should track what has been accessed through e.g. getattr or other ops. I would find it confusing if the access_log would contain information that was never needed, independently if it was loaded or not. Also, I think it's nice to be able to pass the access_log in a consecutive run to preload and thus not load data that wasn't used. Maybe I'm misinterpreting the idea of this log though?

[ikrommyd]

I think the access log is meant to keep track of which branches were read overall. The executors also report back the bytesread via the following:

                    if isinstance(file, uproot.ReadOnlyDirectory):
                        metrics["bytesread"] = file.file.source.num_requested_bytes

I may be miss-interpreting something though.

[lgray]

I'm going to enable auto-merge on this one and we can fix the args/kwargs things before the next release. I've lodged this to-do in an Issue.