Is form_mapping required for uproot.dask?

[tomeichlersmith]

When I try to use uproot.dask out-of-the-box (so to speak), I see an error at .compute() time. This error seems to allude to the form mapping and so I am curious if that part of uproot.dask is required (or at least required for this type of input file).

I am using uproot 5.2.0rc1 and awkward 2.5.0.

In [1]: import uproot

In [2]: events = uproot.dask('test.root:LDMX_Events')

In [3]: events['EcalRecHits_eat.energy_'].compute()
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[3], line 1
----> 1 events['EcalRecHits_eat.energy_'].compute()

... removed functions from call stack internal to dask ...

File ~/.local/lib/python3.11/site-packages/uproot/_dask.py:1039, in _UprootRead.__call__(self, i_start_stop)
   1036 def __call__(self, i_start_stop) -> AwkArray:
   1037     i, start, stop = i_start_stop
-> 1039     return self.read_tree(self.ttrees[i], start, stop)

File ~/.local/lib/python3.11/site-packages/uproot/_dask.py:912, in UprootReadMixin.read_tree(self, tree, start, stop)
    906 container = {}
    907 for buffer_key, dtype in self.expected_form.expected_from_buffers(
    908     buffer_key=self.form_mapping_info.buffer_key
    909 ).items():
    910     # Which key(s) does this buffer require. This code permits the caller
    911     # to require multiple keys to compute a single buffer.
--> 912     keys_for_buffer = self.form_mapping_info.keys_for_buffer_keys(
    913         frozenset({buffer_key})
    914     )
    915     # If reading this buffer loads a permitted key, read from the tree
    916     # We might not have _all_ keys if e.g. buffer A requires one
    917     # but not two of the keys required for buffer B
    918     if all(k in self.common_keys for k in keys_for_buffer):

File ~/.local/lib/python3.11/site-packages/uproot/_dask.py:826, in TrivialFormMappingInfo.keys_for_buffer_keys(self, buffer_keys)
    824     form_key, attribute = buffer_key.rsplit("-", maxsplit=1)
    825     # Identify key from form_key
--> 826     keys.add(self._form_key_to_key[form_key])
    827 return frozenset(keys)

KeyError: 'None'

A non-dask read looks normal to me

In [1]: import uproot

In [2]: events = uproot.open('test.root:LDMX_Events')['EcalRecHits_eat.energy_'].array()

In [3]: events
Out[3]: <Array [[4.31, 5.47, 42.7, ..., 3.55, 46], ...] type='608027 * var * float32'>

And I get a different error if I try to use dask without awkward.

In [1]: import uproot

In [2]: events = uproot.dask('test.root:LDMX_Events', library='np')
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[2], line 1
----> 1 events = uproot.dask('test.root:LDMX_Events', library='np')

File ~/.local/lib/python3.11/site-packages/uproot/_dask.py:213, in dask(files, filter_name, filter_typename, filter_branch, recursive, full_paths, step_size, steps_per_file, library, ak_add_doc, custom_classes, allow_missing, open_files, form_mapping, **options)
    211 if library.name == "np":
    212     if open_files:
--> 213         return _get_dask_array(
    214             files,
    215             filter_name,
    216             filter_typename,
    217             filter_branch,
    218             recursive,
    219             full_paths,
    220             step_size,
    221             custom_classes,
    222             allow_missing,
    223             real_options,
    224             interp_options,
    225             steps_per_file,
    226         )
    227     else:
    228         return _get_dask_array_delay_open(
    229             files,
    230             filter_name,
   (...)
    239             steps_per_file,
    240         )

File ~/.local/lib/python3.11/site-packages/uproot/_dask.py:611, in _get_dask_array(files, filter_name, filter_typename, filter_branch, recursive, full_paths, step_size, custom_classes, allow_missing, real_options, interp_options, steps_per_file)
    608 assert entry_step >= 1
    610 for key in common_keys:
--> 611     dt = ttrees[0][key].interpretation.numpy_dtype
    612     if dt.subdtype is None:
    613         inner_shape = ()

File ~/.local/lib/python3.11/site-packages/uproot/interpretation/__init__.py:57, in Interpretation.numpy_dtype(self)
     52 @property
     53 def numpy_dtype(self):
     54     """
     55     The ``numpy.dtype`` to use to put objects of this type in a NumPy array.
     56     """
---> 57     raise AssertionError

AssertionError: 

[agoose77]

No, form remapping is not required. It might be that we have a bug in the trivial case — I'll take a look.

[tomeichlersmith]

I'm unsure if this information is necessary, but here is more detail on the branch I am trying to access.

I'm trying to access the member of a class that was serialized within a std::vector in each entry of the tree.

In [1]: import uproot

In [2]: tree = uproot.open('test.root:LDMX_Events')

In [3]: tree['EcalRecHits_eat'].show()
name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
EcalRecHits_eat      | vector<ldmx::EcalHit>    | AsGroup(<TBranchElement 'EcalR
EcalRecHits_eat.id_  | int32_t[]                | AsJagged(AsDtype('>i4'))
EcalRecHits_eat.a... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.e... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.t... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.x... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.y... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.z... | float[]                  | AsJagged(AsDtype('>f4'))
EcalRecHits_eat.i... | bool[]                   | AsJagged(AsDtype('bool'))

[tomeichlersmith]

I also should add that I can currently load (and .compute) the same file using coffea's wrapper of uproot.dask which I think does do some form_mapping but I cannot confirm that yet (perhaps you are more familiar).

https://github.com/CoffeaTeam/coffea/blob/5f7651aff69cf4fb19e5c7cbcb18be5fe9c89a91/src/coffea/nanoevents/factory.py#L322-L353