fix(rntuple): preserve optional records when writing and reading (#1713)

src/uproot/models/RNTuple.py

@@ -726,7 +726,10 @@ def _field_form(self, this_id, key_prefixes, ak_add_doc=False):
             idx_type = (
                 "i32" if self._column_records_dict[cfid][0].nbits == 32 else "i64"
             )
-            if self._all_fields[cfid].record.type_name.startswith("std::optional"):
+            if (
+                self._all_fields[cfid].record.type_name.startswith("std::optional")
+                or self._all_fields[cfid].record.type_alias == "std::optional"
+            ):
                 keyname = keyname + "-optional"
                 return ak.forms.IndexedOptionForm(
                     idx_type, inner, form_key=keyname, parameters=parameters

src/uproot/writing/_cascadentuple.py

@@ -88,33 +88,40 @@ def _cpp_typename(akform, subcall=False):
         ak_primitive = akform.primitive
         typename = _ak_primitive_to_typename_dict[ak_primitive]
     elif isinstance(akform, awkward.forms.ListOffsetForm):
+        if akform.parameters.get("__array__", "") == "string":
+            return "std::string"
         content_typename = _cpp_typename(akform.content, subcall=True)
+        if content_typename is None:
+            return None if subcall else ""
         typename = f"std::vector<{content_typename}>"
-        # Check if it contains strings and fix the type
-        if akform.parameters.get("__array__", "") == "string":
-            typename = "std::string"
     elif isinstance(akform, awkward.forms.RecordForm):
         if akform.is_tuple:
             field_typenames = [_cpp_typename(t, subcall=True) for t in akform.contents]
+            if any(t is None for t in field_typenames):
+                return None if subcall else ""
             typename = f"std::tuple<{','.join(field_typenames)}>"
         else:
-            typename = "UntypedRecord"
+            return None if subcall else ""
     elif isinstance(akform, awkward.forms.RegularForm):
         content_typename = _cpp_typename(akform.content, subcall=True)
+        if content_typename is None:
+            return None if subcall else ""
         typename = f"std::array<{content_typename},{akform.size}>"
     elif isinstance(akform, awkward.forms.IndexedOptionForm):
         content_typename = _cpp_typename(akform.content, subcall=True)
+        if content_typename is None:
+            return None if subcall else ""
         typename = f"std::optional<{content_typename}>"
     elif isinstance(akform, awkward.forms.UnionForm):
         field_typenames = [_cpp_typename(t, subcall=True) for t in akform.contents]
+        if any(t is None for t in field_typenames):
+            return None if subcall else ""
         typename = f"std::variant<{','.join(field_typenames)}>"
     elif isinstance(akform, awkward.forms.UnmaskedForm):
-        return _cpp_typename(akform.content, subcall=True)
+        return _cpp_typename(akform.content, subcall=subcall)
     else:
         msg = "This should not have happened. Please report this issue to the Uproot developers."
         raise AssertionError(msg)
-    if not subcall and "UntypedRecord" in typename:
-        typename = ""  # empty types for anything that contains UntypedRecord
     return typename
 
 
@@ -424,11 +431,13 @@ def _build_field_col_records(
             )
         elif isinstance(akform, awkward.forms.IndexedOptionForm):
             type_name = _cpp_typename(akform)
+            type_alias = "" if type_name else "std::optional"
             field = NTuple_Field_Description(
                 parent_fid,
                 uproot.const.RNTupleFieldRole.COLLECTION,
                 field_name,
                 type_name,
+                type_alias=type_alias,
                 field_description=description,
             )
             self._field_records.append(field)

tests/test_1496_rntuple_write_index_arrays.py

@@ -79,13 +79,5 @@ def test_writing_and_reading(tmp_path):
     arrays = obj.arrays()
 
     for f in data.fields:
-        if f in ("record_index64_option_array", "record_index32_option_array"):
-            assert [
-                t[0].tolist() if len(t) > 0 else None for t in arrays[f][:4]
-            ] == data[f].tolist()
-            assert [
-                t[0].tolist() if len(t) > 0 else None for t in arrays[f][4:]
-            ] == data[f].tolist()
-        else:
-            assert arrays[f][:4].tolist() == data[f].tolist()
-            assert arrays[f][4:].tolist() == data[f].tolist()
+        assert arrays[f][:4].tolist() == data[f].tolist()
+        assert arrays[f][4:].tolist() == data[f].tolist()

tests/test_1713_rntuple_optional_records.py

@@ -0,0 +1,140 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/uproot5/blob/main/LICENSE
+
+from __future__ import annotations
+
+import awkward as ak
+import pytest
+
+import uproot
+
+
+def test_optional_record_round_trip(tmp_path):
+    arr = ak.Array([{"o": {"a": 1.0}}, {"o": None}])
+    path = tmp_path / "opt_record.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr
+
+    with uproot.open(path) as f:
+        nt = f["t"]
+        assert nt.arrays().tolist() == arr.tolist()
+        assert nt["o"].array().tolist() == arr["o"].tolist()
+        assert nt["o.a"].array().tolist() == [1.0, None]
+        assert nt.arrays(filter_name="o")["o"].tolist() == arr["o"].tolist()
+        assert nt["o"].typename == ""
+        assert nt["o"].record.type_alias == "std::optional"
+        assert [field.typename for field in nt.all_fields] == ["", "", "double"]
+
+
+def test_optional_record_with_multiple_fields(tmp_path):
+    arr = ak.Array([{"o": {"a": 1.0, "b": 2}}, {"o": None}])
+    path = tmp_path / "opt_multi.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr
+
+    with uproot.open(path) as f:
+        nt = f["t"]
+        assert nt.arrays().tolist() == arr.tolist()
+        assert nt["o.a"].array().tolist() == [1.0, None]
+        assert nt["o.b"].array().tolist() == [2, None]
+        assert nt["o"].typename == ""
+        assert nt["o"].record.type_alias == "std::optional"
+        assert [field.typename for field in nt.all_fields] == [
+            "",
+            "",
+            "double",
+            "std::int64_t",
+        ]
+
+
+def test_optional_record_nested_tuple(tmp_path):
+    arr = ak.Array([{"o": (1.0, {"x": 2})}, {"o": None}])
+    path = tmp_path / "opt_tuple.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr
+
+    with uproot.open(path) as f:
+        nt = f["t"]
+        assert nt.arrays().tolist() == arr.tolist()
+        assert nt["o"].array().tolist() == arr["o"].tolist()
+        assert nt["o"].typename == ""
+        assert nt["o"].record.type_alias == "std::optional"
+        assert [field.typename for field in nt.all_fields] == [
+            "",
+            "",
+            "double",
+            "",
+            "std::int64_t",
+        ]
+
+
+def test_optional_record_nested_list(tmp_path):
+    arr = ak.Array([{"o": [{"a": 1.0}]}, {"o": None}])
+    path = tmp_path / "opt_list.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr
+
+    with uproot.open(path) as f:
+        nt = f["t"]
+        assert nt.arrays().tolist() == arr.tolist()
+        assert nt["o"].array().tolist() == arr["o"].tolist()
+        assert nt["o"].typename == ""
+        assert nt["o"].record.type_alias == "std::optional"
+        assert [field.typename for field in nt.all_fields] == ["", "", "", "double"]
+
+
+def test_optional_record_extend(tmp_path):
+    arr1 = ak.Array([{"o": {"a": 1.0}}, {"o": None}])
+    arr2 = ak.Array([{"o": None}, {"o": {"a": 2.0}}])
+    path = tmp_path / "opt_extend.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr1
+
+    with uproot.update(path) as f:
+        f["t"].extend(arr2)
+
+    with uproot.open(path) as f:
+        expected = ak.concatenate([arr1, arr2])
+        assert f["t"].arrays().tolist() == expected.tolist()
+
+
+def test_optional_record_dask(tmp_path):
+    pytest.importorskip("dask_awkward")
+
+    arr1 = ak.Array([{"o": {"a": 1.0}}, {"o": None}])
+    arr2 = ak.Array([{"o": None}, {"o": {"a": 2.0}}])
+    path = tmp_path / "opt_dask.root"
+
+    with uproot.recreate(path) as f:
+        f["t"] = arr1
+
+    with uproot.update(path) as f:
+        f["t"].extend(arr2)
+
+    expected = ak.concatenate([arr1, arr2])
+    dask_arr = uproot.dask(str(path) + ":t")
+    assert dask_arr.compute().tolist() == expected.tolist()
+
+
+def test_optional_record_ROOT(tmp_path, capfd):
+    ROOT = pytest.importorskip("ROOT")
+    if ROOT.gROOT.GetVersionInt() < 63400:
+        pytest.skip("ROOT version does not support RNTuple v1.0.0.0")
+
+    filepath = str(tmp_path / "opt_record_root.root")
+    arr = ak.Array([{"o": {"a": 1.0}}, {"o": None}])
+
+    with uproot.recreate(filepath) as f:
+        f["ntuple"] = arr
+
+    if ROOT.gROOT.GetVersionInt() < 63600:
+        RT = ROOT.Experimental.RNTupleReader.Open("ntuple", filepath)
+    else:
+        RT = ROOT.RNTupleReader.Open("ntuple", filepath)
+    RT.PrintInfo()
+    out = capfd.readouterr().out
+    assert "ntuple" in out