fix: preserve structure of jagged subfields when accessing the arrays directly (#1476)

* Fix jagged subfields

* Rename tuple subfields to match Awkward

* Added test

* Minor refactoring

Author: ariostas

src/uproot/_dask.py

@@ -1543,7 +1543,7 @@ def real_filter_branch(branch):
         entry_stop = ttree.num_entries
 
         if isinstance(ttree, HasFields):
-            akform = ttree.to_akform(filter_name=common_keys)
+            akform, _ = ttree.to_akform(filter_name=common_keys)
             ttree_step = _RNTuple_regularize_step_size(
                 ttree, akform, step_size, entry_start, entry_stop
             )
@@ -1593,7 +1593,7 @@ def real_filter_branch(branch):
                     partition_args.append((i, start, stop))
 
     if isinstance(ttrees[0], HasFields):
-        base_form = ttrees[0].to_akform(filter_name=common_keys)
+        base_form, _ = ttrees[0].to_akform(filter_name=common_keys)
     else:
         base_form = _get_ttree_form(
             awkward, ttrees[0], common_keys, interp_options.get("ak_add_doc")

src/uproot/behaviors/RNTuple.py

@@ -525,7 +525,9 @@ def to_akform(
                 for compatibility with software that was used for :doc:`uproot.behaviors.TBranch.TBranch`. This argument should not be used
                 and will be removed in a future version.
 
-        Returns the an Awkward Form with the structure of the data in the ``RNTuple`` or ``RField``.
+        Returns a 2-tuple where the first entry is the Awkward Form with the structure of the data in the ``RNTuple`` or ``RField``,
+        and the second entry is the relative path of the requested RField. The second entry is needed in cases where the requested RField
+        is a subfield of a collection, which requires constructing the form with information about the parent field.
         """
         ak = uproot.extras.awkward()
 
@@ -540,6 +542,7 @@ def to_akform(
         top_names = []
         record_list = []
         if self is rntuple:
+            field_path = None
             for field in self.fields:
                 # the field needs to be in the keys or be a parent of a field in the keys
                 if any(
@@ -551,14 +554,39 @@ def to_akform(
                         rntuple.field_form(field.field_id, keys, ak_add_doc=ak_add_doc)
                     )
         else:
+            # If it is a subfield of a collection, we need to include the collection in the keys
+            path_keys = self.path.split(".")
+            top_collection = None
+            tmp_field = self.ntuple
+            field_path = [self.name]
+            for i, key in enumerate(path_keys):
+                tmp_field = tmp_field[key]
+                if (
+                    tmp_field.record.struct_role
+                    == uproot.const.RNTupleFieldRole.COLLECTION
+                ):
+                    top_collection = tmp_field
+                    field_path = path_keys[i:]
+                    break
             # Always use the full path for keys
             # Also include the field itself
             keys = [self.path] + [f"{self.path}.{k}" for k in keys]
-            # The field needs to be in the keys or be a parent of a field in the keys
-            if any(key.startswith(f"{self.path}.") or key == self.path for key in keys):
-                top_names.append(self.name)
+            if top_collection is None:
+                # The field needs to be in the keys or be a parent of a field in the keys
+                if any(
+                    key.startswith(f"{self.path}.") or key == self.path for key in keys
+                ):
+                    top_names.append(self.name)
+                    record_list.append(
+                        rntuple.field_form(self.field_id, keys, ak_add_doc=ak_add_doc)
+                    )
+            else:
+                keys += [top_collection.path]
+                top_names.append(top_collection.name)
                 record_list.append(
-                    rntuple.field_form(self.field_id, keys, ak_add_doc=ak_add_doc)
+                    rntuple.field_form(
+                        top_collection.field_id, keys, ak_add_doc=ak_add_doc
+                    )
                 )
 
         parameters = None
@@ -572,7 +600,7 @@ def to_akform(
         form = ak.forms.RecordForm(
             record_list, top_names, form_key="toplevel", parameters=parameters
         )
-        return form
+        return (form, field_path)
 
     def arrays(
         self,
@@ -697,7 +725,7 @@ def arrays(
             [c.num_entries for c in clusters[start_cluster_idx:stop_cluster_idx]]
         )
 
-        form = self.to_akform(
+        form, field_path = self.to_akform(
             filter_name=filter_name,
             filter_typename=filter_typename,
             filter_field=filter_field,
@@ -755,6 +783,20 @@ def arrays(
         # no longer needed; save memory
         del container_dict
 
+        # If we constructed some parent fields, we need to get back to the requested field
+        if field_path is not None:
+            for field in field_path[:-1]:
+                if field in arrays.fields:
+                    arrays = arrays[field]
+                # tuples are a trickier since indices no longer match
+                else:
+                    if field.isdigit() and arrays.fields == ["0"]:
+                        arrays = arrays["0"]
+                    else:
+                        raise AssertionError(
+                            "The array was not constructed correctly. Please report this issue."
+                        )
+
         # FIXME: This is not right, but it might temporarily work
         if library.name == "np":
             return arrays.to_numpy()
@@ -896,7 +938,7 @@ def iterate(
             )
         )
 
-        akform = self.to_akform(
+        akform, _ = self.to_akform(
             filter_name=filter_name,
             filter_typename=filter_typename,
             filter_field=filter_field,
@@ -1408,7 +1450,7 @@ def num_entries_for(
             )
         )
 
-        akform = self.to_akform(
+        akform, _ = self.to_akform(
             filter_name=filter_name,
             filter_typename=filter_typename,
             filter_field=filter_field,

src/uproot/models/RNTuple.py

@@ -562,7 +562,7 @@ def field_form(self, this_id, keys, ak_add_doc=False):
                 ):
                     recordlist.append(self.field_form(i, keys, ak_add_doc=ak_add_doc))
                     namelist.append(field_records[i].field_name)
-            if all(name == f"_{i}" for i, name in enumerate(namelist)):
+            if all(re.fullmatch(r"_[0-9]+", name) is not None for name in namelist):
                 namelist = None
             return ak.forms.RecordForm(
                 recordlist, namelist, form_key="whatever", parameters=parameters
@@ -1485,7 +1485,15 @@ def name(self):
         """
         Name of the ``RField``.
         """
-        return self._ntuple.field_records[self._fid].field_name
+        # We rename subfields of tuples to match Awkward
+        name = self._ntuple.field_records[self._fid].field_name
+        if (
+            not self.top_level
+            and self.parent.record.struct_role == uproot.const.RNTupleFieldRole.RECORD
+            and re.fullmatch(r"_[0-9]+", name) is not None
+        ):
+            name = name[1:]
+        return name
 
     @property
     def description(self):
@@ -1637,14 +1645,25 @@ def array(
         See also :ref:`uproot.behaviors.RNTuple.HasFields.arrays` to read
         multiple ``RFields`` into a group of arrays or an array-group.
         """
-        return self.arrays(
+        arrays = self.arrays(
             entry_start=entry_start,
             entry_stop=entry_stop,
             library=library,
             interpreter=interpreter,
             backend=backend,
             ak_add_doc=ak_add_doc,
-        )[self.name]
+        )
+        if self.name in arrays.fields:
+            arrays = arrays[self.name]
+        # tuples are a trickier since indices no longer match
+        else:
+            if self.name.isdigit() and arrays.fields == ["0"]:
+                arrays = arrays["0"]
+            else:
+                raise AssertionError(
+                    "The array was not constructed correctly. Please report this issue."
+                )
+        return arrays
 
 
 # No cupy version of numpy.insert() provided

tests/test_1406_improved_rntuple_methods.py

@@ -80,10 +80,10 @@ def test_getitem(tmp_path):
     assert obj["struct1"]["x"] is obj[r"struct1\x"]
 
     # Make sure it accesses the grandchildren field instead of the "real" _0
-    assert obj["struct5._0"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
-    assert obj["struct5._1"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
-    assert obj["struct5._2"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
-    assert obj["struct6._0"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
+    assert obj["struct5.0"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
+    assert obj["struct5.1"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
+    assert obj["struct5.2"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
+    assert obj["struct6.0"].record.struct_role == uproot.const.RNTupleFieldRole.LEAF
 
 
 def test_to_akform(tmp_path):
@@ -94,21 +94,22 @@ def test_to_akform(tmp_path):
 
     obj = uproot.open(filepath)["ntuple"]
 
-    akform = obj.to_akform()
+    akform, field_path = obj.to_akform()
     assert akform == data.layout.form
+    assert field_path is None
 
-    assert obj["struct1"].to_akform() == akform.select_columns("struct1")
-    assert obj["struct2"].to_akform() == akform.select_columns("struct2")
-    assert obj["struct3"].to_akform() == akform.select_columns("struct3")
-    assert obj["struct4"].to_akform() == akform.select_columns("struct4")
-    assert obj["struct5"].to_akform() == akform.select_columns("struct5")
+    assert obj["struct1"].to_akform() == (akform.select_columns("struct1"), ["struct1"])
+    assert obj["struct2"].to_akform() == (akform.select_columns("struct2"), ["struct2"])
+    assert obj["struct3"].to_akform() == (akform.select_columns("struct3"), ["struct3"])
+    assert obj["struct4"].to_akform() == (akform.select_columns("struct4"), ["struct4"])
+    assert obj["struct5"].to_akform() == (akform.select_columns("struct5"), ["struct5"])
 
-    assert obj["struct1"].to_akform(filter_name="x") == akform.select_columns(
+    assert obj["struct1"].to_akform(filter_name="x")[0] == akform.select_columns(
         ["struct1.x"]
     )
-    assert obj["struct3"].to_akform(filter_typename="double") == akform.select_columns(
-        ["struct3.t"]
-    )
+    assert obj["struct3"].to_akform(filter_typename="double")[
+        0
+    ] == akform.select_columns(["struct3.t"])
 
 
 def test_iterate_and_concatenate(tmp_path):
@@ -144,5 +145,5 @@ def test_array(tmp_path):
 
     obj = uproot.open(filepath)["ntuple"]
 
-    assert obj["struct5._0"].array().tolist() == [1, 4]
-    # assert obj["struct6._0"].array().tolist() == [[1, 4], [7]] # TODO: Need to fix this
+    assert obj["struct5.0"].array().tolist() == [1, 4]
+    assert obj["struct6.0"].array().tolist() == [[1, 4], [7]]

tests/test_1476_rntuple_jagged_subfields.py

@@ -0,0 +1,120 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
+
+import os
+
+import pytest
+
+import uproot
+
+ak = pytest.importorskip("awkward")
+
+data = ak.Array(
+    {
+        "field1": [[(0, 1), (2, 3)], [], [(4, 5)]],
+        "field2": [[{"x": 0, "y": 1}, {"x": 2, "y": 3}], [], [{"x": 4, "y": 5}]],
+        "field3": [
+            {"x": [1, 2, 3], "y": [(1, 2), (2, 3), (3, 4)]},
+            {"x": [], "y": [(7, 8)]},
+            {"x": [9, 10], "y": [(11, 12), (13, 14)]},
+        ],
+        "field4": [(1, [(1, 2), (3, 4)]), (2, []), (3, [(5, 6)])],
+        "field5": [
+            {
+                "x": [
+                    {"up": [(0, 2), (1, 2)], "down": [[1, 2, 3], []]},
+                    {"up": [], "down": [[4]]},
+                ],
+                "y": [
+                    (
+                        {"left": [[0, "hi", 2.3]], "right": 6},
+                        {"left": [[], [""]], "right": 8.0},
+                    )
+                ],
+            },
+            {
+                "x": [
+                    {"up": [(10, 2), (12, 2)], "down": []},
+                    {"up": [(1, 2)], "down": [[], [4, 2, 3]]},
+                ],
+                "y": [
+                    (
+                        {"left": [[23, 4.1, "hello"]], "right": 14},
+                        {"left": [], "right": 16.0},
+                    )
+                ],
+            },
+            {
+                "x": [
+                    {"up": [], "down": [[4, 5], [1, 2, 3], []]},
+                    {"up": [(0, 2), (1, 2)], "down": []},
+                ],
+                "y": [
+                    (
+                        {"left": [[]], "right": 14},
+                        {"left": [[2, 3], ["bye", ""]], "right": 16.0},
+                    )
+                ],
+            },
+        ],
+        "field6": [(1, 2, 3, 4), (2, 5, 6, 7), ("hello", 2.3, 8, "")],
+    }
+)
+
+
+def test_jagged_subfields(tmp_path):
+    filepath = os.path.join(tmp_path, "test.root")
+
+    with uproot.recreate(filepath) as file:
+        obj = file.mkrntuple("ntuple", data)
+
+    obj = uproot.open(filepath)["ntuple"]
+
+    assert ak.array_equal(obj["field1"].array(), data["field1"])
+    assert ak.array_equal(obj["field1.0"].array(), data["field1"]["0"])
+    assert ak.array_equal(obj["field1.1"].array(), data["field1"]["1"])
+
+    assert ak.array_equal(obj["field2"].array(), data["field2"])
+    assert ak.array_equal(obj["field2.x"].array(), data["field2"]["x"])
+    assert ak.array_equal(obj["field2.y"].array(), data["field2"]["y"])
+
+    assert ak.array_equal(obj["field3"].array(), data["field3"])
+    assert ak.array_equal(obj["field3.x"].array(), data["field3"]["x"])
+    assert ak.array_equal(obj["field3.y"].array(), data["field3"]["y"])
+    assert ak.array_equal(obj["field3.y.0"].array(), data["field3"]["y"]["0"])
+    assert ak.array_equal(obj["field3.y.1"].array(), data["field3"]["y"]["1"])
+
+    assert ak.array_equal(obj["field4"].array(), data["field4"])
+    assert ak.array_equal(obj["field4.0"].array(), data["field4"]["0"])
+    assert ak.array_equal(obj["field4.1"].array(), data["field4"]["1"])
+    assert ak.array_equal(obj["field4.1.0"].array(), data["field4"]["1"]["0"])
+    assert ak.array_equal(obj["field4.1.1"].array(), data["field4"]["1"]["1"])
+
+    assert obj["field5"].array().tolist() == data["field5"].tolist()
+    assert obj["field5.x"].array().tolist() == data["field5"]["x"].tolist()
+    assert obj["field5.x.up"].array().tolist() == data["field5"]["x"]["up"].tolist()
+    assert obj["field5.x.down"].array().tolist() == data["field5"]["x"]["down"].tolist()
+    assert obj["field5.y"].array().tolist() == data["field5"]["y"].tolist()
+    assert obj["field5.y.0"].array().tolist() == data["field5"]["y"]["0"].tolist()
+    assert (
+        obj["field5.y.0.left"].array().tolist()
+        == data["field5"]["y"]["0"]["left"].tolist()
+    )
+    assert (
+        obj["field5.y.0.right"].array().tolist()
+        == data["field5"]["y"]["0"]["right"].tolist()
+    )
+    assert obj["field5.y.1"].array().tolist() == data["field5"]["y"]["1"].tolist()
+    assert (
+        obj["field5.y.1.left"].array().tolist()
+        == data["field5"]["y"]["1"]["left"].tolist()
+    )
+    assert (
+        obj["field5.y.1.right"].array().tolist()
+        == data["field5"]["y"]["1"]["right"].tolist()
+    )
+
+    assert obj["field6"].array().tolist() == data["field6"].tolist()
+    assert obj["field6.0"].array().tolist() == data["field6"]["0"].tolist()
+    assert obj["field6.1"].array().tolist() == data["field6"]["1"].tolist()
+    assert obj["field6.2"].array().tolist() == data["field6"]["2"].tolist()
+    assert obj["field6.3"].array().tolist() == data["field6"]["3"].tolist()