feat: supply a pre-calculated base form to avoid file opening

[lgray]

Add known_base_form option (awkward.forms.Form | None) to uproot.dask so that opening root files at the client can be avoided for mature analyses where the base forms of all files are well known. This is pertinent for root tree with large numbers of keys like NanoAOD.

Right now it is mutually exclusive with open_files=True, since it brings no advantage in that case.

An initial implementation is included.

I tested the speedup with the following code:

import uproot
import json
import gzip
import awkward

import pyinstrument

if __name__ == "__main__":

    prof = pyinstrument.Profiler()
    prof.start()

    with gzip.open("nano_form.json.gz", "r") as in_form:
	known_form = awkward.forms.from_json(in_form.read())

    for _ in range(20):
	events = uproot.dask({"tests/samples/nano_dy.root": {"object_path": "Events", "steps": [[0,40]]}}, open_files=False, known_base_form=known_form)

    prof.stop()
    print(prof.output_text(unicode=True, color=True, show_all=True))

    prof = pyinstrument.Profiler()
    prof.start()

    for _ in range(20):
        events = uproot.dask({"tests/samples/nano_dy.root": {"object_path": "Events", "steps": [[0,40]]}}, open_files=False)

    prof.stop()
    print(prof.output_text(unicode=True, color=True, show_all=True))

Which on a arm64 macbook pro with SSD gives about a 6x speedup. The gain will be enormous when using xrootd.

[jpivarski]

I agree. At the last minute, I dithered on this, thinking that we could solve the underlying problem by optionally incorporating Numba into the TBranch and TLeaf-deserialization process, since these are the two classes that are especially numerous for TTree metadata. In the end, though, I saw that it would have to go through read_object_any, which builds a generic object graph by reference. It would be a bigger project than the timescale on which you need this.

I'll merge this now and release 5.2.1.