`uproot.dask` cannot handle grouped `TObject` properly

[mrzimu]

Hi,
I'm trying to read a simple object derived from TObject through uproot.dask. It seems that uproot.dask cannot handle branch relationship between TObject and TObject.fUniqueID, TObject.fBits.

When I read the whole tree:

import uproot
dak_arr = uproot.dask({"demo-data.root" : "/my_tree/basic_types"})
dak_arr.compute()

The exception is:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mrli/work/uproot5/.venv/lib/python3.12/site-packages/dask/base.py", line 370, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/.venv/lib/python3.12/site-packages/dask/base.py", line 656, in compute
    results = schedule(dsk, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1261, in __call__
    result, _ = self._call_impl(i, start, stop)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1265, in _call_impl
    return self.read_tree(
           ^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1016, in read_tree
    mapping = self.form_mapping_info.load_buffers(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 969, in load_buffers
    assert src_dtype == dst_dtype  # Sanity check!
           ^^^^^^^^^^^^^^^^^^^^^^
AssertionError

I simply checked the loop this assertion in:

uproot5/src/uproot/_dask.py

Lines 957 to 970 in 65b4af6

	for key, array in zip(keys, arrays):
	# First, convert the sub-array into buffers
	ttree_subform, length, ttree_container = awkward.to_buffers(array)
	# Load the associated projection subform
	projection_subform = self._form.content(key)
	# Correlate each TTree form key with the projection form key
	for (src, src_dtype), (dst, dst_dtype) in zip(
	ttree_subform.expected_from_buffers().items(),
	projection_subform.expected_from_buffers(self.buffer_key).items(),
	):
	assert src_dtype == dst_dtype # Sanity check!
	container[dst] = ttree_container[src]

and found the key is TObject (not TObject.fUniqueID or TObject.fBits), but the array is a tuple containing the array of fUniqueID and fBits, so the ttree_subform and projection_subform are different, causing the AssersionError.

If I directly read the whole file:

dak_arr = uproot.dask("demo-data.root")
dak_arr.compute()

The exception is different from above:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/mrli/work/uproot5/.venv/lib/python3.12/site-packages/dask/base.py", line 370, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/.venv/lib/python3.12/site-packages/dask/base.py", line 656, in compute
    results = schedule(dsk, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1261, in __call__
    result, _ = self._call_impl(i, start, stop)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1265, in _call_impl
    return self.read_tree(
           ^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 1016, in read_tree
    mapping = self.form_mapping_info.load_buffers(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/_dask.py", line 932, in load_buffers
    arrays = tree.arrays(
             ^^^^^^^^^^^^
  File "/home/mrli/work/uproot5/src/uproot/behaviors/TBranch.py", line 906, in arrays
    _fix_asgrouped(
  File "/home/mrli/work/uproot5/src/uproot/behaviors/TBranch.py", line 3221, in _fix_asgrouped
    assert arrays[branch.cache_key] is None
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

I'm not sure whether this error is caused by the same reason.

---

The file is:

demo-data.zip

And the version of uproot is 5.6.6

Hope this report can help :)

[ariostas]

Thank you, @mrzimu! I can indeed reproduce the bug.

We're thinking of possibly deprecating the dask part of Uproot in favor of using virtual arrays with Awkward. So I think it's likely that we won't work on fixing this, but instead, I'll get back to you once there is an option to read the contents into virtual arrays.

[mrzimu]

OK, looking forward to this new feature!

[ikrommyd]

Thank you, @mrzimu! I can indeed reproduce the bug.

We're thinking of possibly deprecating the dask part of Uproot in favor of using virtual arrays with Awkward. So I think it's likely that we won't work on fixing this, but instead, I'll get back to you once there is an option to read the contents into virtual arrays.

I hope deprecating does not mean removing because quite a few people have analyses running with dask awkward. Leaving it stale and “just working” should be the worse case.

[ariostas]

I'll leave the issue open to make sure I don't forget. And at some point I'll check if it's easy to fix.

[ikrommyd]

I got a bit interested in this and I found out that tree.arrays() just takes the subbranches out and puts them in the outer record array while uproot.dask errors.

In [51]: file = uproot.open("/Users/iason/Downloads/demo-data.root")

In [52]: print(file["my_tree;1"].arrays().fields)
['fUniqueID', 'fBits', 'm_bool', 'm_char', 'm_schar', 'm_uchar', 'm_short', 'm_sshort', 'm_ushort', 'm_int', 'm_sint', 'm_uint', 'm_long', 'm_slong', 'm_ulong', 'm_llong', 'm_sllong', 'm_ullong', 'm_float', 'm_double', 'm_ldouble', 'm_int8', 'm_uint8', 'm_int16', 'm_uint16', 'm_int32', 'm_uint32', 'm_int64', 'm_uint64', 'm_rt_bool', 'm_rt_char', 'm_rt_uchar', 'm_rt_short', 'm_rt_ushort', 'm_rt_int', 'm_rt_uint', 'm_rt_long', 'm_rt_ulong', 'm_rt_llong', 'm_rt_ullong', 'm_rt_float', 'm_rt_double', 'm_rt_ldouble']

In [53]: array = uproot.dask("/Users/iason/Downloads/demo-data.root")

In [54]: print(array.fields)
['basic_types', 'TObject', 'fUniqueID', 'fBits', 'm_bool', 'm_char', 'm_schar', 'm_uchar', 'm_short', 'm_sshort', 'm_ushort', 'm_int', 'm_sint', 'm_uint', 'm_long', 'm_slong', 'm_ulong', 'm_llong', 'm_sllong', 'm_ullong', 'm_float', 'm_double', 'm_ldouble', 'm_int8', 'm_uint8', 'm_int16', 'm_uint16', 'm_int32', 'm_uint32', 'm_int64', 'm_uint64', 'm_rt_bool', 'm_rt_char', 'm_rt_uchar', 'm_rt_short', 'm_rt_ushort', 'm_rt_int', 'm_rt_uint', 'm_rt_long', 'm_rt_ulong', 'm_rt_llong', 'm_rt_ullong', 'm_rt_float', 'm_rt_double', 'm_rt_ldouble']

In [55]: for field in array.fields:
    ...:     try:
    ...:         array[field].compute()
    ...:     except Exception as e:
    ...:         print(field, type(e))
    ...:
basic_types <class 'AssertionError'>
TObject <class 'AssertionError'>

Notice that in tree.arrays() case, the subbranches fUniqueID and fBits are put in the outer record array.
Also notice how uproot.dask will keep basic_types and TObject as fields of the array which is nonsense.
Also this is totally fine

In [66]: uproot.dask({"/Users/iason/Downloads/demo-data.root" : "my_tree/basic_types/TObject"}).compute()
Out[66]: <Array [{fUniqueID: 0, fBits: 0}, ..., {...}] type='10 * {fUniqueID: uint32...'>

I think that the problem is that uproot.dask just doesn't structure those sub branches properly. I don't know if it's even meant to.

[mrzimu]

This problem is a little bit complicated for me. Just making a note:

1. TrivialFormMappingInfo.load_buffers always passes how=tuple to tree.arrays:

   uproot5/src/uproot/_dask.py

   Lines 917 to 936 in 097b017

   	def load_buffers(
   	self,
   	tree: HasBranches,
   	keys: frozenset[str],
   	start: int,
   	stop: int,
   	decompression_executor,
   	interpretation_executor,
   	options: Any,
   	) -> Mapping[str, AwkArray]:
   	# First, let's read the arrays as a tuple (to associate with each key)
   	arrays = tree.arrays(
   	keys,
   	entry_start=start,
   	entry_stop=stop,
   	ak_add_doc=options["ak_add_doc"],
   	decompression_executor=decompression_executor,
   	interpretation_executor=interpretation_executor,
   	how=tuple,
   	)

2. The how=tuple is passed to the standalone _fix_asgrouped method:

   uproot5/src/uproot/behaviors/TBranch.py

   Lines 3529 to 3543 in 097b017

   	if isinstance(interpretation, uproot.interpretation.grouped.AsGrouped):
   	assert arrays[branch.cache_key] is None
   	limited_context = dict(expression_context[index_start:index_stop])
   	subarrays = {}
   	subcontext = []
   	for subname in interpretation.subbranches:
   	subbranch = branch[subname]
   	subarrays[subname] = arrays[subbranch.cache_key]
   	subcontext.append((subname, limited_context[subname]))
   	arrays[branch.cache_key] = _ak_add_doc(
   	library.group(subarrays, subcontext, how), branch, ak_add_doc
   	)

   resulting to arrays[<key of TObject branch>] is a tuple of awkward array (that contains fUniqueID and fBits arrays)

3. Back to the TrivialFormMappingInfo.load_buffers method, when looping the keys and arrays, the tuple of fUniqueID and fBits is passed to awkward.to_buffers (as variable array in the loop):

   uproot5/src/uproot/_dask.py

   Lines 953 to 966 in 097b017

   	for key, array in zip(keys, arrays):
   	# First, convert the sub-array into buffers
   	ttree_subform, _length, ttree_container = awkward.to_buffers(array)
   	# Load the associated projection subform
   	projection_subform = self._form.content(key)
   	# Correlate each TTree form key with the projection form key
   	for (src, src_dtype), (dst, dst_dtype) in zip(
   	ttree_subform.expected_from_buffers().items(),
   	projection_subform.expected_from_buffers(self.buffer_key).items(),
   	):
   	assert src_dtype == dst_dtype # Sanity check!
   	container[dst] = ttree_container[src]

   but awkward.to_buffers will transform a tuple to ListOffsetForm:

   >>> import awkward as ak
   >>> tpl = (ak.Array([0, 1, 2, 3]), ak.Array([0, 1, 2, 3]))
   >>> ak.to_buffers(tpl)[0]
   ListOffsetForm('i64', NumpyForm('int64', form_key='node1'), form_key='node0')

   which is a wrong form for grouped TObject.

Simply replace how with None in _fix_asgrouped, i.e. from

uproot5/src/uproot/behaviors/TBranch.py

Lines 3541 to 3543 in 097b017

	arrays[branch.cache_key] = _ak_add_doc(
	library.group(subarrays, subcontext, how), branch, ak_add_doc
	)

to

arrays[branch.cache_key] = _ak_add_doc(
    library.group(subarrays, subcontext, None), branch, ak_add_doc
)

can solve the problem posted on this issue, but I'm not sure whether it will break in other cases.

[mrzimu]

Besides, such replacement can let uproot.dask read the TObject in a group, but it still keeps fUniqueID and fBits as standalone field:

>>> import uproot
>>> dak_arr = uproot.dask({"demo-data.root": "/my_tree/basic_types"})
>>> dak_arr.compute().show()
[{TObject: {fUniqueID: 0, fBits: 0}, fUniqueID: 0, fBits: 0, m_bool: ..., ...},
 {TObject: {fUniqueID: 0, fBits: 0}, fUniqueID: 0, fBits: 0, m_bool: ..., ...},
 ...
 {TObject: {fUniqueID: 0, fBits: 0}, fUniqueID: 0, fBits: 0, m_bool: ..., ...}]

I guess this could be the issue of keys generation in _get_dask_array:

uproot5/src/uproot/_dask.py

Lines 619 to 632 in 097b017

	new_keys = obj.keys(
	recursive=recursive,
	filter_name=filter_name,
	filter_typename=filter_typename,
	**{
	(
	"filter_field"
	if isinstance(obj, HasFields)
	else "filter_branch"
	): real_filter_branch
	},
	full_paths=full_paths,
	ignore_duplicates=True,
	)