Reading doubly-ragged data in coffea

[tomeichlersmith]

Hi coffea drinkers ☕ ! I really appreciate coffea and its developments, I have a question that may have an easy answer but I cannot find it in prior discussions or issues.

One of the branches of the TTree of events I'm reading has a doubly-ragged structure: each entry in an event is a variable-length array and each entry in that array has its own variable-length arrays as members. When I try to read one of these doubly-ragged member arrays within a coffea processor, I get an error about offsets not being defined.

Traceback (most recent call last):
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/coffea/processor/executor.py", line 1640, in _work_function
    out = processor_instance.process(events)
  File "unbiased-study.py", line 187, in process
    print(events['HcalSimHits_eat/HcalSimHits_eat.pdgCodeContribs_'])
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/highlevel.py", line 1283, in __str__
    return self._str()
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/highlevel.py", line 1304, in _str
    limit_value, self.layout, self._behavior
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/_util.py", line 1694, in minimally_touching_string
    lft = next(leftgen)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/_util.py", line 1682, in forever
    yield from iterable
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/_util.py", line 1576, in forward
    for token in forward(x[i], sp):
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/awkward/operations/convert.py", line 4713, in _form_to_layout
    container[key_format(form_key=fk, attribute="offsets", partition=partnum)]
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/coffea/nanoevents/mapping/base.py", line 97, in __getitem__
    getattr(transforms, tname)(stack)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/venv/lib64/python3.6/site-packages/coffea/nanoevents/transforms.py", line 28, in offsets
    stack.append(to_layout(stack.pop()).offsets)
AttributeError: 'awkward._ext.ListArray64' object has no attribute 'offsets'

The other branches in this TTree are readable like normal (i.e. I have been only looking at single-entry or singly-ragged arrays so far). I am using the BaseSchema and this occurs for both the FuturesExecutor and the IterativeExecutor.

I have confirmed that this isn't an issue with my file by using uproot to load the file into awkward arrays directly. Below is the output of pip freeze in my python venv for version reference. My main questions are:

• Can coffea's lazy-loading technique handle doubly-ragged arrays?
• If so, do I need to write my own schema?
• If so, are there simpler schema examples than the ones currently in coffea?

pip freeze output

awkward==1.9.0
awkward0==0.15.5
boost-histogram==1.1.0
cachetools==4.2.4
certifi==2023.5.7
charset-normalizer==2.0.12
click==8.0.4
cloudpickle==2.2.1
coffea==0.7.20
commonmark==0.9.1
correctionlib==2.2.2
cycler==0.11.0
dataclasses==0.8
fsspec==2022.1.0
hist==2.4.0
histoprint==2.4.0
idna==3.4
importlib-metadata==4.8.3
importlib-resources==5.4.0
kiwisolver==1.3.1
llvmlite==0.36.0
lz4==3.1.10
matplotlib==3.3.4
mplhep==0.2.18
numba==0.53.1
numpy==1.19.5
packaging==21.3
pandas==1.1.5
Pillow==8.4.0
pyarrow==6.0.1
pydantic==1.9.2
Pygments==2.14.0
pyparsing==3.0.9
python-dateutil==2.8.2
pytz==2023.3
requests==2.27.1
rich==12.6.0
scipy==1.5.4
six==1.16.0
toml==0.10.2
tqdm==4.64.1
typing_extensions==4.1.1
uhi==0.3.3
uproot==4.3.7
uproot3==3.14.4
uproot3-methods==0.10.1
urllib3==1.26.16
zipp==3.6.0

[lgray]

Are you able to update uproot4 and awkward1? These are particularly old versions, even for awkward1/uproot4 and that may be the origin of your problem.

[tomeichlersmith]

Alright, I switched to a python3.10 container and upgraded my coffea, uproot, and awkward installations. pip warns me about this though

coffea 0.7.22 requires awkward<2,>=1.10.3, but you have awkward 2.5.0 which is incompatible.
coffea 0.7.22 requires uproot!=4.2.4,!=4.3.0,!=4.3.1,==4.*,>=4.1.6, but you have uproot 5.1.2 which is incompatible.

which (if I"m remembering correctly) is why I stayed on 0.7.20 with python 3.6.

This leads to a different (but seemingly earlier) error - it looks like coffea expects a version of uproot that has a certain function.

Traceback (most recent call last):
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/unbiased-study.py", line 406, in <module>
    run(args.output_name, nfilelimit = args.file_limit, ncores = args.ncores, quiet = args.quiet)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/unbiased-study.py", line 364, in run
    out = runner(
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 1700, in __call__
    wrapped_out = self.run(fileset, processor_instance, treename)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 1848, in run
    wrapped_out, e = executor(chunks, closure, None)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 672, in __call__
    accumulate(
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/accumulator.py", line 95, in accumulate
    accum = next(gen)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/accumulator.py", line 92, in <genexpr>
    gen = (x for x in items if x is not None)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/rich/progress.py", line 1215, in track
    for value in sequence:
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 1367, in automatic_retries
    raise e
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 1336, in automatic_retries
    return func(*args, **kwargs)
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/processor/executor.py", line 1615, in _work_function
    factory = NanoEventsFactory.from_root(
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/nanoevents/factory.py", line 124, in from_root
    base_form = mapping._extract_base_form(
  File "/export/scratch/users/eichl008/ldmx/eat/v14/analysis/coffea/.local/lib/python3.10/site-packages/coffea/nanoevents/mapping/uproot.py", line 116, in _extract_base_form
    form = uproot._util.awkward_form_remove_uproot(awkward, form)
AttributeError: module 'uproot._util' has no attribute 'awkward_form_remove_uproot'

Upgrading coffea to the latest prerelease pip install --pre --upgrade coffea removes this error in favor of a NotImplemented error since the newer prerelease coffea provides a dask-awkward array rather than awkward. Adding in a bunch of .compute() where I need the values of the arrays seems to resolve this error, I'm going to write a shorter summary into a separate reply so that future readers can know what to do. Thank you!

P.S. Does coffea plan to make a stable release in the future? Or will it stay as pre-releases to reflect how it is being developed?

[lgray]

The pre-releases are for coffea 2023 (which has some major differences to coffea 0.7)- which we're tying up and about to ship in the next couple weeks.

You may use the pre-releases if you like! Just realize things can be quite different.

Unfortunately coffea 0.7 and awkward 2 / uproot5 are incompatible. You'll have to cook your own environment.

[lgray]

Similarly, can you show the code you're using to access the data which is causing this error?

[tomeichlersmith]

Updating to a more recent (prerelease) version of coffea and updating my access pattern to use dask-awkward resolved the issue, but for future readers:

I was initially going to use the nested values to select parent values but I noticed that some issue was happening. I was able to narrow it down to the following lines of code inside my def process(self, events) function

        print(events.layout)
        print('id')
        print(events['HcalSimHits_eat/HcalSimHits_eat.id_'])
        print('pdg')
        print(events['HcalSimHits_eat/HcalSimHits_eat.pdgCodeContribs_'])

The last line in the code snippet is never printed since that is where the no attribute 'offsets' error is thrown. HcalSimHits is a branch of a variable-length array of objects. Each object has some single-element members (like id_) and some variable-length members (like pdgCodeContribs_).

[tomeichlersmith]

Thank you @lgray for helping debug this issue, I'm just going to summarize the resolution here for any future reader.

Update to a more recent coffea version

The last stable versions of coffea (v0.7.2[0-2]) use a version of uproot/awkward that is unable¹ to implement lazy-loading of doubly-ragged arrays. The easiest way around this issue is to update to a pre-release version of coffea pip install --pre --upgrade coffea remembering to also update the processor(s) to reflect the change that the object provided to def process(self, events) will be a dask-awkward array rather than a simple awkward array. This basically amounts to adding a .compute() call where-ever you need the values (e.g. like filling histograms). https://dask-awkward.readthedocs.io/en/stable/getting-started/intro.html

Footnotes

1. I'm not sure if those versions of uproot/awkward are literally unable to implement lazy-loading of doubly-ragged arrays or if its just coffea v0.7's implementation that does not work for doubly-ragged arrays. In any case, the easiest way forward is still to simply upgrade coffea and update to dask-awkward. ↩

[tomeichlersmith]

There is a dask-awkward demo which highlights the existence of dask-histogram.

I guess I'm curious to see if there is a better option than using dask-awkward and dask-histogram and then doing

for key in histograms:
    histograms[key] = histograms[key].compute()

before returning the histograms to be accumulated.

[tomeichlersmith]

Oh neat, sorry I missed your reply while typing my own 👍

[lgray]

Ah yeah, the concept of accumulating histograms is ... gone... in coffea 2023 it's handled in the background since what you write down operates on a whole dataset.

You can get a flavor of it from query 6 here:
https://github.com/CoffeaTeam/coffea-benchmarks/blob/coffea2023/coffea-adl-benchmarks.ipynb

[tomeichlersmith]

So the move towards dask would mean I'd use some combination of uproot.dask, dask-awkward, and hist.dask if I don't need a special schema from coffea? Or am I oversimplifying things?

I just noticed that when I accidentally returned by dict of hist.dask.Hist objects, the resulting output dict (post accumulation) has the fully accumulated hist.Hist objects. By "handled in the background", does this mean the dask objects I return from a processor will be computed during accumulation? (Is this the intended behavior and something I can rely on?)

Thank you for answering all these questions 🙌 Your support and knowledge is very appreciated.

[lgray]

Yes - if you're not using NanoEventsFactory.from_root you can start from uproot.dask and go from there. In fact if you use latest dask-awkward and uproot you can just use ak.some_function instead of dak.some_function and easily port any existing code you have. The py3.10 container you're using will let you do this.

Handled in the background -> the histogram object itself appears like a single object to the user but gets represented in the dask taskgraph as a bunch of iterated fills and a corresponding tree-reduce to yield the final filled histogram. Since the fill operation is operating on a whole dataset instead of a single chunk within the processor the accumulation is in effect "in the background".