ATLAS `ElementLinks` used to work, NOT anymore

[ekourlit]

Version of Awkward Array

2.6.2

Description and code to reproduce

I'm writing this on behalf of a team within the ATLAS Collaboration who tries to read ATLAS-specific data (PHYSLITE) with uproot and awkward. The problem is related to ElementLinks.

We use coffea and the PHYSLITESchema to load and structure our data in awkward arrays. ElementLinks are also treated within the PHYSLITESchema. This is something that it used to work fine in the past with:

coffea version: 2023.7.0rc1.dev83+g52950d1
uproot version: 5.0.13
awkward version: 2.4.4

but it is not working anymore with:

coffea version: 2024.3.0
uproot version: 5.3.2
awkward version: 2.6.2

When we try to retrieve the linked collection from the parent one we get the following error:

Traceback (most recent call last):
  File "/venv/lib/python3.9/site-packages/awkward/_dispatch.py", line 56, in dispatch
    raise NotImplementedError
NotImplementedError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ekourlit/failing_example.py", line 67, in <module>
    clusters = events.Electrons.caloClusters
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 1535, in __getattr__
    return self.map_partitions(
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 1588, in map_partitions
    return map_partitions(func, self, *args, traverse=traverse, **kwargs)
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 2107, in map_partitions
    return _map_partitions(
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 1943, in _map_partitions
    meta = map_meta(fn, *args, **kwargs)
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 2467, in map_meta
    meta = fn(*to_meta(deps))
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 1903, in __call__
    return self.fn(*args, **kwargs)
  File "/venv/lib/python3.9/site-packages/dask_awkward/lib/core.py", line 2575, in __call__
    return getattr(coll, self.attr)
  File "/home/ekourlit/failing_example.py", line 29, in caloClusters
    return _element_link_method(
  File "/venv/lib/python3.9/site-packages/coffea/nanoevents/methods/physlite.py", line 55, in _element_link_method
    return _element_link(
  File "/venv/lib/python3.9/site-packages/coffea/nanoevents/methods/physlite.py", line 39, in _element_link
    global_index = _get_global_index(target_collection, eventindex, index)
  File "/venv/lib/python3.9/site-packages/coffea/nanoevents/methods/physlite.py", line 132, in _get_global_index
    target_offsets = _get_target_offsets(load_column, eventindex)
  File "/venv/lib/python3.9/site-packages/coffea/nanoevents/methods/physlite.py", line 109, in _get_target_offsets
    if awkward.backend(load_column) == "typetracer":
  File "/venv/lib/python3.9/site-packages/awkward/_dispatch.py", line 70, in dispatch
    return gen_or_result
  File "/venv/lib/python3.9/site-packages/awkward/_errors.py", line 85, in __exit__
    self.handle_exception(exception_type, exception_value)
  File "/venv/lib/python3.9/site-packages/awkward/_errors.py", line 95, in handle_exception
    raise self.decorate_exception(cls, exception)
NotImplementedError: 

See if this has been reported at https://github.com/scikit-hep/awkward/issues

We put together a reproducer script, which downloads one of our public files for testing: failing_example.py.zip

I also tried to factorise out coffea by linking collections without using the PHYSLITESchema, just by using awkward. This is working with pure awkward (2.6.2), although I was not able to make it work using dask_awkward (2024.3.0). Please see the reproducer script: ElementLinks_awkward.py.zip.

We would appreciate your insights. Please let us know if you need further information from our side.

Tagging @alexander-held and @matthewfeickert.

Thank you in advance!

[jpivarski]

I tried to run ElementLinks_awkward.py because I wanted to get more of the stack trace. However, I see a different error from the NotImplemented handling you see above:

% python ElementLinks_awkward.py 
Try to read the events with awkward
Events read!
Try ElementLinks and print the clusters via electrons
[[[{calE: 1.38e+05, calEta: 1.55, calPhi: -1.85, calM: 0}], ..., [{...}]], ...]

Now trying to read the events with dask_awkward
Traceback (most recent call last):
  File "/home/jpivarski/Downloads/ElementLinks_awkward.py", line 91, in <module>
    events = read_events_dak(columns)
  File "/home/jpivarski/Downloads/ElementLinks_awkward.py", line 40, in read_events_dak
    events[cname] = dak.zip(events[cname], depth_limit=2)
  File "/home/jpivarski/irishep/dask-awkward/src/dask_awkward/lib/structure.py", line 1247, in zip
    metadict[k] = coll._meta
  File "/home/jpivarski/irishep/awkward/src/awkward/highlevel.py", line 1239, in __getattr__
    raise AttributeError(f"no field named {where!r}")
AttributeError: no field named '_meta'

Of course, that's very weird: these are supposed to be dak.Arrays, and one of the fundamental parts of a dak.Array is its _meta (which is an ak.Array of typetracers, to propagate type information through the growing DAG).

Right before the failing line, I added

        print("--------------------")
        for k, v in events[cname].items():
            print(f"{k = } {type(v) = } {v = } {getattr(v, '_meta', None) = }")
        print("--------------------")

and see

--------------------
k = 'pt' type(v) = <class 'awkward.highlevel.Array'> v = <Array [[5.5e+04, 3.38e+04, 4.25e+03], ..., []] type='90000 * var * float32'> getattr(v, '_meta', None) = None
k = 'eta' type(v) = <class 'awkward.highlevel.Array'> v = <Array [[1.55, 1.84, 0.717], [1.59], ..., []] type='90000 * var * float32'> getattr(v, '_meta', None) = None
k = 'phi' type(v) = <class 'awkward.highlevel.Array'> v = <Array [[-1.85, 1.66, -2.33], [2.76], ..., []] type='90000 * var * float32'> getattr(v, '_meta', None) = None
k = 'm' type(v) = <class 'awkward.highlevel.Array'> v = <Array [[0.511, 0.511, 0.511], [0.511], ..., []] type='90000 * var * float32'> getattr(v, '_meta', None) = None
k = 'caloClusterLinks' type(v) = <class 'awkward.highlevel.Array'> v = <Array [[[{m_persKey: ..., ...}], ...], ...] type='90000 * var * var * {m_p...'> getattr(v, '_meta', None) = None
--------------------

Okay, so that's the problem with this example; it's trying to dak.zip (lazy) some ak.Arrays (eager). Oh, but that's because tree is an open ROOT file and you pull a TBranch from that and call uproot.TBranch.array to get an array eagerly. I think this demonstrator is not what you intended: in trying to show the difference between eager and lazy, this script is mixing eager and lazy and encountering new problems because of that.

I wanted to start with a pure-Awkward example first, then add dask-awkward, then add Coffea, in order to narrow in on where the error actually is. If pure-Awkward works fine, then we'll probably need to move this issue to dask-awkward or Coffea, unless the actual error is dask-awkward or Coffea calling an Awkward routine and the Awkward routine is breaking.

---

I also ran failing_example.py, and the stack trace is not any deeper than the NotImplementedError. I don't know why we're not seeing more stack trace—we don't need to see the code that handles errors, we need to see the error!

That line,

awkward/src/awkward/_dispatch.py

Lines 53 to 58 in 87d2191

	# Future proof the implementation by permitting the `__awkward_function__` to return `NotImplemented`
	# This may later be used to signal that another overload should be used.
	if result is NotImplemented:
	raise NotImplementedError
	else:
	return result

is to handle a function call like ak.whatever(actually_a_dak_Array) and pass it to dak.whatever. The NotImplemented handling is so that dask-awkward can progressively add coverage of Awkward functions; it's saying that this one hasn't been implemented yet. (An odd thing to say for a function that used to work.)

At least we know we're looking for an __awkward_function__ magic method that returns NotImplemented (returns NotImplemented, not raises NotImplementedError). There's no such method in https://github.com/CoffeaTeam/coffea/blob/master/src/coffea/nanoevents/methods/physlite.py, though a lot of the second stack trace seems to be concentrated there. That stack trace ends here:

https://github.com/CoffeaTeam/coffea/blob/bef54a2771a4ec034b58cae928fb59cc02e65603/src/coffea/nanoevents/methods/physlite.py#L107-L110

which is just touching columns, to let Dask know that they'll be needed in the computation on the Dask workers. I don't see why that would go to any dispatching; this is not a high-level Awkward function. Well, it does define a HighLevelContext:

awkward/src/awkward/typetracer.py

Lines 133 to 163 in 87d2191

	def touch_data(
	array: Any,
	*,
	highlevel: bool = True,
	behavior: Mapping | None = None,
	attrs: Mapping[str, Any] | None = None,
	) -> Array | Record:
	"""
	Args:
	array: Array-like data (anything #ak.to_layout recognizes).
	highlevel (bool): If True, return an #ak.Array; otherwise, return
	a low-level #ak.contents.Content subclass.
	behavior (None or dict): Custom #ak.behavior for the output array, if
	high-level.
	attrs (None or dict): Custom attributes for the output array, if
	high-level.
	Recursively touches the data and returns a shall copy of the given array.
	"""
	with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
	layout = ctx.unwrap(
	array,
	allow_unknown=False,
	allow_record=True,
	primitive_policy="error",
	none_policy="error",
	string_policy="as-characters",
	)
	_touch_data(layout)
	return ctx.wrap(layout, highlevel=highlevel)

which is for decorating errors. (I'd say that's unnecessary for a function like this. ak.typetracer.touch_data is public for downstream libraries, but not high-level for data analysts.) But that HighLevelContext is not the cause of this error. I commented it out and the error is unchanged.

I need more of the stack trace in the place where NotImplemented is handled in the dispatch mechanism—like, which __awkward_function__ is being called? Even adding

                        import traceback
                        try:
                            raise Exception("fake-o")
                        except Exception:
                            print(traceback.format_exc())

doesn't do it. Surely the Python stack is deeper than one level here. @martindurant, do you know why Python isn't showing more of the stack trace? Are we doing something that hides it?

---

I made a guess and got a bit more progress: in the Coffea code, the predicate of the if statement is awkward.backend(load_column), which is a high-level function. load_column is a dak.Array, and dask-awkward doesn't define a dak.backend function. So that must be what's happening: ak.backend(load_column) is trying to dispatch to dak.backend(load_column), which doesn't exist, so we handle NotImplemented. I'm realizing now that I should have concluded it from

  File "/venv/lib/python3.9/site-packages/coffea/nanoevents/methods/physlite.py", line 109, in _get_target_offsets
    if awkward.backend(load_column) == "typetracer":
  File "/venv/lib/python3.9/site-packages/awkward/_dispatch.py", line 70, in dispatch

but I was looking at the stack trace that ended with NotImplementedError, since it's more obvious that something is going wrong there. In lieu of actually getting a deeper stack trace at that point in the code (I still don't know why we don't see more stack trace there), I added text to the NotImplementedError message (#3078), so that now it reads:

NotImplementedError: no overload for ak.backend with arguments (dask.awkward<calE, npartitions=1>,)

That would have made this debugging a little easier.

Ultimately, the error, then, is in dask-awkward not having an overload for ak.backend, which would be trivial to write:

# 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)

Request this feature in dask-awkward, referencing this issue.

[jpivarski]

I made that feature request: dask-contrib/dask-awkward#497, so I'll close this issue in anticipation of the dask-awkward feature fixing the problem. If it doesn't, we'll reopen.

[jpivarski]

With PR dask-contrib/dask-awkward#498, which fixes issue dask-contrib/dask-awkward#497, failing_example.py gets a little further:

TypeError: cannot convert typetracer arrays to Python lists

This error occurred while calling

    ak.from_iter(
        Array-instance
        highlevel = False
        allow_record = True
    )

and that's true: typetracer arrays (which are the _meta of dak.Array) have no data and therefore can't be iterated over. But it's already an indication of a problem if anything is being iterated over in the first place.

  File "/home/jpivarski/irishep/coffea/src/coffea/nanoevents/methods/physlite.py", line 110, in _get_target_offsets
    awkward.typetracer.touch_data(load_column)
  File "/home/jpivarski/irishep/awkward/src/awkward/typetracer.py", line 153, in touch_data
    layout = ctx.unwrap(
  File "/home/jpivarski/irishep/awkward/src/awkward/_layout.py", line 145, in unwrap
    return to_layout_impl(
  File "/home/jpivarski/irishep/awkward/src/awkward/operations/ak_to_layout.py", line 263, in _impl
    return ak.operations.from_iter(

touch_data is trying to coerce its argument into an Awkward Array (with a layout)? That doesn't make sense. If its argument is not already an Awkward Array backed by typetracers, there's no point in calling touch_data at all. Instead of coercing, this function should halt with an exception if its argument is not the right type. (It's not a high-level data-analyst-facing function; it's a utility for downstream dependencies like Coffea.)

[agoose77]

@jpivarski I am on the move at the moment, so I'll hold off engaging with the part of the question pertaining to HighLevelContext.

However, if we have a known array type, e.g. Typetracer, we should not be iterating it. I am curious as to why that happens at all!

[jpivarski]

I've been digging a little deeper, and the HighLevelContext code replaced to_layout in #2806, so although that's too-loose argument handling for a high-level function, it's no looser than it was before. That can't be regarded as a cause of the present issue because the old code would have had the same issue.

The issue is that touch_data is being called on a dak.Array, rather than a dak.Array's _meta. When I pass a dak.Array into touch_data, the HighLevelContext goes to to_layout, which goes to from_iter.

Potential solutions:

1. make HighLevelContext.unwrap or to_layout recognize dak.Arrays and unwrap them
2. say this is another instance of dask-awkward not having the appropriate overloads

Solution (1) would fix this case and any similar cases in the future, but it has the danger of hiding the fact that some ak.* function needs a Dask-specific dak.* overload by attempting to operate on its _meta. Some dak.* functions have to do special things with the partitions, and a blanket solution like (1) would blithely ignore that. It's better to be stopped with an error message if there's some decision (about partitions) to be made. Unfortunately, the error message is useful to developers, but presented to users!

Solution (2) might not be possible here because these are deliberately hidden public functions, ak.typetracer.touch_data, rather than ak.*. I saw in dask-awkward that the overloading is not an explicit mapping, and it might assume that everything is in the top-level namespace. We don't want this low-level, yet public, function to be in the top-level namespace.

Perhaps

3. add type(array).__module__.startswith("dask_awkward.") or hasattr(array, "_meta") checks to these hidden-but-public functions.

We know that these functions don't have to do special things with partitions, because they operate on metadata. This is something that we can do on a case-by-case basis, to avoid opening floodgates to "accidental" dak.* overloads.

I'm going to start a PR on solution (3) and tag you as a reviewer, @agoose77, for whenever you're ready. (Not now, I know.)

[jpivarski]

With #3079, the final error message will be

TypeError: ak.typetracer.touch_data cannot take a dask-awkward Array; pass its '_meta' instead

which passes the problem on to Coffea, which should be sending load_column._meta to this function, rather than load_column. I'll make an issue there. It's #1075.