fix: correctly write all Awkward array types to RNTuple (#1496)

* Fixed writing indexed (option) arrays

* Fixed writing of ListArrays

* Updated test filename

* Minor tweaks

* Fixed issue with 32-bit indices

* Test with 32-bit indices

* Fix 32-bit indices for string arrays

* Implemented custom carry to make sure that form is preserved

* Fixed writing of trimmed ListOffsetArrays

* Switched to using to_packed

* Simpified code a bit

* Read std::optional as IndexedOptionArray

* Fixed warning message

* Fixed typo

Author: ariostas

src/uproot/behaviors/RNTuple.py

@@ -1750,6 +1750,27 @@ def _recursive_find(form, res):
         _recursive_find(form.content, res)
 
 
+def _cupy_insert(arr, obj, value):
+    # obj is assumed to be sorted
+    # both arr and obj are assumed to be flat arrays
+    cupy = uproot.extras.cupy()
+    out_size = arr.size + obj.size
+    out = cupy.empty(out_size, dtype=arr.dtype)
+    src_i = 0
+    dst_i = 0
+    for idx in obj.get():
+        n = idx - src_i
+        if n > 0:
+            out[dst_i : dst_i + n] = arr[src_i : src_i + n]
+            dst_i += n
+            src_i += n
+        out[dst_i] = value
+        dst_i += 1
+    if src_i < arr.size:
+        out[dst_i:] = arr[src_i:]
+    return out
+
+
 def _fill_container_dict(container_dict, content, key, dtype_byte):
     array_library_string = uproot._util.get_array_library(content)
 
@@ -1758,7 +1779,19 @@ def _fill_container_dict(container_dict, content, key, dtype_byte):
     if "cardinality" in key:
         content = library.diff(content)
 
-    if dtype_byte == uproot.const.rntuple_col_type_to_num_dict["switch"]:
+    if "optional" in key:
+        # We need to convert from a ListOffsetArray to an IndexedOptionArray
+        diff = library.diff(content)
+        missing = library.nonzero(diff == 0)[0]
+        missing -= library.arange(len(missing), dtype=missing.dtype)
+        dtype = "int64" if content.dtype == library.uint64 else "int32"
+        indices = library.arange(len(content) - len(missing), dtype=dtype)
+        if array_library_string == "numpy":
+            indices = numpy.insert(indices, missing, -1)
+        else:
+            indices = _cupy_insert(indices, missing, -1)
+        container_dict[f"{key}-index"] = indices
+    elif dtype_byte == uproot.const.rntuple_col_type_to_num_dict["switch"]:
         kindex, tags = uproot.models.RNTuple._split_switch_bits(content)
         # Find invalid variants and adjust buffers accordingly
         invalid = numpy.flatnonzero(tags == -1)

src/uproot/models/RNTuple.py

@@ -441,8 +441,9 @@ def col_form(self, field_id, extra_parameters=None, is_cardinality=False):
             parameters = {"__array__": "string"}
             if extra_parameters is not None:
                 parameters.update(extra_parameters)
+            idx_type = "i32" if rel_crs[0].nbits == 32 else "i64"
             return ak.forms.ListOffsetForm(
-                "i64", inner, form_key=form_key, parameters=parameters
+                idx_type, inner, form_key=form_key, parameters=parameters
             )
         else:
             raise (RuntimeError(f"Missing special case: {field_id}"))
@@ -547,6 +548,11 @@ def field_form(self, this_id, keys, 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"):
+                keyname = keyname + "-optional"
+                return ak.forms.IndexedOptionForm(
+                    idx_type, inner, form_key=keyname, parameters=parameters
+                )
             return ak.forms.ListOffsetForm(
                 idx_type, inner, form_key=keyname, parameters=parameters
             )